import h5py
import numpy as np
import inspect
import fabio
#from pyFAI.azimuthalIntegrator import AzimuthalIntegrator
from pyFAI.integrator.azimuthal import AzimuthalIntegrator
from pySAXS.LS.QIE import *
from pySAXS.LS import absolute
import pySAXS.LS.SAXSparametersXML as SAXSparameters
from pySAXS.guisaxs import pySaxsColors
from collections.abc import Iterable
import matplotlib.pyplot as plt
from pySAXS.mcsas import MCtools
import pandas
from pySAXS.guisaxs.dataset import *



EDF_TO_GEOMETRY={'directDist':'SampleDistance','centerX':'Center_1','centerY':'Center_2',\
                 'pixelX':'PSize_1','pixelY':'PSize_1'}#, 'wavelength':'Wavelength'}
EDF_TO_GEOMETRY_CONVERSION={'directDist':1000,'pixelX':1e6,'pixelY':1e6}

INTEGRATION_NPT=1000
INTEGRATION_PARAMETERS_REQUIRED={'npt':INTEGRATION_NPT,'error_model':"poisson",'unit':"q_A^-1"}

EDF_TO_ABSOLUTE={'SampleDistance':'D','TransmittedFlux':'TransmittedFlux','ExposureTime':'time','Wavelength':'wavelength','PSize_1':'pixel_size'}
EDF_TO_ABSOLUTE_CONVERSION={'SampleDistance':100,'TransmittedFlux':1,'Wavelength':1e10,'PSize_1':100}



class Node:
    def __init__(self, hdf5_file, node_path):
        """
        Represents a node in an HDF5 file.

        :param hdf5_file: Path to the HDF5 file.
        :param node_path: Path to the node within the file.
        """
        self.hdf5_file = hdf5_file
        self.node_path = node_path

    def add_child(self, child_name, value):
        """
        Adds a child node with a value.

        :param child_name: Name of the child node.
        :param value: Value to store. If a NumPy array, it's stored as a dataset; otherwise, as an attribute.
        :return: New Node object representing the child.
        """
        with h5py.File(self.hdf5_file, "a") as hdf:
            parent = hdf[self.node_path]

            if child_name in parent:
                del parent[child_name]  # Delete existing node

            if isinstance(value, np.ndarray):
                parent.create_dataset(child_name, data=value)
            else:
                child = parent.create_group(child_name)
                try:
                    child.attrs["value"] = value
                except TypeError:
                    child.attrs["value"] = str(value)

        return Node(self.hdf5_file, f"{self.node_path}/{child_name}")

    def add_child_dict(self, data_dict):
        """
        Adds multiple child nodes from a dictionary.

        :param data_dict: Dictionary where keys are child names and values are the data to store.
        """
        for key, value in data_dict.items():
            self.add_child(key, value)

    def get_children_with_values(self):
        """
        Retrieves children and their values as a dictionary.

        :return: Dictionary with children names as keys and their values.
        """
        children = {}

        with h5py.File(self.hdf5_file, "r") as hdf:
            parent = hdf[self.node_path]

            for child_name in parent:
                child = parent[child_name]

                if isinstance(child, h5py.Dataset):
                    children[child_name] = child[()]
                elif "value" in child.attrs:
                    children[child_name] = child.attrs["value"]
                else:
                    children[child_name] = None

        return children

    def __repr__(self):
        return f"Node(path='{self.node_path}')"

