#!/usr/bin/env python3
from __future__ import annotations

import argparse
import traceback
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.animation import PillowWriter, FuncAnimation
from scipy.optimize import minimize

import tklsq.tkdataio as tkdataio
import tklsq.tkparamio as tkparamio
import tklsq.tkminfit as tkminfit
import tklsq.tknlsq as tknlsq
import tklsq.tkfitdiag as tkfitdiag


DEFAULT_PARAMS = [
    {"varname": "amp", "optid": 1, "optid_lin": 1, "p0": 100.0, "pscale": "linear", "dp": 5.0, "pmin": 0.0, "pmax": "", "kpenalty": 1.0e8},
    {"varname": "x0", "optid": 1, "optid_lin": 0, "p0": 0.0, "pscale": "linear", "dp": 0.01, "pmin": "", "pmax": "", "kpenalty": 1.0e8},
    {"varname": "fwhm", "optid": 1, "optid_lin": 0, "p0": 1.0, "pscale": "log", "dp": 0.02, "pmin": 1.0e-12, "pmax": "", "kpenalty": 1.0e8},
    {"varname": "eta", "optid": 1, "optid_lin": 0, "p0": 0.5, "pscale": "linear", "dp": 0.02, "pmin": 0.0, "pmax": 1.0, "kpenalty": 1.0e8},
    {"varname": "b0", "optid": 1, "optid_lin": 1, "p0": 0.0, "pscale": "linear", "dp": 1.0, "pmin": "", "pmax": "", "kpenalty": 1.0e8},
    {"varname": "b1", "optid": 1, "optid_lin": 1, "p0": 0.0, "pscale": "linear", "dp": 0.1, "pmin": "", "pmax": "", "kpenalty": 1.0e8},
]


def auto_stem(infile: str | Path) -> str:
    return Path(infile).with_suffix("").name


def pseudo_voigt_unit(x: np.ndarray, x0: float, fwhm: float, eta: float) -> np.ndarray:
    fwhm = max(float(fwhm), 1.0e-300)
    eta = float(eta)
    z = (np.asarray(x, dtype=float) - float(x0)) / fwhm
    gaussian = np.exp(-4.0 * np.log(2.0) * z**2)
    lorentzian = 1.0 / (1.0 + 4.0 * z**2)
    return eta * lorentzian + (1.0 - eta) * gaussian


def model(x: np.ndarray, p: dict[str, float]) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    peak = float(p["amp"]) * pseudo_voigt_unit(x, p["x0"], p["fwhm"], p["eta"])
    baseline = float(p.get("b0", 0.0)) + float(p.get("b1", 0.0)) * x
    return peak + baseline


def read_params(args):
    params = tkparamio.read_param_csv(args.paramfile, defaults=DEFAULT_PARAMS, create_if_missing=True)
    tkparamio.validate_param_scales(params)
    return params


def load_data(args):
    return tkdataio.read_xy(
        args.infile,
        x=args.xcol,
        y=args.ycol,
        sheet_name=args.sheet,
        xmin=args.xmin,
        xmax=args.xmax,
    )


def free_names_all_optid1(params) -> list[str]:
    # 仕様: optid=1 は optid_lin=0/1 の両方を必ず q に含める。
    return [name for name, row in params.items() if int(row.get("optid", 0)) == 1]


def make_progress_figure():
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8.5, 6.2), sharex=True, gridspec_kw={"height_ratios": [3, 1]})
    return fig, ax1, ax2


def update_progress_plot(fig, ax1, ax2, x, y, y_fit, residual, p, iteration, obj, xlabel, ylabel):
    ax1.clear()
    ax2.clear()
    ax1.scatter(x, y, s=20, c="black", alpha=0.75, label="data")
    ax1.plot(x, y_fit, c="tab:blue", lw=2.0, label="current fit")
    ax1.set_ylabel(ylabel)
    ax1.set_title(f"pseudo-Voigt fit: iteration={iteration}, objective={obj:.6g}")
    ax1.grid(True, alpha=0.3)
    ax1.legend(loc="best")
    ax2.axhline(0.0, c="gray", lw=1.0)
    ax2.plot(x, residual, c="tab:red", lw=1.2)
    ax2.set_xlabel(xlabel)
    ax2.set_ylabel("residual")
    ax2.grid(True, alpha=0.3)
    text = ", ".join(f"{k}={v:.5g}" for k, v in p.items())
    ax1.text(0.02, 0.98, text, transform=ax1.transAxes, va="top", fontsize=8)
    fig.tight_layout()


