%pip install torchimize

n = 10

import torch
import numpy as np

sqrt2π = torch.sqrt(torch.tensor(2*torch.pi))
def lm_fx(p, x):
    A, μ, σ = p.T
    x, A, μ, σ = x[None, :], A[:, None], μ[:, None], σ[:, None]
    fx = A*torch.exp(((x - μ)/σ)**2/-2)
    return fx

def lm_fun(p, x, y, dy=1):
    fx = lm_fx(p, x)
    cost = ((fx - y)/dy)**2
    return torch.unsqueeze(cost, 1)  # Add an empty multi-cost dimension

def lm_jac(p, x, y, dy=1):
    A, μ, σ = p.T
    x, A, μ, σ = x[None, :], A[:, None], μ[:, None], σ[:, None]
    shift = x - μ
    scale = shift/σ
    scalesq = scale**2
    G = torch.exp(scalesq/-2)
    fx = A * G
    dfdA = G
    dfdμ = fx * (scale/σ)
    dfdσ = fx * (scalesq/σ)

    Δ = (fx - y)/dy
    #cost = Δ**2
    partials = torch.stack([dfdA, dfdμ, dfdσ], dim=2)
    #print(partials.shape, Δ.shape)
    jac = (2*Δ[..., None])*partials
    return torch.unsqueeze(jac, 1) # Add an empty multi-cost dimension

def gen(n, noise=0):
    x = torch.linspace(0, 50, 30)
    mid = (x[0] + x[-1])/2
    width = (x[-1] - x[0])
    μ, σ, A = torch.randn(n)*width/3 + mid, torch.rand(n)*width/5, 10**(1+torch.rand(n))
    minwidth = 5
    σ[σ < minwidth] = minwidth
    #A = torch.ones_like(A)
    pars = torch.stack([A, μ, σ]).T
    y = lm_fx(pars, x)
    if noise > 0:
        dy = noise*y
        dy[dy < noise] = noise
        y += torch.randn(*y.shape)*dy
    else:
        dy = 1
    return x, y, dy, pars

x, y, dy, target = gen(n, noise=0.05)

# Generate initial guess from the data
np_y, np_x = y.numpy(), x.numpy()
μ = np_x[np_y.argmax(axis=1)]
# Cheat for now and use target value rather than finding FWHM/2.35
σ = target[:, 2].numpy()
A = np_y.max(axis=1)
p0 = torch.tensor(np.vstack([A, μ, σ]).T)

# Bigger cheat: Use the target values as the intial guess
#p0 = target

from torchimize.functions.parallel.lma_fun_parallel import lsq_lma_parallel as opt
extra = dict(meth = "mar", ftol = 0, ptol = 1e-8, gtol = 0)

coeffs = opt(
    p = p0,
    function = lm_fun,
    jac_function = lm_jac,
    args = (x.to(dtype=p0.dtype), y.to(dtype=p0.dtype), dy),
    max_iter = 99,
    **extra,
)
print("num steps", len(coeffs))
best = coeffs[-1]

from matplotlib import pyplot as plt
%matplotlib inline
fx = lm_fx(best, x).detach()
for k in range(min(y.shape[0], 10)):
    h, = plt.plot(x, y[k], '.')
    plt.plot(x, fx[k], '-', color=h.get_color())