from nhf.build import Model, TargetKind, target, assembly, submodel
from nhf.materials import Role, Material
import nhf.utils

import math
from dataclasses import dataclass
import cadquery as Cq

@dataclass
class Onbashira(Model):

    n_side: int = 6
    # Dimensions of each side panel
    side_width: float = 200.0
    side_length: float = 600.0

    side_thickness: float = 25.4 / 8

    # Dimensions of gun barrels
    barrel_diam: float = 25.4 * 2
    barrel_length: float = 300.0
    # Radius from barrel centre to axis
    rotation_radius: float = 75.0
    n_bearing_balls: int = 24
    # Size of ball bearings
    bearing_ball_diam: float = 25.4 * 1/2
    bearing_ball_gap: float = .5
    bearing_height: float = 40.0
    bearing_thickness: float = 20.0
    bearing_track_radius: float = 120.0
    # Gap between the inner and outer bearing disks
    bearing_gap: float = 10.0
    bearing_disk_thickness: float = 25.4 / 8

    material_side: Material = Material.WOOD_BIRCH
    material_bearing: Material = Material.PLASTIC_PLA
    material_bearing_ball: Material = Material.ACRYLIC_TRANSPARENT
    material_brace: Material = Material.METAL_AL

    def __post_init__(self):
        assert self.n_side >= 3
        # Bulk must be large enough for the barrel + bearing to rotate
        assert self.bulk_radius - self.side_thickness - self.bearing_thickness - self.bearing_diam > self.rotation_radius + self.barrel_diam / 2
        assert self.bearing_height > self.bearing_diam
        assert self.bearing_gap < 0.95 * self.bearing_ball_diam

    @property
    def angle_side(self) -> float:
        return 360 / self.n_side
    @property
    def angle_dihedral(self) -> float:
        return 180 - self.angle_side
    @property
    def bulk_radius(self) -> float:
        """
        Calculate radius of the bulk to the centre
        """
        return self.side_width / 2 / math.tan(math.radians(self.angle_side / 2))
    @property
    def bearing_diam(self) -> float:
        return self.bearing_ball_diam + self.bearing_ball_gap

    @property
    def bearing_disk_gap(self) -> float:
        diag = self.bearing_ball_diam
        dx = self.bearing_gap
        return math.sqrt(diag ** 2 - dx ** 2)

    @target(name="bearing-stator", kind=TargetKind.DXF)
    def profile_bearing_stator(self) -> Cq.Sketch:
        return (
            Cq.Sketch()
            .regularPolygon(self.side_width, self.n_side)
            .circle(self.bearing_track_radius + self.bearing_gap/2, mode="s")
        )
    def bearing_stator(self) -> Cq.Workplane:
        return (
            Cq.Workplane()
            .placeSketch(self.profile_bearing_stator())
            .extrude(self.bearing_disk_thickness)
        )
    @target(name="bearing-rotor", kind=TargetKind.DXF)
    def profile_bearing_rotor(self) -> Cq.Sketch:
        return (
            Cq.Sketch()
            .circle(self.bearing_track_radius - self.bearing_gap/2)
            .regularPolygon(self.rotation_radius, self.n_side)
            .vertices()
            .circle(self.barrel_diam/2, mode="s")
        )
    def bearing_rotor(self) -> Cq.Workplane:
        return (
            Cq.Workplane()
            .placeSketch(self.profile_bearing_rotor())
            .extrude(self.bearing_disk_thickness)
        )
    @target(name="bearing-gasket", kind=TargetKind.DXF)
    def profile_bearing_gasket(self) -> Cq.Sketch:
        pass


    @target(name="pipe", kind=TargetKind.DXF)
    def pipe(self) -> Cq.Sketch:
        pass

    @target(name="side-panel", kind=TargetKind.DXF)
    def profile_side_panel(self) -> Cq.Sketch:
        return (
            Cq.Sketch()
            .rect(self.side_width, self.side_length)
        )

    def side_panel(self) -> Cq.Workplane:
        w = self.side_width
        l = self.side_length
        result = (
            Cq.Workplane()
            .placeSketch(self.profile_side_panel())
            .extrude(self.side_thickness)
        )
        intersector = (
            Cq.Workplane('XZ')
            .polyline([
                (-w/2, 0),
                (w/2, 0),
                (0, self.bulk_radius),
            ])
            .close()
            .extrude(l)
            .translate(Cq.Vector(0, l/2,0))
        )
        # Intersect the side panel
        return result * intersector



    def bearing_ball(self) -> Cq.Solid:
        return Cq.Solid.makeSphere(radius=self.bearing_ball_diam/2, angleDegrees1=-90)

    def assembly_rotor(self) -> Cq.Assembly:
        z_lower = -self.bearing_disk_gap - self.bearing_disk_thickness
        a = (
            Cq.Assembly()
            .addS(
                self.bearing_stator(),
                name="stator1",
                material=self.material_bearing,
                role=Role.STATOR,
                loc=Cq.Location(0, 0, self.bearing_disk_gap/2)
            )
            .addS(
                self.bearing_rotor(),
                name="rotor1",
                material=self.material_bearing,
                role=Role.ROTOR,
                loc=Cq.Location(0, 0, self.bearing_disk_gap/2)
            )
            .addS(
                self.bearing_stator(),
                name="stator2",
                material=self.material_bearing,
                role=Role.STATOR,
                loc=Cq.Location(0, 0, z_lower)
            )
            .addS(
                self.bearing_rotor(),
                name="rotor2",
                material=self.material_bearing,
                role=Role.ROTOR,
                loc=Cq.Location(0, 0, z_lower)
            )
        )
        for i in range(self.n_bearing_balls):
            ball = self.bearing_ball()
            loc = Cq.Location.rot2d(i * 360/self.n_bearing_balls) * Cq.Location(self.bearing_track_radius, 0, 0)
            a = a.addS(
                ball,
                name=f"bearing_ball{i}",
                material=self.material_bearing_ball,
                role=Role.BEARING,
                loc=loc,
            )
        return a

    def assembly(self) -> Cq.Assembly:
        a = Cq.Assembly()
        side = self.side_panel()
        r = self.bulk_radius
        a = a.add(self.assembly_bearing(), name="bearing")
        for i in range(self.n_side):
            a = a.addS(
                side,
                name=f"side{i}",
                material=self.material_side,
                role=Role.STRUCTURE | Role.DECORATION,
                loc=Cq.Location.rot2d(i*360/self.n_side) * Cq.Location(-r,0,0,90,0,90),
            )
        return a