#!/usr/bin/env python3
"""
I3 - A Kramers-Kronig-consistent impedance spectrum is not evidence of linearity.

Claim under test
    Electrochemical impedance spectra are screened with a Kramers-Kronig
    (Lin-KK) residual test and accepted when that residual is small. We show,
    exactly and numerically, that the small-amplitude describing function of a
    Butler-Volmer interface is KK-consistent at *every* amplitude, while the
    charge-transfer resistance it returns is strongly amplitude-dependent.
    The standard validity test cannot see the dominant systematic.

Method
    Full harmonic balance. eta(t) is expanded in Nh harmonics on N samples; the
    current is the exact Fourier transform of the memoryless Butler-Volmer law
    plus a linear double-layer term; a linear external impedance (series
    resistance, Warburg) acts harmonic by harmonic. Solve
        eta_k = V_k - Zext(j w_k) I_k ,  V_0 = Vdc,  V_1 = -iA/2,  V_{k>=2} = 0.
    Exact for this model; no time stepping, no describing-function approximation.
    Z1 = V1/I1 is then tested with Lin-KK and fitted with a Randles model.

Exact identity (not a new prediction)
    Symmetric Butler-Volmer, no series element, no C_dl:
        i = 2 i0 sinh(eta/b), eta = A sin(wt)
        fundamental(i) = 4 i0 I1(A/b); G_app/G_true = 2 I1(x)/x, x = A/b
        I3/I1 = I3(x)/I1(x)          <-- amplitude meter

No physical experiment was performed. All numbers are model computations.
Run: python3 investigations/I3_kk_blindspot.py
"""
import numpy as np, os, json, csv
from scipy.optimize import least_squares
from scipy.special import iv
import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.abspath(__file__)); ROOT = os.path.dirname(HERE)
DATA, FIGS = os.path.join(ROOT,'data'), os.path.join(ROOT,'figs')
os.makedirs(DATA, exist_ok=True); os.makedirs(FIGS, exist_ok=True)

FAR, RGAS, TEMP = 96485.332, 8.314462, 298.15
BT = RGAS*TEMP/FAR            # 25.69 mV thermal voltage
ALPHA = 0.5
BTAF = BT/ALPHA               # 51.39 mV Tafel parameter b
I0, CDL, RS = 1.0e-6, 2.0e-5, 10.0
NH, N = 10, 128

def zext(w, sw=0.0):
    z = np.full(np.shape(w), RS, dtype=complex)
    if sw > 0:
        z = z + sw/np.sqrt(1j*np.asarray(w, float))
    return z

def ifar(eta):
    return I0*(np.exp(eta/BTAF) - np.exp(-eta/BTAF))

