"""HDP30 RevC endurance runs (2026-10-18 evening): one fresh unit, one PLA flexspline.

Run from anywhere:  python endurance_analysis.py
Reads the endurance*.csv files one folder up, writes en*.png + numbers_endurance.txt next to this file.

Runs, in order, all on the same unit:
  static hold   0.5 A into CH1 for 21 s, then disarmed
  0.4 A         20 min of reversing strokes, 150 RPM, 4 s hold at each cell
  0.6 / 0.8 / 1.0 / 1.2 A   the same stroke pattern, shorter runs
The flexspline wall cracked during the 1.2 A run (cracking was heard during 1.0 A).
"""
import glob
import os
import sys
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, RATIO, SURF, INK, INK2, GRID, TILE,
                      C_CH1, C_CH2, C_IN, C_SKIP, C_ALT)

REVB = os.environ.get("REVB_DATA") or os.path.join(
    A.DATA, "..", "..", "260830_Harmonic_Pancake_RevB_30_to_1", "Test Data")
if not os.path.isdir(REVB):
    REVB = A.DATA
sys.path.append(os.path.join(REVB, "analysis"))
import key_metrics as K  # noqa: E402  (stop/stiffness extraction)

KT = 0.151                       # N*m/A, motor alone (leftbeef baseline)
STATIC = "endurancetest-20261018-191254-659912.csv"
RUN_COL = {"0.4": "#7d8fa8", "0.6": C_IN, "0.8": C_CH1, "1.0": C_ALT, "1.2": C_CH2}


def load(path):
    d = pd.read_csv(path).interpolate().bfill()
    d["in_turns"] = np.unwrap(np.radians(d.input_deg)) / (2 * np.pi)
    d["mt_turns"] = np.unwrap(np.radians(d.output_deg)) / (2 * np.pi)
    d["T_nm"] = (d.ch1_n - d.ch2_n) * ARM_M
    return d


def run_files():
    fs = glob.glob(os.path.join(A.DATA, "endurancetestreversing*.csv"))
    return sorted(fs, key=lambda p: os.path.basename(p).split("-", 1)[1])


def per_contact():
    """One row per held contact across every reversing run, in time order."""
    rows, cum_in, t_off, runs = [], 0.0, 0.0, {}
    for f in run_files():
        d = load(f)
        tag = os.path.basename(f).split("-")[0].replace("endurancetestreversing", "").replace("amps", "") or "0.4"
        c = A.contacts(d)
        s = A.skips(d)
        t, inp, mt, T = d.time_s.values, d.in_turns.values, d.mt_turns.values, d.T_nm.values
        trav = np.concatenate([[0], np.cumsum(np.abs(np.diff(inp)))])
        st = K.stops("x", d, c, s)
        held = c[c.cmd > 0]
        for r, sto in zip(held.itertuples(), st):
            sg = 1 if r.ch == "ch1" else -1
            a = np.searchsorted(t, r.t0)
            tq = sg * T
            ipk = a + np.argmax(tq[a:a + 60])
            i2 = np.searchsorted(t, t[ipk] + 0.15)
            k = ipk + np.argmin(tq[ipk:i2 + 1])
            rows.append(dict(run=tag, ch=r.ch, t=r.t0, tg=t_off + r.t0, cum_in=cum_in + trav[a],
                             peak=tq[ipk], hold=r.F_hold * ARM_M, iq=r.iq, cmd=r.cmd, stop=r.in_pos,
                             k=sto["k_u"] * 57.3, unwind=sto["u_tw"].max(),
                             give=tq[ipk] - tq[k], give_motor=sg * (inp[k] - inp[ipk]) * 12))
        v = np.gradient(inp, t) * 60
        m = (np.abs(np.abs(v) - 150) < 8) & (d.current_cmd_a.values > 0)
        for r in c.itertuples():
            m &= ~((t > r.t0 - 0.4) & (t < r.t1 + 0.6))
        runs[tag] = dict(dur=t[-1], turns=trav[-1], skips=len(s), drag=np.abs(d.iq_a.values[m]).mean(),
                         t_off=t_off, file=os.path.basename(f))
        cum_in += trav[-1]
        t_off += t[-1]
    R = pd.DataFrame(rows)
    # yield: the torque falls by more than 0.1 N*m within 0.15 s of the peak while the motor keeps
    # advancing into the stop. A healthy hold relaxes in place; a yielding one gives way.
    R["yield"] = (R.give > 0.1) & (R.give_motor > 0.15)
    return R, runs, cum_in


