# @author Arnaud Morin <arnaud.gfpv@mailops.fr>
# SPDX-License-Identifier: Apache-2.0
#

import os

import cadquery as cq
from cadquery import exporters

import frame_params as P
from arm import result as ARM
from motor_base import result as MOTOR_BASE
from plate import result as PLATE
from spar import result as SPAR
from camera_mount import result as CAMERA_MOUNT
import standoff as SO
from standoff import result as STANDOFF
from lollipop import result as LOLLIPOP
from fc import result as FC, PCB_DENSITY

OUT = "build"

TPU_DENSITY = 1.21e-3    # g/mm3, for the printed camera_mount only

COLORS = {
    "plate":        cq.Color(0.16, 0.17, 0.20),   # near-black carbon
    "arm":          cq.Color(0.85, 0.24, 0.16),   # red
    "motor_base":   cq.Color(0.20, 0.45, 0.80),   # blue
    "spar":         cq.Color(0.22, 0.62, 0.35),   # green
    "camera_mount": cq.Color(0.95, 0.78, 0.09),   # yellow TPU
    "standoff":     cq.Color(0.75, 0.76, 0.78),   # aluminium
    "lollipop":     cq.Color(0.95, 0.78, 0.09),   # yellow TPU
    "fc":           cq.Color(0.42, 0.28, 0.62),   # PCB purple
    "prop":         cq.Color(0.55, 0.55, 0.60),   # opaque: a ring hides
}

DENSITY = {
    "plate": P.CF_DENSITY, "arm": P.CF_DENSITY,
    "motor_base": P.CF_DENSITY, "spar": P.CF_DENSITY,
    "camera_mount": TPU_DENSITY,
    "standoff": SO.ALU_DENSITY,
    "lollipop": TPU_DENSITY,
    "fc": PCB_DENSITY,
}

# --- propellers -------------------------------------------------------------
# Shown, never made.  A prop is not a part of this frame: it is the volume the
# frame has to stay out of, so it goes into the assembly for the picture and
# is kept out of PARTS, which is what drives the BOM and the exports.
PROP_D = 78.0            # 3 inch.  Same number as the motor-to-motor gap, so
                         # on a true X neighbouring discs just touch -- that
                         # is the whole point of drawing them
PROP_T = 1.0             # token thickness; this is a swept circle, not a blade
PROP_RING = 1.5          # radial width.  A ring, not a filled disc: the tip
                         # circle is the whole point and a disc just buries
                         # the frame under it
PROP_MOTOR_H = 15.0      # motors are not modelled either, so this is only how
                         # far above the motor base the disc floats -- about a
                         # 1404's height
# The prop's colour lives in COLORS with the others, deliberately.  A bare
# module-level cq.Color breaks CQ-editor: it walks the module's globals after
# running the script and compares them against the object being shown, and
# Color.__eq__ does self.toTuple() == other.toTuple() -- which blows up on an
# Assembly, since Assembly has no toTuple.  Inside a dict it is never
# compared, which is why the four part colours never caused this.
SHOW_PROPS = False        # the discs are for looking at, so if a viewer will
                         # not show the assembly with them in it, turn them
                         # off here rather than unpicking build()

# name, solid, thickness, flat-pattern face selector ("flat" lies in XY already,
# "edge" is a vertical plate and has to be tipped down for the cutter, "solid"
# is a printed part with no flat pattern -- STEP only, no DXF)
PARTS = [
    ("plate",        PLATE,        P.THICKNESS, "flat"),
    ("arm",          ARM,          P.THICKNESS, "edge"),
    ("motor_base",   MOTOR_BASE,   P.THICKNESS, "flat"),
    ("spar",         SPAR,         P.THICKNESS, "edge"),
    ("camera_mount", CAMERA_MOUNT, P.THICKNESS,   "solid"),
    ("standoff",     STANDOFF,     SO.OD,        "solid"),
    ("lollipop",     LOLLIPOP,     3.5,          "solid"),
    ("fc",           FC,           11.0,         "solid"),
]


