"""The one implementation of Johnny 6's control laws.

Five copies of the balance law used to exist - three in serve.py, one in
learn_balance.py's episodes, one in assist_lab.py - and they drifted: the live
demo negated the foot-shift reflex the trainer had measured and fixed, the demo
dropped recover's brace slots, and drive mode re-derived balance with enough
small differences to produce a 1.9 Nm limit cycle while standing still. The fix
recorded in AUDIT.md and NEXT.md is structural: one module owns the law, and
the trainer, the scorer, the live demo and the drive mode all call it. A
divergence is now an import error, not a Saturday.

Everything here is deliberately plain: readable gains in, one command out,
state carried explicitly. These are the functions that would be transcribed
into firmware.
"""

import math

import numpy as np

L_SEG = 0.200          # thigh and shank, both 200mm
R_WHEEL = 0.0625


# --------------------------------------------------------------- kinematics

def stance_to_leg(h, dx=0.0, branch=-1.0):
    """Two-link leg: put the axle at height h, and dx ahead of the hip.

    dx=0 reproduces the simple case - hip +q/2, knee -q, axle directly under
    the hip. A non-zero dx moves the wheel fore or aft WITHOUT changing height,
    which is how the robot shifts its mass over the contact patch: sliding the
    feet forward under a forward dive is the same as throwing the body back.

    branch picks which way the knee folds. Both postures reach the same foot
    position, but each can only recover a fall one way, and swapping is a
    large, slow movement - choose by direction of travel, not mid-fall.
    """
    H = max(0.02, h - R_WHEEL)
    r = math.hypot(dx, H)
    r = min(r, 2.0 * L_SEG - 1e-4)              # keep it reachable
    c = max(-1.0, min(1.0, (r * r - 2.0 * L_SEG ** 2) / (2.0 * L_SEG ** 2)))
    knee = branch * math.acos(c)
    hip = math.atan2(dx, H) - math.atan2(L_SEG * math.sin(knee),
                                         L_SEG * (1.0 + math.cos(knee)))
    return hip, knee


def lean_trim_deg(stance_mm):
    """The lean the robot must hold at a given stance height to stay put.

    REFITTED FOR THE R2 BODY (13 Aug): the r1 polynomial was not just stale,
    it was BACKWARDS - at a 380 crouch it commanded -6.1 degrees where the
    r2 body truly needs +7.0, opposite in sign at every height, because the
    knee-mounted drive motors flipped the crouch CoM geometry. That 13-degree
    error is why bob wandered (a wrong trim can only be "balanced" by
    accelerating), why gate changes rocked, and why look measured +50mm of
    feet-forward compensation (~6.6 degrees - the same error, absorbed
    geometrically). Solved kinematically on the r2 model: pose the body at
    each stance, find the pitch that puts the total CoM over the axle.
    """
    s = float(stance_mm)
    # sign settled empirically (the derivation frame was flipped): this
    # orientation measured 139/98/75mm of gate-change rocking against the
    # stale r1 curve's 184/155/131 - and decaying, not lingering
    kin = -(-3.10426766e-06 * s ** 3 + 3.58345452e-03 * s ** 2
            - 1.42838739e+00 * s + 2.02662285e+02)
    # DRIFT-NULL CORRECTION (15 Aug, operator caught the crouch bench
    # creeping backwards): the kinematic pose puts the CoM over the
    # axle, but the dynamic equilibrium includes tyre-contact
    # migration - measured by holding each bench's baked law at its
    # stance and finding the trim that nulls the drift. The error is
    # systematic and grows with crouch: -0.75 deg at 340, -0.45 at
    # 400, +0.21 at 460. Interpolated, clamped at the measured ends.
    _fix = ((340.0, -0.75), (400.0, -0.45), (460.0, 0.21))
    if s <= _fix[0][0]:
        c = _fix[0][1]
    elif s >= _fix[-1][0]:
        c = _fix[-1][1]
    else:
        c = _fix[0][1]
        for (s0, c0), (s1, c1) in zip(_fix, _fix[1:]):
            if s0 <= s <= s1:
                f = (s - s0) / (s1 - s0)
                c = c0 + f * (c1 - c0)
                break
    return kin + c


# ------------------------------------------------------------- balance law

def lean_reference(station, speed, odo, clamp):
    """The lean target: speed and odometry pick a small lean, clamped.

    The clamp is load-bearing, not incidental - flattening the law into a
    plain weighted sum was measured at 0.33, worse than all-zeros, because a
    0.66 m/s roll then asks for a 19 degree lean.
    """
    return max(-abs(clamp), min(abs(clamp),
               station * (0.5 * speed + 0.1 * odo)))


def wheel_command(kp, kd, lean, rate, ref=0.0, trim=0.0):
    """The inner balance law. Returns a normalised command in [-1, 1];
    the actuator maps 1.0 to its full 1.4 Nm."""
    return float(np.clip(-kp * (lean - ref - trim) + kd * rate, -1.0, 1.0))


def slew(cur, want, rate, dt):
    """Move all leg joints TOGETHER: one scale factor for the whole move so
    hip and knee arrive in ratio and the axle stays under the hip."""
    d = want - cur
    far = float(np.max(np.abs(d)))
    if far > 1e-9:
        cur = cur + d * (min(far, rate * dt) / far)
    return cur


