"""Size the legs and the drivetrain from a height budget and real parts.

Answers three questions before any simulation happens:
  1. How long should each leg segment be, given a target standing height?
  2. What joint torque does that geometry demand?
  3. What reduction, if any, do the chosen motors need to supply it?

  python sizing.py                       # the current build
  python sizing.py --height 0.6 --wheel 0.25

The geometry result is worth stating up front. With the hip and the wheel axle
on one vertical line and two equal segments bent by theta from vertical:

    drop      D = 2 L cos(theta)          so   L = D / (2 cos theta)
    knee load T = (W/2) L sin(theta)      so   T = (W D / 4) tan(theta)

Knee torque depends on the bend angle and NOT on the segment length, once the
height is fixed. Longer legs held straighter cost nothing extra in torque. The
bend angle is the whole design variable, and it is bounded from below by the
control authority and suspension travel you need, not by strength.
"""

import argparse
import math

MOTORS = {
    "DD-5050Pro": {
        "kv": 300, "v_max": 36, "v_nom": 24,
        "rated_a": 33, "max_a": 67,
        "rated_torque": 1.2, "max_torque": 2.5,
        "rated_w": 900, "max_w": 1800,
        "mass": 0.430, "size": "50x50mm",
    },
    "DD-4125": {
        "kv": 300, "v_max": 25.2, "v_nom": 24,
        "rated_a": 12, "max_a": 18.5,
        "rated_torque": 8.27 / 300 * 12, "max_torque": 8.27 / 300 * 18.5,
        "rated_w": 178, "max_w": 178,
        "mass": 0.200, "size": "41x25mm",
    },
}

# Carbon fibre tube, unidirectional-dominant layup.
CF_DENSITY = 1600.0     # kg/m3
CF_E = 70e9             # Pa, conservative for a filament-wound tube


def tube(od_mm, wall_mm=2.0):
    od, wall = od_mm / 1000, wall_mm / 1000
    idm = od - 2 * wall
    area = math.pi / 4 * (od ** 2 - idm ** 2)
    I = math.pi / 64 * (od ** 4 - idm ** 4)
    return {"od_mm": od_mm, "wall_mm": wall_mm, "kg_per_m": area * CF_DENSITY,
            "area": area, "I": I, "EI": CF_E * I}


def solve(height=0.60, wheel_d=0.25, chassis_h=0.12, theta_deg=25.0,
          tube_od=40.0, joint_motor="DD-5050Pro", wheel_motor="DD-4125",
          wheel_ratio=1.0, extra_mass=2.2, gearbox_mass=0.25,
          dynamic_factor=2.0):
    wheel_r = wheel_d / 2
    theta = math.radians(theta_deg)

    # Hips sit on the chassis centre line, so the drop the legs must cover is
    # from there to the wheel axle.
    chassis_z = height - chassis_h / 2
    drop = chassis_z - wheel_r
    if drop <= 0:
        raise SystemExit("wheels alone exceed the height budget")
    L = drop / (2 * math.cos(theta))

    t = tube(tube_od)
    jm, wm = MOTORS[joint_motor], MOTORS[wheel_motor]

    seg_mass = t["kg_per_m"] * L
    mass = (4 * jm["mass"] + 2 * wm["mass"] + 4 * seg_mass
            + 2 * 0.35 + extra_mass)

    W = mass * 9.81
    knee_static = (W * drop / 4) * math.tan(theta)
    hip_static = knee_static                      # symmetric two-link stance
    knee_dyn = knee_static * dynamic_factor

    # Wheels: what tilt can it catch? An inverted pendulum needs base
    # acceleration g*tan(lean) to hold a lean, so available traction sets the
    # largest disturbance the machine can recover from at all.
    wheel_tau = wm["max_torque"] * wheel_ratio
    traction = 2 * wheel_tau / wheel_r
    accel = traction / mass
    max_lean = math.degrees(math.atan(accel / 9.81))
    no_load_rad = jm["kv"] * jm["v_nom"] * 2 * math.pi / 60
    wheel_no_load = wm["kv"] * min(wm["v_max"], 24) * 2 * math.pi / 60 / wheel_ratio
    top_speed = wheel_no_load * wheel_r

    return {
        "height": height, "chassis_z": chassis_z, "wheel_r": wheel_r,
        "drop": drop, "L": L, "theta_deg": theta_deg, "mass": mass,
        "seg_mass": seg_mass, "tube": t,
        "knee_static": knee_static, "hip_static": hip_static, "knee_dyn": knee_dyn,
        "joint_motor": joint_motor, "wheel_motor": wheel_motor,
        "ratio_cont": knee_static / jm["rated_torque"],
        "ratio_dyn": knee_dyn / jm["max_torque"],
        "joint_no_load_rad": no_load_rad,
        "traction_n": traction, "accel": accel, "max_lean_deg": max_lean,
        "top_speed": top_speed, "wheel_ratio": wheel_ratio,
        "extend_z": wheel_r + 2 * L + chassis_h / 2,
        "crouch_z": wheel_r + 2 * L * math.cos(math.radians(60)) + chassis_h / 2,
    }


