import xara
from collections import defaultdict


class ParaModel:
    """
    A wrapper around xara.Model that allows us to use JAX's vjp/jvp
    with the model's methods.
    """
    def __init__(self, *args, parameters=0, **kwargs):
        self._model = xara.Model(*args, **kwargs)
        for i in range(parameters):
            self._model.parameter(i)

        self._p = defaultdict(lambda: {
            "material": set(),
            "section":  set(),
            "element":  set()
        })


    def __getattr__(self, name):
        return getattr(self._model, name)


    def analyze(self, *args, **kwargs):
        self.sensitivityAlgorithm("-computeAtEachStep")
        for (index,name),param in self._p.items():
            for level in "element",:
                for tag in param[level]:
                    self._model.addToParameter(index, level, tag, name)

        return self._model.analyze(*args, **kwargs)


    def _pull_parameters(self, level, tag, args, kwds):
        new_args = []
        new_kwds = {}
        for arg in args:
            if isinstance(arg, _Parameter):
                self._p[(arg.index,arg.name)][level].add(tag)
                new_args.append(float(arg))
            else:
                new_args.append(arg)

        for key, arg in kwds.items():
            if isinstance(arg, _Parameter):
                self._p[(arg.index,arg.name)][level].add(tag)
                new_kwds[key] = float(arg)
            else:
                new_kwds[key] = arg

        if "material" in kwds:
            for param in self._p.values():
                if kwds["material"] in param["material"]:
                    param[level].add(tag)

        if "section" in kwds:
            for param in self._p.values():
                if kwds["section"] in param["section"]:
                    param[level].add(tag)

        return new_args, new_kwds


    def fiber(self, y, z, area, material, section):
        self._pull_parameters("fiber", )


    def section(self, *args, **kwds):
        tag = args[1]
        args, kwds = self._pull_parameters("section", tag, args, kwds)
        return self._model.section(*args, **kwds)

    def material(self, *args, **kwds):
        tag = args[1]
        args, kwds = self._pull_parameters("material", tag, args, kwds)
        return self._model.material(*args, **kwds)
    
    def element(self, *args, **kwds):
        tag = args[1]
        args, kwds = self._pull_parameters("element", tag, args, kwds)
        return self._model.element(*args, **kwds)


def _element_sections(model, element):
    sections = model.eleResponse(element, "sectionTags")
    if isinstance(sections, (int,float)):
        sections = [int(sections)]
    elif isinstance(sections, (list, tuple)):
        pass
    else:
        raise ValueError(f"Unexpected section response: {sections}")
    return sections


class _Parameter(float):
    __slots__ = ("index", "name")

    def __new__(cls, value, index=0, name=None):
        obj = super().__new__(cls, value)
        obj.index = int(index)
        obj.name  = str(name)
        return obj

    def __repr__(self):
        return f"{{Parameter {self.index} {super().__repr__()}}}"


class Chain:
    def __init__(self,type,tag,parameter,material=None, parent=None):
        self.type = type
        self.tag = tag
        self.parameter = parameter
        self.parent = parent

        # fiber chain
        self._material = material

        
        self.elements = set()
        self.sections = defaultdict(lambda: {"fibers": set()})

    def __repr__(self):
        return f"Chain(type={self.type}, tag={self.tag}, parameter={self.parameter}, sections={self.sections}, elements={self.elements})"
    
    def activate(self, model):
        if self.type == "material":
            for element in self.elements:
                sections = _element_sections(model, element)
                for i,section in enumerate(sections):
                    if section in self.sections:
                        # for fiber in self.sections[section]["fibers"]:
                        model.addToParameter(self.parameter.index, 
                                            "element", element, 
                                            "section", i+1,
                                            # "fiber", fiber,
                                            "material",
                                            self.tag,
                                            self.parameter.name
                                            )
                    
        elif self.type == "fiber":
            for element in self.elements:
                sections = _element_sections(model, element)
                for i,section in enumerate(sections):
                    if section == self.parent[1]:
                        model.addToParameter(self.parameter.index, 
                                            "element", element, 
                                            "section", i+1,
                                            "fiber", self.tag,
                                            self.parameter.name
                                            )
        elif self.type == "section":
            for element in self.elements:
                sections = _element_sections(model, element)
                for i,section in enumerate(sections):
                    # if section == self.tag:
                    model.addToParameter(self.parameter.index, 
                                        "element", element, 
                                        "section", i+1,
                                        self.parameter.name
                                        )
        else:
            raise NotImplementedError(f"Chain type '{self.type}' is not implemented.")


