import xara
from xcsi.csi import find_row, find_rows
from collections import defaultdict
from dataclasses import dataclass

class _Pattern:
    @dataclass
    class JointLoad:
        joint: int
        force: list

    @dataclass
    class FrameLoad:
        pass

    def __init__(self):
        self.joint_forces = []
        self.frame_forces = []
        self.shell_forces = []


class LoadPattern:
    """
    "LOAD CASE DEFINITIONS"
    "LOAD PATTERN DEFINITIONS"

    "JOINT LOADS - FORCE"
    "FRAME LOADS - DISTRIBUTED" OK
    "FRAME LOADS - GRAVITY"
    "FRAME LOADS - POINT" TODO

    "CABLE LOADS - DISTRIBUTED"
    """
    def __init__(self, name, assembly):
        self.name = name
        csi = assembly.tables
        self._assembly = assembly

        joint_sp = defaultdict(dict)
        for joint in find_rows(csi.get("JOINT LOADS - GROUND DISPLACEMENT", []),
                               LoadPat=name):
            restraints = find_row(csi.get("JOINT RESTRAINT ASSIGNMENTS", []),
                                 Joint=joint["Joint"])
            for key in "U1", "U2", "U3", "R1", "R2", "R3":
                if key in joint:
                    if joint[key] == 0.0 and (restraints and not restraints[key]):
                        continue
                    joint_sp[joint["Joint"]][key] = joint[key]

        self._joint_sp = joint_sp

        self._read_joint_loads(csi)
        self._read_frame_loads(csi)


    def _read_joint_loads(self, csi):

        joint_force = defaultdict(lambda: [0.0 for _ in range(6)])
        for joint in find_rows(csi.get("JOINT LOADS - FORCE", []),
                               LoadPat=self.name):
            for i,key in enumerate(("F1", "F2", "F3", "M1", "M2", "M3")):
                if key in joint:
                    joint_force[joint["Joint"]][i] = joint[key]

        self._joint_force = joint_force


    def _apply_joint_loads(self, asm, ptag):
        model: xara.Model = asm.model
        names = asm.names

        for joint, force in self._joint_force.items():
            if not any(force):
                continue
            model.load(
                names.identify("JointBase", "node", joint),
                tuple(force),
                pattern=ptag
            )

    def _read_frame_loads(self, csi):
        name = self.name 

        # Frame loads
        self._frame_loads = defaultdict(list)
        for assign in find_rows(csi.get("FRAME LOADS - DISTRIBUTED", []),
                               LoadPat=name):
            
            if assign.get("RelDistA", 0) != 0:
                raise ValueError("Relative distance not supported for distributed frame loads")
            if assign.get("RelDistB", 1) != 1:
                raise ValueError("Relative distance not supported for distributed frame loads")
            if assign.get("DistType", "RelDist") != "RelDist":
                raise ValueError("Only relative distance supported for distributed frame loads")
            
            sign = 1.0
            direction = None
            try:
                direction = int(assign["Dir"])-1
            except ValueError:
                if assign["Dir"].lower() == "gravity":
                    direction = 2  # Z direction
                    sign = -1.0
                # elif assign["Dir"].lower() == "proj z":
                #     # Test 1-001, Case 4
                #     direction = 2  # Z direction
                #     frm = self._assembly.create_frame(
                #         find_row(
                #             csi.get("CONNECTIVITY - FRAME", []),
                #             Frame=assign["Frame"]
                #         )
                #     )
                #     sign = None
                else:
                    raise ValueError(f"Unsupported direction {assign['Dir']} for frame load {assign['Frame']}")

            force = [0,0,0]
            moment = [0,0,0]
            if assign.get("Type", None) == "Force":
                if assign["FOverLA"] != assign["FOverLB"]:
                    raise ValueError("FOverLA and FOverLB must be equal for distributed frame loads")

                force[direction] = sign * assign["FOverLA"]

            elif assign.get("Type", None) == "Moment":
                if assign["MOverLA"] != assign["MOverLB"]:
                    raise ValueError("MOverLA and MOverLB must be equal for distributed frame loads")

                raise NotImplementedError("Moment loads not implemented for frames")
            else:
                raise ValueError(f"Unsupported frame load type: {assign.get('Type', None)}")

            self._frame_loads[assign["Frame"]].append({
                "type": "Uniform",
                "basis": assign["CoordSys"].lower(),
                "force": force,
                "moment": moment,
                "offset": [0.0, 0.0, 0.0],
            })

        for assign in find_rows(csi.get("FRAME LOADS - GRAVITY", []),
                               LoadPat=name):
            sign = 1.0
            direction = None
            multipliers = [
                assign.get("MultiplierX", 0.0),
                assign.get("MultiplierY", 0.0),
                assign.get("MultiplierZ", 0.0),
            ]
            if sum(int(m!=0) for m in multipliers if m != 0.0) != 1.0:
                raise ValueError("Only one multiplier can be non-zero for distributed frame loads")
            direction = multipliers.index(next(m for m in multipliers if m != 0.0))
            sign = multipliers[direction]
            

            force = [0,0,0]
            moment = [0,0,0]
            frame = self._assembly.create_frame(
                find_row(
                    csi.get("CONNECTIVITY - FRAME", []),
                    Frame=assign["Frame"]
                )
            )
            mass = frame.mass
            gravity = self._assembly.units.gravity
            force[direction] = float(gravity*sign*mass/frame.length)

            self._frame_loads[assign["Frame"]].append({
                "type": "Uniform",
                "basis": assign["CoordSys"].lower(),
                "force": force,
                "moment": moment,
                "offset": [0.0, 0.0, 0.0],
            })

        for assign in find_rows(csi.get("FRAME LOADS - POINT", []),
                               LoadPat=name):
            if "RelDist" not in assign:
                raise ValueError("RelDist must be specified for point frame loads")

            xa = assign["RelDist"]
            sign = 1.0
            if assign["Dir"].lower() == "gravity":
                direction = 2  # Z direction
                sign = -1.0
            else:
                try:
                    direction = int(assign["Dir"])-1
                except ValueError:
                    raise ValueError(f"Unsupported direction {assign['Dir']} for frame load {assign['Frame']}")

            force = [0,0,0]
            moment = [0,0,0]
            if assign["Type"] == "Force":
                force[direction] = sign * assign["Force"]
            else:
                raise NotImplementedError("Moment loads not implemented for frames")

            self._frame_loads[assign["Frame"]].append({
                "type":  "Point",
                "basis":  assign["CoordSys"].lower(),
                "force":  force,
                "moment": moment,
                "offset": [xa, 0.0, 0.0],
            })

    def _apply_frame_loads(self, asm, ptag):
        model: xara.Model = asm.model
        names = asm.names
        for frame, loads in self._frame_loads.items():
            etag = names.identify("Frame", "element", frame)
            for load in loads:
                if load["type"] in {"Uniform", "Point"}:
                    model.eleLoad(
                        "Frame",
                        load["type"],
                        element=[etag],
                        offset=load["offset"],
                        basis=load["basis"],
                        force=load["force"],
                        couple=load["moment"],
                        pattern=ptag
                    )
                else:
                    raise NotImplementedError(f"Unsupported frame load type: {load['type']}")


    def apply(self, asm):
        """
        Apply the load pattern to the model.
        """
        model: xara.Model = asm.model
        dofs = ["U1", "U2", "U3", "R1", "R2", "R3"]

        ptag = asm.names.define("LoadPattern", "pattern", self.name)

        model.pattern("Plain", ptag, "Linear")
        for joint, sp in self._joint_sp.items():
            if not sp:
                continue
            for dof, value in sp.items():
                model.sp(
                    asm.names.identify("JointBase", "node", joint),
                    dofs.index(dof)+1,
                    value,
                    pattern=ptag
                )
        
        self._apply_frame_loads(asm, ptag)
        self._apply_joint_loads(asm, ptag)


    def __str__(self):
        return f"LoadPattern(name={self.name})"



class AmbientAcceleration:
    def __init__(self, name, csi):
        self.name = name
        self._csi = csi

        self.direction = None
        self.accel = None

        for row in find_rows(csi.get("LOAD PATTERN DEFINITIONS", []),
                             LoadPat=name):
            self.direction = row["Dir"]
            self.accel = row["Accel"]

        if self.direction is None or self.accel is None:
            raise ValueError(f"Ambient acceleration pattern {name} not found")

    def apply(self, asm):
        """
        Apply the ambient acceleration to the model.
        """
        model = asm.model
        dofs = ["U1", "U2", "U3"]

        ptag = asm.names.define("LoadPattern", "pattern", self.name)

        model.pattern("UniformExcitation", ptag, dofs.index(self.direction))
        model.ambient(self.direction, self.accel, pattern=ptag)

    def __str__(self):
        return f"AmbientAcceleration(name={self.name}, direction={self.direction}, accel={self.accel})"