def leg_rate(leg_speed, floor=0.1, unlimited_at=None):
    """Leg slew rate in rad/s from a stage's leg-speed gain (sign-blind,
    floored - a negative half-line was an inescapable dead zone that froze
    the legs). Stages that may use the servos flat out (level, on the BNO
    reference) pass unlimited_at: above it the ramp opens to servo speed.

    Base is 150 deg/s per unit. The old 60 was the accidental-stabiliser
    era: the accel-polluted roll reference punished fast legs for their own
    motion, and the crawl was load-bearing. On the BNO reference with the
    r2 peak-torque legs it is only a brake - the searches now choose their
    speed on a scale where 1.0 means genuinely quick."""
    ls = max(floor, abs(leg_speed))
    if unlimited_at is not None and ls >= unlimited_at:
        return math.radians(1e5)
    return math.radians(150.0 * ls)


def gaze_pose(stance_m, dx_w, phi, branch=-1.0):
    """Leg IK for a PITCHED chassis: hips and knees move together so the
    wheel stays planted at dx_w (world) while the chassis rotates by phi.

    The naive version - adding the gaze to the hip joint alone - swings the
    whole leg and translates the wheel (the leg is the long lever, the
    chassis the heavy part); measured at 5m of drift. Rotating the leg
    TARGET into the pitched chassis frame is the coordinated pose the
    operator specified: gaze through hips AND knees, wheels balancing only.
    """
    Hh = max(0.05, stance_m - R_WHEEL)
    dx_c = dx_w * math.cos(phi) + Hh * math.sin(phi)
    Hh_c = Hh * math.cos(phi) - dx_w * math.sin(phi)
    return stance_to_leg(Hh_c + R_WHEEL, dx_c, branch)


# ---------------------------------------------------------------- reflexes

def foot_shift(shift_k, lean, shift_max):
    """Feet slide TOWARD the fall - dx has the SAME sign as lean. Measured in
    the assist lab: toward the fall cuts a 2 m/s shove from 24.9 deg peak lean
    to 6.7; the inverted sign (feet fleeing) closed an unstable pitch loop
    that fell in 5s unshoved. That inversion has now happened twice, once in
    the trainer and once in serve.py's copies - which is why the sign lives
    here and nowhere else."""
    return float(np.clip(shift_k * lean, -abs(shift_max), abs(shift_max)))


class ShoveBrace:
    """Detect a shove and lean into it, fading afterwards.

    Fires on acceleration the robot did not command for itself:
        shove_felt = filtered_accel - (commanded_speed_change / dt)
    Raw wheel-speed derivative carries the balance loop's own activity
    (measured 17 m/s^2 peak just standing, against a threshold of 3 - the
    brace became a random 12 degree injection). Through a 0.10s low-pass the
    same standing noise peaks at 3.7 while a real 2.5 m/s shove reads 10.5,
    so 5.0 separates cleanly. A driving stage must feed its own v_cmd or it
    will brace against its own throttle.
    """

    TAU = 0.10
    THRESHOLD = 5.0

    def __init__(self, gain, fade):
        self.gain = gain
        self.fade = fade
        self.v_filt = 0.0
        self.prev_speed = 0.0
        self.prev_v_cmd = 0.0
        self.brace = 0.0

    def update(self, speed, v_cmd, dt):
        self.v_filt += (dt / max(dt, self.TAU)) * (speed - self.v_filt)
        accel = (self.v_filt - self.prev_speed) / dt
        self.prev_speed = self.v_filt
        expected = (v_cmd - self.prev_v_cmd) / dt
        self.prev_v_cmd = v_cmd
        felt = accel - expected
        if abs(felt) > self.THRESHOLD:
            self.brace = self.gain * felt
        self.brace *= math.exp(-dt / max(0.05, abs(self.fade)))
        return self.brace


class SpeedFilter:
    """Low-pass on wheel speed before it drives the lean target.

    Raw wheel speed contains the limit cycle itself, so feeding it back
    unfiltered closes a positive loop: the lean target saturates at its
    clamp, flips sign every step, and the wheels chatter at full torque while
    the robot stands perfectly still - measured at 2.22 Nm mean where holding
    still needs about 0.05. tau=0 disables (raw)."""

    def __init__(self, tau=0.10):
        self.tau = tau
        self.v = 0.0

    def update(self, v_raw, dt):
        a = dt / max(dt, self.tau)
        self.v += a * (v_raw - self.v)
        return self.v


class ComplementaryTilt:
    """Gyro-fast tilt with slow accel correction, for the REAL robot.

    In the sim the fused estimator (est_up) already plays the BNO086 and is
    the reference every stage controls and scores on; this class is the
    firmware-portable equivalent for hardware without a fused IMU. The gyro
    carries the fast path (immune to the linear acceleration the legs
    inject - the loop that made accel-only references punish fast legs);
    the accelerometer corrects drift with time constant tau."""

    def __init__(self, tau=1.0):
        self.tau = tau
        self.angle = 0.0

    def update(self, gyro_rate, accel_angle, dt):
        a = self.tau / (self.tau + dt)
        self.angle = a * (self.angle + gyro_rate * dt) + (1.0 - a) * accel_angle
        return self.angle