class ChainModel:
    """
    A wrapper around xara.Model that allows us to use JAX's vjp/jvp
    with the model's methods.
    """
    def __init__(self, *args, parameters=0, **kwargs):
        self._model = xara.Model(*args, **kwargs)
        for i in range(parameters):
            self._model.parameter(i)

        self._chains = {}


    def __getattr__(self, name):
        return getattr(self._model, name)


    def analyze(self, *args, **kwargs):
        self.sensitivityAlgorithm("-computeAtEachStep")
        for chain in self._chains.values():
            chain.activate(self._model)

        return self._model.analyze(*args, **kwargs)


    def _pull_parameters(self, chain_type, tag, args, kwds, parent=None):
        new_args = []
        new_kwds = {}
        for arg in args:
            if isinstance(arg, _Parameter):
                if arg.index not in self._chains:
                    self._chains[arg.index] = Chain(chain_type, tag, arg, 
                                                                 material=kwds.get("material", None),
                                                                 parent=parent
                                                            )
                new_args.append(float(arg))
            else:
                new_args.append(arg)

        for key, arg in kwds.items():
            if isinstance(arg, _Parameter):
                if arg.index not in self._chains:
                    self._chains[arg.index] = Chain(chain_type, tag, arg, 
                                                    material=kwds.get("material", None),
                                                    parent=parent
                                               )
                new_kwds[key] = float(arg)
            else:
                new_kwds[key] = arg

        return new_args, new_kwds


    def fiber(self, y, z, area, material:int, section:int):
        i = self._model.fiber(y=float(y), z=float(z), area=float(area), material=material, section=section)
        self._pull_parameters("fiber", i, [y,z,area], {"material": material}, parent=("section", section))

        for chain in self._chains.values():
            if chain.type == "material" and material == chain.tag:
                chain.sections[section]["fibers"].add(i)


    def section(self, *args, **kwds):
        tag = args[1]
        args, kwds = self._pull_parameters("section", tag, args, kwds)

        return self._model.section(*args, **kwds)

    def material(self, *args, **kwds):
        tag = args[1]
        args, kwds = self._pull_parameters("material", tag, args, kwds)
        return self._model.material(*args, **kwds)
    
    def element(self, *args, **kwds):
        tag = args[1]
        args, kwds = self._pull_parameters("element", tag, args, kwds)

        if "section" in kwds:
            for chain in self._chains.values():
                if chain.type in {"material"} : #chain.parent is None:
                    if kwds["section"] in chain.sections:
                        chain.elements.add(tag)
                
                elif chain.type == "section" and kwds["section"] == chain.tag:
                    chain.elements.add(tag)

                # type == "fiber"
                elif chain.parent and chain.parent[0] == "section" and kwds["section"] == chain.parent[1]:
                    chain.elements.add(tag)       

        return self._model.element(*args, **kwds)

try:
    from jax import tree_util
    def _flatten_model(model):
        return (), model

    def _unflatten_model(aux_data, _children):
        return aux_data

    tree_util.register_pytree_node(ParaModel, _flatten_model, _unflatten_model)
    tree_util.register_pytree_node(ChainModel, _flatten_model, _unflatten_model)
except ImportError:
    # JAX is not available, so we won't register the pytree node
    pass
