#===----------------------------------------------------------------------===#
#
#         STAIRLab -- STructural Artificial Intelligence Laboratory
#
#===----------------------------------------------------------------------===#
#
import sys
import numpy as np
import warnings

from ..utility import UnimplementedInstance, find_row
from xcsi.csi._material import find_material
from shps.rotor import exp



def _is_truss(frame, csi):
    if "FRAME RELEASE ASSIGNMENTS 1 - GENERAL" in csi:
        release = find_row(csi["FRAME RELEASE ASSIGNMENTS 1 - GENERAL"],
                           Frame=frame["Frame"])
    else:
        return False

    return release and all(release[i] for i in ("TI", "M2I", "M3I", "M2J", "M3J"))



def add_frames(instance):
    csi = instance._tree # Parsed input file
    names = instance.names # class mapping from CSI names to Xara tags
    model = instance.model # the xara.Model
    asm = instance.assembly

    tags = {}
    transforms = _FrameAxes(instance)

    for frame in csi.get("CONNECTIVITY - FRAME",[]):

        if _is_truss(frame, csi):
            instance._log(UnimplementedInstance("Truss", frame))
            continue

        if "IsCurved" in frame and frame["IsCurved"]:
            instance._log(UnimplementedInstance("Frame.Curve", frame))

        


        automesh = find_row(
            csi.get("FRAME AUTO MESH ASSIGNMENTS", []),
            Frame=frame["Frame"]
        )

        #
        # 1) Nodes
        #
        springs = find_row(csi.get("FRAME RELEASE ASSIGNMENTS 2 - PARTIAL FIXITY", []), Frame=frame["Frame"])
        if springs:
            instance._log(
                UnimplementedInstance(f"Frame with springs", 
                                      frame,
                                      table="FRAME RELEASE ASSIGNMENTS 2 - PARTIAL FIXITY"))

        if False: #springs:
            inode = names.define("Joint", "node")
            jnode = names.define("Joint", "node")
            model.node(inode, *asm.create_node(
                find_row(
                    csi["JOINT COORDINATES"],
                    Joint=frame["JointI"]
                )
            ).location)
            model.node(jnode, *asm.create_node(
                find_row(
                    csi["JOINT COORDINATES"],
                    Joint=frame["JointJ"]
                )
            ).location)
            for i,dof in enumerate(["_", "V2", "V3", "_", "M2", "M3"]):
                for n in "IJ":
                    key = f"{dof}{n}"
                    if key in springs and springs[key] > 0.0:
                        mat = names.define("Material", "material")
                        model.uniaxialMaterial("Elastic", mat, springs[key])
                        model.spring(
                            inode if n == "I" else jnode,
                            i,
                            k=springs[key]
                        )
            nodes = (inode, jnode)
        else:
            nodes = (
                names.identify("Joint", "node", frame["JointI"]),
                names.identify("Joint", "node", frame["JointJ"])
            )

        xi = np.array(model.nodeCoord(nodes[0]))
        xj = np.array(model.nodeCoord(nodes[1]))
        for nodes in _mesh_nodes((xi, xj), nodes, automesh, instance):

            tags[frame["Frame"]] = create_frame(frame, nodes,
                                                instance=instance,
                                                transforms=transforms)

    # TODO: remove this
    names._library["frame_tags"] = tags

    return []


