import math
from dataclasses import dataclass, field
from typing import Optional, Tuple
import cadquery as Cq
from nhf import Material, Role
from nhf.build import Model, target, assembly
from nhf.parts.springs import TorsionSpring
from nhf.parts.joints import TorsionJoint
from nhf.parts.box import Hole, MountingBox, box_with_centre_holes
import nhf.utils

TOL = 1e-6

@dataclass
class ShoulderJoint(Model):

    height: float = 60.0
    torsion_joint: TorsionJoint = field(default_factory=lambda: TorsionJoint(
        radius_track=18,
        radius_rider=18,
        groove_radius_outer=16,
        groove_radius_inner=13,
        track_disk_height=5.0,
        rider_disk_height=5.0,
        # M8 Axle
        radius_axle=3.0,
        spring=TorsionSpring(
            # inner diameter = 9
            radius=9/2 + 1.2,
            thickness=1.3,
            height=7.5,
        ),
        rider_slot_begin=0,
        rider_n_slots=2,
        rider_slot_span=15,
    ))

    # On the parent side, drill vertical holes

    parent_conn_hole_diam: float = 6.0
    # Position of the holes relative
    parent_conn_hole_pos: list[Tuple[float, float]] = field(default_factory=lambda: [
        (15, 8),
        (15, -8),
    ])

    parent_lip_length: float = 25.0
    parent_lip_width: float = 30.0
    parent_lip_thickness: float = 5.0
    parent_lip_ext: float = 40.0
    parent_lip_guard_height: float = 8.0

    # Measured from centre of axle
    child_lip_length: float = 45.0
    child_lip_width: float = 20.0
    child_conn_hole_diam: float = 6.0
    # Measured from centre of axle
    child_conn_hole_pos: list[float] = field(default_factory=lambda: [25, 35])
    child_core_thickness: float = 3.0

    # Rotates the torsion joint to avoid collisions or for some other purpose
    axis_rotate_bot: float = 225.0
    axis_rotate_top: float = -225.0

    directrix_id: int = 0

    def __post_init__(self):
        assert self.parent_lip_length * 2 < self.height

    def parent(self, top: bool = False) -> Cq.Assembly:
        joint = self.torsion_joint
        # Thickness of the lip connecting this joint to the wing root
        assert self.parent_lip_width <= joint.radius_track * 2
        assert self.parent_lip_ext > joint.radius_track

        lip_guard = (
            Cq.Solid.makeBox(
                self.parent_lip_ext,
                self.parent_lip_width,
                self.parent_lip_guard_height)
            .located(Cq.Location((0, -self.parent_lip_width/2 , 0)))
            .cut(Cq.Solid.makeCylinder(joint.radius_track, self.parent_lip_guard_height))
        )
        lip = MountingBox(
            length=self.parent_lip_length,
            width=self.parent_lip_width,
            thickness=self.parent_lip_thickness,
            holes=[
                Hole(x=self.height / 2 - x, y=y)
                for x, y in self.parent_conn_hole_pos
            ],
            hole_diam=self.parent_conn_hole_diam,
            generate_side_tags=False,
        )
        # Flip so the lip's holes point to -X
        loc_axis = Cq.Location((0,0,0), (0, 1, 0), -90)
        # so they point to +X
        loc_dir = Cq.Location((0,0,0), (0, 0, 1), 180)
        loc_pos = Cq.Location((self.parent_lip_ext - self.parent_lip_thickness, 0, 0))

        rot = -self.axis_rotate_top if top else self.axis_rotate_bot

        result = (
            Cq.Assembly()
            .add(joint.track(), name="track",
                 loc=Cq.Location((0, 0, 0), (0, 0, 1), rot))
            .add(lip_guard, name="lip_guard")
            .add(lip.generate(), name="lip", loc=loc_pos * loc_dir * loc_axis)
        )
        return result

    @target(name="parent-bot")
    def parent_bot(self) -> Cq.Assembly:
        return self.parent(top=False)
    @target(name="parent-top")
    def parent_top(self) -> Cq.Assembly:
        return self.parent(top=True)

    @property
    def child_height(self) -> float:
        """
        Calculates the y distance between two joint surfaces on the child side
        of the shoulder joint.
        """
        joint = self.torsion_joint
        return self.height - 2 * joint.total_height + 2 * joint.rider_disk_height

    @target(name="child")
    def child(self) -> Cq.Assembly:
        """
        Creates the top/bottom shoulder child joint
        """

        joint = self.torsion_joint
        assert all(r > joint.radius_rider for r in self.child_conn_hole_pos)
        assert all(r < self.child_lip_length for r in self.child_conn_hole_pos)

        # Half of the height of the bridging cylinder
        dh = self.height / 2 - joint.total_height
        core_start_angle = 30
        core_end_angle1 = 90
        core_end_angle2 = 180

        radius_core_inner = joint.radius_rider - self.child_core_thickness
        core_profile1 = (
            Cq.Sketch()
            .arc((0, 0), joint.radius_rider, core_start_angle, core_end_angle1-core_start_angle)
            .segment((0, 0))
            .close()
            .assemble()
            .circle(radius_core_inner, mode='s')
        )
        core_profile2 = (
            Cq.Sketch()
            .arc((0, 0), joint.radius_rider, -core_start_angle, -(core_end_angle2-core_start_angle))
            .segment((0, 0))
            .close()
            .assemble()
            .circle(radius_core_inner, mode='s')
        )
        core = (
            Cq.Workplane('XY')
            .placeSketch(core_profile1)
            .toPending()
            .extrude(dh * 2)
            .copyWorkplane(Cq.Workplane('XY'))
            .placeSketch(core_profile2)
            .toPending()
            .extrude(dh * 2)
            .translate(Cq.Vector(0, 0, -dh))
        )
        assert self.child_lip_width / 2 <= joint.radius_rider
        lip_thickness = joint.rider_disk_height
        lip = box_with_centre_holes(
            length=self.child_lip_length,
            width=self.child_lip_width,
            height=lip_thickness,
            hole_loc=self.child_conn_hole_pos,
            hole_diam=self.child_conn_hole_diam,
        )
        lip = (
            lip
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                radius=joint.radius_rider,
                height=lip_thickness,
                centered=(True, True, False),
                combine='cut')
        )
        theta = self.torsion_joint.spring.angle_neutral - self.torsion_joint.rider_slot_span
        loc_rotate = Cq.Location((0, 0, 0), (1, 0, 0), 180)
        loc_axis_rotate_bot = Cq.Location((0, 0, 0), (0, 0, 1), self.axis_rotate_bot)
        loc_axis_rotate_top = Cq.Location((0, 0, 0), (0, 0, 1), self.axis_rotate_top)
        result = (
            Cq.Assembly()
            .add(core, name="core", loc=Cq.Location())
            .add(joint.rider(rider_slot_begin=-90, reverse_directrix_label=True), name="rider_top",
                 loc=loc_axis_rotate_top * Cq.Location((0, 0, dh), (0, 0, 1), -90) * Cq.Location((0, 0, 0), (0, 0, 1), theta))
            .add(joint.rider(rider_slot_begin=180), name="rider_bot",
                 loc=loc_axis_rotate_bot * Cq.Location((0, 0, -dh), (0, 0, 1), -90) * loc_rotate)
            .add(lip, name="lip_top",
                 loc=Cq.Location((0, 0, dh)))
            .add(lip, name="lip_bot",
                 loc=Cq.Location((0, 0, -dh)) * loc_rotate)
        )
        return result

    @assembly()
    def assembly(self, deflection: float = 0) -> Cq.Assembly:
        directrix = self.directrix_id
        mat = Material.RESIN_TRANSPERENT
        mat_spring = Material.STEEL_SPRING
        result = (
            Cq.Assembly()
            .addS(self.child(), name="child",
                  role=Role.CHILD, material=mat)
            .constrain("child/core", "Fixed")
            .addS(self.torsion_joint.spring.generate(deflection=-deflection), name="spring_top",
                  role=Role.DAMPING, material=mat_spring)
            .addS(self.parent_top(),
                  name="parent_top",
                  role=Role.PARENT, material=mat)
            .addS(self.torsion_joint.spring.generate(deflection=deflection), name="spring_bot",
                  role=Role.DAMPING, material=mat_spring)
            .addS(self.parent_bot(),
                  name="parent_bot",
                  role=Role.PARENT, material=mat)
        )
        TorsionJoint.add_constraints(
            result,
            rider="child/rider_top",
            track="parent_top/track",
            spring="spring_top",
            directrix=directrix)
        TorsionJoint.add_constraints(
            result,
            rider="child/rider_bot",
            track="parent_bot/track",
            spring="spring_bot",
            directrix=directrix)
        return result.solve()