class FairSAXSDataHDF5:
    """
    Manages a SAXS data workflow using HDF5.
    """

    def __init__(self, filename="workflow.hdf5", group="root", rwa='w'):
        """
        Initializes the HDF5 file and root group.

        :param filename: Name of the HDF5 file.
        :param group: Name of the root group.
        :param rwa: File access mode.
        """
        self.filename = filename
        with h5py.File(self.filename, rwa) as hdf:
            if group not in hdf:
                hdf.create_group(group)
        self.order = []
        self.node_order = []
        self.current = None
        self.group = group

    def get_group(self):
        return self.group

    def change_group(self, group_name):
        self.group = group_name

    def add_node(self, node_name, function_name=None, values=None):
        """
        Adds a node to the root and returns a Node object.

        :param node_name: Name of the node.
        :param function_name: Name of the function associated with the node.
        :param values: Values to store in the node.
        :return: Node object representing the new node.
        """
        self.order.append(function_name)
        self.node_order.append(node_name)
        self.current = function_name

        with h5py.File(self.filename, "a") as hdf:
            root = hdf[self.group]
            if node_name in root:
                del root[node_name]

            node = root.create_group(node_name)
            if values is not None:
                node.attrs["values"] = values

            if 'metadata_order' in root:
                del root['metadata_order']
            metadata = root.create_group('metadata_order')
            metadata.attrs["execution_order"] = ",".join(self.node_order)
            metadata.attrs["execution_function"] = ",".join(self.order)

        return Node(self.filename, f"{self.group}/{node_name}")

    def get_node(self, node_name):
        """
        Retrieves an existing node as a Node object.

        :param node_name: Name of the node.
        :return: Node object.
        """
        with h5py.File(self.filename, "r") as hdf:
            if f"{self.group}/{node_name}" not in hdf:
                raise ValueError(f"Node '{node_name}' does not exist.")
        return Node(self.filename, f"{self.group}/{node_name}")

    def get_node_values(self, node_name):
        """
        Retrieves the values of a node.

        :param node_name: Name of the node.
        :return: Dictionary of node values.
        """
        node = self.get_node(node_name)
        return node.get_children_with_values()

    def get_info(self, parent_group, child_group=None):
        """
        Retrieves attributes, datasets, or child nodes from a group.

        :param parent_group: Main group name.
        :param child_group: Optional subgroup name.
        :return: Dictionary with attributes, dataset values, or sub-nodes.
        """
        with h5py.File(self.filename, "r") as hdf:
            if parent_group not in hdf:
                raise ValueError(f"Group '{parent_group}' does not exist.")

            node = hdf[parent_group]
            if child_group:
                if child_group not in node:
                    raise ValueError(f"Subgroup '{child_group}' not found in '{parent_group}'.")
                node = node[child_group]

            result = {}
            if len(node.attrs) > 0:
                result["attributes"] = {attr: node.attrs[attr] for attr in node.attrs}

            if isinstance(node, h5py.Dataset):
                result["data"] = node[()]

            if isinstance(node, h5py.Group):
                result["children"] = list(node.keys())

            return result

    def display_tree(self):
        """
        Prints the tree structure and values stored in HDF5.
        """
        with h5py.File(self.filename, "r") as hdf:
            def print_group(group, indent=0):
                for key, item in group.items():
                    if isinstance(item, h5py.Group):
                        attrs = item.attrs
                        attr_str = f": '{attrs['value']}'" if "value" in attrs else ""
                        print("  " * indent + f"Node: {key}{attr_str}")
                        print_group(item, indent + 1)
                    elif isinstance(item, h5py.Dataset):
                        if item.size > 10:
                            print("  " * indent + f"Dataset: {key}, Shape: {item.shape}, Dtype: {item.dtype}")
                        else:
                            print("  " * indent + f"Dataset: {key}, Value: {item[()]}")

            print("Current data treatment:", self.current)
            print("Order:", self.order)
            print_group(hdf[self.group])

