from dataclasses import dataclass, field
from typing import Optional, Tuple
import cadquery as Cq
from nhf import Role
from nhf.build import Model, target
import nhf.parts.springs as springs
import nhf.utils

TOL = 1e-6

@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 beam(self) -> Cq.Assembly:
        beam = (
            Cq.Workplane('XZ')
            .box(self.spine_length, self.spine_thickness, self.spine_height)
        )
        h = self.spine_height / 2 + self.foot_height
        result = (
            Cq.Assembly()
            .add(beam, name="beam")
            .add(self.foot(), name="top",
                 loc=Cq.Location((0, h, 0)))
            .add(self.foot(), name="bot",
                 loc=Cq.Location((0, -h, 0), (0, 0, 1), 180))
        )
        return result

@dataclass
class DiskJoint(Model):
    """
    Sandwiched disk joint for the wrist and elbow
    """

    radius_housing: float = 22.0
    radius_disk: float = 20.0
    radius_spring: float = 9 / 2
    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_thickness: float = 1.3
    spring_height: float = 6.5
    spring_tail_length: float = 45.0

    # Spring angle at 0 degrees of movement
    spring_angle: float = 30.0
    # Angle at which the spring exerts no torque
    spring_angle_neutral: float = 90.0
    spring_angle_shift: float = 30
    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

    def spring(self):
        return springs.torsion_spring(
            radius=self.radius_spring,
            height=self.spring_height,
            thickness=self.spring_thickness,
            tail_length=self.spring_tail_length,
            right_handed=False,
        )

    @property
    def neutral_movement_angle(self) -> Optional[float]:
        a = self.spring_angle_neutral - self.spring_angle
        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

    @property
    def radius_spring_internal(self):
        return self.radius_spring - self.spring_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.radius_spring_internal, 0)))
            .rotate((0, 0, 0), (0, 0, 1), self.spring_angle_shift)
        )
        result = (
            Cq.Workplane('XY')
            .cylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
                centered=(True, True, False)
            )
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                height=self.disk_thickness,
                radius=self.radius_spring,
                centered=(True, True, False),
                combine='cut',
            )
            .cut(cut)
        )
        plane = result.copyWorkplane(Cq.Workplane('XY'))
        plane.tagPlane("dir", direction="+X")
        plane.workplane(offset=self.disk_thickness).tagPlane("mate_top")
        result.copyWorkplane(Cq.Workplane('YX')).tagPlane("mate_bot")

        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 = result.union(tongue, tol=TOL)
        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("dir", 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 = (
            Cq.Solid.makeCylinder(
                radius=self.radius_spring,
                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.radius_spring, 0))))
        ).rotate((0, 0, 0), (0, 0, 1), self.spring_angle - self.spring_angle_shift)
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
        )
        result.faces("<Z").tag("mate")
        result.copyWorkplane(Cq.Workplane('XY')).tagPlane("dir", direction="-X")
        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

    def add_constraints(self,
                        assembly: Cq.Assembly,
                        housing_lower: str,
                        housing_upper: str,
                        disk: str,
                        angle: Tuple[float, float, float] = (0, 0, 0),
                        ) -> Cq.Assembly:
        """
        The angle supplied must be perpendicular to the disk normal.
        """
        (
            assembly
            .constrain(f"{disk}?mate_bot", f"{housing_lower}?mate", "Plane")
            .constrain(f"{disk}?mate_top", f"{housing_upper}?mate", "Plane")
            .constrain(f"{housing_lower}?dir", f"{housing_upper}?dir", "Axis")
            .constrain(f"{disk}?dir", "FixedRotation", angle)
        )


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

@dataclass
class ElbowJoint:
    """
    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

    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_joint(self) -> Cq.Assembly:
        angle = -self.disk_joint.tongue_span / 2
        dz = self.disk_joint.disk_thickness / 2
        result = (
            self.child_beam.beam()
            .add(self.disk_joint.disk(), name="disk",
                 loc=Cq.Location((-self.child_arm_radius, 0, -dz), (0, 0, 1), angle))
        )
        return result

    def parent_joint_bot(self) -> Cq.Workplane:
        return self.disk_joint.housing_lower()

    def parent_joint_top(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()
        result = (
            Cq.Assembly()
            .add(housing, name="housing",
                 loc=axial_offset * Cq.Location((0, 0, housing_dz)))
            .add(connector, name="connector",
                 loc=axial_offset)
            .add(self.parent_beam.beam(), name="beam")
        )
        return result

    def assembly(self, angle: float = 0) -> Cq.Assembly:
        da = self.disk_joint.tongue_span / 2
        result = (
            Cq.Assembly()
            .add(self.child_joint(), name="child", color=Role.CHILD.color)
            .add(self.parent_joint_bot(), name="parent_bot", color=Role.CASING.color)
            .add(self.parent_joint_top(), name="parent_top", color=Role.PARENT.color)
            .constrain("parent_bot", "Fixed")
        )
        self.disk_joint.add_constraints(
            result,
            housing_lower="parent_bot",
            housing_upper="parent_top/housing",
            disk="child/disk",
            angle=(0, 0, angle + da),
        )
        return result.solve()

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