#!/usr/bin/env python
"""Characterise the detector dark as a per-pixel line: dark(t) = offset + rate * t_us.

Block ALL light (sources off, aperture capped), then this captures dark frames
at several integration times and least-squares fits, for each pixel, a fixed
electronic offset plus a dark-current rate. Writes dark_model.csv
(wavelength_nm, offset_counts, rate_counts_per_us), which reconstructs the dark
at ANY integration time, so grab / live / calibrate can subtract the right dark
without a fresh capture per measurement.

Usage: python capture_dark.py [--times US [US ...]] [--frames N] [--out dark_model.csv]
"""
import argparse

import numpy as np
from seabreeze.spectrometers import Spectrometer

p = argparse.ArgumentParser(description=__doc__,
                            formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--times", type=int, nargs="+",
               default=[50000, 500000, 1000000, 1500000],
               help="integration times in us to sample (default: 0.05 .. 1.5 s)")
p.add_argument("--frames", type=int, default=10, help="frames to average per time")
p.add_argument("--out", default="dark_model.csv")
args = p.parse_args()

spec = Spectrometer.from_first_available()
wl = spec.wavelengths()
times = np.array(sorted(args.times), dtype=float)

stack = np.empty((times.size, wl.size))
for i, us in enumerate(times):
    spec.integration_time_micros(int(us))
    spec.intensities()  # flush one frame after changing integration time
    acc = np.zeros_like(wl)
    for _ in range(args.frames):
        acc += spec.intensities(correct_nonlinearity=True)
    stack[i] = acc / args.frames
    print(f"  {us / 1e6:.3f}s: median {np.median(stack[i]):.0f} counts")
spec.close()

# Per-pixel linear fit: counts = offset + rate * t_us.
rate, offset = np.polyfit(times, stack, 1)
resid = stack - (offset + np.outer(times, rate))
print(f"fit residual rms: {np.sqrt(np.mean(resid ** 2)):.1f} counts")
print(f"offset median {np.median(offset):.0f} counts, "
      f"rate median {np.median(rate) * 1e6:.1f} counts/s")

np.savetxt(args.out, np.column_stack([wl, offset, rate]), delimiter=",",
           header="wavelength_nm,offset_counts,rate_counts_per_us", comments="")
print(f"wrote {args.out}")