def compute_fit_band(x, p_fit, free_names, cov_free):
    if cov_free is None:
        y_fit = model(x, p_fit)
        sigma = np.zeros_like(y_fit)
        return y_fit, y_fit, y_fit, sigma

    q_fit = tkparamio.pack_optim_values(p_fit, {k: {"pscale": "linear"} for k in p_fit}, free_names)
    # 共分散は実パラメータ空間で推定しているので、ここでは q=実値としてデルタ法を使う。
    def output_func(v):
        p = dict(p_fit)
        for n, vv in zip(free_names, np.asarray(v, dtype=float)):
            p[n] = float(vv)
        return model(x, p)

    pvec = np.array([p_fit[n] for n in free_names], dtype=float)
    y0, var_y, _ = tknlsq.delta_method_variance(output_func, pvec, cov_free)
    sigma = np.sqrt(np.maximum(var_y, 0.0))
    return y0, y0 - sigma, y0 + sigma, sigma


def save_final_plot(path, x, y, y_fit, y_lower, y_upper, residual, xlabel, ylabel, show):
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8.5, 6.2), sharex=True, gridspec_kw={"height_ratios": [3, 1]})
    ax1.scatter(x, y, s=24, c="black", alpha=0.75, label="data")
    ax1.plot(x, y_fit, c="tab:blue", lw=2.0, label="fit")
    ax1.fill_between(x, y_lower, y_upper, color="tab:blue", alpha=0.18, label="±1σ")
    ax1.set_ylabel(ylabel)
    ax1.set_title("pseudo-Voigt peak fit")
    ax1.grid(True, alpha=0.3)
    ax1.legend(loc="best")
    ax2.axhline(0.0, c="gray", lw=1.0)
    ax2.plot(x, residual, c="tab:red", lw=1.2)
    ax2.set_xlabel(xlabel)
    ax2.set_ylabel("residual")
    ax2.grid(True, alpha=0.3)
    fig.tight_layout()
    if path is not None:
        fig.savefig(path, dpi=180)
    if show:
        plt.show()
    plt.close(fig)


def save_animation(path, x, y, frames, xlabel, ylabel):
    if not frames:
        return
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8.5, 6.2), sharex=True, gridspec_kw={"height_ratios": [3, 1]})

    def draw(i):
        fr = frames[i]
        update_progress_plot(fig, ax1, ax2, x, y, fr["y_fit"], fr["residual"], fr["params"], fr["iteration"], fr["objective"], xlabel, ylabel)
        return []

    ani = FuncAnimation(fig, draw, frames=len(frames), interval=300, blit=False)
    ani.save(path, writer=PillowWriter(fps=3))
    plt.close(fig)


def run_read(args):
    x, y, xlabel, ylabel = load_data(args)
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.scatter(x, y, s=24, c="black")
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title("read")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    if args.save:
        fig.savefig(f"{args.outfile}_read.png", dpi=180)
    if args.show:
        plt.show()
    plt.close(fig)


def run_sim(args):
    params = read_params(args)
    p = tkparamio.values_from_params(params)
    x, y, xlabel, ylabel = load_data(args)
    y_sim = model(x, p)
    residual = y - y_sim
    save_final_plot(f"{args.outfile}_sim.png" if args.save else None, x, y, y_sim, y_sim, y_sim, residual, xlabel, ylabel, bool(args.show))


def run_lfit(args):
    # pseudo-Voigt形状を固定し、amp/b0/b1 など optid_lin=1 だけ線形最小二乗。
    params = read_params(args)
    p0 = tkparamio.values_from_params(params)
    x, y, xlabel, ylabel = load_data(args)
    lin_names = tkparamio.linear_param_names(params)

    def design_matrix_func(p, names):
        basis = {
            "amp": pseudo_voigt_unit(x, p["x0"], p["fwhm"], p["eta"]),
            "b0": np.ones_like(x),
            "b1": x,
        }
        return np.vstack([basis[n] for n in names]).T

    p_fit, lin_res = tkminfit.solve_linear_block(y, params, p0, design_matrix_func, lin_names=lin_names)
    stderr = {}
    if lin_res is not None and lin_res.beta_std is not None:
        stderr = {n: float(se) for n, se in zip(lin_names, lin_res.beta_std)}
    print(tkparamio.format_params(p_fit, stderr=stderr, title="[lfit result]"))
    if args.save:
        tkparamio.write_param_csv(args.paramfile, params, values=p_fit, stderr=stderr)
    y_fit = model(x, p_fit)
    residual = y - y_fit
    save_final_plot(f"{args.outfile}_lfit.png" if args.save else None, x, y, y_fit, y_fit, y_fit, residual, xlabel, ylabel, bool(args.show))


