from ._material import Material, _Material

class ShellSection:
    def __init__(self, type, material, shape, **kwds):
        if not isinstance(material, dict):
            material = {
                None: material
            }
        self.materials = material
        self.shape = shape
        self.kwds = kwds
        self.type = type

    def _add_to_model(self, model, tag):
        if "fiber" in self.type.lower():
            if isinstance(self.shape, float):
                nip = 5
                mtag = self.materials[None]._mtag
                if mtag is None:
                    raise ValueError("Material must be added to model before section.")
                thickness = self.shape
                fibers = [(mtag, thickness/nip) for _ in range(nip)]
                fibers = [x for pair in fibers for x in pair]
                model.section("LayeredShell", tag, nip, *fibers)
            else:
                raise NotImplementedError("Only uniform fiber sections are implemented.")
        else:
            raise NotImplementedError(f"Shell section type '{self.type}' is not implemented.")



def Section(type, shape, material=None, mixed_type=None, mixed=True, nu=None):
    if not isinstance(material, (dict, _Material,None.__class__)):
        shape, material = material, shape
    
    if material is None:
        material = shape.materials
    elif not isinstance(material, dict):
        material = {
            None: material
        }

    if type == "Elastic":
        return _ElasticSection(type, shape, material, mixed_type=mixed_type, mixed=mixed, nu=nu)
    elif type in {"ShearFiber", "Multiaxial"}:#, "AxialFiber"}:
        return _MultiaxialSection(type, shape, material, mixed_type=mixed_type, mixed=mixed, nu=nu)
    elif type in {"AxialFiber"}:
        return _UniaxialSection(type, shape, material, mixed_type=mixed_type, mixed=mixed, nu=nu)
    else:
        raise NotImplementedError(f"Section type '{type}' is not implemented.")




class _Section:
    def __init__(self, type, shape, material, mixed_type=None, nu=None, mixed=True):
        self.type = type
        self._tag = None

        self._shape = shape
        if material is None:
            material = shape.materials
        self._material = material

        if mixed_type == "energetic":
            mixed_type = "UE"
        elif mixed_type == "geometric":
            mixed_type = "UG"

        self.__mixed_type = mixed_type
        self._mixed = mixed

        from xsection.analysis import SaintVenantSectionAnalysis
        self._shape_model = shape.model
        sv = SaintVenantSectionAnalysis(shape, nu=nu, model=self._shape_model)

        self.GJ = sv.twist_rigidity()
        self.EA = sv.axial_rigidity()
        self.rc = sv.centroid()
        self.sv = sv


    def _mixed_type(self, model):
        if self.__mixed_type is not None:
            return self.__mixed_type
        
        ndm = model._call("getNDM")
        ndf = model._call("getNDF")
        if ndm == 2 or not self._mixed:
            return None 
        elif ndf == 6:
            return "UE"
        elif ndf == 7:
            return "NR"

        
    def _warp_shapes(self, model, mixed_type):
        shape = self._shape

        warp_shapes = []

        ndf = model._call("getNDF")

        if mixed_type in {"UE", "UG", "UT", "NR"} or ndf >= 7:
            sv = self.sv
            warp_shapes.append(sv.solve_twist())
            if mixed_type in {"UE", "UG"} or ndf >= 8:
                warp_shapes.extend(sv.solve_shear())
            elif mixed_type in {"NR"}:
                warp_shapes.append(sv.solve_mixed_twist())

        return warp_shapes

    @property
    def materials(self)->dict:
        return self._material

    def _add_to_model(self, model, tag):
        raise NotImplementedError("Subclasses must implement _add_to_model method.")