@dataclass
class Beam:
    """
    A I-shaped spine with two feet
    """

    foot_length: float = 40.0
    foot_width: float = 20.0
    foot_height: float = 5.0
    spine_thickness: float = 4.0
    spine_length: float = 10.0
    total_height: float = 50.0

    hole_diam: float = 8.0
    # distance between the centres of the two holes
    hole_dist: float = 24.0

    def __post_init__(self):
        assert self.spine_height > 0
        assert self.hole_diam + self.hole_dist < self.foot_length
        assert self.hole_dist - self.hole_diam >= self.spine_length

    @property
    def spine_height(self):
        return self.total_height - self.foot_height * 2

    def foot(self) -> Cq.Workplane:
        """
        just one foot
        """
        dx = self.hole_dist / 2
        result = (
            Cq.Workplane('XZ')
            .box(self.foot_length, self.foot_width, self.foot_height,
                 centered=(True, True, False))
            .faces(">Y")
            .workplane()
            .pushPoints([(dx, 0), (-dx, 0)])
            .hole(self.hole_diam)
        )
        plane = result.faces(">Y").workplane()
        plane.moveTo(dx, 0).tagPlane("conn1")
        plane.moveTo(-dx, 0).tagPlane("conn0")
        return result

    def generate(self, flip: bool = False) -> Cq.Assembly:
        beam = (
            Cq.Workplane('XZ')
            .box(self.spine_length, self.spine_thickness, self.spine_height)
        )
        h = self.spine_height / 2 + self.foot_height

        tag_p, tag_n = "top", "bot"
        if flip:
            tag_p, tag_n = tag_n, tag_p
        result = (
            Cq.Assembly()
            .add(beam, name="beam")
            .add(self.foot(), name=tag_p,
                 loc=Cq.Location((0, h, 0)))
            .add(self.foot(), name=tag_n,
                 loc=Cq.Location((0, -h, 0), (1, 0, 0), 180))
        )
        return result


