"""HDP30 (Rinnegan drive) RevC, gearbox 2: slow backlash runs (2026-10-23).

Run from anywhere:  python backlash_analysis.py
Reads ../Second PLA Gearbox/*.csv, writes bl*.png + numbers_backlash.txt next to this file.

Each run: start at rest in free travel, drive at 5 input RPM (1 deg/s at the output) into one
load cell (CCW -> CH1, CW -> CH2), stall at the current limit (0.4 or 0.6 A), then reverse and
drive back through free travel. At 1 deg/s one 5 ms sample is 0.005 deg of output travel, so
encoder timing skew (a few ms) no longer matters.

Backlash is estimated three ways, all in output degrees:
  hysteresis gap   motor position at the same output reading, push-direction pass minus return pass,
                   in free travel. Comparing at the same output reading cancels the magnet-ring scale error.
  departure        motor travel after the cell reaches zero torque until the arm has moved 0.03 deg.
  start lag        motor travel from the start of motion until the arm has moved 0.03 deg (only for
                   runs logged from rest; the drive's last move was the other way, so slack must close).
The 0.03 deg threshold is about 2 sigma of the raw MT6701 noise, so departure and start lag carry
that much bias upward.
"""
import glob
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

import analysis as A
from analysis import HERE, OUT, DPI, ARM_M, INK, INK2, GRID, C_CH1, C_CH2, C_SKIP, C_ALT

DATA = next((os.path.join(A.DATA, n) for n in ("Second PLA Gearbox", "gearbox2-backlash")   # lab / website folder
             if os.path.isdir(os.path.join(A.DATA, n))), os.path.join(A.DATA, "gearbox2-backlash"))
THR = 0.03   # deg of arm motion that counts as "moving"


def sm(x, n=9):
    return np.convolve(x, np.ones(n) / n, mode="same")


def analyse(path):
    name = os.path.basename(path)
    d = pd.read_csv(path).interpolate().bfill()
    t = d.time_s.values
    sg = 1 if name.startswith("CCW") else -1           # push direction
    m = sg * np.unwrap(np.radians(d.input_deg)) / (2 * np.pi) * 12     # motor, output-equivalent deg
    o_raw = sg * np.unwrap(np.radians(d.output_deg)) / (2 * np.pi) * 30
    o = sm(o_raw)
    tq = sg * (d.ch1_n - d.ch2_n).values * ARM_M
    v = np.gradient(m, t)
    i_start = int(np.argmax(np.abs(v) > 0.3))
    i_contact = int(np.argmax(tq > 0.02))
    ipk = int(np.argmax(m))                              # motor reversal
    i_free = ipk + int(np.argmax(tq[ipk:] < 0.01))       # cell back at zero
    from_rest = d.current_cmd_a.values[0] == 0

    # hysteresis gap in free travel
    app = np.arange(i_start + 60, i_contact - 20)
    back = np.where((np.arange(len(t)) > i_free + 20) & (v < -0.3))[0][:-10]
    gap, overlap = np.nan, 0.0
    if len(app) > 20 and len(back) > 20:
        lo = max(o[app].min(), o[back].min()) + 0.05
        hi = min(o[app].max(), o[back].max()) - 0.05
        if hi > lo:
            g = np.linspace(lo, hi, 50)
            a, b = np.argsort(o[app]), np.argsort(o[back])
            gap = np.median(np.interp(g, o[app][a], m[app][a]) - np.interp(g, o[back][b], m[back][b]))
            overlap = hi - lo

    rest = np.median(o[i_free:i_free + 3])
    k = i_free + int(np.argmax(o[i_free:] < rest - THR))
    departure = m[i_free] - m[k]
    start_lag = np.nan
    if from_rest:
        o0 = np.median(o[max(i_start - 40, 0):i_start])
        ks = i_start + int(np.argmax(o[i_start:] > o0 + THR))
        start_lag = m[ks] - m[i_start]

    # quasi-static loading branch: twist across the gearbox against torque
    hold = np.median(tq[ipk - 200:ipk - 20])
    load = np.arange(i_contact - 40, ipk)
    tw = (m - o)[load] - (m - o)[i_contact - 40]
    mm = (tq[load] > 0.1 * hold) & (tq[load] < 0.9 * hold)
    k_load = np.polyfit(tw[mm], tq[load][mm], 1)[0] if mm.sum() > 10 else np.nan
    unl = np.arange(ipk, i_free + 1)
    tw_u = (m - o)[unl] - (m - o)[i_contact - 40]
    w50 = (np.interp(0.5 * hold, tq[unl][::-1], tw_u[::-1]) -
           np.interp(0.5 * hold, np.maximum.accumulate(tq[load]), tw))
    amps = float(name.split("superslow")[1][:3].replace("a", "")) if name.split("superslow")[1][0] == "." else np.nan
    return dict(run=name.split("-2026")[0], dir="CCW" if sg > 0 else "CW", amps=amps, from_rest=from_rest,
                hold=hold, eff=100 * hold / (0.151 * 30 * amps), k_load=k_load * 57.3, w50=w50,
                gap=gap, overlap=overlap, departure=departure, start_lag=start_lag,
                curve=(m[i_free - 80:i_free + 160] - m[i_free], o[i_free - 80:i_free + 160] - rest,
                       tq[i_free - 80:i_free + 160]),
                loop=(np.concatenate([tw, tw_u]), np.concatenate([tq[load], tq[unl]])))


