# The pitch-shift reflex as a delayed, leaky feedback loop (a sketch, not a fitted model).
# y: sung pitch (cents re target). Headphones return h = y + s, s = +100 cents, 0.5-0.7 s.
# Correction acts on what was heard tau earlier: c' = -G h(t - tau) - c / T_leak; y = c +
# small motor noise; tau = 120 ms (assumed, inside the 100-150 ms range). Latency is read
# as experimenters do: when the mean of 40 trials leaves +/- 2 SD of the pre-shift
# baseline and stays out for 20 ms.
import math, random

DT, TAU, T_LEAK = 0.001, 0.120, 0.5

def trial(G, rnd, shift=100.0):
    n, lag = 1300, int(TAU / DT)
    y, h, c = [0.0] * n, [0.0] * n, 0.0
    for i in range(n):
        t = i * DT
        c += (-G * (h[i - lag] if i >= lag else 0.0) - c / T_LEAK) * DT
        y[i] = c + rnd.gauss(0, 2.0)
        h[i] = y[i] + (shift if 0.5 <= t < 0.7 else 0.0)
    return y

def average(G, k=40, seed=7):
    rnd = random.Random(seed)
    runs = [trial(G, rnd) for _ in range(k)]
    return [sum(r[i] for r in runs) / k for i in range(len(runs[0]))]

print("gain G (1/s) | latency (ms after shift) | largest opposing change | share")
for G in (1.0, 2.0, 4.0):
    m = average(G)
    base = m[300:500]
    mu = sum(base) / len(base)
    sd = math.sqrt(sum((v - mu) ** 2 for v in base) / (len(base) - 1))
    onset = next(i for i in range(500, 1300) if all(v < mu - 2 * sd for v in m[i:i + 20]))
    peak = min(m[500:1000])
    print(f"    {G:4.1f}     | {onset - 500:24d} | {peak:+17.1f} cents |"
          f" {-peak / 100:4.0%}")
print(f"loop delay assumed: {TAU * 1000:.0f} ms; responses oppose the shift,"
      " as in experiments")
# gain G (1/s) | latency (ms after shift) | largest opposing change | share
#      1.0     |                      129 |             -16.5 cents |  17%
#      2.0     |                      122 |             -32.1 cents |  32%
#      4.0     |                      122 |             -61.5 cents |  62%
# loop delay assumed: 120 ms; responses oppose the shift, as in experiments

# Holding a note with a 120 ms auditory delay, 10 ms steps. Pitch error y (cents) drifts:
#   y[k+1] = y[k] + b u[k] + w, w ~ N(0, 1.5^2);  hearing z[k] = y[k-D] + v, v ~ N(0, 8^2)
# (a) direct feedback u = -g z.  (b) forward model + Kalman: estimate y[k-D] from z, add
# the commands sent since (efference copy) to predict y[k], u = -L x prediction
# (Smith-predictor form, after state-feedback accounts of speech). Noise levels and gains
# are illustrative.
import math, random

D, SW, SV = 12, 1.5, 8.0

def simulate(mode, gain, b_true=1.0, steps=2000, seed=0, trace=False):
    rnd = random.Random(seed)
    y, ys, us, e, P, err2, out = 0.0, [0.0] * (D + 1), [0.0] * (D + 1), 0.0, 25.0, 0.0, []
    for k in range(steps):
        z = ys[-D - 1] + rnd.gauss(0, SV)                  # heard now: y from D steps ago
        if mode == "direct":
            u = -gain * z
        else:
            # internal-model step for y[k-D]
            e, P = e + us[-D - 1], P + SW ** 2
            K = P / (P + SV ** 2)
            # Kalman update with the late sample
            e, P = e + K * (z - e), (1 - K) * P
            u = -gain * (e + sum(us[-D:]))                 # + commands sent since then
        y = y + b_true * u + rnd.gauss(0, SW)
        ys, us = ys[1:] + [y], us[1:] + [u]
        err2 += y * y
        out.append(y)
    return (math.sqrt(err2 / steps), out) if trace else math.sqrt(err2 / steps)

def rms(mode, gain, **kw):                       # average over 10 noise realisations
    return sum(simulate(mode, gain, seed=s, **kw) for s in range(10)) / 10

if __name__ == "__main__":
    print(f"no correction (feedforward only): RMS error {rms('direct', 0.0):6.1f} cents")
    print("direct delayed feedback, gain per 10 ms step (G = gain / 0.01 s):")
    for g in (0.02, 0.05, 0.08, 0.11, 0.13, 0.15):
        r = rms("direct", g)
        txt = (f"{r:6.1f} cents" if r < 60 else
               f"{r:6.0f} cents  <- oscillates" if r < 1000 else
               "  > 1000 cents  <- oscillation grows")
        print(f"   g = {g:.2f} (G = {g * 100:4.0f}/s, G*tau = {g * D:.2f}): RMS {txt}")
    print("forward model + Kalman estimate, then the same kind of correction:")
    for L in (0.2, 0.5, 0.9):
        print(f"   L = {L:.1f}: RMS {rms('kalman', L):6.1f} cents")
    print("forward model wrong by 40 % (true command effect 1.4 x the model's), L = 0.9:")
    print(f"   RMS {rms('kalman', 0.9, b_true=1.4):6.1f} cents")
    print(f"stability limit, direct feedback: G*tau below about pi/2 = {math.pi / 2:.2f}")
    print(f"floor no controller can beat: drift over the delay and one step, sqrt(D + 1) x 1.5"
          f" = {math.sqrt(D + 1) * SW:.1f} cents")
# no correction (feedforward only): RMS error   30.4 cents
# direct delayed feedback, gain per 10 ms step (G = gain / 0.01 s):
#    g = 0.02 (G =    2/s, G*tau = 0.24): RMS    8.6 cents
#    g = 0.05 (G =    5/s, G*tau = 0.60): RMS    7.0 cents
#    g = 0.08 (G =    8/s, G*tau = 0.96): RMS    7.6 cents
#    g = 0.11 (G =   11/s, G*tau = 1.32): RMS   11.9 cents
#    g = 0.13 (G =   13/s, G*tau = 1.56): RMS    274 cents  <- oscillates
#    g = 0.15 (G =   15/s, G*tau = 1.80): RMS   > 1000 cents  <- oscillation grows
# forward model + Kalman estimate, then the same kind of correction:
#    L = 0.2: RMS    6.8 cents
#    L = 0.5: RMS    6.5 cents
#    L = 0.9: RMS    6.4 cents
# forward model wrong by 40 % (true command effect 1.4 x the model's), L = 0.9:
#    RMS    6.8 cents
# stability limit, direct feedback: G*tau below about pi/2 = 1.57
# floor no controller can beat: drift over the delay and one step, sqrt(D + 1) x 1.5 = 5.4 cents