@dataclass
class DiskJoint(Model):
    """
    Sandwiched disk joint for the wrist and elbow
    """
    spring: TorsionSpring = field(default_factory=lambda: TorsionSpring(
        radius=9 / 2,
        thickness=1.3,
        height=6.5,
        tail_length=45.0,
        right_handed=False,
    ))

    radius_housing: float = 22.0
    radius_disk: float = 20.0
    radius_axle: float = 3.0

    housing_thickness: float = 5.0
    disk_thickness: float = 5.0
    # Gap between disk and the housing
    #disk_thickness_gap: float = 0.1

    # Spring angle at 0 degrees of movement
    spring_angle_at_0: float = 60.0
    spring_slot_offset: float = 15.0

    wall_inset: float = 2.0

    # Angular span of movement
    movement_angle: float = 120.0
    # Angular span of tongue on disk
    tongue_span: float = 30.0
    tongue_length: float = 10.0

    generate_inner_wall: bool = False


    def __post_init__(self):
        super().__init__(name="disk-joint")
        assert self.housing_thickness > self.wall_inset
        assert self.radius_housing > self.radius_disk
        assert self.radius_disk > self.radius_axle
        assert self.housing_upper_carve_offset > 0

    @property
    def neutral_movement_angle(self) -> Optional[float]:
        a = self.spring.angle_neutral - self.spring_angle_at_0
        if 0 <= a and a <= self.movement_angle:
            return a
        return None

    @property
    def total_thickness(self) -> float:
        return self.housing_thickness * 2 + self.disk_thickness

    @property
    def opening_span(self) -> float:
        return self.movement_angle + self.tongue_span

    @property
    def housing_upper_carve_offset(self) -> float:
        """
        Distance between the spring track and the outside of the upper housing
        """
        return self.housing_thickness + self.disk_thickness - self.spring.height

    @property
    def housing_upper_dz(self) -> float:
        """
        Distance between the default upper housing location and the median line
        """
        return self.total_thickness / 2 - self.housing_thickness

    @target(name="disk")
    def disk(self) -> Cq.Workplane:
        cut = (
            Cq.Solid.makeBox(
                length=self.spring.tail_length,
                width=self.spring.thickness,
                height=self.disk_thickness,
            )
            .located(Cq.Location((0, self.spring.radius_inner, 0)))
            .rotate((0, 0, 0), (0, 0, 1), self.spring_slot_offset)
        )
        radius_tongue = self.radius_disk + self.tongue_length
        tongue = (
            Cq.Solid.makeCylinder(
                height=self.disk_thickness,
                radius=radius_tongue,
                angleDegrees=self.tongue_span,
            ).cut(Cq.Solid.makeCylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
            ))
        )
        result = (
            Cq.Workplane('XY')
            .cylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
                centered=(True, True, False)
            )
            .union(tongue, tol=TOL)
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                height=self.disk_thickness,
                radius=self.spring.radius,
                centered=(True, True, False),
                combine='cut',
            )
            .cut(cut)
        )
        plane = result.copyWorkplane(Cq.Workplane('XY'))
        theta = math.radians(self.spring_slot_offset)
        plane.tagPlane("dir", direction=(math.cos(theta), math.sin(theta), 0))
        plane.workplane(offset=self.disk_thickness).tagPlane("mate_top")
        result.copyWorkplane(Cq.Workplane('YX')).tagPlane("mate_bot")

        return result

    def wall(self) -> Cq.Compound:
        height = self.disk_thickness + self.wall_inset
        wall = Cq.Solid.makeCylinder(
            radius=self.radius_housing,
            height=height,
            angleDegrees=360 - self.opening_span,
        ).cut(Cq.Solid.makeCylinder(
            radius=self.radius_disk,
            height=height,
        )).rotate((0, 0, 0), (0, 0, 1), self.opening_span)
        return wall

    @target(name="housing-lower")
    def housing_lower(self) -> Cq.Workplane:
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
            .cut(Cq.Solid.makeCylinder(
                radius=self.radius_axle,
                height=self.housing_thickness,
            ))
        )
        result.faces(">Z").tag("mate")
        result.faces(">Z").workplane().tagPlane("dirX", direction="+X")
        result = result.cut(
            self
            .wall()
            .located(Cq.Location((0, 0, self.disk_thickness - self.wall_inset)))
            #.rotate((0, 0, 0), (1, 0, 0), 180)
            #.located(Cq.Location((0, 0, self.disk_thickness + self.housing_thickness)))
        )
        return result

    @target(name="housing-upper")
    def housing_upper(self) -> Cq.Workplane:
        carve_angle = -(self.spring_angle_at_0 - self.spring_slot_offset)
        carve = (
            Cq.Solid.makeCylinder(
                radius=self.spring.radius,
                height=self.housing_thickness
            ).fuse(Cq.Solid.makeBox(
                length=self.spring.tail_length,
                width=self.spring.thickness,
                height=self.housing_thickness
            ).located(Cq.Location((0, -self.spring.radius, 0))))
        ).rotate((0, 0, 0), (0, 0, 1), carve_angle)
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
        )
        theta = math.radians(carve_angle)
        result.faces("<Z").tag("mate")
        p_xy = result.copyWorkplane(Cq.Workplane('XY'))
        p_xy.tagPlane("dirX", direction="+X")
        p_xy.tagPlane("dir", direction=(math.cos(theta), math.sin(theta), 0))
        result = result.faces(">Z").hole(self.radius_axle * 2)

        # tube which holds the spring interior
        if self.generate_inner_wall:
            tube = (
                Cq.Solid.makeCylinder(
                    radius=self.radius_spring_internal,
                    height=self.disk_thickness + self.housing_thickness,
                ).cut(Cq.Solid.makeCylinder(
                    radius=self.radius_axle,
                    height=self.disk_thickness + self.housing_thickness,
                ))
            )
            result = result.union(tube)
        wall = (
            self.wall()
            .located(Cq.Location((0, 0, -self.disk_thickness-self.wall_inset)))
        )
        result = (
            result
            .cut(carve.located(Cq.Location((0, 0, -self.housing_upper_carve_offset))))
            .union(wall, tol=TOL)
        )
        return result.clean()

    def add_constraints(self,
                        assembly: Cq.Assembly,
                        housing_lower: str,
                        housing_upper: str,
                        disk: str,
                        angle: float = 0.0,
                        ) -> Cq.Assembly:
        assert 0 <= angle <= self.movement_angle
        deflection = angle - self.neutral_movement_angle
        spring_name = disk.replace("/", "__Z") + "_spring"
        (
            assembly
            .addS(
                self.spring.generate(deflection=-deflection),
                name=spring_name,
                role=Role.DAMPING,
                material=Material.STEEL_SPRING)
            .constrain(f"{disk}?mate_bot", f"{housing_lower}?mate", "Plane")
            .constrain(f"{disk}?mate_top", f"{housing_upper}?mate", "Plane")
            .constrain(f"{housing_lower}?dirX", f"{housing_upper}?dirX", "Axis", param=0)
            .constrain(f"{housing_upper}?dir", f"{spring_name}?dir_top", "Axis", param=0)
            .constrain(f"{spring_name}?dir_bot", f"{disk}?dir", "Axis", param=0)
            .constrain(f"{disk}?mate_bot", f"{spring_name}?bot", "Plane", param=0)
            #.constrain(f"{housing_lower}?dirX", f"{housing_upper}?dir", "Axis", param=0)
            #.constrain(f"{housing_lower}?dirX", f"{disk}?dir", "Axis", param=angle)
            #.constrain(f"{housing_lower}?dirY", f"{disk}?dir", "Axis", param=angle - 90)
        )
        return (
            assembly
        )


    def assembly(self, angle: Optional[float] = 0) -> Cq.Assembly:
        if angle is None:
            angle = self.neutral_movement_angle
            if angle is None:
                angle = 0
        else:
            assert 0 <= angle <= self.movement_angle
        result = (
            Cq.Assembly()
            .addS(self.disk(), name="disk", role=Role.CHILD)
            .addS(self.housing_lower(), name="housing_lower", role=Role.PARENT)
            .addS(self.housing_upper(), name="housing_upper", role=Role.CASING)
            .constrain("housing_lower", "Fixed")
        )
        result = self.add_constraints(
            result,
            housing_lower="housing_lower",
            housing_upper="housing_upper",
            disk="disk",
            angle=angle,
        )
        return result.solve()


