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

Purpose
-------
Generate (or read) data that need not be exactly representable by the analysis
model, and extract three *representative* relaxation times in a reproducible
way.

Synthetic physical model (--source model)
-----------------------------------------
The data-generating model is fixed to FOUR exponential terms:

    y_true(x) = y0 + sum_{i=1}^4 A_i exp(-x/tau_i)

Analysis model
--------------
The analysis deliberately uses only THREE exponential terms:

    y_fit(x) = c + B_fast exp(-x/tau_fast)
                 + B_mid  exp(-x/tau_mid)
                 + B_slow exp(-x/tau_slow)

Workflow
--------
1) Determine tau_fast from a short-time window with a local 1-exp + constant fit.
2) Determine tau_slow from a long-time window with a local 1-exp + constant fit.
3) Fix tau_fast and tau_slow.
4) Determine only tau_mid by nonlinear minimization over the full analysis range.
   For each trial tau_mid, c and B_fast/B_mid/B_slow are solved by linear least
   squares.  Thus the final nonlinear optimization is only one-dimensional.
5) Plot profile SSE(tau_mid), which directly shows how sharply tau_mid is
   determined.

This is intended for comparative parameterization when a finite exponential
sum is NOT assumed to be the exact microscopic model.

Examples
--------
Synthetic 4-exp data, analyzed by 3 representative time constants:

    python fit_4exp_to_3tau_windowed.py ^
      --source model ^
      --true-A 0.40,0.30,0.20,0.10 ^
      --true-tau 1,2,4,12 ^
      --xmin 0 --xmax 30 --noise 0.01 ^
      --fast-xmin 0 --fast-xmax 1.5 ^
      --slow-xmin 10 --slow-xmax 30 ^
      --method Nelder-Mead ^
      -o test_4to3

Excel input:

    python fit_4exp_to_3tau_windowed.py ^
      --source excel -i decay.xlsx ^
      --xcol x --ycol y ^
      --fast-xmin 0 --fast-xmax 2 ^
      --slow-xmin 20 --slow-xmax 100 ^
      --fit-xmin 0 --fit-xmax 100 ^
      --method Powell ^
      -o sampleA

Manual fast/slow relaxation times can also be supplied:

    python fit_4exp_to_3tau_windowed.py --source excel -i decay.xlsx ^
      --tau-fast-fixed 0.8 --tau-slow-fixed 60

Dependencies
------------
    pip install numpy scipy pandas matplotlib openpyxl
