#===----------------------------------------------------------------------===#
#
#         STAIRLab -- STructural Artificial Intelligence Laboratory
#
#===----------------------------------------------------------------------===#
#
import xara
import numpy as np
import warnings
from ..csi.utility import UnimplementedInstance, find_row, find_rows


class _LinkMaterial:
    _link_tables = {
    #   LinkType                                    TableName
        "Linear"               : "LINK PROPERTY DEFINITIONS 02 - LINEAR",
        "??                  " : "LINK PROPERTY DEFINITIONS 03 - MULTILINEAR",
        "Damper - Exponential" : "LINK PROPERTY DEFINITIONS 04 - DAMPER",
        "Gap"                  : "LINK PROPERTY DEFINITIONS 05 - GAP",
        "????                " : "LINK PROPERTY DEFINITIONS 06 - HOOK",
        "?????               " : "LINK PROPERTY DEFINITIONS 07 - RUBBER ISOLATOR",
        "??????              " : "LINK PROPERTY DEFINITIONS 08 - SLIDING ISOLATOR",
        "Plastic (Wen)"        : "LINK PROPERTY DEFINITIONS 10 - PLASTIC (WEN)",
        "MultiLinear Plastic"  : "LINK PROPERTY DEFINITIONS 11 - MULTILINEAR PLASTIC",
    }
    def __init__(self, prop_table, instance):
        self.name = prop_table["LinkProp"]
        self._gen_prop_table = prop_table
        self._dofs = []
        self._tags = []
        self._instance = instance
        self.create(instance)

    @property 
    def property_table(self):
        return self._link_tables.get(self._gen_prop_table["LinkType"], None)

    @property 
    def material_dofs(self):
        return tuple(["U1", "U2", "U3", "R1", "R2", "R3"].index(i)+1 for i in self._dofs)
    
    @property
    def material_tags(self):
        return tuple(self._tags)


class _LinkMaterial02(_LinkMaterial):
    def create(self, instance):
        csi = instance._tree
        model: xara.Model = instance.model
        names  = instance.names

        for dof in ["U1", "U2", "U3", "R1", "R2", "R3"]:
            table = find_row(csi["LINK PROPERTY DEFINITIONS 02 - LINEAR"],
                            Link=self.name, DOF=dof)
            if table is None:
                continue

            if dof in {"U1", "U2", "U3"}:
                stiff = table["TransKE"]
                damp  = table["TransCE"]
            elif dof in {"R1", "R2", "R3"}:
                stiff = table["RotKE"]
                damp  = table["RotCE"]
            
            if damp != 0:
                instance._log(UnimplementedInstance(f"LinkMaterial02.Damping", table))
            
            tag = names.define("LinkProp", "material", f"Link/{self.name}.{dof}")
            model.uniaxialMaterial(
                "Elastic",
                tag,
                stiff
            )
            self._dofs.append(dof)
            self._tags.append(tag)

        return

class _LinkMaterial05(_LinkMaterial):
    def create(self, instance):
        csi = instance._tree
        model  = instance.model
        names  = instance.names

        for dof in ["U1", "U2", "U3", "R1", "R2", "R3"]:
            table = find_row(csi["LINK PROPERTY DEFINITIONS 05 - GAP"],
                            Link=self.name, DOF=dof)
            if table is None:
                continue

            if dof in {"U1", "U2", "U3"}:
                stiff = table.get("TransKE", None)
                damp  = table.get("TransCE", None)
            
            if damp != 0:
                instance._log(UnimplementedInstance(f"LinkMaterial02.Damping", table))

            if not table["Fixed"] and stiff is not None:
                assert stiff is not None, f"Stiffness must be defined for non-fixed gap link {self.name} DOF {dof}"
                tag = names.define("LinkProp", "material", f"Link/{self.name}.{dof}")
                model.uniaxialMaterial(
                    "Elastic",
                    tag,
                    stiff
                )
            else:
                tag = None

            self._dofs.append(dof)
            self._tags.append(tag)

        return

