import cadquery as cq
from ocp_vscode import *
from utils import *

import localselectors

tube_w = inch(2.0)
tube_h = inch(1.5)
tube_wall = inch(0.075)  # 1.519 #inch(0.062)
column_spacing = 560
column_h = 400
plate_z = 2.54
plate_z2 = 10

x_mount_z = 5

ballscrew_x_offset = 4

ballscrew_len = 500
column_spacing = (
    ballscrew_len - 15 - 39 + 25 - 10 + 7 / 2 + 20 / 2 + 2 * x_mount_z + tube_w
)
print("column_spacing = ", column_spacing)

x_rail_spacing = 240
x_carriage_spacing = 120


class Gantry:
    def __init__(self, w):
        self.left = tubing(
            column_h,
            tube_w,
            tube_h,
            tube_wall,
            w.transformed(offset=(-0.5 * column_spacing, 0, 0)),
        )
        self.right = tubing(
            column_h,
            tube_w,
            tube_h,
            tube_wall,
            w.transformed(offset=(0.5 * column_spacing, 0, 0)),
        )
        upper_center = column_h - 0.5 * tube_w
        self.upper = tubing(
            column_spacing - tube_w,
            tube_h,
            tube_w,
            tube_wall,
            w.transformed(
                offset=(-0.5 * column_spacing + 0.5 * tube_w, 0, upper_center),
                rotate=(90, 90, 0),
            ),
        )
        lower_center = upper_center - x_rail_spacing
        self.lower = tubing(
            column_spacing - tube_w,
            tube_h,
            tube_w,
            tube_wall,
            w.transformed(
                offset=(-0.5 * column_spacing + 0.5 * tube_w, 0, lower_center),
                rotate=(90, 90, 0),
            ),
        )
        self.upper_rail = HGR20(
            w.transformed(
                offset=(0, 0.5 * tube_h + plate_z, upper_center), rotate=(-90, 0, 0)
            ),
            column_spacing + tube_w,
        )
        self.lower_rail = HGR20(
            w.transformed(
                offset=(0, 0.5 * tube_h + plate_z, lower_center), rotate=(-90, 0, 0)
            ),
            column_spacing + tube_w,
        )
        self.upper_front_plate = (
            w.transformed(offset=(0, tube_h / 2, upper_center), rotate=(-90, 0, 0))
            .rect(column_spacing + tube_w, tube_w)
            .extrude(plate_z)
        )
        self.lower_front_plate = (
            w.transformed(offset=(0, tube_h / 2, lower_center), rotate=(-90, 0, 0))
            .rect(column_spacing + tube_w, tube_w)
            .extrude(plate_z)
        )
        self.upper_back_plate = (
            w.transformed(offset=(0, -tube_h / 2, upper_center), rotate=(90, 0, 0))
            .rect(column_spacing + tube_w, tube_w)
            .extrude(plate_z)
        )
        self.lower_back_plate = (
            w.transformed(offset=(0, -tube_h / 2, lower_center), rotate=(90, 0, 0))
            .rect(column_spacing + tube_w, tube_w)
            .extrude(plate_z)
        )
        self.ball_screw_w = w.transformed(
            offset=(
                ballscrew_len / 2 - (-15 - 39 + 25 + 10 - 7 / 2 - 20 / 2) / 2,
                ballscrew_x_offset,
                (upper_center + lower_center) / 2,
            ),
            rotate=(0, 0, 90),
        ).transformed(rotate=(0, 90, 0))
        self.ball_screw_model = SFU1204(self.ball_screw_w, ballscrew_len)
        self.floating_mount_w = self.ball_screw_w.transformed(
            offset=(0, ballscrew_len + (-10 + 7 / 2 + 20 / 2), 0), rotate=(-90, 0, 0)
        )

        self.ball_screw = self.ball_screw_model.ballscrew()
        self.fixed_bearing = self.ball_screw_model.fixedBearing()
        self.floating_bearing = self.ball_screw_model.floatingBearing()

        self.left = self.ball_screw_model.floatingBearingHoles(self.left)
        self.right = self.ball_screw_model.fixedBearingHoles(self.right)
        self.right = (
            self.right
            # holes for the ball screw
            .copyWorkplane(self.ball_screw_w.transformed(rotate=(90, 0, 0)))
            .hole(20)
            .copyWorkplane(self.ball_screw_w.transformed(rotate=(-90, 0, 0)))
            .hole(20)
        )

        self.upper_front_plate = self.upper_rail.holes(self.upper_front_plate)
        self.upper_back_plate = self.upper_rail.holes(self.upper_back_plate)

        self.lower_front_plate = self.lower_rail.holes(self.lower_front_plate)
        self.lower_back_plate = self.lower_rail.holes(self.lower_back_plate)

        right_bearing_base = (
            self.ball_screw_w.transformed(offset=(0, 0, -ballscrew_x_offset))
            .transformed(rotate=(-90, 0, 0))
            .transformed(offset=(0, 0, 15 + 39 - 25 - x_mount_z))
            .transformed(rotate=(0, 0, 90))
        )

        self.right_bearing_mount = (
            w.copyWorkplane(right_bearing_base)
            .rect(tube_h, x_rail_spacing - tube_w)
            .extrude(x_mount_z)
            .center(-ballscrew_x_offset, 0)
            .tag("ballscrew_center")
            .circle(12 / 2)
            .extrude(20, combine="s")
            .workplaneFromTagged("ballscrew_center")
            .center(35 / 2 - 4 / 2, 0)
            .rect(4, 35)
            .extrude(20, combine="s")
            .workplaneFromTagged("ballscrew_center")
            .center(-35 / 2 + 4 / 2, 0)
            .rect(4, 35)
            .extrude(20, combine="s")
            .workplaneFromTagged("ballscrew_center")
            .center(0, -35 / 2 + 4 / 2)
            .rect(25, 4)
            .extrude(20, combine="s")
            .workplaneFromTagged("ballscrew_center")
            .center(0, 35 / 2 - 4 / 2)
            .rect(25, 4)
            .extrude(20, combine="s")
            # .rect(35, 35)
            # .extrude(20, combine="s")
        )
        self.right_bearing_mount = self.ball_screw_model.fixedBearingHoles(
            self.right_bearing_mount
        )
        self.left_bearing_mount = (
            w.copyWorkplane(self.floating_mount_w)
            .transformed(offset=(0, ballscrew_x_offset, 0), rotate=(0, 0, 90))
            .rect(tube_h, x_rail_spacing - tube_w)
            .extrude(x_mount_z)
        )
        # this doesn't work for some reason, i guess just use the fixed bearing mount flipped
        self.left_bearing_mount = self.ball_screw_model.floatingBearingHoles(
            self.left_bearing_mount
        )

        points = [(40 * (i - 5.5), y) for i in range(12) for y in [-18, 18]]
        points += [
            (8 * i + j, y)
            for i in range(-2, 3)
            for y in [-18, 18]
            for j in [column_spacing / 2, -column_spacing / 2]
        ]
        top_w = cq.Workplane().copyWorkplane(self.upper_front_plate)
        bot_w = cq.Workplane().copyWorkplane(self.lower_front_plate)

        def holes(w, ref):
            return w.copyWorkplane(ref).pushPoints(points).hole(3)

        self.upper_front_plate = holes(self.upper_front_plate, top_w)
        self.upper_back_plate = holes(self.upper_back_plate, top_w)
        # self.upper = top_holes(self.upper)
        self.lower_front_plate = holes(self.lower_front_plate, bot_w)
        self.lower_back_plate = holes(self.lower_back_plate, bot_w)
        # self.lower = top_holes(self.lower)

        stepper_mount_w = w.transformed(
            offset=(
                column_spacing / 2 + tube_w / 2,
                ballscrew_x_offset,
                (upper_center + lower_center) / 2,
            ),
            rotate=(0, 90, 0),
        )

        self.stepper_mount = (
            cq.Workplane()
            .copyWorkplane(
                stepper_mount_w.transformed(offset=(0, 0, 0), rotate=(0, 0, 45))
            )
            .rect(58.4, 58.4)
            .extrude(5)
            .copyWorkplane(
                stepper_mount_w.transformed(offset=(0, -ballscrew_x_offset, 0))
            )
            .rect(88, tube_h)
            .extrude(5)
            .copyWorkplane(
                stepper_mount_w.transformed(offset=(0, 0, 0), rotate=(180, 0, 45))
            )
            .rect(47.41, 47.41, forConstruction=True)
            .vertices()
            .hole(5)
            .copyWorkplane(
                stepper_mount_w.transformed(offset=(0, 0, 0), rotate=(180, 0, 0))
            )
            .hole(40)
            .copyWorkplane(
                stepper_mount_w.transformed(offset=(0, -ballscrew_x_offset, 5))
            )
            .rect(78, 30)
            .vertices()
            .hole(5)
            # .copyWorkplane(stepper_mount_w)
            # .circle(40)
            # .extrude(10)
        )
        self.stepper_mount = self.ball_screw_model.fixedBearingHoles(self.stepper_mount)

    def assembly(self):
        return (
            cq.Assembly(name="gantry")
            .add(self.left, name="left column")
            .add(self.right, name="right column")
            .add(self.upper, name="upper platform")
            .add(self.lower, name="lower platform")
            .add(self.upper_rail.rail(), name="upper rail")
            .add(self.lower_rail.rail(), name="lower rail")
            .add(self.upper_front_plate, name="upper front plate")
            .add(self.lower_front_plate, name="lower front plate")
            .add(self.upper_back_plate, name="upper back plate")
            .add(self.lower_back_plate, name="lower back plate")
            .add(self.ball_screw, name="ball screw")
            .add(self.fixed_bearing, name="fixed bearing")
            .add(self.floating_bearing, name="floating bearing")
            .add(self.right_bearing_mount, name="fixed bearing mount")
            .add(self.left_bearing_mount, name="floating bearing mount")
            .add(self.stepper_mount, name="stepper mount")
        )


