#!/usr/bin/env python3
"""Small demo for tklib.tknlsq + tklib.tkfitdiag."""

import numpy as np
from tknlsq import nonlinear_lsq, delta_method_band
from tkfitdiag import diagnose_covariance, propose_fix_candidates_from_diagnostics, format_fix_candidates


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


def main():
    rng = np.random.default_rng(0)
    x = np.linspace(0, 5, 40)
    y = model(x, {"a": 2.0, "tau": 1.3, "c": 0.2}) + rng.normal(0, 0.05, size=x.size)

    def residuals(p):
        return y - model(x, p)

    res = nonlinear_lsq(
        residuals,
        p0={"a": 1.0, "tau": 1.0, "c": 0.0},
        bounds={"a": (0, np.inf), "tau": (1e-9, np.inf), "c": (-np.inf, np.inf)},
    )

    print("params:", res.params)
    print("stderr:", res.stderr)
    print("warning:", res.warning)

    if res.cov_free is not None:
        vals = [res.params[n] for n in res.free_names]
        diag = diagnose_covariance(res.free_names, vals, res.cov_free, jacobian=res.jacobian)
        cand = propose_fix_candidates_from_diagnostics(diag)
        print("\nfix/constrain candidates:")
        print(format_fix_candidates(cand))

        def output_func(pfree):
            pp = dict(res.params)
            for n, v in zip(res.free_names, pfree):
                pp[n] = v
            return model(x, pp)

        yfit, ylo, yhi, sig = delta_method_band(output_func, res.params_free, res.cov_free)
        print("\nfirst 5 bands:")
        for row in zip(x[:5], yfit[:5], ylo[:5], yhi[:5], sig[:5]):
            print("x={:.3g}, y={:.5g}, lo={:.5g}, hi={:.5g}, sig={:.3g}".format(*row))


if __name__ == "__main__":
    main()