# _LinkMaterial04 - Damper (Exponential)
    # for damper in csi.get("LINK PROPERTY DEFINITIONS 04 - DAMPER", []):
    #     # TODO: implement dampers
    #     conv.log(UnimplementedInstance("Link.Damper", damper))
    #     continue
    #     name = damper["Link"]
    #     stiff = damper["TransK"]
    #     dampcoeff = damper["TransC"]
    #     exp = damper["CExp"]
    #     model.eval(f"uniaxialMaterial ViscousDamper {mat_total} {stiff} {dampcoeff}' {exp}\n")

    #     dof = damper["DOF"]
    #     library["link_materials"][name][dof] = mat_total
    #     mat_total += 1

class _LinkMaterial10(_LinkMaterial):
    def create(self, instance):
        csi = instance._tree
        model = instance.model
        names  = instance.names

        for dof in ["U1", "U2", "U3", "R1", "R2", "R3"]:

            table = find_row(csi["LINK PROPERTY DEFINITIONS 10 - PLASTIC (WEN)"],
                             Link=self.name, DOF=dof)

            # Skip if property not defined for this DOF
            if table is None:
                continue

            fy = table["YieldForce"]
            ratio = table["ratio"]
            stiff = table["TransK"]
            tag = names.define("LinkProp", "material", f"Link/{self.name}.{dof}")
            # TODO: use the actual Bouc-Wen model here
            model.uniaxialMaterial(
                "Steel01",
                tag,
                fy,
                stiff,
                ratio
            )
            self._dofs.append(dof)
            self._tags.append(tag)
        return
    
class _LinkMaterial11(_LinkMaterial):
    def create(self, instance):
        csi = instance._tree
        model = instance.model
        names  = instance.names
        for dof in ["U1", "U2", "U3", "R1", "R2", "R3"]:

            # tables = find_rows(csi["LINK PROPERTY DEFINITIONS 11 - MULTILINEAR PLASTIC"],
            #                   Link=self.name, DOF=dof)
            
            tables = find_row(csi["LINK PROPERTY DEFINITIONS 11 - MULTILINEAR PLASTIC"],
                              Link=self.name, DOF=dof, Point=1)
            if tables is None:
                continue
            
            K = tables["TransKE"]
            tag = names.define("LinkProp", "material", f"Link/{self.name}.{dof}")
            model.uniaxialMaterial(
                "Elastic",
                tag,
                K
            )
            self._dofs.append(dof)
            self._tags.append(tag)
        return
        points = []
        table = find_row(csi["LINK PROPERTY DEFINITIONS 11 - MULTILINEAR PLASTIC"],
                        Link=self.name, DOF=self.dof)
        
        if len(points) < 2:
            instance._log(UnimplementedInstance(f"LinkMaterial11.InsufficientPoints", table))
            return

        tag = names.define("LinkProp", "material", self.key)

        # https://opensees.berkeley.edu/wiki/index.php/MultiLinear_Material


class _LinkMaterialNA(_LinkMaterial):
    def create(self, instance):
        return


_link_tables = {
#   LinkType                                    TableName
    "Linear"               : _LinkMaterial02, # "LINK PROPERTY DEFINITIONS 02 - LINEAR",
    "??                  " : _LinkMaterialNA, # "LINK PROPERTY DEFINITIONS 03 - MULTILINEAR",
    "Damper - Exponential" : _LinkMaterialNA, # "LINK PROPERTY DEFINITIONS 04 - DAMPER",
    "Gap"                  : _LinkMaterial05, # "LINK PROPERTY DEFINITIONS 05 - GAP",
    "????                " : _LinkMaterialNA, # "LINK PROPERTY DEFINITIONS 06 - HOOK",
    "?????               " : _LinkMaterialNA, # "LINK PROPERTY DEFINITIONS 07 - RUBBER ISOLATOR",
    "??????              " : _LinkMaterialNA, # "LINK PROPERTY DEFINITIONS 08 - SLIDING ISOLATOR",
    "Plastic (Wen)"        : _LinkMaterial10, # "LINK PROPERTY DEFINITIONS 10 - PLASTIC (WEN)",
    "MultiLinear Plastic"  : _LinkMaterial11, # "LINK PROPERTY DEFINITIONS 11 - MULTILINEAR PLASTIC",
}


def create_links(instance):
    csi = instance._tree
    model: xara.Model = instance.model
    names  = instance.names

    link_props = _create_link_materials(instance)

    # Create Single Joint Links
    _create_single_links(instance, link_props)

    # Create Double Joint Links
    _create_double_links(instance, link_props)