def hb(A, w, Vdc=0.0, sw=0.0, nh=NH, n=N, guess=None):
    """Return eta(t), X_I, X_V, max residual, solution vector."""
    th = 2*np.pi*np.arange(n)/n
    kk = np.arange(n//2+1)
    def eta_of(x):
        e = np.full(n, x[0])
        for h in range(1, nh+1):
            e = e + x[1+2*(h-1)]*np.cos(h*th) + x[2+2*(h-1)]*np.sin(h*th)
        return e
    def deriv(e):
        # exact spectral derivative on the periodic grid: rfft/irfft are exact inverses
        return np.fft.irfft(1j*kk*np.fft.rfft(e), n)
    def resid(x):
        e = eta_of(x)
        i = ifar(e) + CDL*w*deriv(e)
        XI = np.fft.rfft(i)/n
        XE = np.fft.rfft(e)/n
        XV = -0.5j*A*np.array([0,1]+[0]*(len(kk)-2), dtype=complex) if len(kk)>=2 else np.zeros(1,complex)
        XV = np.zeros(len(kk), complex)
        if len(kk) > 1: XV[1] = -0.5j*A
        XV[0] = Vdc
        r = []
        for h in range(0, nh+1):
            r.append((XE[h] - (XV[h] - zext(np.array([h*w]), sw)[0]*XI[h])))
        r = np.array(r)
        return np.concatenate([r.real, r.imag])
    x0 = np.zeros(1+2*nh)
    if guess is not None: x0 = guess.copy()
    sol = least_squares(resid, x0, xtol=1e-14, ftol=1e-14, gtol=1e-14, max_nfev=20000)
    e = eta_of(sol.x)
    i = ifar(e) + CDL*w*deriv(e)
    XI = np.fft.rfft(i)/n
    XV = np.zeros(len(kk), complex); XV[0]=Vdc
    if len(kk)>1: XV[1] = -0.5j*A
    return e, XI, XV, float(np.max(np.abs(resid(sol.x)))), sol.x

def Z1(A, w, **kw):
    e, XI, XV, res, x = hb(A, w, **kw)
    return XV[1]/XI[1], res, x

# ---------- validate: R_s=0, C=0 must reproduce the Bessel identity ----------
print('--- validation 1: identity check (R_s=0, C_dl=0) ---')
for x in [0.25, 0.5, 1.0]:
    A = x*BTAF
    z, res, _ = Z1(A, 2*np.pi*10.0, sw=0.0)
    e, XI, XV, _, _ = hb(A, 2*np.pi*10.0)
    pred = 2*iv(1, x)/x
    # with R_s=10 the current is ~A/(Rct) ... instead compute conductance ratio
    print(f'  x={x}: residue={res:.2e}')

# a cleaner validation: set R_s to a negligible value by monkey-patching
globals()['RS'] = 0.0; globals()['CDL'] = 0.0
print('--- validation 2: R_s = 0, C_dl = 0, compare to Bessel ---')
for x in [0.25, 0.5, 1.0, 2.0]:
    A = x*BTAF
    e, XI, XV, res, _ = hb(A, 2*np.pi*10.0)
    # fundamental conductance = |I1| / A
    G = abs(XI[1])/abs(XV[1])
    Gtrue = 2*I0/BTAF
    print(f'  x={x:4.2f}  G_app/G_true={G/Gtrue:.6f}  2I1(x)/x={2*iv(1,x)/x:.6f}  '
          f'I3/I1={abs(XI[3])/abs(XI[1]):.6f}  analytic={iv(3,x)/iv(1,x):.6f}  res={res:.1e}')
globals()['RS'] = 10.0; globals()['CDL'] = 2.0e-5

# ---------- the real model ----------
freqs = np.logspace(-1, 4.5, 24)
amps = [0.002, 0.005, 0.010, 0.020, 0.030, 0.050, 0.070, 0.100]
RD = {}
for A in amps:
    zs, guess = [], None
    for w in 2*np.pi*freqs:
        z, res, guess = Z1(A, w)
        zs.append(z)
    RD[A] = np.array(zs)
print('--- Randles fits ---')

def randles(p, w):
    Rs, Rct, C = p
    return Rs + 1.0/(1.0/Rct + 1j*w*C)

def fit_randles(w, Z, p0=(10., 2.6e4, 2e-5)):
    def res(p): 
        r = randles(p, w) - Z
        return np.concatenate([r.real, r.imag])
    o = least_squares(res, p0, method='trf', x_scale='jac', max_nfev=50000)
    return o.x, o.cost

fits = {A: fit_randles(2*np.pi*freqs, Z)[0] for A, Z in RD.items()}
p_true = fits[0.002]
for A in sorted(fits):
    print(f'  A={A*1000:5.1f} mV  Rs={fits[A][0]:7.2f}  Rct={fits[A][1]:10.1f}  '
          f'Cdl={fits[A][2]*1e6:6.3f} uF   Rct/true={fits[A][1]/p_true[1]:.4f}')

# ---------- Lin-KK ----------
def lin_kk(w, Z, M=14, pad=1.0):
    """Lin-KK residual: bounded Voigt fit with non-negative R_m (Boukamp/Schonle)."""
    from scipy.optimize import lsq_linear
    tau = np.logspace(np.log10(1/w.max())-pad, np.log10(1/w.min())+pad, M)
    X = np.concatenate([np.ones((len(w),1)),
                        1.0/(1.0+1j*w[:,None]*tau[None,:])], axis=1)
    Xr = np.concatenate([X.real, X.imag]); Zr = np.concatenate([Z.real, Z.imag])
    lo = np.r_[-np.inf, np.zeros(M)]
    o = lsq_linear(Xr, Zr, bounds=(lo, np.full(M+1, np.inf)), tol=1e-12, max_iter=4000)
    Zf = X @ o.x
    return float(np.sqrt(np.mean(np.abs(Z-Zf)**2))/np.mean(np.abs(Z)))

print('--- is the finite-amplitude spectrum distinguishable from a linear Debye? ---')
# With R_s = 0 the describing function is EXACTLY Z1 = 1/(G_nl + j w C_dl): a Debye,
# and a Debye satisfies Kramers-Kronig identically. R_s = 10 ohm is 0.04% of R_ct,
# so the measured spectrum must be a Debye to ~1e-4 at every amplitude. We verify
# this, and separately report the classical Lin-KK number with its own floor.
deb = {}
for A, Z in RD.items():
    Zf = randles(fits[A], 2*np.pi*freqs)
    deb[A] = float(np.max(np.abs(Z-Zf))/np.max(np.abs(Z)))
    print(f'  A={A*1000:5.1f} mV  max|Z_Zdeb|/max|Z| = {deb[A]:.2e}   Lin-KK = {lin_kk(2*np.pi*freqs, Z):.3f}')
Zs = randles(p_true, 2*np.pi*freqs)
print('  exactly linear reference (A->0):  max dev = %.2e   Lin-KK = %.3f'
      % (0.0, lin_kk(2*np.pi*freqs, Zs)))
kk = {A: lin_kk(2*np.pi*freqs, Z) for A, Z in RD.items()}

# ---------- harmonic amplitude meter ----------
print('--- amplitude meter: I3/I1 vs x = A/b ---')
meter = []
for x in [0.1,0.2,0.3,0.5,0.8,1.0,1.5,2.0]:
    A = x*BTAF
    e, XI, XV, res, _ = hb(A, 0.02)   # w*C << 1/Rct : faradaic current dominates
    r = abs(XI[3])/abs(XI[1])
    meter.append((x, r, iv(3,x)/iv(1,x)))
    print(f'  x={x:4.2f}  num={r:.6f}  analytic={iv(3,x)/iv(1,x):.6f}')

# ---------- amplitude extrapolation under noise ----------
rng = np.random.default_rng(7)
def protocol(A_list, snr=0.005, nrep=400):
    out = []
    w0 = 2*np.pi*freqs[0]
    for _ in range(nrep):
        re = []
        for A in A_list:
            z, _, _ = Z1(A, w0)
            re.append(z.real*(1+snr*rng.normal()))
        c = np.polyfit(np.array(A_list)**2, re, 1)
        out.append(c[1])
    return float(np.mean(out)), float(np.std(out))
m3, s3 = protocol([0.002,0.005,0.010])
m4, s4 = protocol([0.002,0.005,0.010,0.020])
z0 = Z1(0.001, 2*np.pi*freqs[0])[0].real
print(f'  true Re(Z1) at A->0: {z0:.1f}')
print(f'  3-point extrapolation: {m3:.1f} +- {s3:.1f}')
print(f'  4-point extrapolation: {m4:.1f} +- {s4:.1f}')

# ---------- figures ----------
fig, ax = plt.subplots(1,3, figsize=(13.6,3.9), dpi=180)
x = np.linspace(0.02, 2.5, 300)
ax[0].plot(x, 2*iv(1,x)/x, lw=2.2, color='#1b3b6f')
ax[0].axhline(1, ls=':', c='gray')
ax[0].set_xlabel(r'$x=A/b$'); ax[0].set_ylabel(r'$G_{app}/G_{true}$')
ax[0].set_title('A. Exact amplitude bias\n(symmetric Butler–Volmer)')
ax[0].grid(alpha=.25)
Aarr = np.array(sorted(amps))
ax[1].semilogx(Aarr*1000, [fits[a][1]/p_true[1] for a in Aarr], 'o-', label=r'$R_{ct}$')
ax[1].semilogx(Aarr*1000, [fits[a][2]/p_true[2] for a in Aarr], 's-', label=r'$C_{dl}$')
ax[1].axhline(1, ls=':', c='gray')
ax[1].set_xlabel('perturbation amplitude (mV)'); ax[1].set_ylabel('reported / true')
ax[1].set_title('B. What the Randles fit returns'); ax[1].legend(frameon=False, fontsize=8)
ax[1].grid(alpha=.25, which='both')
ax[2].semilogx(Aarr*1000, [kk[a] for a in Aarr], 'o-', color='#c1553b')
ax[2].axhline(lin_kk(2*np.pi*freqs, Zs), ls='--', c='gray', label='perfect Randles floor')
ax[2].set_xlabel('perturbation amplitude (mV)'); ax[2].set_ylabel('relative Lin-KK residual')
ax[2].set_title('C. The validity test is blind'); ax[2].legend(frameon=False, fontsize=8)
ax[2].grid(alpha=.25, which='both')
plt.tight_layout(); plt.savefig(os.path.join(FIGS,'i3_kk.png'))

json.dump(dict(btaf=BTAF, fits={f'{a*1000:.0f}mV': list(map(float,fits[a])) for a in sorted(fits)},
               kk={f'{a*1000:.0f}mV': kk[a] for a in sorted(kk)},
               kk_floor=lin_kk(2*np.pi*freqs, Zs),
               meter=meter, extrap=[m3,s3,m4,s4], z0=float(z0),
               note='model computation only, no experiment'),
          open(os.path.join(DATA,'i3_summary.json'),'w'), indent=1)