def _loc(x, y, z, rot_z, flip=0.0):
    """Placement: flip about the part's own X axis first, then rot_z about the
    world Z axis, then translate.  Matches the helpers' documented order."""
    L = cq.Location(cq.Vector(x, y, z), cq.Vector(0, 0, 1), rot_z)
    if flip:
        L = L * cq.Location(cq.Vector(0, 0, 0), cq.Vector(1, 0, 0), flip)
    return L


def pieces():
    """[(name, tag, solid, Location)] -- the 25 bodies, straight off the
    placement helpers: 2 plates, 4 arms, 8 motor bases, 4 spars, 4 standoffs,
    1 camera mount, 1 antenna holder, 1 fc stack."""
    out = []
    for i, (x, y, z, rz) in enumerate(P.plate_placements()):
        out.append(("plate", "plate_%s" % ("bottom", "top")[i],
                    PLATE, _loc(x, y, z, rz)))
    for i, (x, y, z, rz) in enumerate(P.arm_placements()):
        out.append(("arm", "arm_%d" % i, ARM, _loc(x, y, z, rz)))
    for i, (x, y, z, rz, fl) in enumerate(P.motor_placements()):
        # first four sit on top of the corners, last four underneath them
        side = "top" if i < 4 else "bot"
        out.append(("motor_base", "motor_base_%s%d" % (side, i % 4),
                    MOTOR_BASE, _loc(x, y, z, rz, fl)))
    for i, (x, y, z, rz, fl) in enumerate(P.spar_placements()):
        out.append(("spar", "spar_%d" % i, SPAR, _loc(x, y, z, rz, fl)))
    # camera_mount: bolts to the top plate's own 25.5x25.5 FC holes, which it
    # matches exactly -- no rotation needed, just sat on the top face
    out.append(("camera_mount", "camera_mount_0", CAMERA_MOUNT,
                _loc(0.0, 0.0, P.Z_TOP_FACE, 0.0)))
    # standoffs: they stand on the bottom plate and their height is the plate
    # gap, so they meet the top plate's underside exactly
    for i, (x, y, z, rz) in enumerate(P.standoff_placements()):
        out.append(("standoff", "standoff_%d" % i, STANDOFF, _loc(x, y, z, rz)))
    # lollipop: clipped over the rear standoff pair, the only pair 28 apart.
    # Its own pads sit at (5, +-14), so they run along Y -- the 90 deg turn
    # is what lays them across the X pair, and the -25 then carries them from
    # y = 5 onto y = -20.  z = 7 puts the pads over the posts and is the one
    # height that clears both plates.
    out.append(("lollipop", "lollipop_0", LOLLIPOP,
                _loc(0.0, -22.0, 7.0, 90.0)))
    # fc: the stack stands on the bottom plate's top face, over the 25.5 square
    # the plate is drilled for.  fc.py puts its own origin at the underside of
    # its posts, so that face is the whole placement.  The 45 deg is not
    # cosmetic: fc.py sets its posts on the diagonals at (+-12.75, +-12.75),
    # while the plate drills the same 25.5 square with its holes on the axes at
    # 25.5/sqrt(2) = 18.03.  Same radius, turned an eighth of a turn, so without
    # this the screws miss.
    out.append(("fc", "fc_0", FC, _loc(0.0, 0.0, P.THICKNESS, 45.0)))
    return out


def propellers():
    """(tag, solid, Location) for the four prop rings, centred on the motor
    axes.  Representation only -- see the note by PROP_D."""
    ring = (cq.Workplane("XY")
            .circle(PROP_D / 2.0)
            .circle(PROP_D / 2.0 - PROP_RING)
            .extrude(PROP_T))
    z = P.Z_MOTOR_BASE + P.THICKNESS + PROP_MOTOR_H
    return [("prop_%d" % i, ring, _loc(x, y, z, 0.0))
            for i, (x, y) in enumerate(P.motor_positions())]