def main():
    files = sorted(glob.glob(os.path.join(DATA, "*.csv")))
    res = [analyse(f) for f in files]
    R = pd.DataFrame([{k: v for k, v in r.items() if k not in ("curve", "loop")} for r in res])
    pd.set_option("display.width", 220)
    lines = [f"{len(R)} runs from {DATA}",
             R.drop(columns=["from_rest"]).round(3).to_string(), ""]
    g = R[R.overlap > 0.5]
    lines.append(f"Hysteresis gap (runs with >0.5 deg free-travel overlap, n={len(g)}): median {g.gap.median():.3f} deg, "
                 f"mean {g.gap.mean():.3f}, sd {g.gap.std():.3f}, range {g.gap.min():.3f} to {g.gap.max():.3f}")
    lines.append(f"Departure: median {R.departure.median():.3f} deg (CCW {R[R.dir == 'CCW'].departure.median():.3f}, "
                 f"CW {R[R.dir == 'CW'].departure.median():.3f}), includes up to {THR} deg of threshold bias")
    s = R.start_lag.dropna()
    lines.append(f"Start lag (runs logged from rest, n={len(s)}): median {s.median():.3f} deg, range {s.min():.3f}-{s.max():.3f}")
    lines.append(f"Slow loading stiffness: median {R.k_load.median():.1f} N*m/rad; loop width at half torque median {R.w50.median():.3f} deg")
    lines.append(f"Hold / static efficiency: 0.4 A {R[R.amps == 0.4].hold.median():.2f} N*m ({R[R.amps == 0.4].eff.median():.0f}%), "
                 f"0.6 A {R[R.amps == 0.6].hold.median():.2f} N*m ({R[R.amps == 0.6].eff.median():.0f}%)")
    txt = "\n".join(lines)
    print(txt)
    open(os.path.join(HERE, "numbers_backlash.txt"), "w", encoding="utf-8").write(txt)

    # ---- BL1: departure off the cell, every run ------------------------------------
    fig, ax = plt.subplots(figsize=(11, 7))
    for r in res:
        x, y, _ = r["curve"]
        jump = np.where(np.abs(np.diff(y)) > 0.2)[0]      # end of the run / disarm glitch
        if len(jump):
            x, y = x[:jump[0] + 1], y[:jump[0] + 1]
        col = C_CH1 if r["dir"] == "CCW" else C_CH2
        ax.plot(-x, -y, color=col, lw=1.4, alpha=0.8)
    xx = np.linspace(0, 1.2, 10)
    ax.plot(xx, xx, color=INK2, ls="--", lw=1.3)
    ax.text(0.62, 0.5, "arm follows motor 1:1\n(zero backlash)", color=INK2, fontsize=11)
    ax.plot([0, 0.1, 1.2], [0, 0, 1.1], color=C_SKIP, ls=":", lw=1.8)
    ax.text(0.3, 0.03, "what 0.1° of backlash would look like", color=C_SKIP, fontsize=11)
    ax.text(0.02, 0.9, "the CCW curves bend upward at the same\narm position every run: magnet-ring\nscale error, not real motion",
            color=INK2, fontsize=10.5)
    ax.axvline(0, color=INK, lw=1)
    ax.axvspan(-0.4, 0, color=GRID, alpha=0.35, lw=0)
    ax.text(-0.38, 1.05, "still unloading\nthe cell", color=INK2, fontsize=11)
    ax.set_xlim(-0.4, 1.2)
    ax.set_ylim(-0.15, 1.3)
    ax.set_xlabel("motor travel back since the cell reached zero  [deg, output-equivalent]")
    ax.set_ylabel("arm motion off the cell  [deg]")
    ax.legend(handles=[Line2D([], [], color=C_CH1, lw=3, label="CCW runs (off CH1)"),
                       Line2D([], [], color=C_CH2, lw=3, label="CW runs (off CH2)")], loc="upper left")
    ax.set_title(f"BACKLASH at 5 RPM: the arm leaves within {R.departure.median():.2f}° of motor travel (12 runs)")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "bl1_departure.png"), dpi=DPI)

    # ---- BL2: three estimates per run ------------------------------------------------
    fig, ax = plt.subplots(figsize=(12, 5.5))
    xs = np.arange(len(R))
    ax.bar(xs - 0.27, R.gap.where(R.overlap > 0.5), width=0.26, color=C_ALT, label="free-travel hysteresis gap")
    ax.bar(xs, R.departure, width=0.26, color=C_CH1, label=f"departure (incl. {THR}° threshold)")
    ax.bar(xs + 0.27, R.start_lag, width=0.26, color=C_CH2, label=f"start lag (incl. {THR}° threshold)")
    ax.axhline(0, color=INK, lw=1)
    ax.axhline(0.1, color=C_SKIP, ls=":", lw=1.4)
    ax.text(len(R) - 0.5, 0.104, "0.1° = 6 arcmin", color=C_SKIP, ha="right", fontsize=11)
    ax.set_xticks(xs)
    ax.set_xticklabels([f"{r.dir} {r.amps:.1f} A\n#{r.run[-1]}" for r in R.itertuples()], fontsize=9)
    ax.set_ylabel("lost motion  [deg at the output]")
    ax.set_ylim(-0.12, 0.16)
    ax.legend(loc="lower left", fontsize=10, ncol=3)
    ax.set_title("BACKLASH, three ways: every run between 0 and 0.1° (blank = not enough free travel)")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "bl2_estimates.png"), dpi=DPI)
    print("\nfigures written to", OUT)


if __name__ == "__main__":
    main()