import numpy as np

import tklsq.tkparamio as tkparamio
import tklsq.tkminfit as tkminfit
import tklsq.tkplot as tkplot
import tklsq.tksynthetic as tksynthetic


DEFAULT_PARAMS = [
    {"varname": "a", "optid": 0, "optid_lin": 1, "p0": 1.0, "pmin": "", "pmax": "", "kpenalty": 1e6},
    {"varname": "tau", "optid": 1, "optid_lin": 0, "p0": 1.0, "pmin": 0.05, "pmax": 10.0, "kpenalty": 1e6},
    {"varname": "b", "optid": 0, "optid_lin": 1, "p0": 0.0, "pmin": "", "pmax": "", "kpenalty": 1e6},
    {"varname": "c", "optid": 0, "optid_lin": 1, "p0": 0.0, "pmin": "", "pmax": "", "kpenalty": 1e6},
]


def model(x, p):
    return p["a"] * np.exp(-x / p["tau"]) + p["b"] * x + p["c"]


def main():
    param_csv = "params_expfit.csv"

    params = tkparamio.read_param_csv(
        param_csv,
        defaults=DEFAULT_PARAMS,
        create_if_missing=True,
    )

    true_p = {"a": 2.4, "tau": 1.25, "b": -0.12, "c": 0.35}
    data = tksynthetic.generate_noisy_data(
        model,
        np.linspace(0.0, 5.0, 60),
        true_p,
        noise_std=0.08,
        seed=123,
    )

    x = data.x
    y = data.y

    p0 = tkparamio.values_from_params(params)

    def model_closure(p):
        return model(x, p)

    def design_matrix_func(p, lin_names):
        tau = max(p["tau"], 1e-12)
        basis = {
            "a": np.exp(-x / tau),
            "b": x,
            "c": np.ones_like(x),
        }
        return np.vstack([basis[name] for name in lin_names]).T

    res = tkminfit.variable_projection_lsq(
        y,
        params,
        model_func=model_closure,
        design_matrix_func=design_matrix_func,
        method="Nelder-Mead",
        print_interval=10,
        options={"maxiter": 500, "xatol": 1e-9, "fatol": 1e-9},
    )

    print(tkparamio.format_params(res.params, stderr=res.stderr, title="[fit result]"))
    if res.warning:
        print(res.warning)

    tkparamio.write_param_csv(
        param_csv,
        params,
        values=res.params,
        stderr=res.stderr,
    )

    tkplot.plot_fit_before_after(
        x,
        y,
        model,
        p0,
        p_after=res.params,
        out_png="fit.png",
        title="variable projection fit",
    )


if __name__ == "__main__":
    main()