def build(ps=None, props=True):
    asy = cq.Assembly(name="frame_3in")
    for name, tag, solid, loc in (ps or pieces()):
        asy.add(solid, name=tag, loc=loc, color=COLORS[name])
    if props:
        for tag, solid, loc in propellers():
            asy.add(solid, name=tag, loc=loc, color=COLORS["prop"])
    return asy


# --- flat profile for the cutter --------------------------------------------

def cut_profile(part, kind):
    """The single face the part is milled out of, laid into XY at z = 0.

    A vertical plate is tipped down about X so that its -Y side face comes up
    normal-up with the part's own +Z running up the page: the outline is then
    the right way round, not mirrored."""
    face = part.faces("<Z" if kind == "flat" else "<Y").val()
    if kind == "edge":
        face = face.rotate((0, 0, 0), (1, 0, 0), -90.0)
    z = min(v.Z for v in face.Vertices())
    return cq.Workplane("XY").add(face.translate((0.0, 0.0, -z)))


# --- measurements ------------------------------------------------------------


def bom(ps):
    print("bill of materials")
    print("  %-12s %4s  %6s  %9s  %9s" % ("part", "qty", "t/mm", "g each", "g"))
    total = 0.0
    for name, solid, thick, _ in PARTS:
        qty = sum(1 for n, _, _, _ in ps if n == name)
        each = solid.val().Volume() * DENSITY[name]
        total += each * qty
        print("  %-12s %4d  %6.1f  %9.3f  %9.3f"
              % (name, qty, thick, each, each * qty))
    print("  %-12s %4d  %6s  %9s  %9.3f" % ("FRAME", len(ps), "", "", total))
    print("  carbon at %.2f g/cm3, TPU parts at %.2f g/cm3, aluminium"
          " standoffs at %.2f g/cm3," % (P.CF_DENSITY * 1e3, TPU_DENSITY * 1e3,
                                         SO.ALU_DENSITY * 1e3))
    print("  fc as a solid block of FR4 at %.2f g/cm3 -- an envelope, not a"
          " board" % (PCB_DENSITY * 1e3,))
    print("  hardware counted is the standoffs and the fc stack -- no screws,"
          " motors, props or receiver")
    return total


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--no-png", action="store_true",
                        help="skip rendering PNGs -- they're slow, step/dxf export is not")
    args = parser.parse_args()

    os.makedirs(OUT, exist_ok=True)
    os.makedirs(OUT + '/step/', exist_ok=True)
    os.makedirs(OUT + '/dxf/', exist_ok=True)
    os.makedirs(OUT + '/stl/', exist_ok=True)

    ps = pieces()
    asy = build(ps)
    bom(ps)

    print()
    # Export each piece in step, stl and dxf
    for name, solid, thick, kind in PARTS:
        step = os.path.join(OUT + '/step', "%s.step" % name)
        print(f'Exporting {step}')
        exporters.export(solid.val(), step)
        # STL for the slicer.  0.01 mm of chordal deviation is far below what
        # any printer resolves, and costs a third of what 0.005 does.
        stl = os.path.join(OUT + '/stl', "%s.stl" % name)
        print(f'Exporting {stl}')
        exporters.export(solid.val(), stl,
                         tolerance=0.01, angularTolerance=0.1)
        if kind == "solid":
            continue      # printed part, no flat pattern to cut
        dxf = os.path.join(OUT + '/dxf', "%s.dxf" % name)
        print(f'Exporting {dxf}')
        exporters.exportDXF(cut_profile(solid, kind), dxf)

    if args.no_png:
        return

    # Render the frame in png
    os.makedirs(OUT + '/png/', exist_ok=True)
    from render import render, render_grid
    for view in ("iso", "top", "front"):
        render(asy, "%s/frame_%s.png" % (OUT + '/png', view), view=view)
    render_grid(asy, "%s/frame_grid.png" % (OUT + '/png'))


result = build(props=SHOW_PROPS)

if "show_object" not in globals():          # running outside CQ-editor
    def show_object(*args, **kwargs):
        pass

show_object(result)

if __name__ == "__main__":
    main()