from math import radians, tan, cos

from javampire.cad.framework.assembly import Assembly
from javampire.cad.gdmutils import ymove, xrot, zmove, zrot, yrot, xmove, xflip_copy
from javampire.cad.parts.annotations.ruler import Ruler
from javampire.cad.parts.cnc.servo.servo_arm import ServoArm
from javampire.solidpython.path import Color, v_diff, v_sum, fit3
from solid import *
from tt.ball_machine.joint_arm import JointArm
from tt.ball_machine.joint_box import JointBox
from tt.ball_machine.joint_box_assembly import JointBoxAssembly
from tt.ball_machine.motor_holder import MotorHolder
from tt.ball_machine.motor_holder_assembly import MotorHolderAssembly
from tt.ball_machine.servo_calculator import ServoCalculator
from tt.ball_machine.top_ball_box import TopBallBox
from tt.ball_machine.top_ball_box_assembly import TopBallBoxAssembly


class BmAssembly(Assembly):

    CONFIG_DESCRIPTOR = dict(
        pipe=dict(
            outer_diameter=50.0,
            wall=1.8,
            length=500.0
        ),
        motor_holder=dict(
            # override here the motor holder settings
            motor_assembly=dict(
                type="a2208",
                wheel=dict(
                    diameter=60
                ),
                motor=dict(
                    stator=dict(
                        visible=False
                    )
                ),
            ),
            vertical_angle=0,
            joint=dict(
                rim=dict(
                    length=12.0,
                    joint_clearance=2.0
                ),
                hole=dict(
                    y_offset=17.0
                )
            ),
            enforcement_plate=dict(
                joint_clearance=1.0
            ),
            vertical_joint=dict(
                wall=6.0,
                hole=dict(
                    diameter=2.5
                )
            ),
        ),
        joint_box_assembly=dict(
            joint_box=dict(
                # override here the joint box settings
                side=dict(
                    wall=4.0
                ),
                bottom=dict(
                    horizontal_width=48.0,
                ),
                vertical_servo_anchor=dict(
                    max_width=6.0
                )
            ),
            servo=dict(
                z_clearance=2.0,
                x_offset=17.0
            ),
            axis="x-z",
            horizontal_angle=0
        ),
        top_ball_box=dict(
            # override here the top ball box settings
        ),
        servo_calculator=dict(
            horizontal=dict(
                max_angle=18.0,
                servo_arm_length=14.0,
                servo_arm_initial_angle=-13,
                show=False,
                calculate_response=False,
                axis_vector=(0, 1),
                font_color="aqua"
            ),
            vertical=dict(
                max_angle=35.0,
                servo_arm_length=20.0,
                servo_arm_initial_angle=27,
                show=False,
                calculate_response=False,
                axis_vector=(-1, 0),
                font_color="aqua",
                axis="-yz"
            ),
        ),
        vertical_servo_fastener=dict(
            type="3F-25T",
            length=20.0,
            wall=3.0,
            wing=dict(
                spacer_width=3.0,
                joint_clearance=0.4
            ),
            knob=dict(
                height=4.0
            ),
            bolt=dict(
                diameter=3.0,
                head=dict(
                    diameter=8.0
                ),
                wall=1.6
            ),
            color="PaleTurquoise"
        ),
        joint_clearance=0.4,
        segments=16
    )

    def __init__(self, *overrides, **kw_overrides):
        super(BmAssembly, self).__init__(*overrides, **kw_overrides)
        if not isinstance(self.joint_box_assembly.joint_box, JointBox):
            self.joint_box_assembly.joint_box = JointBoxAssembly(self.joint_box_assembly).joint_box
        if not isinstance(self.motor_holder, MotorHolder):
            self.motor_holder = MotorHolder(
                self.motor_holder,
                joint=dict(
                    hole=dict(
                        diameter=self.joint_box_assembly.joint_box.vertical_joint.bearing.inner_diameter
                    )
                )
            )
        self.motor_holder_assembly = MotorHolderAssembly(self.motor_holder)
        if not isinstance(self.top_ball_box, TopBallBoxAssembly):
            self.top_ball_box = TopBallBoxAssembly(
                self.top_ball_box,
                guide=dict(
                    end_width=dict(
                        top=self.joint_box_assembly.joint_box.side.wall
                    )
                )
            )
        sch = self.servo_calculator.horizontal
        sch.max_rangle = radians(sch.max_angle)
        sch.box_side = (
            0.5 * self.joint_box_assembly.joint_box.bottom.horizontal_width
            + self.joint_box_assembly.joint_box.side.wall
        )
        sch.max_side_incursion = (
            tan(sch.max_rangle) * sch.box_side + self.top_ball_box.joint.bearing.center_offset / cos(sch.max_rangle)
        )
        sch.origin = (0, self.top_ball_box.joint.bearing.y_center_position)
        self.servo_calculator.horizontal = ServoCalculator(
            sch,
            joint_vector=(
                self.top_ball_box.horizontal_servo_position[0] + 2,
                sch.max_side_incursion + 3.0
            ),
            gear_knob_vector=v_sum(
                (0, self.top_ball_box.horizontal_servo.axle_knob.x_position),
                v_diff(
                    self.top_ball_box.horizontal_servo_position[:-1],
                    sch.origin
                )
            ),
            joint_angle=self.joint_box_assembly.horizontal_angle,
        )
        self.horizontal_servo_gear_knob_z_elevation = (
            self.top_ball_box.horizontal_servo_position[2] + self.top_ball_box.horizontal_servo.z_width
        )
        self.horizontal_servo_fastener = ServoArm(
            knob=dict(type="SG90-21T"),
            length=14.0,
            wing=dict(spacer_width=3.0),
            head=dict(
                tooth_proportion=1.5
            ),
            color="PaleTurquoise", axis="-y-x"
        )
        self.vertical_servo_fastener = ServoArm(
            self.vertical_servo_fastener,
            head=dict(
                tooth_proportion=2.0
            ),
            axis="x"
        )
        self.joint_box_assembly = JointBoxAssembly(
            self.joint_box_assembly,
            joint_box=dict(
                motor_mount_brace=dict(
                    brace_gap=self.motor_holder.joint.horizontal_width + 2 * self.joint_clearance
                ),
                horizontal_joint=dict(
                    hole=dict(
                        diameter=self.top_ball_box.joint.bearing.inner_diameter
                    ),
                    arm_wing=dict(
                        joint_vector=(
                            self.servo_calculator.horizontal.joint_vector[0],
                            - self.horizontal_servo_fastener.wing.z_elevation
                            - self.horizontal_servo_gear_knob_z_elevation,
                            self.servo_calculator.horizontal.joint_vector[1]
                        )
                    )
                )
            ),
            servo=dict(
                z_offset=self.top_ball_box.cap.wall + self.joint_box_assembly.servo.z_clearance,
                y_offset=self.vertical_servo_fastener.wing.z_elevation
            ),
        )
        self.joint_box = self.joint_box_assembly.joint_box
        if self.servo_calculator.horizontal.calculate_response:
            self.servo_calculator.horizontal.plot_angle_response(
                title="Horizontal servo - joint angle to servo arm angle map"
            )
        scv = self.servo_calculator.vertical
        scv.origin = (0, 0)
        scv.max_rangle = radians(scv.max_angle)
        self.vertical_joint_arm = JointArm()
        self.servo_calculator.vertical = ServoCalculator(
            scv,
            joint_vector=(
                - self.motor_holder.joint.arm.y_offset,
                self.motor_holder.joint.arm.z_offset
            ),
            gear_knob_vector=(
                self.joint_box.vertical_joint.z_offset + 0.5 * self.joint_box_assembly.servo.body.x_width
                - self.joint_box_assembly.servo.axle_knob.x_position - self.joint_box_assembly.servo.x_offset,
                0.5 * (self.joint_box.bottom.vertical_width + self.joint_box_assembly.servo.body.y_width)
                + self.joint_box.side.wall + self.joint_box_assembly.servo.z_offset
            ),
            joint_angle=-self.motor_holder.vertical_angle,
        )
        if self.servo_calculator.vertical.calculate_response:
            self.servo_calculator.vertical.plot_angle_response(
                title="Vertical servo - joint angle to servo arm angle map"
            )
        self.horizontal_joint_arm = JointArm(
            length=self.servo_calculator.horizontal.joint_arm_length
        )
        self.vertical_joint_arm = JointArm(
            length=self.servo_calculator.vertical.joint_arm_length,
            axis="yz"
        )
        self.horizontal_joint_arm_length = self.servo_calculator.horizontal.joint_arm_length
        self.vertical_joint_arm_length = self.servo_calculator.vertical.joint_arm_length
        self.horizontal_servo_arm_length = self.servo_calculator.horizontal.servo_arm_length
        self.vertical_servo_arm_length = self.servo_calculator.vertical.servo_arm_length

    def build(self):
        result = union()(
            ymove(self.joint_box.vertical_joint.z_offset)(
                xrot(self.motor_holder.vertical_angle)(
                    ymove(self.motor_holder.joint.y_position)(
                        self.motor_holder_assembly.build()
                    )
                ),
                xmove(-100)(
                    self.servo_calculator.vertical.wrap() if self.servo_calculator.vertical.show else union()
                ),
                translate((
                    self.vertical_servo_fastener.wing.z_elevation,
                    -self.servo_calculator.vertical.gear_knob_vector[0],
                    self.servo_calculator.vertical.gear_knob_vector[1]
                ))(
                    xrot(
                        -(
                            self.servo_calculator.vertical.servo_angle
                            + self.servo_calculator.vertical.servo_arm_initial_angle
                        )
                    )(
                        self.vertical_servo_fastener.wrap()
                    ),
                    translate((
                        -0.5 * self.vertical_joint_arm.delta_bearing.shaft.diameter
                        - self.vertical_servo_fastener.wing.z_elevation
                        - self.vertical_servo_fastener.wing.joint_clearance,
                        - self.servo_calculator.vertical.rotated_servo_arm_vector[0],
                        self.servo_calculator.vertical.rotated_servo_arm_vector[1]
                    ))(
                        xrot(
                            - self.servo_calculator.vertical.servo_arm_initial_angle
                            - self.servo_calculator.vertical.servo_angle
                            - self.servo_calculator.vertical.rotated_servo_arm_to_joint_arm_angle
                        )(
                            self.vertical_joint_arm.wrap()
                        )
                    )
                ),
                xflip_copy()(
                    xmove(
                        0.5 * (
                            self.joint_box.vertical_joint.bearing.height + self.joint_box.motor_mount_brace.brace_gap)
                        + 0.1
                    )(
                        yrot(-90)(
                            self.joint_box.vertical_joint.bearing.wrap()
                        )
                    )
                )
            ),
            self.joint_box_assembly.wrap()
        )
        result = union()(
            zmove(0.1)(
                ymove(self.top_ball_box.joint.bearing.y_center_position)(
                    zrot(self.joint_box_assembly.horizontal_angle)(
                        result
                    ),
                    Color("lightgrey")(
                        # sphere(d=self.motor_holder.ball.diameter),
                    )
                )
            ),
            self.top_ball_box.wrap(),
            translate(
                v_sum(
                    fit3(self.servo_calculator.horizontal.origin),
                    fit3(self.servo_calculator.horizontal.gear_knob_vector),
                    fit3(self.servo_calculator.horizontal.rotated_servo_arm_vector),
                    (
                        0, 0,
                        self.horizontal_servo_gear_knob_z_elevation
                        + 0.5 * self.horizontal_joint_arm.delta_bearing.shaft.diameter
                        + self.horizontal_servo_fastener.wing.z_elevation + 0.4
                    )
                )
            )(
                zrot(
                    90 + self.servo_calculator.horizontal.servo_arm_initial_angle
                    + self.servo_calculator.horizontal.servo_angle
                    + self.servo_calculator.horizontal.rotated_servo_arm_to_joint_arm_angle
                )(
                    self.horizontal_joint_arm.wrap()
                )
            ),
            translate(
                v_sum(
                    fit3(self.servo_calculator.horizontal.origin),
                    fit3(self.servo_calculator.horizontal.gear_knob_vector),
                    (
                        0, 0, self.horizontal_servo_gear_knob_z_elevation + 0.2
                    )
                )
            )(
                zrot(
                    self.servo_calculator.horizontal.servo_arm_initial_angle
                    + self.servo_calculator.horizontal.servo_angle
                )(
                    self.horizontal_servo_fastener.wrap()
                )
            ),
        )
        if self.servo_calculator.horizontal.show:
            result = union()(
                result,
                zmove(50)(
                    self.servo_calculator.horizontal.build()
                )
            )
        measure = False
        if measure:
            result = zrot(-90)(
                result,
                Ruler(axis="y-x", origin=(1, -37, 52), body=dict(color="pink:0.5")).wrap()
            )
        return result

    def iter_components(self):
        yield "motor_holder", self.motor_holder
        yield "joint_box", self.joint_box
        yield "top_ball_box", TopBallBox(self.top_ball_box)
        yield "top_ball_box_bracket", self.top_ball_box.bracket
        yield "top_ball_box_cap", self.top_ball_box.top_ball_box_cap
        yield "top_ball_box_guide", self.top_ball_box.guide
        yield "top_ball_box_guide_bracket_left", self.top_ball_box.guide_bracket
        yield "top_ball_box_guide_bracket_right", self.top_ball_box.guide_bracket.override(mirror=True)
        yield "horizontal_servo_fastener", self.horizontal_servo_fastener.override(axis="z")
        yield "vertical_servo_fastener", self.vertical_servo_fastener.override(axis="z")
        yield "wheel_rim", self.motor_holder.motor_assembly.wheel.rim
        yield "wheel_tire", self.motor_holder.motor_assembly.wheel.tire
        yield "vertical_servo_anchor", self.joint_box_assembly.joint_box.vertical_servo_anchor
        yield "vertical_servo_mount_anchor", self.joint_box_assembly.joint_box.vertical_servo_mount_anchor
        yield "vertical_servo_mount_arm", self.joint_box_assembly.vertical_servo_mount_arm
        yield (
            "motor_placement_helper",
            self.motor_holder_assembly.motor_placement_helper_assembly.motor_placement_helper
        )
        yield (
            "motor_placement_helper_insert",
            self.motor_holder_assembly.motor_placement_helper_assembly.motor_placement_helper_insert
        )
        yield "cable_clip_0", self.motor_holder_assembly.clip_holder.hook