class _UniaxialSection(_Section):
    def __init__(self, type, shape, material, **kwds):
        super().__init__(type, shape, material, **kwds)


    def _add_to_model(self, model, tag):
        shape = self._shape

        mixed_type = self._mixed_type(model)

        sdata = {
            "mixed_type": mixed_type,
        }
        if mixed_type in {"UE", "UG", "UT", "NV", "NR", "NT"}:
            GJ = self.GJ
        elif mixed_type is None:
            GJ = 1e4 # TODO

        mixed_type = self._mixed_type(model)
        if False: #mixed_type in {"UE", "UG"}:
            trace = shape.create_trace(form=mixed_type)
            align = trace.shift_shear_gamma()
            align = list(map(float, [align[0,0], align[0,1], align[1,0], align[1,1]]))
            sdata["align"] = align
            sdata["shiftTwist"] = trace.shift_twist_gamma()
            sdata["shiftAxial"] = trace.shift_twist_axial()

        if self.type == "AxialFiber":
            # and ndm >= 3:
            sdata["GJ"] = GJ

        type = "Fiber" # self.type
        model.section(type, tag, **sdata)

        for group in self._material:
            mtag = self._material[group]._mtag
            if mtag is None:
                raise ValueError("Material must be added to model before section.")
            for fiber in shape.create_fibers(warp_type=None, group=group):
                model.fiber(y=fiber["y"], 
                            z=fiber["z"],
                            area=fiber["area"],
                            material=mtag, 
                            section=tag)


class _MultiaxialSection(_Section):
    def __init__(self, type, shape, material, **kwds):
        if not isinstance(material, dict):
            material = {
                None: material
            }
        super().__init__(type, shape, material, **kwds)

        self._shape_model = shape.create_model(materials=material)


    def _add_to_model(self, model, tag):
        shape = self._shape

        mixed_type = self._mixed_type(model)

        sdata = {
            "mixed_type": mixed_type,
        }
        if mixed_type in {"UE", "UG"}:
            trace = self.sv.create_trace(form=mixed_type)
            align = trace.shift_shear_gamma()
            align = list(map(float, [align[0,0], align[0,1], align[1,0], align[1,1]]))
            sdata["align"] = align
            sdata["shiftTwist"] = trace.shift_twist_gamma()
            sdata["shiftAxial"] = trace.shift_twist_axial()


        model.section(self.type, tag, **sdata)

        warp_shapes = self._warp_shapes(model, mixed_type)
        for group in self._material:
            mtag = self._material[group]._mtag
            if mtag is None:
                raise ValueError("Material must be added to model before section.")

            # for fiber in shape.create_fibers(warp_type=self._warp, group=group):
            for fiber in shape._create_fibers(self._shape_model, warp_shapes, group=group):
                model.fiber(**fiber, material=mtag, section=tag)



class _ElasticSection(_Section):
    def __init__(self, type, shape, material, warp=None, **kwds):
        super().__init__("Elastic", shape, material, warp)

        if shape._is_composite:
            self.G = 1 
            self.E = 1
        else:
            self.E = self._material[None]["E"]
            self.G = self._material[None]["G"]

    def _add_to_model(self, model, tag):
        shape = self._shape

        # sdata = {
        #     "mixed_type": self._warp,
        # }

        # if self._warp in {"UE", "UG"}:
        #     trace = shape.create_trace(form=self._warp)
        #     align = trace.shift_shear_gamma()
        #     align = list(map(float, [align[0,0], align[0,1], align[1,0], align[1,1]]))
            # sdata["align"] = align
            # sdata["shiftTwist"] = trace.shift_twist_gamma()
            # sdata["shiftAxial"] = trace.shift_twist_axial()

        cmm = shape.cmm()
        cnn = shape.cnn()
        cnv = shape.cnv()
        cnm = shape.cnm()
        cmw = shape.cmw()
        cvv = shape.cvv()
        cww = shape.cww()
        A = self.EA/self.E
        model.section("ElasticFrame", 
                        tag,
                        E=self.E,
                        G=self.G,
                        A=A,
                        Ay=1*A,
                        Az=1*A,
                        Qy=cnm[0,1],
                        Qz=cnm[2,0],
                        Iy=cmm[1,1]/self.E,
                        Iz=cmm[2,2]/self.E,
                        J = self.GJ/self.G,
                        Cw= cww[0,0],
                        Ry= cnv[1,0],
                        Rz= cnv[2,0],
                        # Sy= cvv[1,1],
                        # Sz= cvv[2,2]
                        Sy= cmw[1,0],#*swch,
                        Sz= cmw[2,0],#*swch
        )