"""
Backlash-induced limit cycles in a 2-DOF manipulator under a fixed PD + gravity
controller.

Model (per joint): two-inertia drive train
    motor side   J_m * qm_ddot = tau_cmd - tau_s - b_m * qm_dot
    link  side   M(q) q_ddot + C(q,q_dot) q_dot + g(q) + b_l * q_dot = tau_s

Transmission torque through a backlash dead zone of half-width `a`
(all motor quantities are expressed in link-side coordinates, i.e. post-gearbox):

    d  = qm - q                      (relative deflection)
    tau_s = k*(d - a) + c*d_dot      if d >  a   (teeth in contact, forward flank)
            k*(d + a) + c*d_dot      if d < -a   (contact, reverse flank)
            0                        otherwise   (free flight inside the gap)

Controller (nominal, never retuned as backlash grows):

    tau_cmd = Kp*(q_d - q_fb) + Kd*(q_d_dot - q_fb_dot) + ghat(q_fb)

`q_fb` is the joint-output encoder by default (--feedback link). With
--feedback motor the loop is closed on the motor encoder instead, which is
stable for any gap but leaves a static droop -- a useful contrast case.

Run:  python3 backlash_sim.py
"""

import argparse
import math
import os

import numpy as np

# ---------------------------------------------------------------- arm & drive

G = 9.81

# link 1 / link 2 planar arm in the vertical plane
L1, L2 = 0.40, 0.40
M1, M2 = 3.0, 2.0
LC1, LC2 = 0.20, 0.20
I1, I2 = M1 * L1**2 / 12.0, M2 * L2**2 / 12.0

JM = (0.50, 0.20)        # reflected motor+gearbox inertia, link-side coords
BM = (0.20, 0.10)        # motor-side viscous friction
BL = (0.05, 0.03)        # link-side viscous friction
KS = (4.0e4, 2.0e4)      # tooth-mesh stiffness [N*m/rad]
CS = (20.0, 10.0)        # tooth-mesh damping, active only while in contact

KP = (400.0, 100.0)      # nominal gains: well damped when the gap is zero
KD = (40.0, 10.0)
TAU_MAX = (120.0, 40.0)  # actuator torque limit; keeps a divergent case bounded
LIVE_PP = 0.2            # hold-phase peak-to-peak [deg] counted as a limit cycle

Q0 = (-math.pi / 2.0, 0.0)              # start: arm hanging
QF = (-math.pi / 2.0 + 0.6, 0.8)        # goal
T_MOVE = 1.5                            # quintic move duration
T_END = 6.0
DT = 5.0e-5


def arm_terms(q1, q2, d1, d2):
    """Mass matrix, Coriolis and gravity terms of the 2-link planar arm."""
    s2, c2 = math.sin(q2), math.cos(q2)
    m11 = M1 * LC1**2 + I1 + M2 * (L1**2 + LC2**2 + 2.0 * L1 * LC2 * c2) + I2
    m12 = M2 * (LC2**2 + L1 * LC2 * c2) + I2
    m22 = M2 * LC2**2 + I2
    h = -M2 * L1 * LC2 * s2
    cor1 = h * d2 * d1 + h * (d1 + d2) * d2
    cor2 = -h * d1 * d1
    g1 = (M1 * LC1 + M2 * L1) * G * math.cos(q1) + M2 * LC2 * G * math.cos(q1 + q2)
    g2 = M2 * LC2 * G * math.cos(q1 + q2)
    return m11, m12, m22, cor1, cor2, g1, g2


def gravity(q1, q2):
    g1 = (M1 * LC1 + M2 * L1) * G * math.cos(q1) + M2 * LC2 * G * math.cos(q1 + q2)
    g2 = M2 * LC2 * G * math.cos(q1 + q2)
    return g1, g2


def traj(t):
    """Quintic point-to-point move, then hold."""
    if t <= 0.0:
        return Q0[0], Q0[1], 0.0, 0.0
    if t >= T_MOVE:
        return QF[0], QF[1], 0.0, 0.0
    tau = t / T_MOVE
    s = 10 * tau**3 - 15 * tau**4 + 6 * tau**5
    sd = (30 * tau**2 - 60 * tau**3 + 30 * tau**4) / T_MOVE
    return (Q0[0] + s * (QF[0] - Q0[0]), Q0[1] + s * (QF[1] - Q0[1]),
            sd * (QF[0] - Q0[0]), sd * (QF[1] - Q0[1]))