def run_fit(args):
    params = read_params(args)
    base_values = tkparamio.values_from_params(params)
    x, y, xlabel, ylabel = load_data(args)

    free_names = free_names_all_optid1(params)
    q0 = tkparamio.pack_optim_values(base_values, params, free_names)
    simplex = tkparamio.initial_simplex_from_dp(base_values, params, free_names)
    options = {"maxiter": args.maxiter, "xatol": args.xatol, "fatol": args.fatol}
    if args.method == "Nelder-Mead":
        options["initial_simplex"] = simplex

    def unpack(q):
        return tkparamio.unpack_optim_values(q, base_values, params, free_names)

    # objective は仕様固定。
    def objective(q):
        p = unpack(q)
        r = y - model(x, p)
        penalty = tkparamio.bounds_penalty(params, p)
        return float(r @ r + penalty)

    fig_rt, ax_rt1, ax_rt2 = make_progress_figure()
    if args.show:
        plt.ion()
        plt.show(block=False)

    iteration_state = {"i": 0, "frames": []}

    # callback は仕様固定 + リアルタイム描画 plt.pause(0.01)。
    def callback(q):
        iteration_state["i"] += 1
        p = unpack(q)
        obj = objective(q)

        print(f"[callback] iteration={iteration_state['i']} objective={obj:.10g}")
        print("  " + ", ".join(f"{k}={v:.8g}" for k, v in p.items()))

        y_fit = model(x, p)
        residual = y - y_fit
        update_progress_plot(fig_rt, ax_rt1, ax_rt2, x, y, y_fit, residual, p, iteration_state["i"], obj, xlabel, ylabel)
        plt.pause(0.01)

        if iteration_state["i"] % args.nplot_interval == 0:
            iteration_state["frames"].append({
                "iteration": iteration_state["i"],
                "params": dict(p),
                "objective": obj,
                "y_fit": y_fit.copy(),
                "residual": residual.copy(),
            })

    # minimize 呼び出しは仕様固定。
    res = minimize(
        objective,
        q0,
        method=args.method,
        options=options,
        callback=callback,
    )

    p_fit = unpack(res.x)
    y_fit = model(x, p_fit)
    residual = y - y_fit
    rss = float(residual @ residual)

    def residual_func_real(v):
        p = dict(p_fit)
        for n, vv in zip(free_names, np.asarray(v, dtype=float)):
            p[n] = float(vv)
        return y - model(x, p)

    pvec = np.array([p_fit[n] for n in free_names], dtype=float)
    J = tknlsq.numerical_jacobian(residual_func_real, pvec)
    dof = len(y) - len(free_names)
    cov_free, stderr_free, sigma2, cov_warning = tknlsq.covariance_from_jacobian(residual, J, dof=dof)
    stderr = {name: None for name in p_fit}
    if stderr_free is not None:
        for name, se in zip(free_names, stderr_free):
            stderr[name] = float(se)

    print("\n[fit result]")
    print(f"success={res.success} message={res.message}")
    print(f"RSS={rss:.10g} objective={objective(res.x):.10g} dof={dof} sigma2={sigma2}")
    print(tkparamio.format_params(p_fit, stderr=stderr))
    if cov_warning:
        print(cov_warning)

    diag = None
    if cov_free is not None:
        diag = tkfitdiag.diagnose_covariance(
            names=free_names,
            values=[p_fit[n] for n in free_names],
            cov=cov_free,
            jacobian=J,
        )
        print("\n[diagnostics]")
        print(f"cond(JTJ) = {diag.cond_jtj}")
        print("eig(JTJ)  =", diag.eig_jtj_values_asc)
        print("corr      =")
        print(diag.corr)

    y_band, y_lower, y_upper, sigma_y = compute_fit_band(x, p_fit, free_names, cov_free)

    if args.save:
        tkparamio.write_param_csv(args.paramfile, params, values=p_fit, stderr=stderr)
        save_final_plot(f"{args.outfile}_fit.png", x, y, y_band, y_lower, y_upper, residual, xlabel, ylabel, False)
        save_animation(f"{args.outfile}_fit_animation.gif", x, y, iteration_state["frames"], xlabel, ylabel)

        fit_df = pd.DataFrame({
            xlabel: x,
            ylabel: y,
            "y_fit": y_band,
            "y_lower": y_lower,
            "y_upper": y_upper,
            "sigma_y": sigma_y,
            "residual": y - y_band,
        })
        param_df = pd.DataFrame({
            "varname": list(p_fit.keys()),
            "value": [p_fit[k] for k in p_fit],
            "stderr": [stderr.get(k) for k in p_fit],
        })
        diag_tables = {"fit": fit_df, "params": param_df}
        if diag is not None:
            diag_tables["corr"] = pd.DataFrame(diag.corr, index=free_names, columns=free_names)
            diag_tables["eig_JTJ"] = pd.DataFrame({"eig_JTJ_asc": diag.eig_jtj_values_asc})
        tkdataio.write_excel_tables(f"{args.outfile}_fit.xlsx", diag_tables)

    if args.show:
        save_final_plot(None, x, y, y_band, y_lower, y_upper, residual, xlabel, ylabel, True)
    plt.close(fig_rt)


