#!/usr/bin/env python
"""Fit a temperature to each spectrum in a sequence and plot T versus time.

For a cooling or heating object: grab a burst of frames, then fit each and watch
the temperature track. Reuses the Planck fit from fit_temperature.py.

Inputs are spectrum CSVs in time order. With --response (and --dark plus
--integration) each is response-corrected first, exactly as grab_spectrum.py
--response would; otherwise they are assumed already corrected. Time comes from
the file modification times (first file = 0), or pass --interval for even spacing.

Note the printed +/- is the formal fit error and understates the truth; the
real scatter is frame-to-frame noise and the gray-body and reference systematics.

Usage: python fit_sequence.py f1.csv f2.csv ... [--response R.csv --dark D.csv
                              --integration US] [--band LO HI] [--interval S]
                              [--out T.csv] [--png fig.png]
"""
import argparse
import os

import numpy as np

from response import add_calibration_args, apply_response, dark_at, load_csv
from fit_temperature import fit_planck

p = argparse.ArgumentParser(description=__doc__,
                            formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("spectra", nargs="+", help="spectrum CSVs, in time order")
p.add_argument("--integration", type=int, metavar="US",
               help="integration time (us) the spectra were captured at; required with --dark")
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("--interval", type=float, default=None,
               help="seconds between frames (default: use file modification times)")
p.add_argument("--out", default=None, help="write a time,temperature_K,dT_K CSV")
p.add_argument("--png", default="temperature_sequence.png")
add_calibration_args(p)
args = p.parse_args()

cal = load_csv(args.response) if args.response else None
dark = None
if args.dark:
    if args.integration is None:
        p.error("--integration US is required with --dark")
    dark = dark_at(args.dark, args.integration)

if args.interval is not None:
    times = np.arange(len(args.spectra)) * args.interval
else:
    t0 = os.path.getmtime(args.spectra[0])
    times = np.array([os.path.getmtime(f) - t0 for f in args.spectra])

rows = []
for t, f in zip(times, args.spectra):
    wl, y = load_csv(f)
    if cal is not None:
        y = apply_response(wl, y, cal, dark)
    result = fit_planck(wl, y, tuple(args.band))
    if result is None:
        print(f"{f}: no fit (too few usable points in band)")
        continue
    temperature, dT, _ = result
    rows.append((t, temperature, dT))
    print(f"{f}: t = {t:7.2f} s   T = {temperature:.2f} +/- {dT:.2f} K")

if not rows:
    p.error("no spectra could be fit; check --band and inputs")
rows = np.array(rows)

if args.out:
    np.savetxt(args.out, rows, delimiter=",",
               header="time_s,temperature_K,dT_K", comments="")
    print(f"wrote {args.out}")

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(9, 5))
ax.errorbar(rows[:, 0], rows[:, 1], yerr=rows[:, 2], fmt="o-", capsize=3)
ax.set_xlabel("time (s)")
ax.set_ylabel("temperature (K)")
ax.set_title(f"Temperature sequence, fit {args.band[0]:.0f}-{args.band[1]:.0f} nm")
ax.grid(alpha=0.3)
fig.tight_layout()
fig.savefig(args.png, dpi=110)
print(f"wrote {args.png}")