z_rail_spacing = 110
z_carriage_spacing = 100
z_rail_len = 400
carriage_w = 250
carriage_h = 300
bolster_w = 30
bolster_z = 25

z_spacer_z = 10
z_spacer_w = 40

z_ballscrew_len = 400
z_ballscrew_offset = 0


class Carriage:
    def __init__(self, gantry, x=0, z=0):
        self.z = z

        self.bul = gantry.upper_rail.carriage(x)
        self.bur = gantry.upper_rail.carriage(x + x_carriage_spacing)
        self.bll = gantry.lower_rail.carriage(x)
        self.blr = gantry.lower_rail.carriage(x + x_carriage_spacing)
        back = cq.Workplane("XZ").copyWorkplane(
            self.bul.workplaneFromTagged("mountingholes")
        )
        # why doesn't this track the carriage without me adding x to it?
        self.plate = (
            cq.Workplane()
            .copyWorkplane(back)
            .center(-x - HGR20.L - x_carriage_spacing / 2, -x_rail_spacing / 2)
            # .rect(HGR20.L+x_carriage_spacing, HGR20.W+x_rail_spacing)
            .rect(carriage_w, carriage_h)
            .extrude(plate_z2)
            .faces(">z")  # localselectors plugin
            .workplane()
            .tag("carriage_plate")
        )
        self.plate = gantry.upper_rail.carriageHoles(self.plate, x)
        self.plate = gantry.upper_rail.carriageHoles(self.plate, x + x_carriage_spacing)
        self.plate = gantry.lower_rail.carriageHoles(self.plate, x)
        self.plate = gantry.lower_rail.carriageHoles(self.plate, x + x_carriage_spacing)

        # creating the workplane the rails of the spindle carrier
        # mount to, because HGR works from that plane forward
        self.carrier_plane = (
            cq.Workplane()
            .copyWorkplane(self.plate.workplaneFromTagged("carriage_plate"))
            .transformed(offset=(0, 0, HGR20.H), rotate=(180, 0, 0))
        )

        self.left_rail = HGR20(
            self.carrier_plane.transformed(
                offset=(z_rail_spacing / 2, z, 0), rotate=(0, 0, 90)
            ),
            z_rail_len,
        )
        self.right_rail = HGR20(
            self.carrier_plane.transformed(
                offset=(-z_rail_spacing / 2, z, 0), rotate=(0, 0, 90)
            ),
            z_rail_len,
        )

        u_zcarriage_pos = z_rail_len / 2 - HGR20.L / 2 - z - z_carriage_spacing / 2
        l_zcarriage_pos = u_zcarriage_pos + z_carriage_spacing
        self.ful = self.left_rail.carriage(u_zcarriage_pos)
        self.fll = self.left_rail.carriage(l_zcarriage_pos)
        self.fur = self.right_rail.carriage(u_zcarriage_pos)
        self.flr = self.right_rail.carriage(l_zcarriage_pos)

        self.plate = self.left_rail.carriageHoles(self.plate, u_zcarriage_pos)
        self.plate = self.left_rail.carriageHoles(self.plate, l_zcarriage_pos)
        self.plate = self.right_rail.carriageHoles(self.plate, u_zcarriage_pos)
        self.plate = self.right_rail.carriageHoles(self.plate, l_zcarriage_pos)

        self.ball_nut = gantry.ball_screw_model.nut(x + carriage_w / 2)

        lbolster_w = self.plate.workplaneFromTagged("carriage_plate").transformed(
            offset=(-(carriage_w - bolster_w) / 2, 0, 0)
        )
        rbolster_w = self.plate.workplaneFromTagged("carriage_plate").transformed(
            offset=((carriage_w - bolster_w) / 2, 0, 0)
        )

        points = [(0, 50 * (i - 2.5)) for i in range(6)]

        self.lbolster = (
            cq.Workplane()
            .copyWorkplane(lbolster_w)
            .rect(bolster_w, carriage_h)
            .extrude(bolster_z)
            .copyWorkplane(lbolster_w.transformed(rotate=(0, 180, 0)))
            .pushPoints(points)
            .hole(3)
        )
        self.rbolster = (
            cq.Workplane()
            .copyWorkplane(rbolster_w)
            .rect(bolster_w, carriage_h)
            .extrude(bolster_z)
            .copyWorkplane(lbolster_w.transformed(rotate=(0, 180, 0)))
            .pushPoints(points)
            .hole(3)
        )

        self.plate = (
            self.plate.copyWorkplane(lbolster_w)
            .pushPoints(points)
            .hole(3)
            .copyWorkplane(rbolster_w)
            .pushPoints(points)
            .hole(3)
        )

        self.wp_carrier = (
            cq.Workplane()
            .copyWorkplane(self.plate.workplaneFromTagged("carriage_plate"))
            .transformed(offset=(0, -z, HGR20.H), rotate=(0, 0, 0))
        )
        self.x_nut_mount = (
            cq.Workplane()
            .copyWorkplane(back)
            .center(-x - HGR20.L - x_carriage_spacing / 2, -x_rail_spacing / 2)
            .transformed(rotate=(180, 0, 0))
            .rect(60, 60)
            .extrude(64)
        )

        # these are filled in by the carrier
        self.z_nut_mount = None
        self.z_nut = None

    def assembly(self):
        assmbly = (
            cq.Assembly(name="x carriage")
            .add(self.bul, name="back ul carriage")
            .add(self.bur, name="back ur carriage")
            .add(self.bll, name="back ll carriage")
            .add(self.blr, name="back lr carriage")
            .add(self.plate, name="carrier plate")
            .add(self.left_rail.rail(), name="left rail")
            .add(self.right_rail.rail(), name="right rail")
            .add(self.ful, name="front ul carriage")
            .add(self.fur, name="front ur carriage")
            .add(self.fll, name="front ll carriage")
            .add(self.flr, name="front lr carriage")
            .add(self.lbolster, name="left bolster")
            .add(self.rbolster, name="right bolster")
            .add(self.ball_nut, name="ball nut")
            .add(self.x_nut_mount, name="x nut mount")
        )
        if self.z_nut_mount is not None:
            assmbly = assmbly.add(self.z_nut_mount, name="z nut mount")
        if self.z_nut is not None:
            assmbly = assmbly.add(self.z_nut, name="z nut")
        return assmbly


