#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
fit_vbo_xps.py

Fit a film/substrate mixed XPS valence-band spectrum as a weighted sum of
reference film and substrate spectra with independent energy shifts.

Model:
    I_mix(E) = A_f * I_film(E - dE_f) + A_s * I_sub(E - dE_s)
               + polynomial_background(E)

The fitted offset is reported mainly as
    dE_shift = dE_f - dE_s
and, when reference VBM positions are supplied,
    VBO_BE = (VBM_film_ref + dE_f) - (VBM_sub_ref + dE_s)

For binding-energy axes, be careful with the sign convention.  The script
prints both the binding-energy difference and its negative.

ex.
python fit_vbo_xps.py ^
  --film Bi2OS2100nm_20260508-1_Bi2OS2_pSi_HAXPES_fit_YBPS100_VBwideY2-deconvoluted_point13.xlsm ^
  --substrate n-Si_20260508-1_Bi2OS2_pSi_HAXPES_fit_YNS_VBnarrowY-Si-deconvoluted_point10.xlsm ^
  --mix TOA90_20260508-1_Bi2OS2_pSi_HAXPES_fit_YBPS25-90_VBwideY2-deconvoluted_point13.xlsm ^
  --fit-xmin -2.5 --fit-xmax 8 ^
  --show 1 --amp-sub0 1.0