def mesh_torque(i, delta, ddelta, a):
    """Dead-zone transmission torque for joint i."""
    if delta > a:
        return KS[i] * (delta - a) + CS[i] * ddelta
    if delta < -a:
        return KS[i] * (delta + a) + CS[i] * ddelta
    return 0.0


def rhs(t, y, a, link_feedback):
    q1, q2, d1, d2, m1, m2, w1, w2 = y
    qd1, qd2, dqd1, dqd2 = traj(t)

    f1, f2 = (q1, q2) if link_feedback else (m1, m2)
    fd1, fd2 = (d1, d2) if link_feedback else (w1, w2)
    gh1, gh2 = gravity(f1, f2)
    u1 = KP[0] * (qd1 - f1) + KD[0] * (dqd1 - fd1) + gh1
    u2 = KP[1] * (qd2 - f2) + KD[1] * (dqd2 - fd2) + gh2
    u1 = max(-TAU_MAX[0], min(TAU_MAX[0], u1))
    u2 = max(-TAU_MAX[1], min(TAU_MAX[1], u2))

    ts1 = mesh_torque(0, m1 - q1, w1 - d1, a)
    ts2 = mesh_torque(1, m2 - q2, w2 - d2, a)

    M11, M12, M22, c1, c2, g1, g2 = arm_terms(q1, q2, d1, d2)
    b1 = ts1 - c1 - g1 - BL[0] * d1
    b2 = ts2 - c2 - g2 - BL[1] * d2
    det = M11 * M22 - M12 * M12
    a1 = (M22 * b1 - M12 * b2) / det
    a2 = (M11 * b2 - M12 * b1) / det

    am1 = (u1 - ts1 - BM[0] * w1) / JM[0]
    am2 = (u2 - ts2 - BM[1] * w2) / JM[1]
    return (d1, d2, a1, a2, w1, w2, am1, am2)


def simulate(backlash_deg, link_feedback=True, dt=DT, t_end=T_END, decim=20):
    """Fixed-step RK4. Returns decimated logs."""
    a = math.radians(backlash_deg)
    n = int(round(t_end / dt))
    y = (Q0[0], Q0[1], 0.0, 0.0, Q0[0], Q0[1], 0.0, 0.0)

    nlog = n // decim + 1
    log = np.empty((nlog, 11))
    j = 0
    for step in range(n + 1):
        t = step * dt
        if step % decim == 0:
            qd1, qd2, _, _ = traj(t)
            log[j] = (t, y[0], y[1], y[2], y[3], y[4], y[5],
                      qd1, qd2, y[4] - y[0], y[5] - y[1])
            j += 1
        if step == n:
            break
        k1 = rhs(t, y, a, link_feedback)
        y2 = tuple(y[i] + 0.5 * dt * k1[i] for i in range(8))
        k2 = rhs(t + 0.5 * dt, y2, a, link_feedback)
        y3 = tuple(y[i] + 0.5 * dt * k2[i] for i in range(8))
        k3 = rhs(t + 0.5 * dt, y3, a, link_feedback)
        y4 = tuple(y[i] + dt * k3[i] for i in range(8))
        k4 = rhs(t + dt, y4, a, link_feedback)
        y = tuple(y[i] + dt / 6.0 * (k1[i] + 2 * k2[i] + 2 * k3[i] + k4[i])
                  for i in range(8))

    t = log[:j, 0]
    return {
        "backlash_deg": backlash_deg,
        "t": t,
        "q": log[:j, 1:3],
        "dq": log[:j, 3:5],
        "qm": log[:j, 5:7],
        "qd": log[:j, 7:9],
        "delta": log[:j, 9:11],
        "e": log[:j, 7:9] - log[:j, 1:3],
    }


