#!/usr/bin/env python3
"""Fit parametric curves (v2): amount_total = A(depth)*S(sens)*H(sharp).
S normalized so S(12)=1, H(10)=1. A→1 as depth→large.
Saturation form: f(x)=1-exp(-(x/k)^m).
"""
import numpy as np
from scipy.optimize import curve_fit

def sat(x, k, m):
    return 1 - np.exp(-(x / k) ** m)

def to_amount(red):
    return 1 - 10 ** (-red / 20.0)

def fit(name, x, y, p0):
    popt, _ = curve_fit(sat, x, y, p0=p0, maxfev=50000)
    pred = sat(x, *popt)
    print(f"{name}: k={popt[0]:.4f} m={popt[1]:.3f}  RMSE={np.sqrt(np.mean((pred-y)**2)):.5f}")
    for xi, yi, pi in zip(x, y, pred):
        print(f"   x={xi:6.2f}  meas={yi:.4f}  pred={pi:.4f}")
    return popt

A_def = to_amount(7.70)

# A: depth sweep at S=H=1 (sens12 sharp10)
x = np.array([0.5, 0.864, 1.0, 2.0, 3.0, 5.0, 10.0, 20.0], float)
y = np.array([to_amount(d) for d in [7.14, 7.70, 7.91, 9.56, 11.29, 14.89, 24.25, 39.41]])
ka = fit('A(depth)', x, y, [3.0, 0.8])

# S: sens at depth0.864 (S saturates -> S(12)=1). amount = A_def * S
# S has a floor at sens=0 (still 2.89dB): S(x) = s0 + (1-s0)*sat(x)
def S(x, s0, k, m):
    return s0 + (1 - s0) * sat(x, k, m)
x = np.array([0.0, 6.0, 12.0, 24.0, 48.0], float)
y = np.array([to_amount(d) for d in [2.89, 5.14, 7.70, 7.70, 7.70]]) / A_def
popt, _ = curve_fit(S, x, y, p0=[0.4, 3.0, 3.0], maxfev=50000)
pred = S(x, *popt)
print(f"S(sens): s0={popt[0]:.4f} k={popt[1]:.4f} m={popt[2]:.3f}  RMSE={np.sqrt(np.mean((pred-y)**2)):.5f}")
for xi, yi, pi in zip(x, y, pred):
    print(f"   x={xi:6.2f}  meas={yi:.4f}  pred={pi:.4f}")
ks = popt

# H: sharp at depth0.864 -> H(10)=1
x = np.array([1.0, 3.0, 5.0, 10.0, 20.0], float)
y = np.array([to_amount(d) for d in [1.11, 3.71, 5.90, 7.70, 7.70]]) / A_def
kh = fit('H(sharp)/norm', x, y, [3.0, 1.5])

print("\ncross-check amount_total = A*S*H at default:")
print(f"  A(0.864)*S(12)*H(10) = {sat(0.864,*ka):.4f}*1*1 -> red {(-20*np.log10(1-sat(0.864,*ka))):.2f} dB (meas 7.70)")
print(f"  depth=10 sens=0:      A(10)*S(0) = {sat(10,*ka):.4f}*{sat(0,*ks):.4f} = {sat(10,*ka)*sat(0,*ks):.4f} -> {( -20*np.log10(1-sat(10,*ka)*sat(0,*ks))):.2f} dB")
np.savez('/home/m/re-tools/curve_fits.npz',
         k_depth=ka, k_sens=ks, k_sharp=kh, A_def=A_def)