#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Bayesian decay-model selection by Replica-Exchange Monte Carlo (parallel tempering)
and thermodynamic integration.

Models
------
  exp1, exp2, ... : y = y0 + sum_i A_i * exp(-x/tau_i)
  stretch         : y = y0 + A * exp(-(x/tau)**s)

Data source
-----------
1) Excel x-y data
   python bayes_decay_model_selection.py --source excel -i data.xlsx \
       --xcol x --ycol y --models exp1,exp2,exp3,stretch

2) Synthetic data generated through model()
   python bayes_decay_model_selection.py --source model \
       --true-model stretch --params "y0=0.05,A=1.0,tau=2.5,s=0.65" \
       --xmin 0 --xmax 10 --npoints 120 --noise 0.02 \
       --models exp1,exp2,exp3,stretch

Output
------
  <prefix>_results.xlsx
  <prefix>_summary.csv
  <prefix>_fit.png
  <prefix>_ti.png

Notes
-----
* Bayesian evidence depends on the prior ranges. The automatic prior ranges are
  convenient defaults, not universally valid physical priors.
* F = -log Z is used as the Bayesian free energy. Smaller F is better.
* Equal prior probabilities are assumed for the candidate models when reporting
  posterior model probabilities.
"""

from __future__ import annotations

import argparse
import math
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


# =============================================================================
# User-editable defaults
# =============================================================================
DEFAULT_MODELS = "exp1,exp2,exp3,stretch"
DEFAULT_N_REPLICAS = 20
DEFAULT_STEPS = 30000
DEFAULT_BURN = 10000
DEFAULT_THIN = 10
DEFAULT_SWAP_INTERVAL = 10
DEFAULT_PROPOSAL_SCALE = 0.06   # in normalized prior coordinates [0, 1]
DEFAULT_BETA_POWER = 4.0        # beta_i = sin(pi*i/2/(R-1))**power; clusters near beta=0 and 1
DEFAULT_SEED = 12345

# Automatic prior ranges. These are multiplied by scales estimated from the data.
AUTO_Y0_MARGIN = 0.75           # y0 in [ymin-margin*yscale, ymax+margin*yscale]
AUTO_AMPLITUDE_MAX = 3.0        # A in [0, AUTO_AMPLITUDE_MAX * yscale]
AUTO_TAU_MIN_FACTOR = 0.1       # tau_min ~ min positive dx * factor
AUTO_TAU_MAX_FACTOR = 20.0      # tau_max ~ x-span * factor
AUTO_SIGMA_MIN_FACTOR = 1e-4    # sigma_min ~ yscale * factor
AUTO_SIGMA_MAX_FACTOR = 1.0     # sigma_max ~ yscale * factor


# =============================================================================
# Decay models
# =============================================================================
def model(x: np.ndarray, model_name: str, params: Dict[str, float]) -> np.ndarray:
    """Evaluate one of the supported decay models.

    Parameters
    ----------
    x : array
        Independent variable, normally time >= 0.
    model_name : str
        'exp1', 'exp2', ... or 'stretch'.
    params : dict
        expK   : y0, A1..AK, tau1..tauK
        stretch: y0, A, tau, s
    """
    x = np.asarray(x, dtype=float)
    name = model_name.lower().strip()

    if name == "stretch":
        y0 = params.get("y0", 0.0)
        A = params["A"]
        tau = params["tau"]
        s = params["s"]
        return y0 + A * np.exp(-np.power(np.clip(x / tau, 0.0, None), s))

    m = re.fullmatch(r"exp(\d+)", name)
    if not m:
        raise ValueError(f"Unknown model: {model_name}")

    k = int(m.group(1))
    y = np.full_like(x, params.get("y0", 0.0), dtype=float)
    for i in range(1, k + 1):
        y += params[f"A{i}"] * np.exp(-x / params[f"tau{i}"])
    return y


def parse_param_string(text: str) -> Dict[str, float]:
    """Parse 'a=1,b=2.5' into {'a':1.0, 'b':2.5}."""
    out: Dict[str, float] = {}
    if not text:
        return out
    for item in text.split(","):
        item = item.strip()
        if not item:
            continue
        if "=" not in item:
            raise ValueError(f"Bad parameter specification: {item!r}")
        key, val = item.split("=", 1)
        out[key.strip()] = float(val)
    return out


def validate_true_params(model_name: str, params: Dict[str, float]) -> None:
    name = model_name.lower()
    required = ["y0"]
    if name == "stretch":
        required += ["A", "tau", "s"]
    else:
        m = re.fullmatch(r"exp(\d+)", name)
        if not m:
            raise ValueError(f"Unknown true model: {model_name}")
        k = int(m.group(1))
        required += [f"A{i}" for i in range(1, k + 1)]
        required += [f"tau{i}" for i in range(1, k + 1)]
    missing = [p for p in required if p not in params]
    if missing:
        raise ValueError(f"Missing parameters for {model_name}: {', '.join(missing)}")


# =============================================================================
# Data input / synthetic generation
# =============================================================================
def read_excel_xy(
    path: str,
    sheet: str | int = 0,
    xcol: str = "x",
    ycol: str = "y",
) -> Tuple[np.ndarray, np.ndarray]:
    df = pd.read_excel(path, sheet_name=sheet, engine="openpyxl")

    # xcol/ycol may be column names or zero-based integer positions.
    def get_col(spec: str) -> pd.Series:
        if spec in df.columns:
            return df[spec]
        try:
            idx = int(spec)
            return df.iloc[:, idx]
        except (ValueError, IndexError):
            raise KeyError(
                f"Column {spec!r} not found. Available columns: {list(df.columns)}"
            )

    d = pd.DataFrame({"x": get_col(xcol), "y": get_col(ycol)}).dropna()
    x = pd.to_numeric(d["x"], errors="coerce").to_numpy(dtype=float)
    y = pd.to_numeric(d["y"], errors="coerce").to_numpy(dtype=float)
    mask = np.isfinite(x) & np.isfinite(y)
    x, y = x[mask], y[mask]
    if len(x) < 3:
        raise ValueError("At least 3 finite x-y points are required.")

    order = np.argsort(x)
    return x[order], y[order]


def generate_data(
    true_model: str,
    params: Dict[str, float],
    xmin: float,
    xmax: float,
    npoints: int,
    noise: float,
    xscale: str,
    rng: np.random.Generator,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    validate_true_params(true_model, params)
    if xscale == "log":
        if xmin <= 0:
            raise ValueError("--xmin must be > 0 when --xscale log is used.")
        x = np.geomspace(xmin, xmax, npoints)
    else:
        x = np.linspace(xmin, xmax, npoints)

    y_true = model(x, true_model, params)
    y = y_true + rng.normal(0.0, noise, size=x.size)
    return x, y, y_true


# =============================================================================
# Prior definition in normalized coordinates u in [0, 1]
# =============================================================================
@dataclass
class PriorSpec:
    name: str
    kind: str            # 'uniform' or 'loguniform'
    lo: float
    hi: float

    def transform(self, u: float) -> float:
        if self.kind == "uniform":
            return self.lo + u * (self.hi - self.lo)
        if self.kind == "loguniform":
            if self.lo <= 0 or self.hi <= 0:
                raise ValueError(f"Log-uniform prior for {self.name} requires lo,hi > 0")
            return math.exp(math.log(self.lo) + u * (math.log(self.hi) - math.log(self.lo)))
        raise ValueError(f"Unknown prior kind: {self.kind}")


@dataclass
class ModelPrior:
    model_name: str
    specs: List[PriorSpec]
    tau_u_indices: List[int]
    estimate_sigma: bool
    fixed_sigma: Optional[float]

    @property
    def ndim(self) -> int:
        return len(self.specs)

    def sample_u(self, rng: np.random.Generator) -> np.ndarray:
        u = rng.random(self.ndim)
        # Keep tau coordinates ordered so expK does not waste sampling on K! label permutations.
        if len(self.tau_u_indices) > 1:
            vals = np.sort(u[self.tau_u_indices])
            u[self.tau_u_indices] = vals
        return u

    def is_valid_u(self, u: np.ndarray) -> bool:
        if np.any(u < 0.0) or np.any(u > 1.0):
            return False
        if len(self.tau_u_indices) > 1:
            vals = u[self.tau_u_indices]
            if np.any(np.diff(vals) <= 0.0):
                return False
        return True

    def transform(self, u: np.ndarray) -> Dict[str, float]:
        return {spec.name: spec.transform(float(ui)) for spec, ui in zip(self.specs, u)}


def automatic_prior_bounds(x: np.ndarray, y: np.ndarray) -> Dict[str, float]:
    ymin, ymax = float(np.min(y)), float(np.max(y))
    yrange = ymax - ymin
    yscale = max(yrange, float(np.std(y)), max(abs(ymin), abs(ymax)) * 0.1, 1e-12)

    xs = np.sort(np.unique(x))
    dx = np.diff(xs)
    pos_dx = dx[dx > 0]
    xspan = max(float(xs[-1] - xs[0]), 1e-12)
    min_dx = float(np.min(pos_dx)) if len(pos_dx) else xspan / max(len(x), 2)

    return {
        "y0_min": ymin - AUTO_Y0_MARGIN * yscale,
        "y0_max": ymax + AUTO_Y0_MARGIN * yscale,
        "A_max": AUTO_AMPLITUDE_MAX * yscale,
        "tau_min": max(AUTO_TAU_MIN_FACTOR * min_dx, 1e-12),
        "tau_max": max(AUTO_TAU_MAX_FACTOR * xspan, 10 * min_dx),
        "sigma_min": max(AUTO_SIGMA_MIN_FACTOR * yscale, 1e-15),
        "sigma_max": max(AUTO_SIGMA_MAX_FACTOR * yscale, AUTO_SIGMA_MIN_FACTOR * yscale * 10),
    }


def build_prior(
    model_name: str,
    x: np.ndarray,
    y: np.ndarray,
    args: argparse.Namespace,
) -> ModelPrior:
    auto = automatic_prior_bounds(x, y)

    y0_min = auto["y0_min"] if args.y0_min is None else args.y0_min
    y0_max = auto["y0_max"] if args.y0_max is None else args.y0_max
    A_max = auto["A_max"] if args.A_max is None else args.A_max
    tau_min = auto["tau_min"] if args.tau_min is None else args.tau_min
    tau_max = auto["tau_max"] if args.tau_max is None else args.tau_max
    sigma_min = auto["sigma_min"] if args.sigma_min is None else args.sigma_min
    sigma_max = auto["sigma_max"] if args.sigma_max is None else args.sigma_max

    if y0_max <= y0_min:
        raise ValueError("y0_max must be > y0_min")
    if A_max <= 0:
        raise ValueError("A_max must be > 0")
    if tau_min <= 0 or tau_max <= tau_min:
        raise ValueError("Require 0 < tau_min < tau_max")

    specs: List[PriorSpec] = [PriorSpec("y0", "uniform", y0_min, y0_max)]
    tau_indices: List[int] = []

    name = model_name.lower()
    if name == "stretch":
        specs += [
            PriorSpec("A", "uniform", 0.0, A_max),
            PriorSpec("tau", "loguniform", tau_min, tau_max),
            PriorSpec("s", "uniform", args.s_min, args.s_max),
        ]
        tau_indices = [2]
    else:
        m = re.fullmatch(r"exp(\d+)", name)
        if not m:
            raise ValueError(f"Unknown model {model_name!r}")
        k = int(m.group(1))
        for i in range(1, k + 1):
            specs.append(PriorSpec(f"A{i}", "uniform", 0.0, A_max))
        for i in range(1, k + 1):
            tau_indices.append(len(specs))
            specs.append(PriorSpec(f"tau{i}", "loguniform", tau_min, tau_max))

    estimate_sigma = args.sigma is None
    if estimate_sigma:
        if sigma_min <= 0 or sigma_max <= sigma_min:
            raise ValueError("Require 0 < sigma_min < sigma_max")
        specs.append(PriorSpec("sigma", "loguniform", sigma_min, sigma_max))
        fixed_sigma = None
    else:
        if args.sigma <= 0:
            raise ValueError("--sigma must be > 0")
        fixed_sigma = float(args.sigma)

    return ModelPrior(name, specs, tau_indices, estimate_sigma, fixed_sigma)


# =============================================================================
# Likelihood
# =============================================================================
def log_likelihood(
    x: np.ndarray,
    y: np.ndarray,
    prior: ModelPrior,
    u: np.ndarray,
) -> float:
    if not prior.is_valid_u(u):
        return -np.inf

    p = prior.transform(u)
    sigma = p.pop("sigma") if prior.estimate_sigma else prior.fixed_sigma
    assert sigma is not None

    yhat = model(x, prior.model_name, p)
    resid = y - yhat
    n = y.size
    return -0.5 * float(np.sum((resid / sigma) ** 2)) - n * math.log(sigma * math.sqrt(2.0 * math.pi))


# =============================================================================
# Replica exchange / parallel tempering
# =============================================================================
@dataclass
class PTResult:
    model_name: str
    betas: np.ndarray
    mean_logl: np.ndarray
    sem_logl: np.ndarray
    log_evidence: float
    log_evidence_se: float
    free_energy: float
    posterior_u: np.ndarray
    posterior_logl: np.ndarray
    best_u: np.ndarray
    best_logl: float
    local_acceptance: np.ndarray
    swap_acceptance: np.ndarray
    prior: ModelPrior


def make_beta_ladder(nreplica: int, power: float) -> np.ndarray:
    if nreplica < 3:
        raise ValueError("--nreplica must be >= 3")
    if power <= 0:
        raise ValueError("--beta-power must be > 0")
    z = np.linspace(0.0, 1.0, nreplica)
    return np.sin(0.5 * np.pi * z) ** power


def batch_sem(values: np.ndarray, max_batches: int = 20) -> float:
    values = np.asarray(values, dtype=float)
    n = len(values)
    if n < 4:
        return float("nan")
    nb = min(max_batches, max(2, n // 20))
    chunks = np.array_split(values, nb)
    means = np.array([np.mean(c) for c in chunks if len(c) > 0])
    if len(means) < 2:
        return float("nan")
    return float(np.std(means, ddof=1) / math.sqrt(len(means)))


def trapezoid_weights(x: np.ndarray) -> np.ndarray:
    """Weights w such that trapz(y,x) == sum_i w_i*y_i."""
    x = np.asarray(x, dtype=float)
    n = len(x)
    w = np.zeros(n, dtype=float)
    if n < 2:
        return w
    w[0] = 0.5 * (x[1] - x[0])
    w[-1] = 0.5 * (x[-1] - x[-2])
    if n > 2:
        w[1:-1] = 0.5 * (x[2:] - x[:-2])
    return w


def run_parallel_tempering(
    x: np.ndarray,
    y: np.ndarray,
    prior: ModelPrior,
    nreplica: int,
    steps: int,
    burn: int,
    thin: int,
    swap_interval: int,
    proposal_scale: float,
    beta_power: float,
    rng: np.random.Generator,
    adapt_interval: int = 500,
) -> PTResult:
    if burn >= steps:
        raise ValueError("--burn must be smaller than --steps")
    if thin < 1 or swap_interval < 1:
        raise ValueError("--thin and --swap-interval must be >= 1")

    betas = make_beta_ladder(nreplica, beta_power)
    ndim = prior.ndim

    states = np.vstack([prior.sample_u(rng) for _ in range(nreplica)])
    logls = np.array([log_likelihood(x, y, prior, u) for u in states])
    scales = np.full(nreplica, proposal_scale, dtype=float)

    accepted = np.zeros(nreplica, dtype=int)
    attempted = np.zeros(nreplica, dtype=int)
    accepted_window = np.zeros(nreplica, dtype=int)
    attempted_window = np.zeros(nreplica, dtype=int)

    swap_acc = np.zeros(nreplica - 1, dtype=int)
    swap_att = np.zeros(nreplica - 1, dtype=int)

    logl_samples: List[List[float]] = [[] for _ in range(nreplica)]
    posterior_u: List[np.ndarray] = []
    posterior_logl: List[float] = []

    best_u = states[-1].copy()
    best_logl = float(logls[-1])

    for step in range(steps):
        # Local random-walk Metropolis update in normalized prior coordinates.
        for r in range(nreplica):
            attempted[r] += 1
            attempted_window[r] += 1
            d = int(rng.integers(0, ndim))
            cand = states[r].copy()
            cand[d] += rng.normal(0.0, scales[r])

            if prior.is_valid_u(cand):
                new_logl = log_likelihood(x, y, prior, cand)
                log_alpha = betas[r] * (new_logl - logls[r])
                if math.log(rng.random()) < min(0.0, log_alpha):
                    states[r] = cand
                    logls[r] = new_logl
                    accepted[r] += 1
                    accepted_window[r] += 1

        # Neighbor swaps. Alternate parity improves mixing without double-using replicas.
        if (step + 1) % swap_interval == 0:
            parity = ((step + 1) // swap_interval) % 2
            for r in range(parity, nreplica - 1, 2):
                swap_att[r] += 1
                log_alpha = (betas[r + 1] - betas[r]) * (logls[r] - logls[r + 1])
                if math.log(rng.random()) < min(0.0, log_alpha):
                    states[[r, r + 1]] = states[[r + 1, r]]
                    logls[[r, r + 1]] = logls[[r + 1, r]]
                    swap_acc[r] += 1

        # Adapt local proposal scales only during burn-in.
        if step < burn and (step + 1) % adapt_interval == 0:
            with np.errstate(divide="ignore", invalid="ignore"):
                rates = np.divide(
                    accepted_window,
                    attempted_window,
                    out=np.zeros_like(scales),
                    where=attempted_window > 0,
                )
            # Gentle multiplicative adaptation toward ~0.30 acceptance.
            scales *= np.exp(0.8 * (rates - 0.30))
            scales = np.clip(scales, 0.003, 0.30)
            accepted_window[:] = 0
            attempted_window[:] = 0

        # Production samples used for TI and posterior summaries.
        if step >= burn and ((step - burn) % thin == 0):
            for r in range(nreplica):
                logl_samples[r].append(float(logls[r]))
            posterior_u.append(states[-1].copy())
            posterior_logl.append(float(logls[-1]))

            if logls[-1] > best_logl:
                best_logl = float(logls[-1])
                best_u = states[-1].copy()

    mean_logl = np.array([np.mean(v) for v in logl_samples], dtype=float)
    sem_logl = np.array([batch_sem(np.asarray(v)) for v in logl_samples], dtype=float)

    logZ = float(np.trapezoid(mean_logl, betas))
    weights = trapezoid_weights(betas)
    safe_sem = np.where(np.isfinite(sem_logl), sem_logl, 0.0)
    logZ_se = float(math.sqrt(np.sum((weights * safe_sem) ** 2)))

    local_acceptance = np.divide(
        accepted,
        attempted,
        out=np.zeros(nreplica, dtype=float),
        where=attempted > 0,
    )
    swap_acceptance = np.divide(
        swap_acc,
        swap_att,
        out=np.full(nreplica - 1, np.nan, dtype=float),
        where=swap_att > 0,
    )

    return PTResult(
        model_name=prior.model_name,
        betas=betas,
        mean_logl=mean_logl,
        sem_logl=sem_logl,
        log_evidence=logZ,
        log_evidence_se=logZ_se,
        free_energy=-logZ,
        posterior_u=np.asarray(posterior_u),
        posterior_logl=np.asarray(posterior_logl),
        best_u=best_u,
        best_logl=best_logl,
        local_acceptance=local_acceptance,
        swap_acceptance=swap_acceptance,
        prior=prior,
    )


# =============================================================================
# Posterior summaries and output
# =============================================================================
def posterior_parameter_table(result: PTResult) -> pd.DataFrame:
    rows = []
    if len(result.posterior_u) == 0:
        return pd.DataFrame()

    samples: Dict[str, List[float]] = {spec.name: [] for spec in result.prior.specs}
    for u in result.posterior_u:
        p = result.prior.transform(u)
        for k, v in p.items():
            samples[k].append(v)

    for name, vals in samples.items():
        a = np.asarray(vals)
        rows.append({
            "parameter": name,
            "median": float(np.median(a)),
            "q16": float(np.quantile(a, 0.16)),
            "q84": float(np.quantile(a, 0.84)),
            "mean": float(np.mean(a)),
            "std": float(np.std(a, ddof=1)) if len(a) > 1 else np.nan,
        })
    return pd.DataFrame(rows)


def representative_params(result: PTResult) -> Dict[str, float]:
    """Use posterior medians as the representative fit parameters."""
    table = posterior_parameter_table(result)
    return {row.parameter: float(row.median) for row in table.itertuples()}


def prediction_from_params(result: PTResult, x: np.ndarray, p: Dict[str, float]) -> np.ndarray:
    pp = dict(p)
    pp.pop("sigma", None)
    return model(x, result.model_name, pp)


def save_outputs(
    prefix: Path,
    x: np.ndarray,
    y: np.ndarray,
    results: List[PTResult],
    y_true: Optional[np.ndarray] = None,
) -> pd.DataFrame:
    logz = np.array([r.log_evidence for r in results])
    max_logz = float(np.max(logz))
    rel = np.exp(logz - max_logz)
    model_prob = rel / np.sum(rel)

    summary_rows = []
    for r, rp, prob in zip(results, rel, model_prob):
        summary_rows.append({
            "model": r.model_name,
            "log_evidence": r.log_evidence,
            "log_evidence_SE": r.log_evidence_se,
            "free_energy_F=-logZ": r.free_energy,
            "delta_F_from_best": r.free_energy - min(rr.free_energy for rr in results),
            "relative_evidence": rp,
            "posterior_model_probability_equal_priors": prob,
            "mean_local_acceptance": float(np.nanmean(r.local_acceptance)),
            "mean_swap_acceptance": float(np.nanmean(r.swap_acceptance)),
            "best_log_likelihood": r.best_logl,
        })
    summary = pd.DataFrame(summary_rows).sort_values("free_energy_F=-logZ").reset_index(drop=True)
    summary.to_csv(prefix.with_name(prefix.name + "_summary.csv"), index=False)

    # Collect fit curves on a dense grid.
    xgrid = np.linspace(float(np.min(x)), float(np.max(x)), 500)
    fit_df = pd.DataFrame({"x": xgrid})
    for r in results:
        p = representative_params(r)
        fit_df[f"{r.model_name}_posterior_median"] = prediction_from_params(r, xgrid, p)

    # Excel workbook with all important results.
    xlsx_path = prefix.with_name(prefix.name + "_results.xlsx")
    with pd.ExcelWriter(xlsx_path, engine="openpyxl") as writer:
        summary.to_excel(writer, sheet_name="Model_summary", index=False)
        data_df = pd.DataFrame({"x": x, "y": y})
        if y_true is not None:
            data_df["y_true"] = y_true
        data_df.to_excel(writer, sheet_name="Data", index=False)
        fit_df.to_excel(writer, sheet_name="Fit_curves", index=False)

        for r in results:
            posterior_parameter_table(r).to_excel(
                writer, sheet_name=f"Params_{r.model_name}"[:31], index=False
            )
            pd.DataFrame({
                "beta": r.betas,
                "mean_log_likelihood": r.mean_logl,
                "SEM_batch": r.sem_logl,
                "local_acceptance": r.local_acceptance,
            }).to_excel(writer, sheet_name=f"TI_{r.model_name}"[:31], index=False)

    # Plot data + representative fits.
    plt.figure(figsize=(8, 5.5))
    plt.scatter(x, y, s=18, label="data")
    if y_true is not None:
        order = np.argsort(x)
        plt.plot(x[order], y_true[order], linewidth=2, label="true model")
    for r in results:
        p = representative_params(r)
        plt.plot(xgrid, prediction_from_params(r, xgrid, p), label=r.model_name)
    plt.xlabel("x")
    plt.ylabel("y")
    plt.title("Decay-model comparison")
    plt.legend()
    plt.tight_layout()
    plt.savefig(prefix.with_name(prefix.name + "_fit.png"), dpi=180)
    plt.close()

    # TI curves: area under each curve is log evidence.
    plt.figure(figsize=(8, 5.5))
    for r in results:
        plt.plot(r.betas, r.mean_logl, marker="o", markersize=3,
                 label=f"{r.model_name}: logZ={r.log_evidence:.2f}")
    plt.xlabel(r"$\beta$")
    plt.ylabel(r"$\langle \log L \rangle_\beta$")
    plt.title("Thermodynamic integration")
    plt.legend()
    plt.tight_layout()
    plt.savefig(prefix.with_name(prefix.name + "_ti.png"), dpi=180)
    plt.close()

    return summary


# =============================================================================
# Command line
# =============================================================================
def parse_sheet(value: str) -> str | int:
    try:
        return int(value)
    except ValueError:
        return value


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Bayesian decay-model selection using replica exchange + thermodynamic integration.",
    )

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

    # Synthetic data options
    p.add_argument("--true-model", default="stretch", help="Model used to generate synthetic data.")
    p.add_argument("--params", default="y0=0.05,A=1.0,tau=2.5,s=0.65",
                   help="Comma-separated true parameters, e.g. 'y0=0,A=1,tau=2,s=0.7'.")
    p.add_argument("--xmin", type=float, default=0.0)
    p.add_argument("--xmax", type=float, default=10.0)
    p.add_argument("--npoints", type=int, default=120)
    p.add_argument("--xscale", choices=["linear", "log"], default="linear")
    p.add_argument("--noise", type=float, default=0.02,
                   help="Gaussian standard deviation for synthetic data.")

    # Candidate models
    p.add_argument("--models", default=DEFAULT_MODELS,
                   help="Comma-separated candidate models, e.g. exp1,exp2,exp3,stretch")

    # Noise model
    p.add_argument("--sigma", type=float, default=None,
                   help="Fix Gaussian sigma. For Excel data, omission means infer sigma; for synthetic data, omission uses --noise unless --infer-sigma is set.")
    p.add_argument("--infer-sigma", action="store_true",
                   help="Infer sigma also for synthetic data instead of fixing it to --noise.")

    # Priors; None means automatic.
    p.add_argument("--y0-min", type=float, default=None)
    p.add_argument("--y0-max", type=float, default=None)
    p.add_argument("--A-max", type=float, default=None)
    p.add_argument("--tau-min", type=float, default=None)
    p.add_argument("--tau-max", type=float, default=None)
    p.add_argument("--s-min", type=float, default=0.1)
    p.add_argument("--s-max", type=float, default=1.0)
    p.add_argument("--sigma-min", type=float, default=None)
    p.add_argument("--sigma-max", type=float, default=None)

    # MCMC / TI
    p.add_argument("--nreplica", type=int, default=DEFAULT_N_REPLICAS)
    p.add_argument("--steps", type=int, default=DEFAULT_STEPS)
    p.add_argument("--burn", type=int, default=DEFAULT_BURN)
    p.add_argument("--thin", type=int, default=DEFAULT_THIN)
    p.add_argument("--swap-interval", type=int, default=DEFAULT_SWAP_INTERVAL)
    p.add_argument("--proposal-scale", type=float, default=DEFAULT_PROPOSAL_SCALE)
    p.add_argument("--beta-power", type=float, default=DEFAULT_BETA_POWER)
    p.add_argument("--seed", type=int, default=DEFAULT_SEED)

    # Output
    p.add_argument("-o", "--output", default="bayes_decay",
                   help="Output prefix (without extension).")
    return p


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()
    rng = np.random.default_rng(args.seed)

    if args.source == "excel":
        if not args.input:
            parser.error("--input is required when --source excel")
        x, y = read_excel_xy(args.input, parse_sheet(args.sheet), args.xcol, args.ycol)
        y_true = None
    else:
        true_params = parse_param_string(args.params)
        x, y, y_true = generate_data(
            args.true_model,
            true_params,
            args.xmin,
            args.xmax,
            args.npoints,
            args.noise,
            args.xscale,
            rng,
        )
        if args.sigma is None and not args.infer_sigma:
            args.sigma = args.noise
            print(f"Synthetic data: fixing likelihood sigma to known noise = {args.noise:g}")

    if np.any(x < 0):
        raise ValueError(
            "This decay implementation assumes x >= 0. Shift the time origin before fitting."
        )

    candidate_models = [m.strip().lower() for m in args.models.split(",") if m.strip()]
    if not candidate_models:
        raise ValueError("No candidate models specified.")

    print("\n=== Data ===")
    print(f"source : {args.source}")
    print(f"points : {len(x)}")
    print(f"x range: {x.min():.6g} .. {x.max():.6g}")
    print(f"y range: {y.min():.6g} .. {y.max():.6g}")
    print(f"models : {', '.join(candidate_models)}")

    results: List[PTResult] = []
    for imodel, name in enumerate(candidate_models, start=1):
        print(f"\n=== [{imodel}/{len(candidate_models)}] {name} ===")
        prior = build_prior(name, x, y, args)
        print("Priors:")
        for spec in prior.specs:
            print(f"  {spec.name:8s} {spec.kind:10s} [{spec.lo:.6g}, {spec.hi:.6g}]")
        if not prior.estimate_sigma:
            print(f"  sigma    fixed      {prior.fixed_sigma:.6g}")

        # Use an independent child RNG stream for each model while preserving reproducibility.
        child_rng = np.random.default_rng(args.seed + 10007 * imodel)
        res = run_parallel_tempering(
            x=x,
            y=y,
            prior=prior,
            nreplica=args.nreplica,
            steps=args.steps,
            burn=args.burn,
            thin=args.thin,
            swap_interval=args.swap_interval,
            proposal_scale=args.proposal_scale,
            beta_power=args.beta_power,
            rng=child_rng,
        )
        results.append(res)
        print(f"log evidence = {res.log_evidence:.6f} +/- {res.log_evidence_se:.4f} (approx.)")
        print(f"free energy = {-res.log_evidence:.6f}")
        print(f"local acceptance mean = {np.nanmean(res.local_acceptance):.3f}")
        print(f"swap acceptance mean  = {np.nanmean(res.swap_acceptance):.3f}")

    prefix = Path(args.output)
    summary = save_outputs(prefix, x, y, results, y_true)

    print("\n=== Model comparison ===")
    with pd.option_context("display.max_columns", None, "display.width", 180):
        print(summary.to_string(index=False, float_format=lambda v: f"{v:.6g}"))

    print("\nFiles written:")
    print(f"  {prefix.with_name(prefix.name + '_results.xlsx')}")
    print(f"  {prefix.with_name(prefix.name + '_summary.csv')}")
    print(f"  {prefix.with_name(prefix.name + '_fit.png')}")
    print(f"  {prefix.with_name(prefix.name + '_ti.png')}")


if __name__ == "__main__":
    main()