def metrics(run, t0=3.0):
    """Steady-state oscillation metrics over the hold phase."""
    m = run["t"] >= t0
    e = np.degrees(run["e"][m])
    t = run["t"][m]
    out = {}
    for i in (0, 1):
        ei = e[:, i]
        pp = ei.max() - ei.min()
        ac = ei - ei.mean()
        rms = float(np.sqrt(np.mean(ac**2)))
        spec = np.abs(np.fft.rfft(ac * np.hanning(len(ac))))
        freqs = np.fft.rfftfreq(len(ac), t[1] - t[0])
        f = float(freqs[np.argmax(spec[1:]) + 1]) if len(spec) > 2 else 0.0
        out[i] = {"pp": float(pp), "rms": rms, "bias": float(ei.mean()), "f": f}
    return out


# ------------------------------------------------------------------- plotting

SURFACE = "#fcfcfb"
INK = "#0b0b0b"
INK2 = "#52514e"
MUTED = "#898781"
GRID = "#e1e0d9"
AXIS = "#c3c2b7"
# ordinal blue ramp, steps 250..700 (nothing lighter than 250 on a light surface)
RAMP = ["#86b6ef", "#6da7ec", "#5598e7", "#2a78d6", "#256abf", "#184f95", "#0d366b"]
CRIT = "#d03b3b"
ORANGE = "#eb6834"


def style(ax, xlabel=None, ylabel=None, title=None):
    ax.set_facecolor(SURFACE)
    ax.grid(True, color=GRID, lw=0.8)
    ax.set_axisbelow(True)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    for s in ("left", "bottom"):
        ax.spines[s].set_color(AXIS)
        ax.spines[s].set_linewidth(1.0)
    ax.tick_params(colors=MUTED, labelsize=9, length=3)
    if xlabel:
        ax.set_xlabel(xlabel, color=INK2, fontsize=10)
    if ylabel:
        ax.set_ylabel(ylabel, color=INK2, fontsize=10)
    if title:
        ax.set_title(title, color=INK, fontsize=11, loc="left", pad=8)


def rolling_pp(x, t, win=0.1):
    """Peak-to-peak of x over a sliding window of `win` seconds."""
    n = max(2, int(round(win / (t[1] - t[0]))))
    v = np.lib.stride_tricks.sliding_window_view(x, n)
    return t[n - 1:], v.max(axis=1) - v.min(axis=1)


def fig_tracking(runs, outdir, tag):
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(3, 1, figsize=(10, 11), facecolor=SURFACE,
                             constrained_layout=True)

    # A -- what the joint actually does, rigid vs worst gap
    ax = axes[0]
    r0, rN = runs[0], runs[-1]
    ax.plot(r0["t"], np.degrees(r0["qd"][:, 0]), color=MUTED, lw=1.5, ls="--")
    ax.plot(r0["t"], np.degrees(r0["q"][:, 0]), color=RAMP[0], lw=2.0)
    ax.plot(rN["t"], np.degrees(rN["q"][:, 0]), color=RAMP[6], lw=1.6)
    k = int(0.8 / (r0["t"][1] - r0["t"][0]))
    ax.annotate("command", (2.6, np.degrees(QF[0])), xytext=(0, -20),
                textcoords="offset points", color=MUTED, fontsize=9)
    ax.annotate("no backlash", (r0["t"][k], np.degrees(r0["q"][k, 0])),
                xytext=(14, -10), textcoords="offset points", color=RAMP[0],
                fontsize=9)
    ax.annotate(f'{rN["backlash_deg"]:.2f}$^\\circ$ backlash',
                (4.6, np.degrees(rN["q"][:, 0]).max()), xytext=(-30, 10),
                textcoords="offset points", color=RAMP[6], fontsize=9)
    ax.axvline(T_MOVE, color=AXIS, lw=1.0, ls=":")
    style(ax, xlabel="time [s]", ylabel="joint 1 position [deg]",
          title="Same nominal PD + gravity controller: the move tracks, the hold does not")

    # B -- oscillation envelope, log scale: decay vs plateau
    ax = axes[1]
    for r, col in zip(runs, RAMP):
        te, pp = rolling_pp(np.degrees(r["e"][:, 0]), r["t"])
        ax.semilogy(te, np.maximum(pp, 1e-6), color=col, lw=2.0,
                    label=f'{r["backlash_deg"]:.2f}$^\\circ$')
    ax.axvline(T_MOVE, color=AXIS, lw=1.0, ls=":")
    ax.set_ylim(1e-5, 1e1)
    style(ax, xlabel="time [s]",
          ylabel="joint 1 error envelope [deg], 0.1 s window",
          title="Small gaps decay away; past ~0.5$^\\circ$ the error plateaus instead")
    leg = ax.legend(title="backlash half-width", frameon=False, ncol=4,
                    fontsize=9, title_fontsize=9, loc="lower right",
                    columnspacing=1.2, handlelength=1.4)
    leg.get_title().set_color(INK2)
    for txt in leg.get_texts():
        txt.set_color(INK2)

    # C -- zoom on the steady state
    ax = axes[2]
    for r, col in zip(runs[3:], RAMP[3:]):
        m = (r["t"] >= 4.0) & (r["t"] <= 4.25)
        ax.plot(r["t"][m], np.degrees(r["e"][m, 0]), color=col, lw=2.0,
                label=f'{r["backlash_deg"]:.2f}$^\\circ$')
    lo, hi = ax.get_ylim()
    ax.set_ylim(lo, hi + 0.45 * (hi - lo))      # headroom for the legend
    style(ax, xlabel="time [s]", ylabel="joint 1 error [deg]",
          title="Hold-phase detail: amplitude set by the gap, not by the command")
    leg = ax.legend(title="backlash half-width", frameon=False, ncol=4,
                    fontsize=9, title_fontsize=9, loc="upper right")
    leg.get_title().set_color(INK2)
    for txt in leg.get_texts():
        txt.set_color(INK2)

    p = os.path.join(outdir, f"fig1_tracking_{tag}.png")
    fig.savefig(p, dpi=150, facecolor=SURFACE)
    plt.close(fig)
    return p