@dataclass(kw_only=True)
class ElbowJoint(Model):
    """
    Creates the elbow and wrist joints.

    This consists of a disk joint, where each side of the joint has mounting
    holes for connection to the exoskeleton. Each side 2 mounting feet on the
    top and bottom, and each foot has 2 holes.

    On the parent side, additional bolts are needed to clamp the two sides of
    the housing together.
    """

    disk_joint: DiskJoint = field(default_factory=lambda: DiskJoint(
        movement_angle=60,
    ))

    # Distance between the child/parent arm to the centre
    child_arm_radius: float = 40.0
    parent_arm_radius: float = 40.0

    child_beam: Beam = field(default_factory=lambda: Beam())
    parent_beam: Beam = field(default_factory=lambda: Beam(
        spine_thickness=8.0,
    ))
    parent_arm_span: float = 40.0
    # Angle of the beginning of the parent arm
    parent_arm_angle: float = 180.0
    parent_binding_hole_radius: float = 30.0

    # Size of the mounting holes
    hole_diam: float = 8.0

    material: Material = Material.RESIN_TRANSPERENT

    # If true, flip the top and bottom tags
    flip: bool = False

    def __post_init__(self):
        assert self.child_arm_radius > self.disk_joint.radius_housing
        assert self.parent_arm_radius > self.disk_joint.radius_housing
        self.disk_joint.tongue_length = self.child_arm_radius - self.disk_joint.radius_disk
        assert self.disk_joint.movement_angle < self.parent_arm_angle < 360 - self.parent_arm_span
        assert self.parent_binding_hole_radius - self.hole_diam / 2 > self.disk_joint.radius_housing

    def child_hole_pos(self) -> list[float]:
        """
        List of hole positions measured from axle
        """
        dx = self.child_beam.hole_dist / 2
        r = self.child_arm_radius
        return [r - dx, r + dx]
    def parent_hole_pos(self) -> list[float]:
        """
        List of hole positions measured from axle
        """
        dx = self.parent_beam.hole_dist / 2
        r = self.parent_arm_radius
        return [r - dx, r + dx]

    @target(name="child")
    def child_joint(self) -> Cq.Assembly:
        angle = -self.disk_joint.tongue_span / 2
        dz = self.disk_joint.disk_thickness / 2
        # We need to ensure the disk is on the "other" side so
        flip_x = Cq.Location((0, 0, 0), (1, 0, 0), 180)
        flip_z = Cq.Location((0, 0, 0), (0, 0, 1), 180)
        result = (
            self.child_beam.generate(flip=self.flip)
            .add(self.disk_joint.disk(), name="disk",
                 loc=flip_x * flip_z * Cq.Location((-self.child_arm_radius, 0, -dz), (0, 0, 1), angle))
            #.constrain("disk", "Fixed")
            #.constrain("top", "Fixed")
            #.constrain("bot", "Fixed")
            #.solve()
        )
        return result

    @target(name="parent-lower")
    def parent_joint_lower(self) -> Cq.Workplane:
        return self.disk_joint.housing_lower()

    @target(name="parent-upper")
    def parent_joint_upper(self):
        axial_offset = Cq.Location((self.parent_arm_radius, 0, 0))
        housing_dz = self.disk_joint.housing_upper_dz
        conn_h = self.parent_beam.spine_thickness
        connector = (
            Cq.Solid.makeCylinder(
                height=conn_h,
                radius=self.parent_arm_radius,
                angleDegrees=self.parent_arm_span)
            .cut(Cq.Solid.makeCylinder(
                height=conn_h,
                radius=self.disk_joint.radius_housing,
            ))
            .located(Cq.Location((0, 0, -conn_h / 2)))
            .rotate((0,0,0), (0,0,1), 180-self.parent_arm_span / 2)
        )
        housing = self.disk_joint.housing_upper()
        housing_loc = Cq.Location(
            (0, 0, housing_dz),
            (0, 0, 1),
            -self.disk_joint.tongue_span / 2
        )
        result = (
            self.parent_beam.generate(flip=self.flip)
            .add(housing, name="housing",
                 loc=axial_offset * housing_loc)
            .add(connector, name="connector",
                 loc=axial_offset)
            #.constrain("housing", "Fixed")
            #.constrain("connector", "Fixed")
            #.solve()
        )
        return result

    @assembly()
    def assembly(self, angle: float = 0) -> Cq.Assembly:
        result = (
            Cq.Assembly()
            .addS(self.child_joint(), name="child",
                  role=Role.CHILD, material=self.material)
            .addS(self.parent_joint_lower(), name="parent_lower",
                  role=Role.CASING, material=self.material)
            .addS(self.parent_joint_upper(), name="parent_upper",
                  role=Role.PARENT, material=self.material)
            #.constrain("child/disk?mate_bot", "Fixed")
        )
        result = self.disk_joint.add_constraints(
            result,
            housing_lower="parent_lower",
            housing_upper="parent_upper/housing",
            disk="child/disk",
            angle=angle,
        )
        return result.solve()

if __name__ == '__main__':
    p = ShoulderJoint()
    p.build_all()
    p = DiskJoint()
    p.build_all()