def main():
    R, runs, total_in = per_contact()
    lines = []
    P = lines.append
    P(f"Motor Kt {KT} N*m/A (direct drive), ideal 30:1 = {KT * RATIO:.2f} N*m/A")
    P(f"Total: {total_in:.0f} input turns = {2 * total_in:.0f} flexspline flex cycles = {total_in / RATIO:.0f} output turns, "
      f"{len(R)} hits into a load cell, {sum(r['skips'] for r in runs.values())} tooth skips")
    for tag, r in runs.items():
        P(f"  {tag} A: {r['dur']:.0f} s, {r['turns']:.0f} input turns, skips {r['skips']}, "
          f"no-load drag at 150 RPM |iq| {r['drag']:.3f} A   ({r['file']})")

    g = R.groupby(["run", "ch"]).agg(n=("peak", "size"), peak=("peak", "median"), hold=("hold", "median"),
                                     k=("k", "median"), unwind=("unwind", "median"),
                                     give=("give", "median"), yields=("yield", "sum")).round(3)
    g["eff_%"] = (100 * g.hold / (KT * RATIO * g.index.get_level_values(0).astype(float))).round(0)
    P("\nPer run and direction (medians):\n" + g.to_string())

    onset = R[R["yield"]].iloc[0] if R["yield"].any() else None
    fail = R.loc[R[R.ch == "ch2"].give.idxmax()]
    big = R[(R.ch == "ch2") & (R.tg < (onset.tg if onset is not None else 1e9))].peak.max()
    first_big = R[(R.ch == "ch2") & (R.peak > 1.65)].iloc[0]
    P(f"\nDamage onset: hit at {first_big.run} A, t={first_big.t:.1f} s, peak {first_big.peak:.3f} N*m (CH2). "
      f"CH2 stop jumped {12 * (R[(R.ch == 'ch2') & (R.run == first_big.run)].stop.diff().min()):.2f} deg (output-equivalent).")
    if onset is not None:
        P(f"First yielding hit: {onset.run} A, t={onset.t:.1f} s (CH2), give {onset.give:.3f} N*m with "
          f"{onset.give_motor:.2f} deg of motor advance. Before that the highest CH2 peak was {big:.3f} N*m.")
    n_y = int(((R.ch == "ch2") & (R.tg >= first_big.tg) & (R.tg <= fail.tg)).sum())
    P(f"Fracture: {fail.run} A, t={fail.t:.1f} s (CH2): peak {fail.peak:.3f} -> {fail.peak - fail.give:.3f} N*m in 30 ms, "
      f"motor advanced {fail.give_motor:.2f} deg. {n_y} CH2 hits from onset to fracture.")
    P(f"Highest CH1 peak in any run: {R[R.ch == 'ch1'].peak.max():.3f} N*m, CH1 never yielded "
      f"({int(R[(R.ch == 'ch1')]['yield'].sum())} yields).")

    # ---- static hold + unpowered hold ------------------------------------------
    d = load(os.path.join(A.DATA, STATIC))
    t, T = d.time_s.values, d.T_nm.values
    on = (d.current_cmd_a.values > 0)
    t_on, t_off = t[on][0], t[on][-1]
    ipk = np.argmax(T * on)
    hm = (t > t[ipk] + 0.05) & (t < t_off - 0.05)
    lt = np.log10(t[hm] - t[ipk])
    b, a0 = np.polyfit(lt, T[hm], 1)
    b = -b  # report creep as a positive loss per decade
    after = T[(t > t_off + 0.5)].mean()
    P(f"\nStatic hold at 0.5 A (CH1): peak {T[ipk]:.3f} N*m -> {T[hm][-1]:.3f} after {t[hm][-1] - t[ipk]:.0f} s "
      f"({100 * (T[ipk] - T[hm][-1]) / T[ipk]:.1f}%), creep {b:.4f} N*m per decade of time "
      f"({100 * b / a0:.2f}%/decade). Static efficiency at the peak {100 * T[ipk] / (KT * RATIO * 0.5):.0f}%.")
    P(f"After disarming (iq = 0): output still pushes {after:.3f} N*m on the cell. "
      f"The drive does not backdrive at that load.")

    txt = "\n".join(lines)
    print(txt)
    open(os.path.join(HERE, "numbers_endurance.txt"), "w", encoding="utf-8").write(txt)

    # ---- EN1: every hit, peak and hold, against cumulative motor turns ------------
    fig, ax = plt.subplots(figsize=(13, 6.5))
    for tag, col in RUN_COL.items():
        for ch, mk in (("ch1", "o"), ("ch2", "s")):
            z = R[(R.run == tag) & (R.ch == ch)]
            ax.scatter(z.cum_in, z.peak, s=26, marker=mk, facecolors="none", edgecolors=col, linewidths=1.3)
            ax.scatter(z.cum_in, z.hold, s=26, marker=mk, color=col)
    for n, (tag, r) in enumerate(runs.items()):
        x0 = R[R.run == tag].cum_in.min()
        ax.text(x0 + 8, 1.86 - 0.07 * (n % 2), f"{tag} A", color=RUN_COL[tag], fontsize=11, fontweight="bold")
        ax.axvline(x0 - 5, color=GRID, lw=1)
    ax.axvline(first_big.cum_in, color=C_SKIP, ls=":", lw=1.5)
    ax.annotate("cracking heard,\nCH2 starts to yield", xy=(first_big.cum_in, first_big.peak), xytext=(1880, 1.62),
                color=C_SKIP, fontsize=11, arrowprops=dict(arrowstyle="->", color=C_SKIP))
    ax.annotate("flexspline wall\ncracks", xy=(fail.cum_in, fail.peak - fail.give), xytext=(2350, 0.6),
                color=C_SKIP, fontsize=11, arrowprops=dict(arrowstyle="->", color=C_SKIP))
    ax.set_ylim(0, 1.95)
    ax.set_xlabel("cumulative motor travel  [input turns]   (2 flex cycles per turn)")
    ax.set_ylabel("output torque on the stop  [N·m]")
    ax.legend(handles=[Line2D([], [], marker="o", ls="", color=INK2, label="hold (CH1 side)"),
                       Line2D([], [], marker="s", ls="", color=INK2, label="hold (CH2 side)"),
                       Line2D([], [], marker="o", ls="", mfc="none", mec=INK2, label="impact peak")],
              loc="lower left", ncol=3)
    ax.set_title(f"ENDURANCE: {2 * total_in:,.0f} flex cycles and {len(R)} hits, then the wall cracked at 1.2 A")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "en1_timeline.png"), dpi=DPI)

    # ---- EN2: torque against current ---------------------------------------------
    fig, ax = plt.subplots(figsize=(11, 6.5))
    ii = np.linspace(0, 1.3, 10)
    ax.plot(ii, KT * RATIO * ii, color=INK2, ls="--", lw=1.4, label=f"ideal 30:1 ({KT * RATIO:.1f} N·m/A)")
    ax.plot(ii, 0.55 * KT * RATIO * ii, color=INK2, ls=":", lw=1.4, label="55 % of ideal")
    gm = R.groupby(["run", "ch"]).agg(hold=("hold", "median"), peak=("peak", "median"), iq=("iq", "median")).reset_index()
    for ch, col, mk in (("ch1", C_CH1, "o"), ("ch2", C_CH2, "s")):
        z = gm[gm.ch == ch]
        ax.plot(z.iq, z.hold, color=col, marker=mk, ms=9, lw=2, label=f"hold, {ch.upper()} side")
        ax.plot(z.iq, z.peak, color=col, marker=mk, ms=8, lw=1.2, ls="--", mfc="none", label=f"impact peak, {ch.upper()} side")
    ax.axhspan(1.5, 1.7, color=C_SKIP, alpha=0.12, lw=0)
    ax.text(0.05, 1.6, "cracking zone (heard at 1.5 yesterday, 1.67 today)", color=C_SKIP, fontsize=11, va="center")
    ax.set_xlim(0, 1.3)
    ax.set_ylim(0, 2.6)
    ax.set_xlabel("motor current limit, measured |iq| during the hold  [A]")
    ax.set_ylabel("output torque  [N·m]")
    ax.legend(loc="upper left", fontsize=10, ncol=2)
    ax.set_title("TORQUE vs CURRENT: holds flatten out near 1.35 N·m past 0.8 A")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "en2_torque_vs_current.png"), dpi=DPI)

    # ---- EN3: damage indicators per hit -------------------------------------------
    fig, axs = plt.subplots(2, 1, figsize=(13, 7.5), sharex=True)
    for ch, col, mk in (("ch1", C_CH1, "o"), ("ch2", C_CH2, "s")):
        z = R[R.ch == ch]
        axs[0].plot(z.cum_in, z.unwind, color=col, marker=mk, ms=4, lw=1, label=f"{ch.upper()} side")
        axs[1].plot(z.cum_in, z.give, color=col, marker=mk, ms=4, lw=1, label=f"{ch.upper()} side")
    for ax in axs:
        ax.axvline(first_big.cum_in, color=C_SKIP, ls=":", lw=1.5)
        for tag in runs:
            ax.axvline(R[R.run == tag].cum_in.min() - 5, color=GRID, lw=1)
    axs[0].set_ylabel("twist to unload  [deg]")
    axs[0].legend(loc="upper left")
    axs[0].set_title("DAMAGE: after one 1.67 N·m hit, the CH2 side softens and gives way on every hit")
    axs[1].set_ylabel("torque lost right\nafter the peak  [N·m]")
    axs[1].set_xlabel("cumulative motor travel  [input turns]")
    axs[1].set_ylim(0, 1.0)
    axs[1].text(first_big.cum_in + 10, 0.8, "1.2 A fracture: 0.89 N·m lost in 30 ms", color=C_SKIP, fontsize=11)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "en3_damage.png"), dpi=DPI)

    # ---- EN4: the fracture, sample by sample ---------------------------------------
    f12 = [f for f in run_files() if "1.2amps" in f][0]
    d = load(f12)
    t, inp, mt, T = d.time_s.values, d.in_turns.values, d.mt_turns.values, d.T_nm.values
    fig, axs = plt.subplots(2, 1, figsize=(12, 7), sharex=True)
    for t0, col, lab in ((fail.t - 20.0, INK2, "previous CH2 hit (20 s earlier)"), (fail.t, C_SKIP, "the fracture")):
        m = (t > t0 + 0.05) & (t < t0 + 0.75)
        i0 = np.where(m)[0][0]
        x = (t[m] - t[i0]) * 1e3
        axs[0].plot(x, -T[m], color=col, lw=2, marker="o", ms=3, label=lab)
        axs[1].plot(x, -(inp[m] - inp[i0]) * 12, color=col, lw=2, label=f"motor, {lab}")
        axs[1].plot(x, -(mt[m] - mt[i0]) * 30, color=col, lw=1.3, ls="--", label=f"arm, {lab}")
    axs[0].set_ylabel("torque on CH2 stop  [N·m]")
    axs[0].legend(loc="upper right")
    axs[0].set_title("FRACTURE: the stop torque halves in 30 ms while the motor keeps pushing")
    axs[1].set_ylabel("travel into the stop\n[deg, output-equivalent]")
    axs[1].set_xlabel("time  [ms]")
    axs[1].legend(loc="lower right", fontsize=9, ncol=2)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "en4_fracture.png"), dpi=DPI)

    # ---- EN5: static hold and the unpowered hold -----------------------------------
    d = load(os.path.join(A.DATA, STATIC))
    t, T = d.time_s.values, d.T_nm.values
    fig, ax = plt.subplots(figsize=(11, 5.8))
    ax.plot(t, T, color=C_CH1, lw=1.8)
    ax.axvspan(t_on, t_off, color=C_IN, alpha=0.08, lw=0)
    ax.text(t_on + 0.5, 0.25, "motor on, 0.5 A", color=INK2, fontsize=11)
    ax.text(t_off + 0.4, 0.25, "motor off\n(iq = 0)", color=INK2, fontsize=11)
    ax.annotate(f"{100 * (T[ipk] - T[hm][-1]) / T[ipk]:.1f}% creep in {t_off - t_on:.0f} s\n({100 * b / a0:.1f}% per decade of time)",
                xy=(15, T[np.searchsorted(t, 15)]), xytext=(9, 0.72), color=INK, fontsize=11,
                arrowprops=dict(arrowstyle="->", color=INK2))
    ax.annotate(f"still holding {after:.2f} N·m\nwith no current", xy=(26, after), xytext=(22.8, 0.55),
                color=INK, fontsize=11, arrowprops=dict(arrowstyle="->", color=INK2))
    ax.set_ylim(0, 1.4)
    ax.set_xlabel("time  [s]")
    ax.set_ylabel("output torque on CH1  [N·m]")
    ax.set_title("STATIC HOLD: slow creep, and it doesn't backdrive at ~1 N·m")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "en5_static_hold.png"), dpi=DPI)
    print("\nfigures written to", OUT)


if __name__ == "__main__":
    main()