def fig_limit_cycle(small, large, outdir, tag):
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 2, figsize=(11, 7.5), facecolor=SURFACE,
                             constrained_layout=True)

    for ax, r, col, name, t0, note in (
            (axes[0, 0], small, RAMP[2], "below onset", 1.0,
             "decaying spiral into the setpoint"),
            (axes[0, 1], large, RAMP[6], "above onset", 2.0,
             "closed orbit: self-sustaining")):
        m = r["t"] >= t0
        ax.plot(np.degrees(r["e"][m, 0]), np.degrees(r["dq"][m, 0]), color=col, lw=1.4)
        ax.plot([0], [0], marker="o", ms=8, color=CRIT, zorder=5)
        ax.annotate("setpoint", (0, 0), xytext=(10, 6), textcoords="offset points",
                    color=CRIT, fontsize=9)
        style(ax, xlabel="joint 1 error [deg]", ylabel="joint 1 velocity [deg/s]",
              title=f'{name}: backlash {r["backlash_deg"]:.2f}$^\\circ$')
        ax.text(0.98, 0.04, note, transform=ax.transAxes, ha="right",
                color=INK2, fontsize=9)

    ax = axes[1, 0]
    m = (large["t"] >= 4.0) & (large["t"] <= 4.6)
    a = large["backlash_deg"]
    ax.fill_between(large["t"][m], -a, a, color=GRID, zorder=0)
    ax.plot(large["t"][m], np.degrees(large["delta"][m, 0]), color=ORANGE, lw=2.0)
    ax.axhline(a, color=AXIS, lw=1.0)
    ax.axhline(-a, color=AXIS, lw=1.0)
    style(ax, xlabel="time [s]", ylabel="motor-link deflection [deg]",
          title="Gap traversal and tooth impacts (high backlash)")
    ax.text(0.02, 0.5, "inside the gap:\nno torque transmitted",
            transform=ax.transAxes, color=INK2, fontsize=9, va="center")

    ax = axes[1, 1]
    m = (large["t"] >= 4.0) & (large["t"] <= 4.6)
    ax.plot(large["t"][m], np.degrees(large["q"][m, 0]), color=RAMP[6], lw=2.0,
            label="link")
    ax.plot(large["t"][m], np.degrees(large["qm"][m, 0]), color=ORANGE, lw=2.0,
            label="motor")
    ax.plot(large["t"][m], np.degrees(large["qd"][m, 0]), color=MUTED, lw=1.5,
            ls="--", label="command")
    style(ax, xlabel="time [s]", ylabel="joint 1 position [deg]",
          title="Motor hunts across the gap, link never settles")
    leg = ax.legend(frameon=False, fontsize=9, loc="best")
    for t in leg.get_texts():
        t.set_color(INK2)

    p = os.path.join(outdir, f"fig2_limit_cycle_{tag}.png")
    fig.savefig(p, dpi=150, facecolor=SURFACE)
    plt.close(fig)
    return p


