"""
To build, execute
```
python3 nhf/touhou/houjuu_nue/__init__.py
```

This cosplay consists of 3 components:

## Trident

The trident is composed of individual segments, made of acrylic, and a 3D
printed head (convention rule prohibits metal) with a metallic paint. To ease
transportation, the trident handle has individual segments with threads and can
be assembled on site.

## Snake

A 3D printed snake with a soft material so it can wrap around and bend

## Wings

This is the crux of the cosplay and the most complex component. The wings mount
on a wearable harness. Each wing consists of 4 segments with 3 joints. Parts of
the wing which demands transluscency are created from 1/16" acrylic panels.
These panels serve double duty as the exoskeleton.

The wings are labeled r1,r2,r3,l1,l2,l3. The segments of the wings are labeled
from root to tip s0 (root),
s1, s2, s3. The joints are named (from root to tip)
shoulder, elbow, wrist in analogy with human anatomy.
"""
from dataclasses import dataclass, field
import cadquery as Cq
from nhf import Material, Role
from nhf.build import Model, TargetKind, target, assembly
from nhf.parts.joints import HirthJoint, TorsionJoint
from nhf.parts.handle import Handle
import nhf.touhou.houjuu_nue.wing as MW
import nhf.touhou.houjuu_nue.trident as MT
import nhf.utils