class Carrier:
    def __init__(self, carriage):
        self.plane = carriage.wp_carrier
        self.left_spacer = (
            cq.Workplane()
            .copyWorkplane(self.plane)
            .transformed(offset=(z_rail_spacing / 2, 0, 0))
            .rect(z_spacer_w, z_rail_len)
            .extrude(z_spacer_z)
        )
        self.right_spacer = (
            cq.Workplane()
            .copyWorkplane(self.plane)
            .transformed(offset=(-z_rail_spacing / 2, 0, 0))
            .rect(z_spacer_w, z_rail_len)
            .extrude(z_spacer_z)
        )

        self.plate = (
            cq.Workplane()
            .copyWorkplane(self.plane)
            .transformed(offset=(0, 0, z_spacer_z))
            .rect(200, z_rail_len)
            .extrude(10)
            .copyWorkplane(self.plane)
            .center(0, z_rail_len / 2 - 40 / 2)
            .rect(68, 40)
            .extrude(20, combine="s")
        )
        self.plate = carriage.left_rail.holes(self.plate)
        self.plate = carriage.right_rail.holes(self.plate)
        self.left_spacer = carriage.left_rail.holes(self.left_spacer)
        self.right_spacer = carriage.right_rail.holes(self.right_spacer)

        self.ballscrew_model = SFU1204(
            self.plane.transformed(
                offset=(
                    0,
                    z_ballscrew_len / 2 + 20,
                    # 25,
                    z_spacer_z - 25 + 5,  # offset into the spindle plate 5mm
                ),  # -(HGR20.H - z_spacer_z) / 2 - 1.2),
                rotate=(180, 0, 0),
            ),
            z_ballscrew_len,
        )
        self.ballscrew = self.ballscrew_model.ballscrew()
        self.fixed_bearing = self.ballscrew_model.fixedBearing()
        self.floating_bearing = self.ballscrew_model.floatingBearing()

        angle_z = 0.188 * 25.4
        angle_w = 2.0 * 25.4
        self.stepper_mount = (
            cq.Workplane()
            .copyWorkplane(self.plane)
            .transformed(offset=(-200 / 2, z_rail_len / 2, 0), rotate=(0, 90, 0))
            .sketch()
            .push([(-10 - 10 - angle_z + angle_w / 2, angle_z / 2 + 5)])
            .rect(angle_w, angle_z)
            .push([(-10 - 10 - angle_z / 2, 5 + angle_z - angle_w / 2)])
            .rect(angle_z, angle_w)
            .finalize()
            .extrude(200)
        )

        carriage.z_nut = self.ballscrew_model.nut(z_ballscrew_len / 2 + carriage.z)
        # TODO add carriage nut mount to carriage

    def assembly(self):
        return (
            cq.Assembly(name="z carrier")
            .add(self.left_spacer, name="left spacer")
            .add(self.right_spacer, name="right spacer")
            .add(self.plate, name="z plate")
            .add(self.ballscrew, name="z ballscrew")
            .add(self.fixed_bearing, name="z fixed bearing")
            .add(self.floating_bearing, name="z floating bearing")
            .add(self.stepper_mount, name="stepper motor/ballscrew mount")
        )


