#!/usr/bin/env python
"""Fit a temperature to a response-corrected emission spectrum.

Takes a spectrum written by `grab_spectrum.py ... --response` (relative
irradiance vs wavelength) and fits a Planck curve over a chosen band, returning
the temperature. The fit is relative: an amplitude absorbs the arbitrary scale
and any flat greybody factor, so only the spectral SHAPE sets T.

Restrict the band with --band LO HI. The default 420 to 800 nm avoids the weak
blue response and the long-wavelength end this spectrometer is not trusted at.
The temperature comes entirely from the slope across the band, so a too-narrow
or too-noisy band gives a poorly constrained T.

Usage: python fit_temperature.py spectrum.csv [--band LO HI] [--png fit.png]
"""
import argparse

import numpy as np
from scipy.optimize import curve_fit

from response import load_csv, planck


def fit_planck(wl, y, band=(420.0, 800.0)):
    """Fit a Planck curve to a response-corrected spectrum over `band`.

    Returns (temperature_K, dT_K, predict), or None if there are too few usable
    points. `predict(wl_nm)` evaluates the fitted curve at any wavelengths, using
    the same normalisation as the fit, so it can be drawn beyond the fit band.
    The model peak is normalised to a free amplitude, so only the spectral slope
    sets T and the fit is well conditioned regardless of the Planck absolute scale.
    """
    lo, hi = band
    keep = np.isfinite(y) & (y > 0) & (wl >= lo) & (wl <= hi)
    if keep.sum() < 10:
        return None
    x, yy = wl[keep], y[keep]

    def model(wl_nm, amplitude, temperature):
        b = planck(wl_nm, temperature)
        return amplitude * b / b.max()

    try:
        popt, pcov = curve_fit(model, x, yy, p0=[yy.max(), 2500.0],
                               bounds=([0, 300], [np.inf, 6000]), maxfev=20000)
    except (RuntimeError, ValueError):
        return None
    amplitude, temperature = popt
    norm = planck(x, temperature).max()  # the scale used over the fit band
    predict = lambda wl_nm: amplitude * planck(wl_nm, temperature) / norm
    return temperature, float(np.sqrt(pcov[1, 1])), predict


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("spectrum", help="response-corrected spectrum CSV")
    p.add_argument("--band", nargs=2, type=float, default=[420.0, 800.0],
                   metavar=("LO", "HI"), help="fit band in nm (default: 420 800)")
    p.add_argument("--png", default=None, help="optional plot of data and fit")
    args = p.parse_args()

    wl, y = load_csv(args.spectrum)
    lo, hi = args.band
    result = fit_planck(wl, y, (lo, hi))
    if result is None:
        p.error(f"too few usable points in {lo:.0f}-{hi:.0f} nm; widen --band "
                "or check the spectrum is response-corrected")
    temperature, dT, predict = result
    print(f"fit band {lo:.0f}-{hi:.0f} nm")
    print(f"T = {temperature:.2f} +/- {dT:.2f} K  (formal fit error; systematics dominate)")

    if args.png:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(9, 5))
        band = (wl >= lo) & (wl <= hi)
        ax.plot(wl, y, lw=0.6, color="0.7", label="corrected spectrum")
        ax.plot(wl, predict(wl), lw=1.2, color="C1", label="Planck fit (extrapolated)")
        ax.plot(wl[band], predict(wl[band]), lw=2.0, color="C3",
                label=f"used {lo:.0f}-{hi:.0f} nm, T = {temperature:.2f} K")
        ax.set_xlabel("wavelength (nm)")
        ax.set_ylabel("relative irradiance")
        ax.grid(True, alpha=0.3)
        ax.legend(fontsize=8)
        fig.tight_layout()
        fig.savefig(args.png, dpi=110)
        print(f"wrote {args.png}")


if __name__ == "__main__":
    main()