def _orient(xi, xj, deg):
    # Calculate the direction vector of the link element
    # Where a is the node number of node i, b is the node number of node j, and degree is the user-specified local axis
    # ------------------------------------------------------------------------------
    d_x, d_y, d_z = xj - xi

    # Local 1-axis points from node I to node J
    l_x = np.array([d_x, d_y, d_z])
    # Global z-axis
    g_z = np.array([0, 0, 1])

    # In SAP2000, if the link is vertical, the local y-axis is the same as the
    # global x-axis, and the local z-axis can be obtained by crossing the local
    # x-axis with the local y-axis
    if d_x == 0 and d_y == 0:
        l_y = np.array([1, 0, 0])
        l_z = np.cross(l_x, l_y)

    # In other cases, the plane formed by the local x-axis and the local y-axis
    # is a vertical plane (i.e., the normal vector is horizontal), and the
    # local z-axis can be obtained by crossing the local x-axis with the global
    # z-axis
    else:
        l_z = np.cross(l_x, g_z)

    # The local axis may also be rotated using the Rodrigues' rotation formula
    angle = deg / 180 * np.pi
    l_z_rot = l_z * np.cos(angle) + np.cross(l_x, l_z) * np.sin(angle)
    # The rotated local y-axis can be obtained by crossing the rotated local z-axis with the local x-axis
    l_y_rot = np.cross(l_z_rot, l_x)
    # Finally, return the normalized local y-axis
    return l_y_rot / np.linalg.norm(l_y_rot)


def _create_link_materials(instance):
    csi = instance._tree

    links = {}

    for link in csi.get("LINK PROPERTY ASSIGNMENTS",[]):
        if link["LinkType"] not in _link_tables:
            instance._log(UnimplementedInstance(f"Link.Prop.LinkType={link['LinkType']}", link))
            continue
        
        if link["Link"] in links:
            continue

        cls = _link_tables[link["LinkType"]](link, instance)
        links[link["LinkProp"]] = cls
    return links


def _create_single_links(instance, link_props):

    csi = instance._tree
    model = instance.model
    names  = instance.names

    for link in csi.get("CONNECTIVITY - LINK",[]):

        assign = find_row(csi["LINK PROPERTY ASSIGNMENTS"],
                          Link=link["Link"])

        if assign.get("LinkJoints", "") != "SingleJoint":
            continue

        nodes = (
            names.identify("JointBase", "node", link["JointI"]),
            names.identify("JointBase", "node", link["JointJ"])
        )

        assert nodes[0] == nodes[1], f"SingleJoint link must connect the same node, got {link['JointI']} and {link['JointJ']} for {link['Link']}"

        # if any(not (isinstance(node, int) or node.isdigit()) for node in nodes):
        #     conv.log(UnimplementedInstance(f"Joint: non-integer nodes", assign))
        #     continue

        if assign["LinkProp"] not in link_props:
            instance._log(UnimplementedInstance(f"Link.Prop.Missing", assign))
            continue
        if isinstance(link_props[assign["LinkProp"]], _LinkMaterialNA):
            instance._log(UnimplementedInstance(f"Link.Prop.LinkType={assign['LinkType']}", assign))
            continue
        
        # Get instance of _LinkMaterial
        link_prop = link_props[assign["LinkProp"]]


        # TODO: Implement soil springs

        flags = [0 for _ in range(6)]
    
        for dof in link_prop.material_dofs:
            flags[dof-1] = 1

        model.fix(names.identify("JointBase", "node", link["JointI"]), tuple(flags))