class FairSAXSDataHDF5_Xeuss3(FairSAXSDataHDF5):
    """
    Extended version of FairSAXSDataHDF5 with additional methods for Xeuss3.
    """

    def __init__(self, filename="workflow_extended.hdf5", group="root", rwa='w',store_images=True):
        super().__init__(filename, group, rwa)
        self.geometry = None
        self.img = None
        self.ai = None
        self.mask = None
        self.store_images=store_images

    def read_edf_image(self, filename):
        """
        Reads an EDF image file and stores metadata.

        :param filename: Path to the EDF file.
        :return: Fabio image object.
        """
        current_function_name = inspect.currentframe().f_code.co_name
        nexus_node = self.add_node("edf", current_function_name, values=filename)
        nexus_node.add_child("filename", filename)

        img = fabio.open(filename)
        image = img.data
        if self.store_images:
            nexus_node.add_child("image", image)

        header = img.header
        nexus_node.add_child("sample name", header['Comment'])
        nexus_node.add_child_dict(header)

        self.img = img
        return img

    def open_mask(self, mask_name):
        """
        Opens a mask EDF file and stores it.

        :param mask_name: Path to the mask file.
        """
        current_function_name = inspect.currentframe().f_code.co_name
        mask_file = fabio.open(mask_name)
        mask_node = self.add_node("mask", current_function_name, values=mask_name)
        mask_node.add_child('mask_name', mask_name)
        if self.store_images:
            mask_node.add_child('mask_edf', mask_file)
        self.mask = mask_file

    def define_geometry(self):
        """
        Defines the geometry from integration parameters.
        """
        current_function_name = inspect.currentframe().f_code.co_name
        self.ai = AzimuthalIntegrator()

        dgeom = {}
        set_geom = {}
        for key, value in EDF_TO_GEOMETRY.items():
            conversion = EDF_TO_GEOMETRY_CONVERSION.get(key, 1)
            dgeom[key] = float(self.img.header[value]) * conversion
            set_geom[f'set_{key}'] = dgeom[key]

        wl = float(self.img.header['Wavelength'])
        self.ai.set_wavelength(wl)
        self.geometry = self.ai.setFit2D(**dgeom)

        pyfai_geometry_node = self.add_node("pyFAI_geometry", current_function_name)
        pyfai_geometry_node.add_child_dict(self.geometry.get_config())
        pyfai_geometry_node.add_child_dict(set_geom)

    def integrate(self):
        """
        Performs radial integration using pyFAI.

        :return: Integration result as a tuple (q, i, e).
        """
        current_function_name = inspect.currentframe().f_code.co_name
        integrate_node = self.add_node("radial_integration", current_function_name)

        dai = INTEGRATION_PARAMETERS_REQUIRED.copy()
        dai['mask'] = self.mask.data

        res1 = self.ai.integrate1d(self.img.data, **dai)
        q, i, e = res1

        integrate_node.add_child('q', q)
        integrate_node.add_child('i', i)
        integrate_node.add_child('e', e)
        integrate_node.add_child('q_unit', INTEGRATION_PARAMETERS_REQUIRED['unit'])

        return res1

    def _append_qie(self, current_method, current_function_name, q, i, e, info_dict=None):
        """
        Appends QIE data to the HDF5 file.

        :param current_method: Name of the method.
        :param current_function_name: Name of the function.
        :param q: Q values.
        :param i: Intensity values.
        :param e: Error values.
        :param info_dict: Additional information dictionary.
        """
        node = self.add_node(current_method, current_function_name)
        node.add_child('q', q)
        node.add_child('i', i)
        node.add_child('e', e)

        if info_dict:
            node.add_child_dict(info_dict)

    def integrate_qie(self):
        """
        Integrates and returns a QIE object.

        :return: QIE object with integration results.
        """
        q, i, e = self.integrate()
        return QIE(q, i, e)

    def read_edf_and_integrate(self,edf_file,mask_file):
        self.read_edf_image(edf_file)
        self.open_mask(mask_file)
        self.define_geometry()
        self.integrate_qie()

    def read_pysaxs_dataset(self,xml_dataset,data_name):
        """
        read data in pysaxs dataset
        """
        #try:
        q, i, e = getQIEFromXMLFile(xml_dataset, data_name)
        current_function_name = inspect.currentframe().f_code.co_name
        self._append_qie("read_pysaxs_data",current_function_name,q,i,e)
        #except:
        #    print(f"error on task read_pysaxs_dataset {xml_dataset} for {data_name}")

    def get_child_value(self, group, child_name):
        """
        Retrieves the value of a child node.

        :param group: Parent group name.
        :param child_name: Child node name.
        :return: Dictionary with attributes or dataset values.
        """
        with h5py.File(self.filename, "r") as hdf:
            if group not in hdf:
                raise ValueError(f"Group '{group}' does not exist.")

            parent_group = hdf[group]
            if child_name not in parent_group:
                raise ValueError(f"Child '{child_name}' not found in group '{group}'.")

            child = parent_group[child_name]
            result = {}

            if len(child.attrs) > 0:
                result["attributes"] = {attr: child.attrs[attr] for attr in child.attrs}

            if isinstance(child, h5py.Dataset):
                result["data"] = child[()]

            return result



    def get_value(self, path):
        """
        Retrieves the value or attributes of a dataset or group.

        :param path: Full hierarchical path.
        :return: Dictionary with attributes and/or dataset values.
        """
        with h5py.File(self.filename, "r") as hdf:
            if path not in hdf:
                raise ValueError(f"Path '{path}' does not exist in the HDF5 file.")

            node = hdf[path]
            result = {}

            if len(node.attrs) > 0:
                result["attributes"] = {attr: node.attrs[attr] for attr in node.attrs}

            if isinstance(node, h5py.Dataset):
                result["data"] = node[()]

            return result

    def get_qie(self, group=None):
        """
        Returns a QIE object associated with a group.

        :param group: Group name.
        :return: QIE object.
        """
        if group is None:
            group = self.group

        r = self.get_info(group, 'metadata_order')
        last = r['attributes']['execution_order'].split(',')[-1]
        r_c = self.get_child_value(group, last)

        q = self.get_value(f"{group}/{last}/q")['data']
        i = self.get_value(f"{group}/{last}/i")['data']
        e = self.get_value(f"{group}/{last}/e")['data']

        return QIE(q, i, e)

    def get_qi(self):
        '''
        return last qi
        '''
        return self.get_qie().get_qi()

    def absolute(self,additional=None,SILENT=True):
        #print(additional)
        #print(SILENT)
        if additional is not None:
            #print(additional_kargs[0])
            self.absolute_from_edf(additional=additional,SILENT=SILENT)
        else:
            self.absolute_from_edf()

    def absolute_from_edf(self, group=None, edf_group='edf',additional=None,SILENT=True):
        """
        Calculates absolute values from EDF headers.

        :param group: Group name.
        :param edf_group: EDF group name.
        """
        if group is None:
            group = self.group

        qie = self.get_qie(group)
        abso_param = SAXSparameters.SAXSparameters()
        self.group = group

        edf_dict = self.get_node_values(edf_group)
        for key, value in EDF_TO_ABSOLUTE.items():
            conversion = EDF_TO_ABSOLUTE_CONVERSION.get(key, 1)
            if key in edf_dict:
                val = float(edf_dict[key])
                abso_param.set(value, conversion * val)
        if additional is not None:
            ADDITIONAL_ABSOLUTE=additional
            for param_key, param_value in ADDITIONAL_ABSOLUTE.items():
                abso_param.set(param_key, param_value)
        #print(f'silent {SILENT}')
        if SILENT:
            set_SILENT = 'SILENT'
        else:
            print(f"------ Absolute intensities from EDF for {self.group} ------------")
            set_SILENT = None
        #print(f'silent {set_SILENT}')
        q, i, e = qie.get_values()
        abso = absolute.absolute(q=q, i=i, ierr=e, parameters=abso_param.parameters,printout=set_SILENT)
        new_i, new_e = abso.calculate()

        current_function_name = inspect.currentframe().f_code.co_name
        dict_absolute = abso_param.get_dict()
        self._append_qie('absolute', current_function_name, q, new_i, new_e, dict_absolute)

    def subtract(self, group2,group1=None):
        """
        Subtracts data from two groups.

        :param group1: First group name.
        :param group2: Second group name.
        :return: Resulting QIE object.
        """
        if group1 is None:
            group1 = self.group
        if group2 is None:
            raise ValueError("Group2 is not provided.")

        qie1 = self.get_qie(group1)
        qie2 = self.get_qie(group2)
        try:
            qie_final = raw_subtract(qie1, qie2)
            q, i, e = qie_final.get_values()
        except:
            #print(f"erreur avec {group1} et {group2}")
            qie_final = smart_subtract(qie1, qie2)
            q, i, e = qie_final.get_values()



        current_function_name = inspect.currentframe().f_code.co_name
        self._append_qie('subtraction', current_function_name, q, i, e, {'what': f'{group1}-{group2}'})

        return qie_final

    def clip_qmax(self,qmax,group=None):
        """
        Clips data above a maximum q value.

        :param group: Group name.
        :param qmax: Maximum q value.
        :return: Clipped QIE object.
        """
        if group is None:
            group = self.group

        qie = self.get_qie(group)
        qie_final = qie.clip_qmax(qmax)
        q, i, e = qie_final.get_values()

        current_function_name = inspect.currentframe().f_code.co_name
        self._append_qie('clip_qmax', current_function_name, q, i, e, {'qmax': qmax})

        return qie_final

    def clip_qmin(self,qmin, group=None, ):
        """
        Clips data below a minimum q value.

        :param group: Group name.
        :param qmin: Minimum q value.
        :return: Clipped QIE object.
        """
        if group is None:
            group = self.group

        qie = self.get_qie(group)
        qie_final = qie.clip_qmin(qmin)
        q, i, e = qie_final.get_values()

        current_function_name = inspect.currentframe().f_code.co_name
        self._append_qie('clip_qmin', current_function_name, q, i, e, {'qmin': qmin})

        return qie_final

    def merge(self, other_group, cut1, cut2):
        """
        Merges data from two groups.

        :param other_group: Second group name.
        :param cut1: First cut value.
        :param cut2: Second cut value.
        :return: Merged QIE object.
        """
        qie0 = self.get_qie()
        qie1 = self.get_qie(other_group)
        qie_final = merge(qie0, qie1, cut1, cut2)
        q, i, e = qie_final.get_values()

        current_function_name = inspect.currentframe().f_code.co_name
        self._append_qie('merging', current_function_name, q, i, e, {'cut1': cut1, 'cut2': cut2, 'with_group': other_group})

        return qie_final

    def mc_sas(self,kwargs):
        """
        do the mc sas optimization

        """
        print(kwargs)
        q,i,e=self.get_qie().get_values()
        res=MCtools.easyMcSas(q,i,e,**kwargs)
        current_function_name = inspect.currentframe().f_code.co_name
        d=kwargs
        d.update(res)
        self._append_qie('mc_sas', current_function_name, q, i, e,d)
        return res