def _mesh_nodes(x, nodes, automesh, instance):
    """ 
    return a list of node pairs defining the meshed frame elements.

    TODO: Get tolerance from, eg,
    
TABLE:  "PREFERENCES - DIMENSIONAL"
   MergeTol=0.008333333333333333   FineGrid=1   Nudge=1   SelectTol=3   SnapTol=12   SLineThick=2   PLineThick=4   MaxFont=8   MinFont=6   AutoZoom=10   ShrinkFact=70   TextFileLen=240
 
    """
    model = instance.model
    names = instance.names
    csi = instance._tree
    xi, xj = x
    ni, nj = nodes
    u = xj - xi
    Lu = float(np.linalg.norm(u))
    if automesh is None:
        return [nodes]
    if not any(value for key, value in automesh.items() if key != "Frame"):
        return [nodes]
    if automesh.get("AtJoints", False):
        # Find all joints (nodes) lying on the segment [xi, xj]
        intersecting = []

        for tag in model.getNodeTags():
            if tag in nodes:
                continue

            xn = np.array(model.nodeCoord(tag), dtype=float)

            # Distance from point to line
            d_line = np.linalg.norm(np.cross(u, xn - xi)) / Lu
            if d_line >= 1e-6:
                continue

            # Ensure it's within the segment (projection parameter t in [0,1])
            t = float(np.dot(xn - xi, u) / (Lu * Lu))
            if t <= 0.0 or t >= 1.0:
                continue

            # Also ensure point is actually close to the segment endpoints numerically
            # (optional; projection test is usually enough)
            intersecting.append((t, tag))

        # Sort by position along the element
        intersecting.sort(key=lambda p: p[0])

        # Build the ordered node list: i, internal..., j
        chain = [ni] + [tag for _, tag in intersecting] + [nj]

        # Convert into consecutive pairs
        pairs = [(chain[k], chain[k + 1]) for k in range(len(chain) - 1)]

        # (Optional) sanity: remove accidental duplicates (same tag repeated)
        pairs = [(a, b) for a, b in pairs if a != b]
        return pairs
    else:
        instance._log(UnimplementedInstance("Frame.AutoMesh.NonJoint", automesh))
    return [nodes]



def create_frame(frame, nodes, instance, transforms):
    csi = instance._tree # Parsed input file
    names = instance.names # class mapping from CSI names to Xara tags
    model = instance.model # the xara.Model
    asm = instance.assembly

    frame_obj = asm.create_frame(frame)

    options = instance.options.find(Frame=frame["Frame"])

    #
    # 1) Nodes
    #
    springs = find_row(csi.get("FRAME RELEASE ASSIGNMENTS 2 - PARTIAL FIXITY", []), Frame=frame["Frame"])
    if springs:
        instance._log(
            UnimplementedInstance(f"Frame with springs", 
                                    frame,
                                    table="FRAME RELEASE ASSIGNMENTS 2 - PARTIAL FIXITY"))

    if False: #springs:
        inode = names.define("Joint", "node")
        jnode = names.define("Joint", "node")
        model.node(inode, *asm.create_node(
            find_row(
                csi["JOINT COORDINATES"],
                Joint=frame["JointI"]
            )
        ).location)
        model.node(jnode, *asm.create_node(
            find_row(
                csi["JOINT COORDINATES"],
                Joint=frame["JointJ"]
            )
        ).location)
        for i,dof in enumerate(["_", "V2", "V3", "_", "M2", "M3"]):
            for n in "IJ":
                key = f"{dof}{n}"
                if key in springs and springs[key] > 0.0:
                    mat = names.define("Material", "material")
                    model.uniaxialMaterial("Elastic", mat, springs[key])
                    model.spring(
                        inode if n == "I" else jnode,
                        i,
                        k=springs[key]
                    )
        nodes = (inode, jnode)

    xi = np.array(model.nodeCoord(nodes[0]))
    xj = np.array(model.nodeCoord(nodes[1]))

    if np.linalg.norm(xj - xi) < 1e-10:
        instance._log(UnimplementedInstance("Frame.ZeroLength", frame))
        return None
    #
    # 2) Section
    #
    sect_asgn = find_row(csi["FRAME SECTION ASSIGNMENTS"],
                            Frame=frame["Frame"])
    
    section = frame_obj.section

    #
    # 3) Geometric Transformation
    #
    transform: int = transforms.find(frame, xi, xj)

    #
    # 4) Mass
    #

    total_mass = frame_obj.mass / frame_obj.length


    #
    # 5) Create the element
    #

    releases = find_row(csi.get("FRAME RELEASE ASSIGNMENTS 1 - GENERAL", []), Frame=frame["Frame"])
    if releases and any(releases[i] for i in ("TI", "M2I", "M3I", "TJ", "M2J", "M3J")):
        names.log(UnimplementedInstance("Frame with releases", 
                                        frame,
                                        table="FRAME RELEASE ASSIGNMENTS 1 - GENERAL"))


    # Case 1: Inelastic frame
    if instance is not None and options.inelastic:

        if section.has_multiaxial:
            stag = names.identify("AnalSect", "section", f"Multiaxial/{section.name}")
            tag = names.define("Frame", "element", frame["Frame"])

            model.element("ForceFrame",
                            tag,
                            nodes,
                            transform=transform,
                            section=stag,
                            mass=total_mass,
                            shear=int(options.use_shear)
            )
            return tag


    # Case 2: Elastic Prismatic

    # TODO: probably dont use "SectionType" from this table; it looks like its 
    # superceded by "Shape" in Section Properties 01 table.
    elif frame_obj.prismatic:

        section = names.identify("AnalSect", "section", sect_asgn["AnalSect"])
        if section is None:
            warnings.warn(f"No section found for {sect_asgn['AnalSect']}")
            return None
        
        tag = names.define("Frame", "element", frame["Frame"])

        model.element("ForceFrame",
                        tag,
                        nodes,
                        section=section,
                        # n=8,
                        # gauss_type="Legendre",
                        transform=transform,
                        mass=total_mass,
                        shear=int(options.use_shear),
                        comment=f"Frame {frame['Frame']}"
        )
        return tag

    # Case 3: Elastic Non-prismatic

    elif sect_asgn["NPSectType"] == "Default" and (
            integr := names.identify("AnalSect", "integration", sect_asgn["AnalSect"])):
        # Non-prismatic sections
        tag = names.define("Frame", "element", frame["Frame"])
        model.element("ForceFrame",
                        tag,
                        nodes,
                        transform,
                        integr,
                        mass=total_mass,
                        shear=int(options.use_shear)
        )
        return tag

    else:
        names.log(UnimplementedInstance("FrameSection.NPSectType", sect_asgn["NPSectType"]))
        return None