@dataclass
class Parameters(Model):
    """
    Defines dimensions for the Houjuu Nue cosplay
    """

    # Thickness of the exoskeleton panel in millimetres
    panel_thickness: float = 25.4 / 16

    # Harness
    harness_thickness: float = 25.4 / 8
    harness_width: float = 300
    harness_height: float = 400
    harness_fillet: float = 10

    harness_wing_base_pos: list[tuple[str, float, float]] = field(default_factory=lambda: [
        ("r1", 70, 150),
        ("l1", -70, 150),
        ("r2", 100, 0),
        ("l2", -100, 0),
        ("r3", 70, -150),
        ("l3", -70, -150),
    ])

    # Holes drilled onto harness for attachment with HS joint
    harness_to_root_conn_diam: float = 6

    hs_hirth_joint: HirthJoint = field(default_factory=lambda: HirthJoint(
        radius=30,
        radius_inner=20,
        tooth_height=10,
        base_height=5
    ))

    # Wing root properties
    #
    # The Houjuu-Scarlett joint mechanism at the base of the wing
    hs_joint_base_width: float = 85
    hs_joint_base_thickness: float = 10
    hs_joint_corner_fillet: float = 5
    hs_joint_corner_cbore_diam: float = 12
    hs_joint_corner_cbore_depth: float = 2
    hs_joint_corner_inset: float = 12

    hs_joint_axis_diam: float = 12
    hs_joint_axis_cbore_diam: float = 20
    hs_joint_axis_cbore_depth: float = 3

    wing_profile: MW.WingProfile = field(default_factory=lambda: MW.WingProfile(
        shoulder_height = 80,
        elbow_height = 100,
    ))

    # Exterior radius of the wing root assembly
    wing_root_radius: float = 40
    wing_root_wall_thickness: float = 8

    shoulder_torsion_joint: TorsionJoint = field(default_factory=lambda: TorsionJoint(
        radius_track=35,
        radius_rider=35,
        groove_radius_outer=32,
        track_disk_height=5.0,
        rider_disk_height=7.0,
        radius_axle=8.0,
    ))

    # Two holes on each side (top and bottom) are used to attach the shoulder
    # joint. This governs the distance between these two holes
    shoulder_attach_dist: float = 25
    shoulder_attach_diam: float = 8

    """
    Heights for various wing joints, where the numbers start from the first
    joint.
    """
    wing_s0_thickness: float = 40

    # Length of the spacer
    wing_s1_thickness: float = 20
    wing_s1_spacer_thickness: float = 25.4 / 8
    wing_s1_spacer_width: float = 20
    wing_s1_spacer_hole_diam: float = 8
    wing_s1_shoulder_spacer_hole_dist: float = 20
    wing_s1_shoulder_spacer_width: float = 60

    trident_handle: Handle = field(default_factory=lambda: Handle(
        diam=38,
        diam_inner=38-2 * 25.4/8,
        # M27-3
        diam_threading=27,
        thread_pitch=3,
        diam_connector_internal=18,
        simplify_geometry=False,
    ))

    material_panel: Material = Material.ACRYLIC_TRANSPARENT
    material_bracket: Material = Material.ACRYLIC_TRANSPARENT

    def __post_init__(self):
        super().__init__(name="houjuu-nue")
        assert self.wing_root_radius > self.hs_hirth_joint.radius,\
            "Wing root must be large enough to accomodate joint"
        assert self.wing_s1_shoulder_spacer_hole_dist > self.wing_s1_spacer_hole_diam, \
            "Spacer holes are too close to each other"

    @target(name="trident/handle-connector")
    def handle_connector(self):
        return self.trident_handle.connector()
    @target(name="trident/handle-insertion")
    def handle_insertion(self):
        return self.trident_handle.insertion()


    def harness_profile(self) -> Cq.Sketch:
        """
        Creates the harness shape
        """
        w, h = self.harness_width / 2, self.harness_height / 2
        sketch = (
            Cq.Sketch()
            .polygon([
                (0.7 * w, h),
                (w, 0),
                (0.7 * w, -h),
                (0.7 * -w, -h),
                (-w, 0),
                (0.7 * -w, h),
            ])
            #.rect(self.harness_width, self.harness_height)
            .vertices()
            .fillet(self.harness_fillet)
        )
        for tag, x, y in self.harness_wing_base_pos:
            conn = [(px + x, py + y) for px, py in self.hs_joint_harness_conn()]
            sketch = (
                sketch
                .push(conn)
                .tag(tag)
                .circle(self.harness_to_root_conn_diam / 2, mode='s')
                .reset()
            )
        return sketch

    @target(name="harness", kind=TargetKind.DXF)
    def harness(self) -> Cq.Shape:
        """
        Creates the harness shape
        """
        result = (
            Cq.Workplane('XZ')
            .placeSketch(self.harness_profile())
            .extrude(self.harness_thickness)
        )
        result.faces(">Y").tag("mount")
        plane = result.faces(">Y").workplane()
        for tag, x, y in self.harness_wing_base_pos:
            conn = [(px + x, py + y) for px, py
                    in self.hs_joint_harness_conn()]
            for i, (px, py) in enumerate(conn):
                (
                    plane
                    .moveTo(px, py)
                    .circle(1, forConstruction='True')
                    .edges()
                    .tag(f"{tag}_{i}")
                )
        return result

    def hs_joint_harness_conn(self) -> list[tuple[int, int]]:
        """
        Generates a set of points corresponding to the connectorss
        """
        dx = self.hs_joint_base_width / 2 - self.hs_joint_corner_inset
        return [
            (dx, dx),
            (dx, -dx),
            (-dx, -dx),
            (-dx, dx),
        ]

    @target(name="hs_joint_parent")
    def hs_joint_parent(self):
        """
        Parent part of the Houjuu-Scarlett joint, which is composed of a Hirth
        coupling, a cylindrical base, and a mounting base.
        """
        hirth = self.hs_hirth_joint.generate()
        conn = self.hs_joint_harness_conn()
        result = (
            Cq.Workplane('XY')
            .box(
                self.hs_joint_base_width,
                self.hs_joint_base_width,
                self.hs_joint_base_thickness,
                centered=(True, True, False))
            .translate((0, 0, -self.hs_joint_base_thickness))
            .edges("|Z")
            .fillet(self.hs_joint_corner_fillet)
            .faces(">Z")
            .workplane()
            .pushPoints(conn)
            .cboreHole(
                diameter=self.harness_to_root_conn_diam,
                cboreDiameter=self.hs_joint_corner_cbore_diam,
                cboreDepth=self.hs_joint_corner_cbore_depth)
        )
        # Creates a plane parallel to the holes but shifted to the base
        plane = result.faces(">Z").workplane(offset=-self.hs_joint_base_thickness)

        for i, (px, py) in enumerate(conn):
            (
                plane
                .pushPoints([(px, py)])
                .circle(1, forConstruction='True')
                .edges()
                .tag(f"h{i}")
            )
        result = (
            result
            .faces(">Z")
            .workplane()
            .union(hirth, tol=0.1)
            .clean()
        )
        result = (
            result.faces("<Z")
            .workplane()
            .cboreHole(
                diameter=self.hs_joint_axis_diam,
                cboreDiameter=self.hs_joint_axis_cbore_diam,
                cboreDepth=self.hs_joint_axis_cbore_depth,
            )
            .clean()
        )
        result.faces("<Z").tag("base")
        return result

    @assembly()
    def harness_assembly(self) -> Cq.Assembly:
        harness = self.harness()
        result = (
            Cq.Assembly()
            .add(harness, name="base", color=Material.WOOD_BIRCH.color)
            .constrain("base", "Fixed")
        )
        for name in ["l1", "l2", "l3", "r1", "r2", "r3"]:
            j = self.hs_joint_parent()
            (
                result
                .add(j, name=name, color=Role.PARENT.color)
                .constrain("base?mount", f"{name}?base", "Axis")
            )
            for i in range(4):
                result.constrain(f"base?{name}_{i}", f"{name}?h{i}", "Point")
        result.solve()
        return result

    #@target(name="wing/joining-plate", kind=TargetKind.DXF)
    #def joining_plate(self) -> Cq.Workplane:
    #    return self.wing_joining_plate.plate()

    @target(name="wing_root")
    def wing_root(self) -> Cq.Assembly:
        """
        Generate the wing root which contains a Hirth joint at its base and a
        rectangular opening on its side, with the necessary interfaces.
        """
        return MW.wing_root(
            joint=self.hs_hirth_joint,
            shoulder_attach_dist=self.shoulder_attach_dist,
            shoulder_attach_diam=self.shoulder_attach_diam,
            wall_thickness=self.wing_root_wall_thickness,
            conn_height=self.wing_profile.shoulder_height,
            conn_thickness=self.wing_s0_thickness,
        )

    @target(name="shoulder_joint_parent")
    def shoulder_joint_parent(self) -> Cq.Workplane:
        joint = self.shoulder_torsion_joint
        # Thickness of the lip connecting this joint to the wing root
        lip_thickness = 10
        lip_width = 25
        lip_guard_ext = 40
        lip_guard_height = self.wing_root_wall_thickness + lip_thickness
        assert lip_guard_ext > joint.radius_track

        lip_guard = (
            Cq.Solid.makeBox(lip_guard_ext, lip_width, lip_guard_height)
            .located(Cq.Location((0, -lip_width/2 , 0)))
            .cut(Cq.Solid.makeCylinder(joint.radius_track, lip_guard_height))
        )
        result = (
            joint.track()
            .union(lip_guard, tol=1e-6)

            # Extrude the handle
            .copyWorkplane(Cq.Workplane(
                'YZ', origin=Cq.Vector((88, 0, self.wing_root_wall_thickness))))
            .rect(lip_width, lip_thickness, centered=(True, False))
            .extrude("next")

            # Connector holes on the lip
            .copyWorkplane(Cq.Workplane(
                'YX', origin=Cq.Vector((57, 0, self.wing_root_wall_thickness))))
            .hole(self.shoulder_attach_diam)
            .moveTo(0, self.shoulder_attach_dist)
            .hole(self.shoulder_attach_diam)
        )
        result.moveTo(0, 0).tagPlane('conn0')
        result.moveTo(0, self.shoulder_attach_dist).tagPlane('conn1')
        return result

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

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

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

        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(joint.radius_rider - core_thickness, 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(joint.radius_rider - core_thickness, 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))
        )
        # Create the upper and lower lips
        lip_height = self.wing_s1_thickness
        lip_thickness = joint.rider_disk_height
        lip_ext = 40 + joint.radius_rider
        hole_dx = self.wing_s1_shoulder_spacer_hole_dist
        assert lip_height / 2 <= joint.radius_rider
        lip = (
            Cq.Workplane('XY')
            .box(lip_ext, lip_height, lip_thickness,
                 centered=(False, True, False))
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(radius=joint.radius_rider, height=lip_thickness,
                      centered=(True, True, False),
                      combine='cut')
            .faces(">Z")
            .workplane()
        )
        hole_x = lip_ext - hole_dx / 2
        for i in range(2):
            x = hole_x - i * hole_dx
            lip = lip.moveTo(x, 0).hole(self.wing_s1_spacer_hole_diam)
        for i in range(2):
            x = hole_x - i * hole_dx
            (
                lip
                .moveTo(x, 0)
                .tagPlane(f"conn{1 - i}")
            )

        loc_rotate = Cq.Location((0, 0, 0), (1, 0, 0), 180)
        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=Cq.Location((0, 0, dh), (0, 0, 1), -90))
            .add(joint.rider(rider_slot_begin=180), name="rider_bot",
                 loc=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 shoulder_assembly(self) -> Cq.Assembly:
        directrix = 0
        result = (
            Cq.Assembly()
            .add(self.shoulder_joint_child(), name="child",
                 color=Role.CHILD.color)
            .constrain("child/core", "Fixed")
            .add(self.shoulder_torsion_joint.spring(), name="spring_top",
                 color=Role.DAMPING.color)
            .add(self.shoulder_joint_parent(), name="parent_top",
                 color=Role.PARENT.color)
            .add(self.shoulder_torsion_joint.spring(), name="spring_bot",
                 color=Role.DAMPING.color)
            .add(self.shoulder_joint_parent(), name="parent_bot",
                 color=Role.PARENT.color)
        )
        TorsionJoint.add_constraints(result,
                                     rider="child/rider_top",
                                     track="parent_top",
                                     spring="spring_top",
                                     directrix=directrix)
        TorsionJoint.add_constraints(result,
                                     rider="child/rider_bot",
                                     track="parent_bot",
                                     spring="spring_bot",
                                     directrix=directrix)
        return result.solve()

    @target(name="wing/s1-spacer", kind=TargetKind.DXF)
    def wing_s1_spacer(self) -> Cq.Workplane:
        result = (
            Cq.Workplane('XZ')
            .sketch()
            .rect(self.wing_s1_spacer_width, self.wing_s1_thickness)
            .finalize()
            .extrude(self.wing_s1_spacer_thickness)
        )
        result.faces("<Z").tag("weld1")
        result.faces(">Z").tag("weld2")
        result.faces(">Y").tag("dir")
        return result

    @target(name="wing/s1-shoulder-spacer", kind=TargetKind.DXF)
    def wing_s1_shoulder_spacer(self) -> Cq.Workplane:
        """
        The mate tags are on the side closer to the holes.
        """
        dx = self.wing_s1_shoulder_spacer_hole_dist
        h = self.wing_s1_spacer_thickness
        result = (
            Cq.Workplane('XZ')
            .sketch()
            .rect(self.wing_s1_shoulder_spacer_width,
                  self.wing_s1_thickness)
            .push([
                (0, 0),
                (dx, 0),
            ])
            .circle(self.wing_s1_spacer_hole_diam / 2, mode='s')
            .finalize()
            .extrude(h)
        )
        # Tag the mating surfaces to be glued
        result.faces("<Z").workplane().moveTo(0, h).tagPlane("weld1")
        result.faces(">Z").workplane().moveTo(0, -h).tagPlane("weld2")

        # Tag the directrix
        result.faces("<Y").tag("dir")

        # Tag the holes
        plane = result.faces(">Y").workplane()
        # Side closer to the parent is 0
        plane.moveTo(-dx, 0).tagPlane("conn0")
        plane.tagPlane("conn1")
        return result

    @target(name="wing/r1s1", kind=TargetKind.DXF)
    def wing_r1s1_profile(self) -> Cq.Sketch:
        return self.wing_profile.wing_r1_profile()

    def wing_r1s1_panel(self, front=True) -> Cq.Workplane:
        profile = self.wing_r1s1_profile()
        w = self.wing_s1_shoulder_spacer_width / 2
        h = (self.wing_profile.shoulder_height - self.shoulder_joint_child_height) / 2
        anchors = [
            ("shoulder_top", w, h + self.shoulder_joint_child_height),
            ("shoulder_bot", w, h),
            ("middle", 50, -20),
            ("tip", 270, 50),
        ]
        result = (
            Cq.Workplane("XY")
            .placeSketch(profile)
            .extrude(self.panel_thickness)
        )
        plane = result.faces(">Z" if front else "<Z").workplane()
        sign = 1 if front else -1
        for name, px, py in anchors:
            plane.moveTo(px, sign * py).tagPlane(name)
        return result

    @assembly()
    def wing_r1s1_assembly(self) -> Cq.Assembly:
        result = (
            Cq.Assembly()
            .add(self.wing_r1s1_panel(front=True), name="panel_front",
                 color=self.material_panel.color)
            .constrain("panel_front", "Fixed")
            .add(self.wing_r1s1_panel(front=False), name="panel_back",
                 color=self.material_panel.color)
            .constrain("panel_front@faces@>Z", "panel_back@faces@<Z", "Point",
                       param=self.wing_s1_thickness)
            .add(self.wing_s1_shoulder_spacer(),
                 name="shoulder_bot_spacer",
                 color=self.material_bracket.color)
            .constrain("panel_front?shoulder_bot", "shoulder_bot_spacer?weld1", "Plane")
            .constrain("panel_back?shoulder_bot", "shoulder_bot_spacer?weld2", "Plane")
            .constrain("shoulder_bot_spacer?dir", "FixedAxis", param=(0, 1, 0))
            .add(self.wing_s1_shoulder_spacer(),
                 name="shoulder_top_spacer",
                 color=self.material_bracket.color)
            .constrain("panel_front?shoulder_top", "shoulder_top_spacer?weld2", "Plane")
            .constrain("panel_back?shoulder_top", "shoulder_top_spacer?weld1", "Plane")
            .constrain("shoulder_top_spacer?dir", "FixedAxis", param=(0, -1, 0))
            # Should be controlled by point value directly
            #.constrain("shoulder_bot_spacer?dir", "shoulder_top_spacer?dir", "Point",
            #           self.shoulder_joint_child_height)
        )
        for tag in ["middle", "tip"]:
            name = f"{tag}_spacer"
            (
                result
                .add(self.wing_s1_spacer(), name=name,
                     color=self.material_bracket.color)
                .constrain(f"panel_front?{tag}", f"{tag}_spacer?weld1", "Plane")
                .constrain(f"panel_back?{tag}", f"{tag}_spacer?weld2", "Plane")
                .constrain(f"{name}?dir", "FixedAxis", param=(0, 1, 0))
            )
        return result.solve()


    @assembly()
    def wing_r1_assembly(self, parts=["root", "shoulder", "s1"]) -> Cq.Assembly:
        result = (
            Cq.Assembly()
        )
        if "root" in parts:
            (
                result
                .add(self.wing_root(), name="root")
                .constrain("root/scaffold", "Fixed")
            )
        if "shoulder" in parts:
            result.add(self.shoulder_assembly(), name="shoulder")

        if "root" in parts and "shoulder" in parts:
            (
                result
                .constrain("root/scaffold?conn_top0", "shoulder/parent_top?conn0", "Plane")
                .constrain("root/scaffold?conn_top1", "shoulder/parent_top?conn1", "Plane")
                .constrain("root/scaffold?conn_bot0", "shoulder/parent_bot?conn0", "Plane")
                .constrain("root/scaffold?conn_bot1", "shoulder/parent_bot?conn1", "Plane")
            )

        if "s1" in parts:
            result.add(self.wing_r1s1_assembly(), name="s1")

        if "s1" in parts and "shoulder" in parts:
            (
                result
                .constrain("shoulder/child/lip_bot?conn0",
                           "s1/shoulder_bot_spacer?conn0",
                           "Plane")
                .constrain("shoulder/child/lip_bot?conn1",
                           "s1/shoulder_bot_spacer?conn1",
                           "Plane")
                .constrain("shoulder/child/lip_top?conn0",
                           "s1/shoulder_top_spacer?conn0",
                           "Plane")
                .constrain("shoulder/child/lip_top?conn1",
                           "s1/shoulder_top_spacer?conn1",
                           "Plane")
            )
        return result.solve()

    @assembly()
    def wings_assembly(self) -> Cq.Assembly:
        """
        Assembly of harness with all the wings
        """
        a_tooth = self.hs_hirth_joint.tooth_angle

        result = (
            Cq.Assembly()
            .add(self.harness_assembly(), name="harness", loc=Cq.Location((0, 0, 0)))
            .add(self.wing_root(), name="w0_r1")
            .add(self.wing_root(), name="w0_l1")
            .constrain("harness/base", "Fixed")
            .constrain("w0_r1/joint?mate", "harness/r1?mate", "Plane")
            .constrain("w0_r1/joint?dir", "harness/r1?dir",
                       "Axis", param=7 * a_tooth)
            .constrain("w0_l1/joint?mate", "harness/l1?mate", "Plane")
            .constrain("w0_l1/joint?dir", "harness/l1?dir",
                       "Axis", param=-1 * a_tooth)
            .solve()
        )
        return result

    @assembly(collision_check=False)
    def trident_assembly(self) -> Cq.Assembly:
        """
        Disable collision check since the threads may not align.
        """
        return MT.trident_assembly(self.trident_handle)



if __name__ == '__main__':
    p = Parameters()
    p.build_all()