def fig_bifurcation(levels, pps, freqs, outdir, tag):
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2), facecolor=SURFACE,
                             constrained_layout=True)
    levels = np.asarray(levels, float)
    pps = np.asarray(pps, float)
    live = pps > LIVE_PP                    # genuinely cycling, not a slow transient
    onset = (levels[~live].max() + levels[live].min()) / 2 if live.any() else None

    ax = axes[0]
    ax.plot(levels, pps, color=RAMP[5], lw=2.0, marker="o", ms=8,
            markerfacecolor=RAMP[5], markeredgecolor=SURFACE, markeredgewidth=2)
    if onset is not None:
        ax.axvline(onset, color=CRIT, lw=1.5, ls="--")
        ax.annotate(f"onset $\\approx${onset:.2f}$^\\circ$", (onset, max(pps) * 0.55),
                    xytext=(-8, 0), textcoords="offset points", ha="right",
                    color=CRIT, fontsize=9)
    style(ax, xlabel="backlash half-width [deg]",
          ylabel="joint 1 error, peak-to-peak [deg]",
          title="Steady-state oscillation amplitude vs backlash")
    for x, y in zip(levels[live], pps[live]):
        ax.annotate(f"{y:.2f}", (x, y), textcoords="offset points",
                    xytext=(0, 11), ha="center", color=INK2, fontsize=9)

    ax = axes[1]
    f = np.where(live, np.asarray(freqs, float), np.nan)
    ax.plot(levels, f, color=ORANGE, lw=2.0, marker="o", ms=8,
            markerfacecolor=ORANGE, markeredgecolor=SURFACE, markeredgewidth=2)
    if onset is not None:
        ax.axvline(onset, color=CRIT, lw=1.5, ls="--")
    ax.set_ylim(0, max(30.0, np.nanmax(f) * 1.3 if live.any() else 30.0))
    ax.set_xlim(*axes[0].get_xlim())
    style(ax, xlabel="backlash half-width [deg]",
          ylabel="dominant frequency [Hz]",
          title="Limit-cycle frequency (no cycle below onset)")
    p = os.path.join(outdir, f"fig3_bifurcation_{tag}.png")
    fig.savefig(p, dpi=150, facecolor=SURFACE)
    plt.close(fig)
    return p


# ----------------------------------------------------------------------- main

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--feedback", choices=("link", "motor"), default="link")
    ap.add_argument("--backlash", type=float, nargs="*",
                    default=[0.0, 0.2, 0.4, 0.6, 0.8, 1.2, 2.0],
                    help="list of backlash half-widths in degrees")
    ap.add_argument("--out", default=os.path.dirname(os.path.abspath(__file__)))
    ap.add_argument("--no-plots", action="store_true")
    args = ap.parse_args()

    link_fb = args.feedback == "link"
    runs, table = [], []
    for b in args.backlash:
        r = simulate(b, link_feedback=link_fb)
        m = metrics(r)
        runs.append(r)
        table.append((b, m))
        print(f"backlash {b:5.2f} deg | "
              f"j1 pp {m[0]['pp']:7.3f} deg  rms {m[0]['rms']:6.3f}  "
              f"bias {m[0]['bias']:7.3f}  f {m[0]['f']:6.2f} Hz | "
              f"j2 pp {m[1]['pp']:7.3f} deg  rms {m[1]['rms']:6.3f}  "
              f"f {m[1]['f']:6.2f} Hz")

    if args.no_plots:
        return
    os.makedirs(args.out, exist_ok=True)
    import matplotlib
    matplotlib.use("Agg")
    tag = args.feedback
    print(fig_tracking(runs, args.out, tag))
    stable = max((r for r in runs if metrics(r)[0]["pp"] < LIVE_PP),
                 key=lambda r: r["backlash_deg"], default=runs[0])
    print(fig_limit_cycle(stable, runs[-1], args.out, tag))
    print(fig_bifurcation([b for b, _ in table],
                          [m[0]["pp"] for _, m in table],
                          [m[0]["f"] for _, m in table], args.out, tag))


if __name__ == "__main__":
    main()