def _find_frame_mass(csi, frame, sect_info, total_length):

    is_nonprismatic = (sect_info["Shape"] == "Nonprismatic")

    if "FRAME ADDED MASS ASSIGNMENTS" in csi:
        row = find_row(csi["FRAME ADDED MASS ASSIGNMENTS"],
                        Frame=frame["Frame"])
        additional_mass = row["MassPerLen"] if row else 0.0
    else:
        additional_mass = 0.0

    if not is_nonprismatic:
        #
        # Prismatic section
        #
        A = sect_info["Area"]  # cross‐sectional area
        mat_info = find_material(csi, sect_info["Material"])
        rho = mat_info["density"]  # material density (mass / volume)

        # self‐weight mass per length = area × density
        self_weight_mpl = A * rho

    else:
        #
        # Nonprismatic section
        #
        # look for total mass and total length in
        #

        # total mass for the entire nonprismatic section
        total_mass = sect_info.get("TotalMass", 0)

        # self‐weight mass per length = total mass / total length
        self_weight_mpl = total_mass / total_length

    # Combine any assigned mass per length with self‐weight mass per length
    total_mass = self_weight_mpl + additional_mass

    return {
        "total_mass": total_mass,
        "mass_type": "lumped"
    }



def _orient(xi, xj, angle):
    """
    Calculate the coordinate transformation vector.
    xi is the location of node I, xj node J,
    and `angle` is the rotation about the local axis in degrees

    By default local axis 2 is always in the 1-Z plane, except if the object
    is vertical and then it is parallel to the global X axis.
    The definition of the local axes follows the right-hand rule.
    """

    # The local 1 axis points from node I to node J
    e1 = xj - xi
    dx, dy, dz = e1
    e1 = e1 / np.linalg.norm(e1)

    # Global z
    E3 = np.array([0, 0, 1])

    # In Sap2000, if the element is vertical, the local y-axis is the same as the
    # global x-axis
    if dx == 0 and dy == 0:
        e2 = np.array([1, 0, 0])
        e3 = np.cross(e1, e2)

    # Otherwise, the plane composed of the local x-axis and the local
    # y-axis is a vertical plane. In this
    # case, the local z-axis is the cross product of the local
    # x-axis and the global z-axis.
    else:
        e3 = -np.cross(E3, e1)
        e2 = np.cross(e3, e1)

    e3 = e3 / np.linalg.norm(e3)
    e2 = e2 / np.linalg.norm(e2)


    # Rotate the local axis using the Rodrigue rotation formula
    # convert from degrees to radians
    angle = angle / 180 * np.pi
    # return exp(angle * e1 / np.linalg.norm(e1)) @ e3
    e3r = e3 * np.cos(angle) + np.cross(e1, e3) * np.sin(angle)
    # Finally, the normalized local z-axis is returned
    return e3r / np.linalg.norm(e3r)


