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

import argparse
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import tklsq.tklsq as tklsq_core
import tklsq.tkdataio as tkdataio
import tklsq.tkparamio as tkparamio
import tklsq.tkminfit as tkminfit
import tklsq.tkplot as tkplot
import tklsq.tkfitdiag as tkfitdiag
import tklsq.tksynthetic as tksynthetic

# デフォルトパラメータ[cite: 5, 7]
DEFAULT_PARAMS_LIST = [
    {"varname": "amplitude", "optid": 0, "optid_lin": 1, "p0": 1e6, "pmin": 0, "pmax": "", "kpenalty": 1e6},
    {"varname": "center",    "optid": 1, "optid_lin": 0, "p0": 0.0, "pmin": "", "pmax": "", "kpenalty": 1e6},
    {"varname": "sigma",     "optid": 1, "optid_lin": 0, "p0": 0.1, "pmin": 1e-6, "pmax": "", "kpenalty": 1e6},
    {"varname": "eta",       "optid": 1, "optid_lin": 0, "p0": 0.5, "pmin": 0.0, "pmax": 1.0, "kpenalty": 1e6},
    {"varname": "offset",    "optid": 0, "optid_lin": 1, "p0": 0.0, "pmin": "", "pmax": "", "kpenalty": 1e6},
]

def pseudo_voigt(x, center, sigma, eta):
    """GaussianとLorentzianの線形結合"""
    g = np.exp(-np.log(2) * ((x - center) / sigma) ** 2)
    l = 1.0 / (1.0 + ((x - center) / sigma) ** 2)
    return (1.0 - eta) * g + eta * l

def model(x, p):
    return p["amplitude"] * pseudo_voigt(x, p["center"], p["sigma"], p["eta"]) + p["offset"]

def make_design_matrix(x, p, lin_names):
    c, s, e = p["center"], p["sigma"], p["eta"]
    basis = {"amplitude": pseudo_voigt(x, c, s, e), "offset": np.ones_like(x)}
    return np.vstack([basis[name] for name in lin_names]).T

def load_data_and_params(args):
    params = tkparamio.read_param_csv(args.paramfile, defaults=DEFAULT_PARAMS_LIST)
    if args.infile and Path(args.infile).exists():
        x, y, xlabel, ylabel = tkdataio.read_xy(args.infile, x=0, y=1)[:4]
    else:
        print("Input file not found. Generating synthetic data...")
        data = tksynthetic.generate_noisy_data(
            model, np.linspace(-0.1, 0.1, 100), 
            {"amplitude": 2e7, "center": 0.008, "sigma": 0.002, "eta": 0.3, "offset": 50},
            noise_std=5.0
        )
        x, y, xlabel, ylabel = data.x, data.y, "x", "y"
    return x, y, xlabel, ylabel, params

def save_results(args, x, y, p_init, p_fit, params_dict, res_obj=None):
    prefix = args.outfile if args.outfile else (Path(args.infile).stem if args.infile else "result")
    if args.save:
        stderr_map = res_obj.stderr if hasattr(res_obj, 'stderr') else None
        tkparamio.write_param_csv(f"{prefix}_params.csv", params_dict, values=p_fit, stderr=stderr_map)
        df = pd.DataFrame({"x": x, "y_obs": y, "y_init": model(x, p_init), "y_fit": model(x, p_fit), "res": y - model(x, p_fit)})
        tkdataio.write_excel_tables(f"{prefix}_output.xlsx", {"FitResult": df})
        tkplot.plot_fit_before_after(x, y, model, p_init, p_after=p_fit, out_png=f"{prefix}_plot.png", show=args.show)
    elif args.show:
        tkplot.plot_fit_before_after(x, y, model, p_init, p_after=p_fit, show=True)

def run_read(args):
    x, y, xl, yl, _ = load_data_and_params(args)
    plt.scatter(x, y, label="Data"); plt.xlabel(xl); plt.ylabel(yl); plt.legend(); plt.grid(True)
    if args.show: plt.show()

def run_sim(args):
    x, y, xl, yl, params = load_data_and_params(args)
    p0 = tkparamio.values_from_params(params)
    tkplot.plot_fit_before_after(x, y, model, p0, title="Simulation", show=args.show)

def run_lfit(args):
    x, y, xl, yl, params = load_data_and_params(args)
    p0 = tkparamio.values_from_params(params)
    p_fit, _ = tkminfit.solve_linear_block(y, params, p0, lambda p, ln: make_design_matrix(x, p, ln))
    save_results(args, x, y, p0, p_fit, params)

def run_fit(args):
    x, y, xl, yl, params = load_data_and_params(args)
    p0 = tkparamio.values_from_params(params)
    res = tkminfit.variable_projection_lsq(y, params, lambda p: model(x, p), lambda p, ln: make_design_matrix(x, p, ln), method=args.method)
    print("\n" + tkparamio.format_params(res.params, stderr=res.stderr, title="[Fit Result]"))
    save_results(args, x, y, p0, res.params, params, res_obj=res)

def run_model_select(args):
    x, y, xl, yl, params = load_data_and_params(args)
    res = tkminfit.variable_projection_lsq(y, params, lambda p: model(x, p), lambda p, ln: make_design_matrix(x, p, ln))
    if res.cov_free is not None:
        diag = tkfitdiag.diagnose_covariance(res.free_names, [res.params[n] for n in res.free_names], res.cov_free, jacobian=res.jacobian)
        print("\n=== Model Selection Diagnostics ===")
        print(tkfitdiag.format_fix_candidates(tkfitdiag.propose_fix_candidates_from_diagnostics(diag)))
        
        # エラー修正箇所: k=res.p_free を明示的に指定
        ic = tklsq_core.information_criteria(res, k=res.p_free)
        print(f"\nAIC: {ic['AIC']:.4f}, BIC: {ic['BIC']:.4f}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["read", "sim", "lfit", "fit", "model_select"], default="fit")
    parser.add_argument("--infile", default=""); parser.add_argument("--paramfile", type=Path, default=Path("params.csv"))
    parser.add_argument("--outfile", default=""); parser.add_argument("--save", type=int, default=1)
    parser.add_argument("--show", type=int, default=1); parser.add_argument("--method", default="nelder-mead")
    args = parser.parse_args()
    modes = {"read": run_read, "sim": run_sim, "lfit": run_lfit, "fit": run_fit, "model_select": run_model_select}
    if args.mode in modes: modes[args.mode](args)

if __name__ == "__main__":
    main()