def report(r):
    jm, wm = MOTORS[r["joint_motor"]], MOTORS[r["wheel_motor"]]
    print(f"\nTARGET  {r['height']*1000:.0f}mm standing, {r['wheel_r']*2000:.0f}mm wheels")
    print("-" * 66)
    print(f"  chassis centre      {r['chassis_z']*1000:6.0f} mm")
    print(f"  hip to wheel axle   {r['drop']*1000:6.0f} mm")
    print(f"  SEGMENT LENGTH      {r['L']*1000:6.0f} mm  each (thigh = shank)")
    print(f"  total leg           {r['L']*2000:6.0f} mm  at {r['theta_deg']:.0f} deg bend")
    print(f"  height range        {r['crouch_z']*1000:.0f} - {r['extend_z']*1000:.0f} mm "
          f"(crouched to straight)")

    t = r["tube"]
    print(f"\n  CF tube {t['od_mm']:.0f}mm OD x {t['wall_mm']:.0f}mm wall: "
          f"{t['kg_per_m']:.3f} kg/m -> {r['seg_mass']*1000:.0f} g per segment")
    print(f"  estimated total mass {r['mass']:.2f} kg")

    print(f"\nJOINTS  {r['joint_motor']}  "
          f"{jm['rated_torque']:.1f} Nm rated / {jm['max_torque']:.1f} Nm peak")
    print(f"  static hold needed  {r['knee_static']:6.2f} Nm at hip and knee")
    print(f"  with {2.0:.0f}x dynamic margin {r['knee_dyn']:6.2f} Nm")
    print(f"  reduction for rated {r['ratio_cont']:6.2f} : 1")
    print(f"  reduction for peak  {r['ratio_dyn']:6.2f} : 1")
    rec = max(1.0, math.ceil(max(r["ratio_cont"], r["ratio_dyn"])))
    verdict = ("DIRECT DRIVE IS ENOUGH" if rec <= 1 else f"NEEDS ABOUT {rec}:1")
    print(f"  ->  {verdict}")
    if rec > 1:
        print(f"      at {rec}:1 the joint gives "
              f"{jm['rated_torque']*rec:.1f} Nm continuous, "
              f"{jm['max_torque']*rec:.1f} Nm peak, "
              f"{r['joint_no_load_rad']/rec:.0f} rad/s no-load")

    print(f"\nWHEELS  {r['wheel_motor']} at {r['wheel_ratio']:.0f}:1  "
          f"{wm['max_torque']:.2f} Nm peak per motor")
    print(f"  traction            {r['traction_n']:6.1f} N total")
    print(f"  base acceleration   {r['accel']:6.2f} m/s2")
    print(f"  RECOVERS LEAN UP TO {r['max_lean_deg']:6.1f} deg")
    print(f"  top speed           {r['top_speed']:6.1f} m/s "
          f"({r['top_speed']*3.6:.0f} km/h)")
    if r["max_lean_deg"] < 12:
        print("  ->  thin. A balancer that can only catch a small lean will")
        print("      survive flat floors and fail on shoves, slopes and thresholds.")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--height", type=float, default=0.60)
    ap.add_argument("--wheel", type=float, default=0.25)
    ap.add_argument("--chassis-h", type=float, default=0.12)
    ap.add_argument("--theta", type=float, default=25.0)
    ap.add_argument("--tube", type=float, default=40.0)
    ap.add_argument("--wheel-ratio", type=float, default=1.0)
    ap.add_argument("--sweep", action="store_true")
    args = ap.parse_args()

    r = solve(args.height, args.wheel, args.chassis_h, args.theta,
              args.tube, wheel_ratio=args.wheel_ratio)
    report(r)

    if args.sweep:
        print("\n\nBEND ANGLE SWEEP (height fixed, so this is the real trade-off)")
        print(f"  {'bend':>5} {'segment':>9} {'knee Nm':>9} {'ratio':>7}   note")
        for th in (10, 15, 20, 25, 30, 35, 40):
            s = solve(args.height, args.wheel, args.chassis_h, th, args.tube,
                      wheel_ratio=args.wheel_ratio)
            note = ("near-singular, poor authority" if th <= 12 else
                    "little travel" if th <= 17 else
                    "good compromise" if th <= 27 else
                    "lots of travel, high torque")
            print(f"  {th:4.0f}d {s['L']*1000:8.0f}mm {s['knee_static']:9.2f} "
                  f"{max(s['ratio_cont'], s['ratio_dyn']):7.2f}   {note}")

        print("\n\nWHEEL REDUCTION SWEEP")
        print(f"  {'ratio':>6} {'traction N':>11} {'m/s2':>7} {'max lean':>9} {'top m/s':>9}")
        for wr in (1, 2, 3, 4, 6, 8):
            s = solve(args.height, args.wheel, args.chassis_h, args.theta,
                      args.tube, wheel_ratio=wr)
            print(f"  {wr:5.0f}:1 {s['traction_n']:11.1f} {s['accel']:7.2f} "
                  f"{s['max_lean_deg']:8.1f}d {s['top_speed']:9.1f}")


if __name__ == "__main__":
    main()