def run_model_select(args):
    # pseudo-Voigtでは、簡易比較として eta 固定値の候補をRSS/AIC/BICで比較。
    x, y, xlabel, ylabel = load_data(args)
    rows = []
    for eta in [0.0, 0.25, 0.5, 0.75, 1.0]:
        params = read_params(args)
        p = tkparamio.values_from_params(params)
        p["eta"] = eta
        for k in ["eta"]:
            if k in params:
                params[k]["optid"] = 0
        free_names = [n for n in free_names_all_optid1(params) if n != "eta"]
        q0 = tkparamio.pack_optim_values(p, params, free_names)

        def unpack(q):
            return tkparamio.unpack_optim_values(q, p, params, free_names)

        def obj(q):
            pp = unpack(q)
            r = y - model(x, pp)
            return float(r @ r + tkparamio.bounds_penalty(params, pp))

        res = minimize(obj, q0, method=args.method, options={"maxiter": args.maxiter})
        pp = unpack(res.x)
        r = y - model(x, pp)
        rss = float(r @ r)
        n = len(y)
        k = len(free_names)
        aic = n * np.log(max(rss / n, 1.0e-300)) + 2 * k
        bic = n * np.log(max(rss / n, 1.0e-300)) + k * np.log(n)
        rows.append({"eta_fixed": eta, "RSS": rss, "AIC": aic, "BIC": bic, "success": res.success})
    df = pd.DataFrame(rows).sort_values(args.criterion.upper())
    print(df.to_string(index=False))
    if args.save:
        tkdataio.write_excel_tables(f"{args.outfile}_model_select.xlsx", {"model_select": df})


def build_parser():
    ap = argparse.ArgumentParser(description="pseudo-Voigt peak fitting program using tklsq")
    ap.add_argument("--mode", required=True, choices=["read", "sim", "lfit", "fit", "model_select"])
    ap.add_argument("--infile", required=True)
    ap.add_argument("--paramfile", default=None)
    ap.add_argument("--outfile", default=None)
    ap.add_argument("--save", type=int, default=1)
    ap.add_argument("--show", type=int, default=1)
    ap.add_argument("--xcol", default=0)
    ap.add_argument("--ycol", default=1)
    ap.add_argument("--sheet", default=0)
    ap.add_argument("--xmin", type=float, default=-1.0e100)
    ap.add_argument("--xmax", type=float, default=1.0e100)
    ap.add_argument("--nplot_interval", type=int, default=10)
    ap.add_argument("--method", default="Nelder-Mead")
    ap.add_argument("--maxiter", type=int, default=1000)
    ap.add_argument("--xatol", type=float, default=1.0e-9)
    ap.add_argument("--fatol", type=float, default=1.0e-9)
    ap.add_argument("--criterion", default="BIC", choices=["AIC", "BIC"])
    return ap


def main():
    try:
        args = build_parser().parse_args()
        stem = auto_stem(args.infile)
        if args.outfile is None:
            args.outfile = stem
        if args.paramfile is None:
            args.paramfile = f"{stem}_params.csv"

        if args.mode == "read":
            run_read(args)
        elif args.mode == "sim":
            run_sim(args)
        elif args.mode == "lfit":
            run_lfit(args)
        elif args.mode == "fit":
            run_fit(args)
        elif args.mode == "model_select":
            run_model_select(args)
        else:
            raise ValueError(f"unknown mode: {args.mode}")
    except Exception:
        traceback.print_exc()
        raise


if __name__ == "__main__":
    main()