def _create_double_links(instance, link_props):
    csi = instance._tree
    model: xara.Model = instance.model
    names  = instance.names

    for link in csi.get("CONNECTIVITY - LINK",[]):

        nodes = (
            names.identify("Joint", "node", link["JointI"]),
            names.identify("Joint", "node", link["JointJ"])
        )

        # if any(not (isinstance(node, int) or node.isdigit()) for node in nodes):
        #     conv.log(UnimplementedInstance(f"Joint: non-integer nodes", assign))
        #     continue

        assign = find_row(csi["LINK PROPERTY ASSIGNMENTS"],
                          Link=link["Link"])

        # Only process double-joint links here
        if assign.get("LinkJoints", "") != "TwoJoint":
            continue

        #
        # Get mats and dofs
        #
        if assign["LinkProp"] not in link_props:
            instance._log(UnimplementedInstance(f"Link.Prop.Missing", assign))
            continue
        if isinstance(link_props[assign["LinkProp"]], _LinkMaterialNA):
            instance._log(UnimplementedInstance(f"Link.Prop.LinkType={assign['LinkType']}", assign))
            continue

        link_prop = link_props[assign["LinkProp"]]

        # if the dof is fixed, then the tag will be None. we need to filter those

        mats = [
            mat for mat in link_prop.material_tags if mat is not None
        ]
        link_dofs = [
            dof for i,dof in enumerate(link_prop.material_dofs)
            if link_prop.material_tags[i] is not None
        ]
        fixed_dofs = [dof for i,dof in enumerate(link_prop.material_dofs)
            if link_prop.material_tags[i] is None
        ]
        for dof in fixed_dofs:
            model.fix(names.identify("JointBase", "node", link["JointI"]), dof=dof)
            model.fix(names.identify("JointBase", "node", link["JointJ"]), dof=dof)

        if len(link_dofs) == 0:
            instance._log(UnimplementedInstance(f"DoubleLink.DOFS", assign))
            continue
        if len(mats) != len(link_dofs):
            instance._log(UnimplementedInstance(f"DoubleLink.MatsVsDofs", assign))
            continue

        # Check whether there are any zero-length link elements
        xi = np.array(model.nodeCoord(nodes[0]))
        xj = np.array(model.nodeCoord(nodes[1]))

        distance = np.linalg.norm(xj - xi)
        zero_length_threshold = 1e-6

        #
        # Get axes and orientation
        #
        axes = find_row(csi.get("LINK LOCAL AXES ASSIGNMENTS 1 - TYPICAL", []),
                        Link=link["Link"])

        orient_vector = None  # Default value

        if axes:
            if axes["AdvanceAxes"]:
                # Handle advanced axes
                axes_advance = find_row(csi.get("LINK LOCAL AXES ASSIGNMENTS 2 - ADVANCED", []),
                                        Link=link["Link"])

                # Common advanced axes setup
                orient_vector = (
                    axes_advance["AxVecX"], axes_advance["AxVecY"], axes_advance["AxVecZ"],
                    axes_advance["PlVecX"], axes_advance["PlVecY"], axes_advance["PlVecZ"],
                )

                # Additional validation for non-zero-length links
                if distance > zero_length_threshold:
                    # Calculate node-based orientation
                    orient_vector_from_nodes = _orient(xi, xj, axes["Angle"])

                    # Calculate y-axis from advanced axes
                    ax_vec = np.array([axes_advance["AxVecX"], axes_advance["AxVecY"], axes_advance["AxVecZ"]])
                    pl_vec = np.array([axes_advance["PlVecX"], axes_advance["PlVecY"], axes_advance["PlVecZ"]])

                    x_axis = ax_vec / np.linalg.norm(ax_vec)
                    pl_projection = pl_vec - np.dot(pl_vec, x_axis) * x_axis
                    y_axis_advanced = pl_projection / np.linalg.norm(pl_projection)

                    # Validate and update orientation
                    if np.allclose(y_axis_advanced, orient_vector_from_nodes, atol=1e-6):
                        orient_vector = tuple(_orient(xi, xj, axes["Angle"]))
                    else:
                        orient_vector = tuple(_orient(xi, xj, axes["Angle"]))
                        warnings.warn(f"Orientation mismatch in link {link['Link']}")
#                       raise ValueError(f"Orientation mismatch in link {link['Link']}")

            else:
                # Handle typical axes (non-advanced)
                orient_vector = tuple(_orient(xi, xj, axes["Angle"]))


        #
        # Create the link
        #
        tag = names.define("Link", "element", link["Link"])
        if orient_vector is not None:
            model.element("TwoNodeLink",
                      tag,
                      nodes,
                      mat=mats,
                      dir=link_dofs,
                      orient=orient_vector
            )
        else:
            model.element("TwoNodeLink",
                      tag,
                      nodes,
                      mat=mats,
                      dir=link_dofs
            )

    return