"""

from __future__ import annotations

import argparse
import os
import sys
import traceback
from dataclasses import dataclass
from typing import Tuple, List

try:
    import numpy as np
    import pandas as pd
    import matplotlib.pyplot as plt
    from scipy.interpolate import interp1d
    from scipy.optimize import minimize
except Exception:
    print("ERROR: required library import failed.")
    traceback.print_exc()
    print("\nInstall example:")
    print("  pip install numpy pandas matplotlib scipy openpyxl")
    input("\nPress ENTER to terminate>>\n")
    sys.exit(1)


@dataclass
class Spectrum:
    name: str
    x: np.ndarray
    y: np.ndarray


def read_spectrum(path: str, sheet: str | int = 0, x_col: int = 0, y_col: int = 2,
                  skip_first_blank: int = 1) -> Spectrum:
    """Read x/y columns from xlsx/xlsm/csv. Column indices are zero-based."""
    ext = os.path.splitext(path)[1].lower()
    if ext in [".csv", ".txt", ".dat"]:
        df = pd.read_csv(path, header=None)
    else:
        df = pd.read_excel(path, sheet_name=sheet, header=None, engine="openpyxl")

    if skip_first_blank:
        # The supplied files have an empty first row.  Numeric conversion below also removes it.
        pass

    x = pd.to_numeric(df.iloc[:, x_col], errors="coerce").to_numpy(dtype=float)
    y = pd.to_numeric(df.iloc[:, y_col], errors="coerce").to_numpy(dtype=float)
    m = np.isfinite(x) & np.isfinite(y)
    x = x[m]
    y = y[m]
    if len(x) < 5:
        raise ValueError(f"too few valid data points in {path}")

    # sort by increasing x and merge duplicate x by averaging y
    order = np.argsort(x)
    x = x[order]
    y = y[order]
    xu, inv = np.unique(x, return_inverse=True)
    if len(xu) != len(x):
        yu = np.zeros_like(xu, dtype=float)
        cnt = np.zeros_like(xu, dtype=float)
        for i, k in enumerate(inv):
            yu[k] += y[i]
            cnt[k] += 1
        x, y = xu, yu / cnt

    return Spectrum(os.path.basename(path), x, y)


def crop(sp: Spectrum, xmin: float | None, xmax: float | None) -> Spectrum:
    m = np.ones_like(sp.x, dtype=bool)
    if xmin is not None:
        m &= sp.x >= xmin
    if xmax is not None:
        m &= sp.x <= xmax
    return Spectrum(sp.name, sp.x[m], sp.y[m])


def make_interp(sp: Spectrum, fill_value: float = 0.0):
    return interp1d(sp.x, sp.y, kind="linear", bounds_error=False, fill_value=fill_value,
                    assume_sorted=True)


def robust_scale(y: np.ndarray) -> float:
    s = np.nanpercentile(y, 95) - np.nanpercentile(y, 5)
    if not np.isfinite(s) or s <= 0:
        s = np.nanmax(np.abs(y))
    return float(s if s > 0 else 1.0)


def poly_background(x: np.ndarray, coeff: np.ndarray, x0: float, xscale: float) -> np.ndarray:
    xr = (x - x0) / xscale
    bg = np.zeros_like(xr)
    for i, c in enumerate(coeff):
        bg += c * xr**i
    return bg


def normalize_intensity(y: np.ndarray, mode: str = "max") -> np.ndarray:
    """Normalize intensity for visual comparison of input spectra."""
    yy = np.asarray(y, dtype=float)
    if len(yy) == 0:
        return yy.copy()

    if mode == "minmax":
        ymin = np.nanmin(yy)
        ymax = np.nanmax(yy)
        scale = ymax - ymin
        if not np.isfinite(scale) or scale <= 0:
            scale = 1.0
        return (yy - ymin) / scale

    if mode == "area":
        # For area normalization, keep the sign of the spectrum and normalize
        # by the integral of abs(y) over the available point index axis.
        scale = np.trapz(np.abs(yy))
        if not np.isfinite(scale) or scale <= 0:
            scale = 1.0
        return yy / scale

    # default: maximum-height normalization without baseline subtraction.
    scale = np.nanmax(np.abs(yy))
    if not np.isfinite(scale) or scale <= 0:
        scale = 1.0
    return yy / scale


def fit_mixture(film: Spectrum, sub: Spectrum, mix: Spectrum, args):
    x = mix.x.copy()
    y = mix.y.copy()
    if args.fit_xmin is not None or args.fit_xmax is not None:
        m = np.ones_like(x, dtype=bool)
        if args.fit_xmin is not None:
            m &= x >= args.fit_xmin
        if args.fit_xmax is not None:
            m &= x <= args.fit_xmax
        x, y = x[m], y[m]
    if len(x) < 10:
        raise ValueError("fit range contains too few points")

    ff = make_interp(film, fill_value=args.outside_fill)
    fs = make_interp(sub, fill_value=args.outside_fill)
    yscale = robust_scale(y)
    x0 = 0.5 * (x.min() + x.max())
    xscale = 0.5 * (x.max() - x.min())
    if xscale <= 0:
        xscale = 1.0

    # Initial parameters are given explicitly.
    # No linear pre-fit is used because this problem has only a few parameters
    # and Nelder-Mead is often more stable for this kind of spectral fitting.
    bg0 = np.zeros(args.bg_order + 1, dtype=float)
    bg0[0] = args.bg0
    if args.bg_order >= 1:
        bg0[1] = args.bg1
    if args.bg_order >= 2:
        bg0[2] = args.bg2

    p0 = np.r_[
        args.shift_film0,
        args.shift_sub0,
        np.log(max(args.amp_film0, args.amp_min)),
        np.log(max(args.amp_sub0, args.amp_min)),
        bg0,
    ]

    def unpack(p):
        d_f, d_s = p[0], p[1]
        a_f, a_s = np.exp(p[2]), np.exp(p[3])
        bgc = p[4:]
        return d_f, d_s, a_f, a_s, bgc

    def model(p, xx=x):
        d_f, d_s, a_f, a_s, bgc = unpack(p)
        return a_f * ff(xx - d_f) + a_s * fs(xx - d_s) + poly_background(xx, bgc, x0, xscale)

    def resid(p):
        yyfit = model(p)
        r = (yyfit - y) / yscale
        if args.loss_on_log:
            yy = np.maximum(y - np.nanmin(y) + 1.0, 1.0)
            mm = np.maximum(yyfit - np.nanmin(y) + 1.0, 1.0)
            r = np.log(mm / yy)
        return r

    eval_counter = {"n": 0, "best": np.inf, "best_p": None}

    def objective(p):
        eval_counter["n"] += 1
        d_f, d_s, a_f, a_s, bgc = unpack(p)
        penalty = 0.0

        # Soft bounds for unconstrained Nelder-Mead.
        if abs(d_f) > args.shift_limit:
            penalty += ((abs(d_f) - args.shift_limit) / max(args.shift_limit, 1.0e-12))**2 * args.penalty
        if abs(d_s) > args.shift_limit:
            penalty += ((abs(d_s) - args.shift_limit) / max(args.shift_limit, 1.0e-12))**2 * args.penalty
        if a_f < args.amp_min:
            penalty += (np.log(args.amp_min / max(a_f, 1.0e-300)))**2 * args.penalty
        if a_s < args.amp_min:
            penalty += (np.log(args.amp_min / max(a_s, 1.0e-300)))**2 * args.penalty
        if a_f > args.amp_max:
            penalty += (np.log(a_f / args.amp_max))**2 * args.penalty
        if a_s > args.amp_max:
            penalty += (np.log(a_s / args.amp_max))**2 * args.penalty

        r = resid(p)
        if args.loss == "linear":
            val = float(np.mean(r * r))
        elif args.loss == "logcosh":
            val = float(np.mean(np.log(np.cosh(r / args.f_scale))) * args.f_scale**2)
        else:
            # Pseudo-Huber-like robust loss.
            z = r / args.f_scale
            val = float(np.mean(2.0 * args.f_scale**2 * (np.sqrt(1.0 + z*z) - 1.0)))
        total = val + penalty
        if total < eval_counter["best"]:
            eval_counter["best"] = total
            eval_counter["best_p"] = np.array(p, dtype=float).copy()
        return total

    callback_counter = {"iter": 0}

    def callback(xk):
        callback_counter["iter"] += 1
        if args.print_every <= 0:
            return
        it = callback_counter["iter"]
        if it == 1 or it % args.print_every == 0:
            val = objective(xk)
            d_f, d_s, a_f, a_s, bgc = unpack(xk)
            print(
                f"Iter {it:5d}  Eval {eval_counter['n']:6d}  "
                f"S={val:.8g}  Best={eval_counter['best']:.8g}  "
                f"dEf={d_f:.6g}  dEs={d_s:.6g}  "
                f"Af={a_f:.6g}  As={a_s:.6g}"
            )

    print("\nNelder-Mead convergence")
    print("  Iter/Eval output is from the callback; S is the objective value including penalties.")
    print("  Parameters: dEf, dEs [eV], Af, As")
    print(
        f"Initial      S={objective(p0):.8g}  "
        f"dEf={args.shift_film0:.6g}  dEs={args.shift_sub0:.6g}  "
        f"Af={args.amp_film0:.6g}  As={args.amp_sub0:.6g}"
    )

    res = minimize(
        objective,
        p0,
        method="Nelder-Mead",
        callback=callback,
        options={
            "maxiter": args.max_iter,
            "maxfev": args.max_nfev,
            "xatol": args.xatol,
            "fatol": args.fatol,
            "disp": bool(args.verbose),
        },
    )
    p = res.x
    d_f, d_s, a_f, a_s, bgc = unpack(p)

    print(
        f"Final        S={float(res.fun):.8g}  "
        f"nit={getattr(res, 'nit', -1)}  nfev={getattr(res, 'nfev', -1)}  "
        f"dEf={d_f:.6g}  dEs={d_s:.6g}  Af={a_f:.6g}  As={a_s:.6g}"
    )

    yfit = model(p)
    film_part = a_f * ff(x - d_f)
    sub_part = a_s * fs(x - d_s)
    bg_part = poly_background(x, bgc, x0, xscale)
    rmse = float(np.sqrt(np.mean((yfit - y)**2)))
    nrmse = rmse / yscale

    return {
        "result": res,
        "x": x, "y": y, "yfit": yfit,
        "film_part": film_part, "sub_part": sub_part, "bg_part": bg_part,
        "params": {"dE_film": d_f, "dE_sub": d_s, "A_film": a_f, "A_sub": a_s,
                   "bg_coeff": bgc, "rmse": rmse, "nrmse": nrmse,
                   "dE_shift_film_minus_sub": d_f - d_s,
                   "objective": float(res.fun),
                   "success": bool(res.success),
                   "message": str(res.message)},
        "interp": (ff, fs),
        "x0": x0, "xscale": xscale,
    }


def save_outputs(film: Spectrum, sub: Spectrum, mix: Spectrum, fit, args):
    stem = args.output_stem
    os.makedirs(os.path.dirname(stem) or ".", exist_ok=True)
    csv_path = stem + "_fit.csv"
    png_path = stem + "_fit.png"
    norm_png_path = stem + "_input_normalized.png"
    summary_path = stem + "_summary.txt"

    x = fit["x"]
    df = pd.DataFrame({
        "E": x,
        "I_mix_obs": fit["y"],
        "I_mix_fit": fit["yfit"],
        "I_film_part": fit["film_part"],
        "I_substrate_part": fit["sub_part"],
        "I_background": fit["bg_part"],
        "residual": fit["yfit"] - fit["y"],
    })
    df.to_csv(csv_path, index=False)

    p = fit["params"]
    vbm_info = []
    if args.vbm_film_ref is not None and args.vbm_sub_ref is not None:
        vbm_f = args.vbm_film_ref + p["dE_film"]
        vbm_s = args.vbm_sub_ref + p["dE_sub"]
        vbo_be = vbm_f - vbm_s
        vbm_info.append(f"VBM_film_fit_BE = {vbm_f:.8g} eV")
        vbm_info.append(f"VBM_sub_fit_BE  = {vbm_s:.8g} eV")
        vbm_info.append(f"VBO_BE = VBM_film_BE - VBM_sub_BE = {vbo_be:.8g} eV")
        vbm_info.append(f"VBO_energy_axis_opposite_sign = {-vbo_be:.8g} eV")

    with open(summary_path, "w", encoding="utf-8") as f:
        f.write("XPS film/substrate mixture fitting summary\n")
        f.write("========================================\n")
        f.write(f"film reference     : {film.name}\n")
        f.write(f"substrate reference: {sub.name}\n")
        f.write(f"mixed spectrum     : {mix.name}\n")
        f.write(f"fit range          : {x.min():.8g} to {x.max():.8g} eV\n")
        f.write(f"dE_film            : {p['dE_film']:.8g} eV\n")
        f.write(f"dE_substrate       : {p['dE_sub']:.8g} eV\n")
        f.write(f"dE_film - dE_sub   : {p['dE_shift_film_minus_sub']:.8g} eV\n")
        f.write(f"A_film             : {p['A_film']:.8g}\n")
        f.write(f"A_substrate        : {p['A_sub']:.8g}\n")
        f.write(f"RMSE               : {p['rmse']:.8g}\n")
        f.write(f"NRMSE              : {p['nrmse']:.8g}\n")
        f.write(f"objective          : {p['objective']:.8g}\n")
        f.write(f"success            : {p['success']}\n")
        f.write(f"message            : {p['message']}\n")
        for i, c in enumerate(p["bg_coeff"]):
            f.write(f"BG_coeff[{i}]       : {c:.8g}\n")
        if vbm_info:
            f.write("\n")
            f.write("\n".join(vbm_info) + "\n")

    fig, ax = plt.subplots(figsize=(8, 5))
    ax.plot(x, fit["y"], label="mixed obs", lw=1.5)
    ax.plot(x, fit["yfit"], label="fit", lw=1.5)
    ax.plot(x, fit["film_part"], label="film part", lw=1.0)
    ax.plot(x, fit["sub_part"], label="substrate part", lw=1.0)
    ax.plot(x, fit["bg_part"], label="background", lw=1.0)
    ax.set_xlabel("Energy / Binding energy (eV)")
    ax.set_ylabel("Intensity")
    ax.legend()
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    if args.save:
        fig.savefig(png_path, dpi=200)
    if args.show:
        plt.show()
    else:
        plt.close(fig)

    fig2, ax2 = plt.subplots(figsize=(8, 5))
    ax2.plot(film.x, normalize_intensity(film.y, args.norm_mode), label="film input normalized", lw=1.2)
    ax2.plot(sub.x, normalize_intensity(sub.y, args.norm_mode), label="substrate input normalized", lw=1.2)
    ax2.plot(mix.x, normalize_intensity(mix.y, args.norm_mode), label="mixed input normalized", lw=1.2)
    if args.fit_xmin is not None:
        ax2.axvline(args.fit_xmin, ls="--", lw=0.8, alpha=0.5)
    if args.fit_xmax is not None:
        ax2.axvline(args.fit_xmax, ls="--", lw=0.8, alpha=0.5)
    ax2.set_xlabel("Energy / Binding energy (eV)")
    ax2.set_ylabel(f"Normalized intensity ({args.norm_mode})")
    ax2.set_title("Input spectra comparison")
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    fig2.tight_layout()
    if args.save:
        fig2.savefig(norm_png_path, dpi=200)
    if args.show:
        plt.show()
    else:
        plt.close(fig2)

    return csv_path, summary_path, png_path, norm_png_path


def build_parser():
    p = argparse.ArgumentParser(description="Fit mixed XPS valence-band spectrum with film/substrate references.")
    p.add_argument("--film", required=True, help="film-only reference spectrum xlsx/xlsm/csv")
    p.add_argument("--substrate", required=True, help="substrate-only reference spectrum xlsx/xlsm/csv")
    p.add_argument("--mix", required=True, help="film/substrate mixed spectrum xlsx/xlsm/csv")
    p.add_argument("--sheet", default=0, help="sheet name or zero-based sheet index [default: 0]")
    p.add_argument("--x-col", type=int, default=0, help="zero-based x column [default: 0]")
    p.add_argument("--y-col", type=int, default=2, help="zero-based intensity column [default: 2; supplied files: deconvoluted data]")
    p.add_argument("--fit-xmin", type=float, default=None, help="minimum x for fitting")
    p.add_argument("--fit-xmax", type=float, default=None, help="maximum x for fitting")
    p.add_argument("--ref-xmin", type=float, default=None, help="minimum x retained in references")
    p.add_argument("--ref-xmax", type=float, default=None, help="maximum x retained in references")
    p.add_argument("--shift-film0", type=float, default=0.0, help="initial film energy shift")
    p.add_argument("--shift-sub0", type=float, default=0.0, help="initial substrate energy shift")
    p.add_argument("--shift-limit", type=float, default=5.0, help="absolute bound for each energy shift [eV]")
    p.add_argument("--bg-order", type=int, default=1, choices=[0, 1, 2], help="polynomial background order")
    p.add_argument("--outside-fill", type=float, default=0.0, help="interpolation value outside reference range")
    p.add_argument("--amp-film0", type=float, default=1.0, help="initial film amplitude")
    p.add_argument("--amp-sub0", type=float, default=0.1, help="initial substrate amplitude")
    p.add_argument("--amp-min", type=float, default=1e-12, help="minimum positive amplitude; enforced by penalty")
    p.add_argument("--amp-max", type=float, default=1e12, help="maximum positive amplitude; enforced by penalty")
    p.add_argument("--bg0", type=float, default=0.0, help="initial constant background")
    p.add_argument("--bg1", type=float, default=0.0, help="initial linear background coefficient")
    p.add_argument("--bg2", type=float, default=0.0, help="initial quadratic background coefficient")
    p.add_argument("--loss", default="soft_l1", choices=["linear", "soft_l1", "logcosh"], help="objective loss for Nelder-Mead")
    p.add_argument("--f-scale", type=float, default=0.05, help="robust loss scale")
    p.add_argument("--loss-on-log", type=int, default=0, choices=[0, 1], help="fit log intensity ratio instead of linear residual")
    p.add_argument("--max-iter", type=int, default=5000, help="maximum Nelder-Mead iterations")
    p.add_argument("--max-nfev", type=int, default=10000, help="maximum function evaluations")
    p.add_argument("--xatol", type=float, default=1e-8, help="Nelder-Mead x tolerance")
    p.add_argument("--fatol", type=float, default=1e-10, help="Nelder-Mead objective tolerance")
    p.add_argument("--penalty", type=float, default=1.0e6, help="soft-bound penalty strength")
    p.add_argument("--print-every", type=int, default=10, help="print Nelder-Mead convergence every N callback iterations; 0 disables")
    p.add_argument("--norm-mode", default="max", choices=["max", "minmax", "area"], help="normalization for input comparison plot")
    p.add_argument("--vbm-film-ref", type=float, default=None, help="VBM of film reference on the same x axis")
    p.add_argument("--vbm-sub-ref", type=float, default=None, help="VBM of substrate reference on the same x axis")
    p.add_argument("--output-stem", default="xps_vbo", help="output file stem")
    p.add_argument("--save", type=int, default=1, choices=[0, 1], help="save png plot")
    p.add_argument("--show", type=int, default=1, choices=[0, 1], help="show plot window")
    p.add_argument("--verbose", type=int, default=0, choices=[0, 1], help="scipy optimizer verbose output")
    return p


def main():
    parser = build_parser()
    args = parser.parse_args()
    # argparse gives sheet as string; convert simple integer strings to int.
    try:
        args.sheet = int(args.sheet)
    except Exception:
        pass

    film = crop(read_spectrum(args.film, args.sheet, args.x_col, args.y_col), args.ref_xmin, args.ref_xmax)
    sub = crop(read_spectrum(args.substrate, args.sheet, args.x_col, args.y_col), args.ref_xmin, args.ref_xmax)
    mix = read_spectrum(args.mix, args.sheet, args.x_col, args.y_col)

    print(f"film     : {film.name}, n={len(film.x)}, x=[{film.x.min():g}, {film.x.max():g}]")
    print(f"substrate: {sub.name}, n={len(sub.x)}, x=[{sub.x.min():g}, {sub.x.max():g}]")
    print(f"mix      : {mix.name}, n={len(mix.x)}, x=[{mix.x.min():g}, {mix.x.max():g}]")

    print("Optimizer: scipy.optimize.minimize(method='Nelder-Mead')")
    fit = fit_mixture(film, sub, mix, args)
    p = fit["params"]
    print("\nOptimized parameters")
    print(f"  dE_film          = {p['dE_film']:.8g} eV")
    print(f"  dE_substrate     = {p['dE_sub']:.8g} eV")
    print(f"  dE_film-dE_sub   = {p['dE_shift_film_minus_sub']:.8g} eV")
    print(f"  A_film           = {p['A_film']:.8g}")
    print(f"  A_substrate      = {p['A_sub']:.8g}")
    print(f"  RMSE             = {p['rmse']:.8g}")
    print(f"  NRMSE            = {p['nrmse']:.8g}")
    print(f"  objective        = {p['objective']:.8g}")
    print(f"  success          = {p['success']}")
    print(f"  message          = {p['message']}")

    if args.vbm_film_ref is not None and args.vbm_sub_ref is not None:
        vbm_f = args.vbm_film_ref + p["dE_film"]
        vbm_s = args.vbm_sub_ref + p["dE_sub"]
        vbo_be = vbm_f - vbm_s
        print("\nVBM / offset")
        print(f"  VBM_film_fit_BE = {vbm_f:.8g} eV")
        print(f"  VBM_sub_fit_BE  = {vbm_s:.8g} eV")
        print(f"  VBO_BE          = {vbo_be:.8g} eV  (= film BE - substrate BE)")
        print(f"  opposite sign   = {-vbo_be:.8g} eV")
    else:
        print("\nNote: supply --vbm-film-ref and --vbm-sub-ref to convert shift difference to VBO.")

    csv_path, summary_path, png_path, norm_png_path = save_outputs(film, sub, mix, fit, args)
    print("\nOutput")
    print(f"  fit csv        : {csv_path}")
    print(f"  summary        : {summary_path}")
    if args.save:
        print(f"  fit plot       : {png_path}")
        print(f"  normalized plot: {norm_png_path}")


if __name__ == "__main__":
    main()