class workflow_exec:
    def __init__(self, hdf5_file,task_args,SILENT=True,store_images=True,colors_inc=None):
        self.task_args = task_args
        self.fairSAXSdata = {}
        self.col = pySaxsColors.pySaxsColors(inc=colors_inc)
        self.colors = {}
        self.SILENT=SILENT
        rwa = 'w'
        for key, values in self.task_args.items():
            if not self.SILENT:
                print(f'Init {key}')
            self.fairSAXSdata[key] = FairSAXSDataHDF5_Xeuss3(hdf5_file, group=key, rwa=rwa,store_images=store_images)
            rwa = 'a'

            if colors_inc is not None:
                self.colors[key] = self.col.getColor()
            else:
                self.colors[key]=None

    def run(self, task_name, task_number):
        for key, value in self.task_args.items():
            if task_number<len(value):
                if value[task_number] is not None:
                    if not self.SILENT:
                        print(f'processing {task_name} for {key} with {value}')
                    # run the task
                    task = getattr(self.fairSAXSdata[key], task_name)
                    if isinstance(value[task_number], dict):
                        # keywords arguments
                        task(**value[task_number])
                    elif isinstance(value[task_number], str):
                        task(value[task_number])
                    elif isinstance(value[task_number], Iterable):
                        task(*value[task_number])
                    else:
                        task(value[task_number])

    def get_color(self, key):
        return self.colors[key]

    def get_data(self, key):
        return self.fairSAXSdata[key]

    def plot(self, task_number=0,everyone=False):
        for key, value in self.task_args.items():
            if task_number < len(value):
                if everyone or value[task_number] is not None:
                    plt.plot(*self.fairSAXSdata[key].get_qi(), color=self.colors[key], label=key)


    def plot_superposed(self,task_number=0,everyone=False):
        k=0
        for key, value in self.task_args.items():
            if task_number < len(value):
                if everyone or value[task_number] is not None:
                    q,i=self.fairSAXSdata[key].get_qi()
                    plt.plot(q,i*10**k, color=self.colors[key], label=key)
                    k+=1
    def plot_iqn(self,task_number=0,n=4,everyone=False):

        for key, value in self.task_args.items():
            if task_number < len(value):
                if everyone or value[task_number] is not None:
                    q,i=self.fairSAXSdata[key].get_qi()
                    plt.plot(q,i*q**n, color=self.colors[key], label=key)
                    


    def plot_fill_between(self,task_number=0,everyone=False):
        for key, value in self.task_args.items():
            if task_number < len(value):
                if everyone or value[task_number] is not None:
                    plt.plot(*self.fairSAXSdata[key].get_qi(), color=self.colors[key], label=key)
                    plt.fill_between(*self.fairSAXSdata[key].get_qie().get_fill_between() , edgecolor=self.colors[key], facecolor=self.colors[key], alpha=0.2, antialiased=True)

    def get_keys(self,task_number=0,everyone=False):
        k=[]
        for key, value in self.task_args.items():
            if task_number < len(value):
                if everyone or value[task_number] is not None:
                    k.append(key)
        return k

    def to_excel(self,filename,task_number=0,everyone=False):
        '''
        save all the datas corresponding to a task number
        '''
        with pandas.ExcelWriter(filename) as writer:
            for key, value in self.task_args.items():
                if task_number < len(value):
                    if everyone or value[task_number] is not None:
                        q,i,e=self.fairSAXSdata[key].get_qie().get_values()
                        df = pandas.DataFrame({'q':q,'i':i,'e':e})
                        df.to_excel(writer,sheet_name=key,index=False)


    def to_separate_csv(self,filename_base,task_number=0,everyone=False):
        '''
        save all the datas corresponding to a task number
        '''
        for key, value in self.task_args.items():
             if task_number < len(value):
                 if everyone or value[task_number] is not None:
                        q,i,e=self.fairSAXSdata[key].get_qie().get_values()
                        df = pandas.DataFrame({'q':q,'i':i,'e':e})
                        valid_key = key.replace(" ", "_")  # Remplace les espaces par des underscores
                        filename=filename_base+'_'+valid_key+'.csv'
                        df.to_csv(filename,index=False)