"""

from __future__ import annotations

import argparse
import math
import warnings
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.optimize import minimize, lsq_linear


# =============================================================================
# Models
# =============================================================================
def physical_model_4exp(x, y0, amplitudes, taus):
    """Four-exponential data-generating model."""
    x = np.asarray(x, dtype=float)
    A = np.asarray(amplitudes, dtype=float)
    tau = np.asarray(taus, dtype=float)
    if len(A) != 4 or len(tau) != 4:
        raise ValueError("physical_model_4exp requires four A values and four tau values")

    y = np.full_like(x, float(y0), dtype=float)
    for Ai, ti in zip(A, tau):
        y += Ai * np.exp(-x / ti)
    return y


def analysis_model_3exp(x, offset, amplitudes, taus):
    """Three-exponential comparative analysis model."""
    x = np.asarray(x, dtype=float)
    B = np.asarray(amplitudes, dtype=float)
    tau = np.asarray(taus, dtype=float)
    if len(B) != 3 or len(tau) != 3:
        raise ValueError("analysis_model_3exp requires three B values and three tau values")

    y = np.full_like(x, float(offset), dtype=float)
    for Bi, ti in zip(B, tau):
        y += Bi * np.exp(-x / ti)
    return y


# =============================================================================
# Utilities
# =============================================================================
def parse_float_list(text, n=None, optname="value"):
    if text is None or str(text).strip() == "":
        return None
    a = np.array([float(v.strip()) for v in str(text).split(",")], dtype=float)
    if n is not None and len(a) != n:
        raise ValueError(f"--{optname} must contain exactly {n} comma-separated values")
    return a


def normalize_method(name):
    table = {
        "nelder-mead": "Nelder-Mead",
        "nelder_mead": "Nelder-Mead",
        "neldermead": "Nelder-Mead",
        "powell": "Powell",
        "bfgs": "BFGS",
        "cg": "CG",
        "l-bfgs-b": "L-BFGS-B",
        "lbfgsb": "L-BFGS-B",
        "slsqp": "SLSQP",
        "cobyla": "COBYLA",
    }
    return table.get(str(name).lower(), name)


def sigmoid_scalar(z):
    z = float(z)
    if z >= 0:
        return 1.0 / (1.0 + math.exp(-z))
    ez = math.exp(z)
    return ez / (1.0 + ez)


def logit_scalar(p):
    p = min(max(float(p), 1e-12), 1.0 - 1e-12)
    return math.log(p / (1.0 - p))


def select_range(x, y, xmin=None, xmax=None):
    mask = np.ones_like(x, dtype=bool)
    if xmin is not None:
        mask &= x >= xmin
    if xmax is not None:
        mask &= x <= xmax
    return x[mask], y[mask], mask


def auto_windows(x):
    """Generic defaults; explicit windows are strongly recommended."""
    xmin = float(np.min(x))
    xmax = float(np.max(x))
    span = xmax - xmin
    return (
        xmin,
        xmin + 0.20 * span,
        xmin + 0.55 * span,
        xmax,
    )


def optimizer_options(method, maxiter, tol):
    method = normalize_method(method)
    options = {"maxiter": maxiter}
    if method == "Nelder-Mead":
        options.update({"xatol": tol, "fatol": tol})
    elif method == "Powell":
        options.update({"xtol": tol, "ftol": tol})
    elif method in ("BFGS", "CG"):
        options.update({"gtol": tol})
    elif method == "L-BFGS-B":
        options.update({"ftol": tol, "gtol": tol})
    elif method == "SLSQP":
        options.update({"ftol": tol})
    return method, options


# =============================================================================
# Data input / synthetic generation
# =============================================================================
def read_excel_xy(filename, sheet, xcol, ycol):
    sheet_arg = int(sheet) if str(sheet).isdigit() else sheet
    df = pd.read_excel(filename, sheet_name=sheet_arg)

    def get_col(spec):
        spec = str(spec)
        if spec in df.columns:
            s = df[spec]
        else:
            try:
                idx = int(spec)
            except ValueError as exc:
                raise ValueError(
                    f"Column '{spec}' was not found. Available columns: {list(df.columns)}"
                ) from exc
            if idx < 0 or idx >= len(df.columns):
                raise IndexError(f"Column index {idx} is outside 0..{len(df.columns)-1}")
            s = df.iloc[:, idx]
        return pd.to_numeric(s, errors="coerce").to_numpy(dtype=float)

    x = get_col(xcol)
    y = get_col(ycol)
    good = np.isfinite(x) & np.isfinite(y)
    x = x[good]
    y = y[good]
    order = np.argsort(x)
    return x[order], y[order]


def generate_model_data(args):
    A = parse_float_list(args.true_A, 4, "true-A")
    tau = parse_float_list(args.true_tau, 4, "true-tau")

    if np.any(tau <= 0):
        raise ValueError("All --true-tau values must be > 0")
    if np.any(np.diff(tau) <= 0):
        raise ValueError("--true-tau must be in increasing order")

    x = np.linspace(args.xmin, args.xmax, args.npoints)
    y_clean = physical_model_4exp(x, args.true_y0, A, tau)

    rng = np.random.default_rng(args.seed)
    y = y_clean + rng.normal(0.0, args.noise, size=len(x))

    truth = {
        "y0": float(args.true_y0),
        "A": A,
        "tau": tau,
        "y_clean": y_clean,
    }
    return x, y, truth


# =============================================================================
# Linear solve at fixed time constants
# =============================================================================
def design_matrix(x, taus):
    x = np.asarray(x, dtype=float)
    cols = [np.ones_like(x)]
    cols.extend(np.exp(-x / float(tau)) for tau in taus)
    return np.column_stack(cols)


def solve_linear_coefficients(x, y, taus, positive_A=False):
    """
    For fixed taus, solve offset and amplitudes.

    positive_A=False:
        ordinary linear least squares.

    positive_A=True:
        offset unrestricted, amplitudes constrained to B_i >= 0.
    """
    X = design_matrix(x, taus)

    if positive_A:
        lo = np.r_[-np.inf, np.zeros(len(taus))]
        hi = np.full(len(taus) + 1, np.inf)
        res = lsq_linear(X, y, bounds=(lo, hi), method="trf", lsmr_tol="auto")
        coef = res.x
    else:
        coef, *_ = np.linalg.lstsq(X, y, rcond=None)

    yfit = X @ coef
    residual = y - yfit
    sse = float(np.dot(residual, residual))
    return coef, yfit, sse


# =============================================================================
# Fast/slow window analysis: 1 exp + constant
# =============================================================================
def one_exp_objective(logtau_vec, x, y, positive_A):
    logtau = float(np.atleast_1d(logtau_vec)[0])
    if not np.isfinite(logtau) or abs(logtau) > 700:
        return 1e300
    tau = math.exp(logtau)
    _, _, sse = solve_linear_coefficients(x, y, [tau], positive_A=positive_A)
    return sse


def estimate_window_tau(x, y, xmin, xmax, method, maxiter, tol, positive_A):
    xw, yw, mask = select_range(x, y, xmin, xmax)
    if len(xw) < 4:
        raise ValueError(
            f"Window [{xmin}, {xmax}] contains only {len(xw)} points; at least 4 are required"
        )

    span = max(float(np.ptp(xw)), np.finfo(float).eps)
    tau0 = max(span / 2.0, np.finfo(float).eps)
    z0 = np.array([math.log(tau0)], dtype=float)

    method, options = optimizer_options(method, maxiter, tol)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        res = minimize(
            one_exp_objective,
            z0,
            args=(xw, yw, positive_A),
            method=method,
            options=options,
        )

    tau = math.exp(float(res.x[0]))
    coef, yfit, sse = solve_linear_coefficients(xw, yw, [tau], positive_A=positive_A)
    return {
        "tau": tau,
        "offset": float(coef[0]),
        "A": float(coef[1]),
        "sse": sse,
        "x": xw,
        "y": yw,
        "yfit": yfit,
        "mask": mask,
        "result": res,
    }


def manual_window_stage(x, y, xmin, xmax, tau, positive_A):
    xw, yw, mask = select_range(x, y, xmin, xmax)
    coef, yfit, sse = solve_linear_coefficients(xw, yw, [tau], positive_A=positive_A)
    return {
        "tau": float(tau),
        "offset": float(coef[0]),
        "A": float(coef[1]),
        "sse": sse,
        "x": xw,
        "y": yw,
        "yfit": yfit,
        "mask": mask,
        "result": None,
    }


# =============================================================================
# Middle relaxation time: one nonlinear variable
# =============================================================================
def tau_mid_bounds(tau_fast, tau_slow, ratio_min):
    lo = tau_fast * ratio_min
    hi = tau_slow / ratio_min
    if not lo < hi:
        raise ValueError(
            "No admissible tau_mid exists. Need tau_slow/tau_fast > tau_ratio_min^2."
        )
    return lo, hi


def decode_tau_mid(u, tau_fast, tau_slow, ratio_min):
    """
    Smoothly map unrestricted u to an allowed tau_mid in logarithmic space:

        tau_fast * ratio_min < tau_mid < tau_slow / ratio_min

    Thus the middle time is always ordered and separated from both fixed ends.
    """
    lo, hi = tau_mid_bounds(tau_fast, tau_slow, ratio_min)
    loglo = math.log(lo)
    loghi = math.log(hi)
    s = sigmoid_scalar(float(np.atleast_1d(u)[0]))
    return math.exp(loglo + s * (loghi - loglo))


def encode_tau_mid(tau_mid, tau_fast, tau_slow, ratio_min):
    lo, hi = tau_mid_bounds(tau_fast, tau_slow, ratio_min)
    loglo = math.log(lo)
    loghi = math.log(hi)
    p = (math.log(tau_mid) - loglo) / (loghi - loglo)
    return logit_scalar(p)


def middle_objective(u, x, y, tau_fast, tau_slow, ratio_min, positive_A):
    try:
        tau_mid = decode_tau_mid(u, tau_fast, tau_slow, ratio_min)
    except (ValueError, OverflowError, FloatingPointError):
        return 1e300

    _, _, sse = solve_linear_coefficients(
        x, y, [tau_fast, tau_mid, tau_slow], positive_A=positive_A
    )
    return sse


def fit_middle_tau(x, y, tau_fast, tau_slow, args):
    lo, hi = tau_mid_bounds(tau_fast, tau_slow, args.tau_ratio_min)

    if args.init_tau_mid is None:
        tau0 = math.sqrt(lo * hi)
    else:
        tau0 = float(args.init_tau_mid)
        tau0 = min(max(tau0, lo * (1 + 1e-9)), hi / (1 + 1e-9))

    u0 = encode_tau_mid(tau0, tau_fast, tau_slow, args.tau_ratio_min)
    method, options = optimizer_options(args.method, args.maxiter, args.tol)

    rng = np.random.default_rng(args.seed + 20260904)
    candidates = []

    for istart in range(args.n_starts):
        if istart == 0:
            start = np.array([u0], dtype=float)
        else:
            start = np.array([rng.normal(0.0, args.start_scale)], dtype=float)

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            res = minimize(
                middle_objective,
                start,
                args=(
                    x, y, tau_fast, tau_slow,
                    args.tau_ratio_min, args.positive_A,
                ),
                method=method,
                options=options,
            )

        tau_mid = decode_tau_mid(
            res.x, tau_fast, tau_slow, args.tau_ratio_min
        )
        taus = np.array([tau_fast, tau_mid, tau_slow], dtype=float)
        coef, yfit, sse = solve_linear_coefficients(
            x, y, taus, positive_A=args.positive_A
        )

        candidates.append({
            "start": istart,
            "result": res,
            "tau_mid": tau_mid,
            "taus": taus,
            "coef": coef,
            "yfit": yfit,
            "sse": sse,
        })

    best = min(candidates, key=lambda c: c["sse"])
    return best, candidates


# =============================================================================
# Profile SSE(tau_mid)
# =============================================================================
def calculate_profile(x, y, tau_fast, tau_slow, args):
    if args.profile_points < 3:
        return None

    lo, hi = tau_mid_bounds(tau_fast, tau_slow, args.tau_ratio_min)
    tau_axis = np.geomspace(lo * (1 + 1e-8), hi / (1 + 1e-8), args.profile_points)
    sse = np.empty_like(tau_axis)

    for i, tm in enumerate(tau_axis):
        _, _, sse[i] = solve_linear_coefficients(
            x, y, [tau_fast, tm, tau_slow], positive_A=args.positive_A
        )

    return tau_axis, sse


def profile_relative_interval(profile, fraction=0.01):
    """Return the tau range with SSE <= (1+fraction)*SSE_min."""
    if profile is None:
        return np.nan, np.nan
    tau_axis, sse = profile
    smin = float(np.min(sse))
    limit = smin * (1.0 + fraction) if smin > 0 else smin + fraction
    mask = sse <= limit
    if not np.any(mask):
        return np.nan, np.nan
    return float(np.min(tau_axis[mask])), float(np.max(tau_axis[mask]))


# =============================================================================
# Statistics / output
# =============================================================================
def fit_statistics(y, yfit, nparam=5):
    """
    Final full-range stage has five fitted quantities:
      offset, B_fast, B_mid, B_slow, tau_mid.
    tau_fast and tau_slow are fixed by the analysis protocol.
    """
    residual = y - yfit
    n = len(y)
    sse = float(np.dot(residual, residual))
    rmse = math.sqrt(sse / n)
    ss_tot = float(np.sum((y - np.mean(y)) ** 2))
    r2 = 1.0 - sse / ss_tot if ss_tot > 0 else np.nan

    sigma2 = max(sse / n, 1e-300)
    logL = -0.5 * n * (math.log(2 * math.pi * sigma2) + 1.0)
    aic = 2 * nparam - 2 * logL
    bic = nparam * math.log(n) - 2 * logL
    return {
        "SSE": sse,
        "RMSE": rmse,
        "R2": r2,
        "logL_hat": logL,
        "AIC": aic,
        "BIC": bic,
    }


def save_window_plot(stage, filename, title, tau_name):
    xd = np.linspace(np.min(stage["x"]), np.max(stage["x"]), 500)
    yd = stage["offset"] + stage["A"] * np.exp(-xd / stage["tau"])

    fig, ax = plt.subplots(figsize=(7.4, 5.0))
    ax.scatter(stage["x"], stage["y"], s=24, label="data in window")
    ax.plot(xd, yd, lw=2, label=f"{tau_name} = {stage['tau']:.6g}")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_title(title)
    ax.grid(alpha=0.25)
    ax.legend()
    fig.tight_layout()
    fig.savefig(filename, dpi=180)
    plt.close(fig)


def save_full_fit_plot(x_all, y_all, xfit, best, truth, args, filename):
    coef = best["coef"]
    taus = best["taus"]
    xd = np.linspace(np.min(xfit), np.max(xfit), 1200)
    yd = analysis_model_3exp(xd, coef[0], coef[1:], taus)

    fig, ax = plt.subplots(figsize=(8.4, 5.7))
    ax.scatter(x_all, y_all, s=20, label="data")
    ax.plot(xd, yd, lw=2.1, label="3-time comparative fit")

    if truth is not None:
        ytrue = physical_model_4exp(xd, truth["y0"], truth["A"], truth["tau"])
        ax.plot(xd, ytrue, "--", lw=1.5, label="true 4-exp model")

    ax.axvspan(args.fast_xmin, args.fast_xmax, alpha=0.08, label="fast window")
    ax.axvspan(args.slow_xmin, args.slow_xmax, alpha=0.08, label="slow window")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_title("4-exp data represented by three characteristic relaxation times")
    ax.grid(alpha=0.25)
    ax.legend()
    fig.tight_layout()
    fig.savefig(filename, dpi=180)
    plt.close(fig)


def save_component_plot(xfit, best, filename):
    coef = best["coef"]
    taus = best["taus"]
    xd = np.linspace(np.min(xfit), np.max(xfit), 1200)

    fig, ax = plt.subplots(figsize=(8.0, 5.4))
    labels = ["fast", "mid", "slow"]
    for i, label in enumerate(labels):
        component = coef[i + 1] * np.exp(-xd / taus[i])
        ax.plot(xd, component, lw=1.8,
                label=f"{label}: B={coef[i+1]:.4g}, tau={taus[i]:.4g}")
    ax.set_xlabel("x")
    ax.set_ylabel("component contribution")
    ax.set_title("Three fitted exponential components")
    ax.grid(alpha=0.25)
    ax.legend()
    fig.tight_layout()
    fig.savefig(filename, dpi=180)
    plt.close(fig)


def save_profile_plot(profile, best, filename):
    if profile is None:
        return
    tau_axis, sse = profile
    smin = np.min(sse)
    rel = sse / smin if smin > 0 else sse - smin

    fig, ax = plt.subplots(figsize=(7.4, 5.0))
    ax.plot(tau_axis, rel, lw=2)
    ax.axvline(best["tau_mid"], ls="--", label=f"best tau_mid={best['tau_mid']:.6g}")
    ax.set_xscale("log")
    ax.set_xlabel(r"$\tau_{mid}$")
    ax.set_ylabel("SSE / SSE_min" if smin > 0 else "SSE - SSE_min")
    ax.set_title("Profile SSE for the middle relaxation time")
    ax.grid(alpha=0.25)
    ax.legend()
    fig.tight_layout()
    fig.savefig(filename, dpi=180)
    plt.close(fig)


def save_outputs(x_all, y_all, xfit, yfit_data, fast, slow,
                 best, candidates, truth, profile, args):
    prefix = Path(args.output)
    prefix.parent.mkdir(parents=True, exist_ok=True)

    coef = best["coef"]
    taus = best["taus"]
    yfit = best["yfit"]
    stats = fit_statistics(yfit_data, yfit, nparam=5)
    profile_1pct_low, profile_1pct_high = profile_relative_interval(profile, 0.01)
    profile_5pct_low, profile_5pct_high = profile_relative_interval(profile, 0.05)

    save_window_plot(
        fast, str(prefix) + "_fast_window.png",
        "Fast-time window: local 1-exp + constant", "tau_fast"
    )
    save_window_plot(
        slow, str(prefix) + "_slow_window.png",
        "Slow-time window: local 1-exp + constant", "tau_slow"
    )
    save_full_fit_plot(
        x_all, y_all, xfit, best, truth, args,
        str(prefix) + "_fit.png"
    )
    save_component_plot(xfit, best, str(prefix) + "_components.png")
    save_profile_plot(profile, best, str(prefix) + "_profile_tau_mid.png")

    summary = {
        "source": args.source,
        "physical_model": "4-exp" if args.source == "model" else "experimental",
        "analysis_model": "3-exp representative times",
        "method": normalize_method(args.method),
        "window_method": normalize_method(args.window_method),
        "positive_A": args.positive_A,
        "tau_ratio_min": args.tau_ratio_min,
        "fit_xmin": args.fit_xmin,
        "fit_xmax": args.fit_xmax,
        "fast_xmin": args.fast_xmin,
        "fast_xmax": args.fast_xmax,
        "slow_xmin": args.slow_xmin,
        "slow_xmax": args.slow_xmax,
        "tau_fast": taus[0],
        "tau_mid": taus[1],
        "tau_slow": taus[2],
        "offset": coef[0],
        "B_fast": coef[1],
        "B_mid": coef[2],
        "B_slow": coef[3],
        "tau_mid_profile_1pct_low": profile_1pct_low,
        "tau_mid_profile_1pct_high": profile_1pct_high,
        "tau_mid_profile_5pct_low": profile_5pct_low,
        "tau_mid_profile_5pct_high": profile_5pct_high,
        "fast_window_SSE": fast["sse"],
        "slow_window_SSE": slow["sse"],
        "final_success": bool(best["result"].success),
        "final_message": str(best["result"].message),
        "nfev": getattr(best["result"], "nfev", np.nan),
        "nit": getattr(best["result"], "nit", np.nan),
        **stats,
    }
    summary_df = pd.DataFrame([summary])
    summary_df.to_csv(str(prefix) + "_summary.csv", index=False)

    data_df = pd.DataFrame({
        "x": xfit,
        "y": yfit_data,
        "y_fit": yfit,
        "residual": yfit_data - yfit,
    })
    if truth is not None:
        data_df["y_true_4exp"] = physical_model_4exp(
            xfit, truth["y0"], truth["A"], truth["tau"]
        )

    parameter_df = pd.DataFrame([
        {"parameter": "tau_fast", "value": taus[0], "role": "fixed from fast window"},
        {"parameter": "tau_mid",  "value": taus[1], "role": "nonlinear full-range fit"},
        {"parameter": "tau_slow", "value": taus[2], "role": "fixed from slow window"},
        {"parameter": "offset",   "value": coef[0], "role": "linear full-range fit"},
        {"parameter": "B_fast",   "value": coef[1], "role": "linear full-range fit"},
        {"parameter": "B_mid",    "value": coef[2], "role": "linear full-range fit"},
        {"parameter": "B_slow",   "value": coef[3], "role": "linear full-range fit"},
    ])

    fast_df = pd.DataFrame({
        "x": fast["x"], "y": fast["y"], "y_fit": fast["yfit"],
        "residual": fast["y"] - fast["yfit"],
    })
    slow_df = pd.DataFrame({
        "x": slow["x"], "y": slow["y"], "y_fit": slow["yfit"],
        "residual": slow["y"] - slow["yfit"],
    })

    starts_rows = []
    for c in candidates:
        starts_rows.append({
            "start": c["start"],
            "success": bool(c["result"].success),
            "message": str(c["result"].message),
            "tau_mid": c["tau_mid"],
            "SSE": c["sse"],
            "offset": c["coef"][0],
            "B_fast": c["coef"][1],
            "B_mid": c["coef"][2],
            "B_slow": c["coef"][3],
            "nfev": getattr(c["result"], "nfev", np.nan),
            "nit": getattr(c["result"], "nit", np.nan),
        })
    starts_df = pd.DataFrame(starts_rows)

    with pd.ExcelWriter(str(prefix) + "_results.xlsx", engine="openpyxl") as writer:
        summary_df.to_excel(writer, sheet_name="Summary", index=False)
        parameter_df.to_excel(writer, sheet_name="Parameters", index=False)
        data_df.to_excel(writer, sheet_name="Full_fit", index=False)
        fast_df.to_excel(writer, sheet_name="Fast_window", index=False)
        slow_df.to_excel(writer, sheet_name="Slow_window", index=False)
        starts_df.to_excel(writer, sheet_name="Multi_start", index=False)

        if profile is not None:
            tau_axis, sse = profile
            profile_df = pd.DataFrame({
                "tau_mid": tau_axis,
                "SSE": sse,
                "SSE_over_min": sse / np.min(sse) if np.min(sse) > 0 else sse - np.min(sse),
            })
            profile_df.to_excel(writer, sheet_name="Profile_tau_mid", index=False)

        if truth is not None:
            rows = [{"parameter": "true_y0", "value": truth["y0"]}]
            for i in range(4):
                rows.append({"parameter": f"true_A{i+1}", "value": truth["A"][i]})
                rows.append({"parameter": f"true_tau{i+1}", "value": truth["tau"][i]})
            pd.DataFrame(rows).to_excel(writer, sheet_name="True_4exp_parameters", index=False)

    return stats


# =============================================================================
# CLI
# =============================================================================
def build_parser():
    p = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description=(
            "Generate/read data and extract three representative relaxation times: "
            "fast and slow from fixed analysis windows, middle from a one-parameter "
            "full-range nonlinear fit. Synthetic data use a four-exponential model."
        ),
    )

    # Data source
    p.add_argument("--source", choices=["excel", "model"], default="model")
    p.add_argument("-i", "--input", help="Input Excel file")
    p.add_argument("--sheet", default="0", help="Excel sheet name or zero-based index")
    p.add_argument("--xcol", default="x", help="x column name or zero-based index")
    p.add_argument("--ycol", default="y", help="y column name or zero-based index")

    # Four-exponential synthetic physical model
    p.add_argument("--true-y0", type=float, default=0.05)
    p.add_argument("--true-A", default="0.40,0.30,0.20,0.10",
                   help="Four amplitudes for synthetic 4-exp data")
    p.add_argument("--true-tau", default="1,2,4,12",
                   help="Four increasing taus for synthetic 4-exp data")
    p.add_argument("--xmin", type=float, default=0.0)
    p.add_argument("--xmax", type=float, default=30.0)
    p.add_argument("--npoints", type=int, default=250)
    p.add_argument("--noise", type=float, default=0.01)

    # Analysis ranges
    p.add_argument("--fit-xmin", type=float, default=None)
    p.add_argument("--fit-xmax", type=float, default=None)
    p.add_argument("--fast-xmin", type=float, default=None)
    p.add_argument("--fast-xmax", type=float, default=None)
    p.add_argument("--slow-xmin", type=float, default=None)
    p.add_argument("--slow-xmax", type=float, default=None)

    # Optional fixed end time constants
    p.add_argument("--tau-fast-fixed", type=float, default=None,
                   help="Skip fast-window tau optimization and use this value")
    p.add_argument("--tau-slow-fixed", type=float, default=None,
                   help="Skip slow-window tau optimization and use this value")

    # Analysis constraints
    p.add_argument("--tau-ratio-min", type=float, default=1.0,
                   help="Require tau_mid/tau_fast and tau_slow/tau_mid >= this ratio")
    p.add_argument(
        "--positive-A",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Require fitted exponential amplitudes >= 0; use --no-positive-A to allow signed components",
    )

    # scipy.optimize.minimize
    p.add_argument("--method", default="Nelder-Mead",
                   help="minimize method for tau_mid")
    p.add_argument("--window-method", default="Nelder-Mead",
                   help="minimize method for fast/slow local fits")
    p.add_argument("--init-tau-mid", type=float, default=None)
    p.add_argument("--n-starts", type=int, default=5)
    p.add_argument("--start-scale", type=float, default=2.0)
    p.add_argument("--maxiter", type=int, default=30000)
    p.add_argument("--tol", type=float, default=1e-10)

    # Profile curve
    p.add_argument("--profile-points", type=int, default=250,
                   help="Number of logarithmic tau_mid values for profile SSE; <3 disables")

    # General
    p.add_argument("--seed", type=int, default=1234)
    p.add_argument("-o", "--output", default="fit_4exp_to_3tau")
    return p


def validate_args(args):
    if args.source == "excel" and not args.input:
        raise ValueError("--input is required for --source excel")
    if args.xmax <= args.xmin:
        raise ValueError("--xmax must be > --xmin")
    if args.npoints < 10:
        raise ValueError("--npoints must be >= 10")
    if args.noise < 0:
        raise ValueError("--noise must be >= 0")
    if args.tau_ratio_min < 1.0:
        raise ValueError("--tau-ratio-min must be >= 1")
    if args.n_starts < 1:
        raise ValueError("--n-starts must be >= 1")


def main():
    args = build_parser().parse_args()
    validate_args(args)

    if args.source == "excel":
        x_all, y_all = read_excel_xy(args.input, args.sheet, args.xcol, args.ycol)
        truth = None
    else:
        x_all, y_all, truth = generate_model_data(args)

    auto_fmin, auto_fmax, auto_smin, auto_smax = auto_windows(x_all)
    if args.fast_xmin is None:
        args.fast_xmin = auto_fmin
    if args.fast_xmax is None:
        args.fast_xmax = auto_fmax
    if args.slow_xmin is None:
        args.slow_xmin = auto_smin
    if args.slow_xmax is None:
        args.slow_xmax = auto_smax
    if args.fit_xmin is None:
        args.fit_xmin = float(np.min(x_all))
    if args.fit_xmax is None:
        args.fit_xmax = float(np.max(x_all))

    xfit, yfit_data, _ = select_range(x_all, y_all, args.fit_xmin, args.fit_xmax)
    if len(xfit) < 8:
        raise ValueError("The full fit range contains too few points")

    if args.tau_fast_fixed is None:
        fast = estimate_window_tau(
            x_all, y_all, args.fast_xmin, args.fast_xmax,
            args.window_method, args.maxiter, args.tol, args.positive_A,
        )
    else:
        if args.tau_fast_fixed <= 0:
            raise ValueError("--tau-fast-fixed must be > 0")
        fast = manual_window_stage(
            x_all, y_all, args.fast_xmin, args.fast_xmax,
            args.tau_fast_fixed, args.positive_A,
        )

    if args.tau_slow_fixed is None:
        slow = estimate_window_tau(
            x_all, y_all, args.slow_xmin, args.slow_xmax,
            args.window_method, args.maxiter, args.tol, args.positive_A,
        )
    else:
        if args.tau_slow_fixed <= 0:
            raise ValueError("--tau-slow-fixed must be > 0")
        slow = manual_window_stage(
            x_all, y_all, args.slow_xmin, args.slow_xmax,
            args.tau_slow_fixed, args.positive_A,
        )

    tau_fast = fast["tau"]
    tau_slow = slow["tau"]

    if tau_slow <= tau_fast:
        raise ValueError(
            f"tau_slow={tau_slow:g} <= tau_fast={tau_fast:g}. "
            "Review the fast/slow windows or supply fixed values."
        )

    tau_mid_bounds(tau_fast, tau_slow, args.tau_ratio_min)

    best, candidates = fit_middle_tau(
        xfit, yfit_data, tau_fast, tau_slow, args
    )
    profile = calculate_profile(
        xfit, yfit_data, tau_fast, tau_slow, args
    )
    stats = save_outputs(
        x_all, y_all, xfit, yfit_data,
        fast, slow, best, candidates, truth, profile, args,
    )

    coef = best["coef"]
    taus = best["taus"]

    print("\n=== 4-exp physical model -> 3 characteristic-time analysis ===")
    print(f"source          : {args.source}")
    print(f"fit range       : {args.fit_xmin:g} .. {args.fit_xmax:g}")
    print(f"fast window     : {args.fast_xmin:g} .. {args.fast_xmax:g}")
    print(f"slow window     : {args.slow_xmin:g} .. {args.slow_xmax:g}")
    print(f"window method   : {normalize_method(args.window_method)}")
    print(f"middle method   : {normalize_method(args.method)}")
    print(f"positive A      : {args.positive_A}")
    print(f"tau ratio min   : {args.tau_ratio_min:g}")
    print()
    print(f"tau_fast [window/fixed] : {taus[0]:.10g}")
    print(f"tau_mid  [nonlinear]    : {taus[1]:.10g}")
    print(f"tau_slow [window/fixed] : {taus[2]:.10g}")
    print(f"offset                  : {coef[0]:.10g}")
    print(f"B_fast                  : {coef[1]:.10g}")
    print(f"B_mid                   : {coef[2]:.10g}")
    print(f"B_slow                  : {coef[3]:.10g}")
    print(f"SSE                     : {stats['SSE']:.10g}")
    print(f"RMSE                    : {stats['RMSE']:.10g}")
    print(f"R^2                     : {stats['R2']:.10g}")
    if profile is not None:
        p1lo, p1hi = profile_relative_interval(profile, 0.01)
        p5lo, p5hi = profile_relative_interval(profile, 0.05)
        print(f"tau_mid profile +1% SSE: {p1lo:.10g} .. {p1hi:.10g}")
        print(f"tau_mid profile +5% SSE: {p5lo:.10g} .. {p5hi:.10g}")
    print(f"AIC*                    : {stats['AIC']:.10g}")
    print(f"BIC*                    : {stats['BIC']:.10g}")
    print("  *Final-stage count: 5 fitted parameters; tau_fast/tau_slow are")
    print("   treated as fixed protocol parameters after the window analyses.")

    if truth is not None:
        print("\nTrue 4-exp model:")
        for i in range(4):
            print(f"  A{i+1}={truth['A'][i]:.8g}, tau{i+1}={truth['tau'][i]:.8g}")
        print("  Note: fitted tau_fast/mid/slow are representative parameters,")
        print("        not estimates of three selected true tau_i values.")

    print("\nOutput:")
    print(f"  {args.output}_fit.png")
    print(f"  {args.output}_fast_window.png")
    print(f"  {args.output}_slow_window.png")
    print(f"  {args.output}_components.png")
    if args.profile_points >= 3:
        print(f"  {args.output}_profile_tau_mid.png")
    print(f"  {args.output}_summary.csv")
    print(f"  {args.output}_results.xlsx")


if __name__ == "__main__":
    main()