class _FrameAxes:
    def __init__(self, instance):
        self._ins = instance
        self._ndm = instance.assembly.ndm
        self._transforms = {}
        self._model = instance.model
        self._ttag = 1
    
    def find(self, frame, xi, xj):
        if self._ndm != 3:
            vecxz = [0, 0, 1]
        else:
            vecxz = self._orient(frame, xi, xj)

        if tuple(vecxz) not in self._transforms:
            # TODO: Use frame options to set linear/pdelta/corotational 
            self._model.geomTransf("Linear", self._ttag, *vecxz)
            self._transforms[tuple(vecxz)] = self._ttag
            self._ttag += 1
        return self._transforms[tuple(vecxz)]
    

    def _orient(self, frame, xi, xj):
        csi = self._ins._tree
        if "FRAME LOCAL AXES ASSIGNMENTS 1 - TYPICAL" not in csi:
            return _orient(xi, xj, 0.0)

        row = find_row(csi["FRAME LOCAL AXES ASSIGNMENTS 1 - TYPICAL"],
                        Frame=frame["Frame"])

        # TODO: implement advanced axes and uncomment this line.
        if row and row.get("AdvanceAxes", False):
            # self._ins._log(UnimplementedInstance("Frame.LocalAxes.AdvancedAxes", frame))
            return self._orient_advanced(frame)

        angle = row["Angle"] if row else 0.0

        return _orient(xi, xj, angle-90)


    def _orient_advanced(self, frame):
        tables = self._ins._tree
        table = find_row(
            tables.get("FRAME LOCAL AXES ASSIGNMENTS 2 - ADVANCED", []),
            Frame=frame["Frame"]
        )
        if table is None:
            raise ValueError("Advanced frame axes requested but no ADVANCED table found.")
        
        R0 = self._ins.assembly._coordinate_system(table["PlCoordSys"]).rotation_matrix

        return R0[:,0]

    def _orient_angle(self, xi, xj, angle):
        """
        Calculate the coordinate transformation vector.
        xi is the location of node I, xj node J,
        and `angle` is the rotation about the local axis in degrees

        By default local axis 2 is always in the 1-Z plane, except if the object
        is vertical and then it is parallel to the global X axis.
        The definition of the local axes follows the right-hand rule.
        """

        # The local 1 axis points from node I to node J
        e1 = xj - xi
        dx, dy, dz = e1
        e1 = e1 / np.linalg.norm(e1)

        # Global z
        E3 = np.array([0, 0, 1])

        # In Sap2000, if the element is vertical, the local y-axis is the same as the
        # global x-axis
        if dx == 0 and dy == 0:
            e2 = np.array([1, 0, 0])
            e3 = np.cross(e1, e2)

        # Otherwise, the plane composed of the local x-axis and the local
        # y-axis is a vertical plane. In this
        # case, the local z-axis is the cross product of the local
        # x-axis and the global z-axis.
        else:
            e3 = -np.cross(E3, e1)
            e2 = np.cross(e3, e1)

        e3 = e3 / np.linalg.norm(e3)
        e2 = e2 / np.linalg.norm(e2)


        # Rotate the local axis using the Rodrigue rotation formula
        # convert from degrees to radians
        angle = angle / 180 * np.pi
        # return exp(angle * e1 / np.linalg.norm(e1)) @ e3
        e3r = e3 * np.cos(angle) + np.cross(e1, e3) * np.sin(angle)
        # Finally, the normalized local z-axis is returned
        return e3r / np.linalg.norm(e3r)
