#!/usr/bin/env python

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Sequence

import numpy as np
import matplotlib.pyplot as plt

import chisurf
from chisurf.experiments.pda import PdaReader
from chisurf.fitting.fit import Fit
from chisurf.models.pda.pdagauss import PdaGaussianDistanceModel


def _parse_channel_list(text: str) -> list[int]:
    parts = [p.strip() for p in (text or "").split(",")]
    out: list[int] = []
    for p in parts:
        if not p:
            continue
        try:
            out.append(int(p))
        except ValueError:
            raise SystemExit(f"Invalid channel entry '{p}' in '{text}'.")
    if not out:
        raise SystemExit(f"Empty channel list in '{text}'.")
    return out


def build_pda_data(
    input_path: Path,
    *,
    file_type: str,
    pattern: str | None,
    channels_1: Sequence[int],
    channels_2: Sequence[int],
    max_photons: int,
    min_photons: int,
    tw_ms: float,
) -> chisurf.data.DataCurve:
    if input_path.is_dir():
        pat = pattern or "*.h5*"
        files = sorted(input_path.glob(pat))
    else:
        files = [input_path]

    if not files:
        raise SystemExit(f"No input files found for '{input_path}'.")

    filenames = [f.as_posix() for f in files]

    reader = PdaReader(
        channels=(list(channels_1), list(channels_2)),
        micro_time_ranges=[(0, 4096)],
        reading_routine=file_type,
        maximum_number_of_photons=int(max_photons),
        minimum_number_of_photons=int(min_photons),
        minimum_time_window_length=float(tw_ms) * 1.0e-3,
    )

    data_group = reader.read(filenames)
    if not len(data_group):
        raise SystemExit("PdaReader.read returned no data.")

    data = data_group[0]
    return data


def build_fit(data: chisurf.data.DataCurve, args: argparse.Namespace) -> Fit:
    fit = Fit(
        model_class=PdaGaussianDistanceModel,
        data=data,
        xmin=0,
        xmax=len(data.y) - 1,
    )
    model: PdaGaussianDistanceModel = fit.model

    n_gauss = max(1, int(args.n_gauss))
    r_min = float(args.r_min)
    r_max = float(args.r_max)
    if r_max <= r_min:
        r_max = r_min + 10.0
    means = np.linspace(r_min, r_max, n_gauss)
    sigma = float(args.sigma)
    amp = 1.0 / float(n_gauss)
    for m in means:
        model.distances.append(mean=float(m), sigma=sigma, amplitude=amp)

    if args.alpha is not None:
        model.nuisance.alpha = float(args.alpha)
    if args.bgG is not None:
        model.nuisance.BG = float(args.bgG)
    if args.bgR is not None:
        model.nuisance.BR = float(args.bgR)
    if args.nph_min is not None:
        model.nuisance.nPh_min = float(args.nph_min)
    if args.nph_max is not None:
        model.nuisance.nPh_max = float(args.nph_max)

    if args.xdonly is not None:
        model.fret_parameters.xDOnly = float(args.xdonly)
    if args.R0 is not None:
        model.fret_parameters.forster_radius = float(args.R0)

    return fit


def export_param_template(model: PdaGaussianDistanceModel, path: Path) -> None:
    state = model.get_state()
    path = path.resolve()
    with path.open("w", encoding="utf-8") as f:
        json.dump(state, f, indent=2, sort_keys=True)
    print(f"Wrote parameter template to '{path}'.")


def load_param_state(model: PdaGaussianDistanceModel, path: Path) -> None:
    path = path.resolve()
    with path.open("r", encoding="utf-8") as f:
        state = json.load(f)
    model.set_state(state)
    print(f"Loaded parameter state from '{path}'.")


def run_fit(fit: Fit) -> None:
    print(f"Initial reduced chi2: {fit.chi2r:.4f}")
    fit.run()
    print(f"Final reduced chi2:   {fit.chi2r:.4f}")


def summarize_parameters(model: PdaGaussianDistanceModel) -> None:
    print(model)