class SpindleMount:
    def __init__(self, carrier, z_offset=0):
        self.plane = (
            carrier.plane.transformed(offset=(0, z_offset, z_spacer_z + 10))
            .transformed(rotate=(-90, 0, 0))
            .transformed(rotate=(0, 0, 180))
        )
        self.spindleaxis = cq.Workplane().copyWorkplane(self.plane).center(0, 54)
        self.mount = (
            cq.Workplane()
            .copyWorkplane(self.plane)
            .hLineTo(184 / 2)
            .vLineTo(17)
            .hLineTo(113 / 2)
            .vLineTo(82)
            .hLineTo(-113 / 2)
            .vLineTo(17)
            .hLineTo(-184 / 2)
            .vLineTo(0)
            .close()
            .extrude(78)
            .copyWorkplane(self.spindleaxis)
            .cylinder(78, 104 / 2, angle=180, centered=[True, True, False])
            .copyWorkplane(self.spindleaxis)
            .cylinder(78, 80 / 2, centered=[True, True, False], combine="s")
        )

    def mounting_holes(self, w):
        # TODO
        pass

    def assembly(self):
        return cq.Assembly(name="spindle motor mount").add(self.mount, name="mount")


g = Gantry(cq.Workplane())
c = Carriage(g, 150, -50)
ca = Carrier(c)
m = SpindleMount(ca, -200)

# TODO add ball screws and associated assembly

full_assembly = (
    cq.Assembly()
    .add(g.assembly(), name="gantry")
    .add(c.assembly(), name="carriage")
    .add(ca.assembly(), name="carrier")
    .add(m.assembly(), name="spindle mount")
)

show_object(full_assembly)
# show_object(ca.assembly())
# show_object(c.assembly())
# show_object(g.assembly())