def plot_result(
    fit: Fit,
    *,
    output_prefix: Path | None = None,
    show: bool = False,
) -> None:
    from chisurf.models.pda.common import pda_1d_residuals_from_s1s2

    model: PdaGaussianDistanceModel = fit.model
    data = fit.data
    pda_meta = getattr(data, "pda", {}) or {}
    s1s2_data = np.asarray(pda_meta.get("s1s2"), dtype=float)
    s1s2_model = np.asarray(model.pda.s1s2, dtype=float)

    shp = pda_meta.get("shape")
    if isinstance(shp, (tuple, list)) and len(shp) == 2:
        ny, nx = int(shp[0]), int(shp[1])
        s1s2_model = s1s2_model[:ny, :nx]

    wres = pda_1d_residuals_from_s1s2(
        fit=fit,
        pda_obj=model.pda,
        nuisance=getattr(model, "nuisance", None),
        kw_hist=model.kw_hist,
    )

    kw_hist = dict(model.kw_hist)
    s1s2_flat_data = s1s2_data.ravel(order="C")
    x_data, y_data = model.pda.get_1dhistogram(s1s2=s1s2_flat_data, **kw_hist)

    s1s2_flat_model = s1s2_model.ravel(order="C")
    x_model, y_model = model.pda.get_1dhistogram(s1s2=s1s2_flat_model, **kw_hist)

    fig, ax = plt.subplots(2, 2, figsize=(8, 6))

    im0 = ax[0, 0].imshow(s1s2_data, origin="lower")
    fig.colorbar(im0, ax=ax[0, 0])
    ax[0, 0].set_title("Experimental S1S2")

    im1 = ax[1, 0].imshow(s1s2_model, origin="lower")
    fig.colorbar(im1, ax=ax[1, 0])
    ax[1, 0].set_title("Model S1S2")

    if kw_hist.get("log_x"):
        ax[0, 1].semilogx(x_data, wres, "k.")
        ax[1, 1].semilogx(x_data, y_data, "k.", label="data")
        ax[1, 1].semilogx(x_model, y_model, "r-", label="model")
    else:
        ax[0, 1].plot(x_data, wres, "k.")
        ax[1, 1].plot(x_data, y_data, "k.", label="data")
        ax[1, 1].plot(x_model, y_model, "r-", label="model")

    ax[0, 1].set_title("1D weighted residuals")
    ax[1, 1].set_title("1D histogram")
    ax[1, 1].legend()

    fig.tight_layout()

    if output_prefix is not None:
        out = output_prefix.with_suffix(".png").resolve()
        fig.savefig(out, dpi=200)
        print(f"Saved figure to '{out}'.")

    if show:
        plt.show()
    else:
        plt.close(fig)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=(
            "PDA FRET Gaussian-distance fit using tttrlib and chisurf. "
            "Builds an S1S2 histogram from TTTR/Photon-HDF5 data, sets up a "
            "two-Gaussian + donor-only PDA model with nuisance parameters, "
            "and optionally runs a bounded least-squares fit."
        )
    )

    p.add_argument(
        "input",
        metavar="INPUT",
        help=(
            "Photon-HDF5/TTTR file or directory containing multiple files. "
            "For a directory, use --pattern to filter files."
        ),
    )

    p.add_argument(
        "--file-type",
        default="PHOTON-HDF5",
        help=(
            "TTTR/Photon file type string understood by tttrlib.TTTR "
            "(e.g. PTU, HT3, SPC-130, PHOTON-HDF5). Default: PHOTON-HDF5."
        ),
    )
    p.add_argument(
        "--pattern",
        default=None,
        help=(
            "Glob pattern used when INPUT is a directory (default: *.h5*)."
        ),
    )

    p.add_argument(
        "--channels1",
        default="0",
        help=(
            "Comma-separated list of routing channels for PDA channel 1 "
            "(donor/green). Default: 0."
        ),
    )
    p.add_argument(
        "--channels2",
        default="1",
        help=(
            "Comma-separated list of routing channels for PDA channel 2 "
            "(acceptor/red). Default: 1."
        ),
    )

    p.add_argument(
        "--tw-ms",
        type=float,
        default=1.0,
        help="Time-window length in milliseconds for PDA bursts (default 1.0).",
    )
    p.add_argument(
        "--min-photons",
        type=int,
        default=30,
        help="Minimum photons per TW for inclusion in PDA (default 30).",
    )
    p.add_argument(
        "--max-photons",
        type=int,
        default=150,
        help="Maximum photons used for S1S2 support (default 150).",
    )

    p.add_argument(
        "--n-gauss",
        type=int,
        default=2,
        help="Number of Gaussian distance components (default 2).",
    )
    p.add_argument(
        "--r-min",
        type=float,
        default=30.0,
        help="Lower bound of initial Gaussian mean range in Å (default 30).",
    )
    p.add_argument(
        "--r-max",
        type=float,
        default=70.0,
        help="Upper bound of initial Gaussian mean range in Å (default 70).",
    )
    p.add_argument(
        "--sigma",
        type=float,
        default=5.0,
        help="Initial sigma (width) for all Gaussians in Å (default 5).",
    )

    p.add_argument("--alpha", type=float, default=None, help="Initial spectral crosstalk alpha (optional).")
    p.add_argument("--bgG", type=float, default=None, help="Initial background in green channel BG (optional).")
    p.add_argument("--bgR", type=float, default=None, help="Initial background in red channel BR (optional).")
    p.add_argument("--nph-min", type=float, default=None, help="Override nuisance nPh_min for photon-number gating.")
    p.add_argument("--nph-max", type=float, default=None, help="Override nuisance nPh_max for photon-number gating.")

    p.add_argument("--xdonly", type=float, default=None, help="Initial donor-only fraction xDOnly (0..1, optional).")
    p.add_argument("--R0", type=float, default=None, help="Override Förster radius R0 in Å (optional).")

    p.add_argument(
        "--export-param-template",
        metavar="JSON",
        default=None,
        help=(
            "Write a JSON parameter template for the constructed model and exit "
            "(unless --fit is also given). The file contains values, bounds, "
            "and fixed flags for all fitting parameters."
        ),
    )
    p.add_argument(
        "--param-state",
        metavar="JSON",
        default=None,
        help=(
            "Load model/parameter state from a JSON file created earlier via "
            "--export-param-template. This allows full control over initial "
            "values, fixed flags, and bounds."
        ),
    )

    p.add_argument(
        "--no-fit",
        dest="do_fit",
        action="store_false",
        help="Build model but do not run the least-squares fit.",
    )
    p.set_defaults(do_fit=True)

    p.add_argument(
        "--plot",
        action="store_true",
        help="Show a matplotlib summary figure with S1S2 and 1D histograms.",
    )
    p.add_argument(
        "--output-prefix",
        metavar="STEM",
        default=None,
        help="If given, save a PNG figure with this filename stem.",
    )

    return p.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)

    input_path = Path(args.input).expanduser()

    channels_1 = _parse_channel_list(args.channels1)
    channels_2 = _parse_channel_list(args.channels2)

    data = build_pda_data(
        input_path,
        file_type=args.file_type,
        pattern=args.pattern,
        channels_1=channels_1,
        channels_2=channels_2,
        max_photons=args.max_photons,
        min_photons=args.min_photons,
        tw_ms=args.tw_ms,
    )

    fit = build_fit(data, args)
    model: PdaGaussianDistanceModel = fit.model

    if args.export_param_template:
        export_param_template(model, Path(args.export_param_template))
        if not args.do_fit:
            return

    if args.param_state:
        load_param_state(model, Path(args.param_state))

    fit.update()

    if args.do_fit:
        run_fit(fit)

    summarize_parameters(model)

    output_prefix: Path | None
    if args.output_prefix is not None:
        output_prefix = Path(args.output_prefix)
    else:
        output_prefix = None

    if args.plot or output_prefix is not None:
        plot_result(fit, output_prefix=output_prefix, show=args.plot)


if __name__ == "__main__":  # pragma: no cover
    main()
