"""The archive's real availability process, as data the sampler can draw from.

`profile_archive.py` classifies every (satellite, band, slot) of the two-satellite
era as GOOD (present and judged good), BAD (present but judged bad) or MISSING (no
usable file), and stores the resulting state sequences here.  This file *is* the
statistical model of the archive: sampling a window of it reproduces the real
run-length distributions, the cross-satellite and cross-band outage correlation, and
the diurnal/eclipse-season structure with no modelling assumptions -- an empirical
joint process rather than a parametric fit.

Per band, because real damage is: of the slots with an unusable band, about four in
five keep one to five good bands (131A alone is the commonest pattern), and the
renderer fills exactly those bands.  A slot-level model that called any partial slot
missing trained the model on outages several times larger than the archive's.

The training sampler overlays these windows onto its stacks (trace bootstrap), so
the model trains against outages shaped like the ones the archive actually
contains: 11% of the dual era has both spacecraft dark simultaneously, in
contiguous runs of up to 26 days -- a regime the earlier hand-tuned injection
constants under-represented, which showed up as a model that only ever needed to
pick the right frame from the right satellite.
"""

from dataclasses import dataclass

import numpy as np

from . import paths

#: States.  Order matters: profile reports index histograms by these values.
GOOD, BAD, MISSING = 0, 1, 2
STATE_NAMES = ("good", "bad", "missing")
#: Bands, in `states` axis-1 order.
BANDS = tuple(paths.WAVELENGTHS)
#: The longest simultaneous two-satellite outage the pipeline will attempt to fill,
#: in slots: 360 slots = 24 hours at the 4-minute cadence.  The joint-outage
#: analogue of ffmpeg_video's per-satellite ``max_frame_interp``: past this, gaps
#: are left as gaps rather than filled.  With a good counterpart frame a fill is
#: sound at any gap length, so no such limit applies to one-sided outages; but when
#: *neither* spacecraft has data, quality is only validated out to ~300 slots (the
#: stratified bench), a day of rotation moves every feature >13 degrees, and the
#: eclipse-season outages that dominate the real archive resolve within one diurnal
#: cycle -- anything longer is a station anomaly whose "fill" would be fabrication.
#: Applied per band, as the renderer applies it.  The training sampler conditions
#: on the same bound (a redrawn trace window), so the model never spends capacity
#: on configurations production refuses to fill.
MAX_DUAL_GAP = 360


@dataclass
class Availability:
    """Per-satellite, per-band state sequences over one contiguous era."""

    #: (satellites, bands, slots) uint8 of GOOD/BAD/MISSING, bands in BANDS order.
    states: np.ndarray
    #: Unix time of slot 0.
    t0: int
    #: Slot spacing in seconds.
    cadence: int
    #: Satellite ids, in `states` row order.
    satellites: tuple

    def __post_init__(self):
        self.states = np.asarray(self.states, dtype=np.uint8)
        if self.states.ndim != 3 or self.states.shape[1] != len(BANDS):
            raise ValueError(f"states must be (satellites, {len(BANDS)} bands, slots), "
                             f"not {self.states.shape}")
        self.satellites = tuple(int(s) for s in self.satellites)
        self._row = {satellite: i for i, satellite in enumerate(self.satellites)}
        self._damaged = {}

    @property
    def slots(self):
        return self.states.shape[2]

    def state(self, satellite, index):
        """The six band states of one slot; all MISSING outside the era or for an
        unknown satellite.

        Unknown satellites read as MISSING rather than GOOD: a trace consulted about
        a spacecraft it never observed has no evidence to offer, and 'assume fine'
        is exactly the optimism this model exists to remove.
        """
        row = self._row.get(int(satellite))
        if row is None or not 0 <= index < self.slots:
            return np.full(len(BANDS), MISSING, dtype=np.uint8)
        return self.states[row, :, index]

    def usable(self, satellite, index):
        """Which of one slot's six bands are GOOD."""
        return self.state(satellite, index) == GOOD

    def pristine(self):
        """(slots,) bool: every band of every satellite GOOD -- the corpus's segments."""
        return (self.states == GOOD).all(axis=(0, 1))

    def damaged_slots(self, satellite):
        """Indices where `satellite` has at least one band not GOOD."""
        if satellite not in self._damaged:
            row = self._row.get(int(satellite))
            self._damaged[satellite] = (
                np.array([], dtype=np.int64) if row is None
                else np.flatnonzero((self.states[row] != GOOD).any(axis=0)))
        return self._damaged[satellite]

    def sample_window(self, rng, satellite=None):
        """A random position in the era, for overlaying onto one training stack.

        With `satellite`, the position is one where that satellite has at least one
        band unusable -- a slot the renderer would fill -- so the target's missing
        bands, and everything around them, are a real fill case.
        """
        if satellite is None:
            return TraceWindow(self, int(rng.integers(0, self.slots)))
        damaged = self.damaged_slots(satellite)
        if damaged.size == 0:
            raise ValueError(f"satellite {satellite} has no damaged slot in the trace")
        return TraceWindow(self, int(damaged[rng.integers(0, damaged.size)]))

    def damage_patterns(self):
        """The empirical distribution of which bands are lost together: (masks (K,
        bands) bool, probabilities (K,)) over every satellite-slot with at least
        one band not GOOD."""
        lost = (self.states != GOOD).transpose(0, 2, 1).reshape(-1, len(BANDS))
        lost = lost[lost.any(axis=1)]
        if lost.size == 0:
            return np.ones((1, len(BANDS)), dtype=bool), np.ones(1)
        codes = lost.astype(np.int64) @ (1 << np.arange(len(BANDS)))
        values, counts = np.unique(codes, return_counts=True)
        masks = ((values[:, None] >> np.arange(len(BANDS))) & 1).astype(bool)
        return masks, counts / counts.sum()

    def fractions(self):
        """{satellite: {state name: fraction of (band, slot)s, "pristine": fraction of
        slots with all six GOOD}} -- the headline rates."""
        out = {}
        for satellite in self.satellites:
            row = self.states[self._row[satellite]]
            rates = {name: float((row == value).mean())
                     for value, name in enumerate(STATE_NAMES)}
            rates["pristine"] = float((row == GOOD).all(axis=0).mean())
            out[satellite] = rates
        return out

    def save(self, path):
        np.savez_compressed(
            path, states=self.states, t0=np.int64(self.t0),
            cadence=np.int64(self.cadence),
            satellites=np.asarray(self.satellites, dtype=np.int64),
            bands=np.asarray(BANDS, dtype=np.int64),
        )

    @classmethod
    def load(cls, path):
        with np.load(path) as data:
            if "bands" not in data or tuple(int(b) for b in data["bands"]) != BANDS:
                raise ValueError(f"{path} is not a per-band availability model of "
                                 f"bands {BANDS}; regenerate it with profile_archive.py")
            return cls(states=data["states"], t0=int(data["t0"]),
                       cadence=int(data["cadence"]),
                       satellites=tuple(int(s) for s in data["satellites"]))


@dataclass
class TraceWindow:
    """One sampled position; states are read relative to it in slot offsets."""

    availability: Availability
    centre: int

    def state(self, satellite, offset_slots):
        return self.availability.state(satellite, self.centre + int(offset_slots))

    def usable(self, satellite, offset_slots):
        """Which bands of the slot `offset_slots` from the centre are GOOD."""
        return self.availability.usable(satellite, self.centre + int(offset_slots))

    def dual_gap(self, limit=MAX_DUAL_GAP):
        """The longest run, over the bands, of slots where no satellite has that band
        GOOD and that contains the centre -- capped at limit + 1.

        Zero when every band is GOOD on some satellite at the centre.  Slots outside
        the era read as MISSING, so a window straddling an edge reports a long gap
        and gets redrawn -- which is right: there is no evidence such a window is
        fillable.  Per band, as the renderer's fill policy is: a band one satellite
        observes is fillable, whatever the other bands are doing.
        """
        model = self.availability
        low, high = self.centre - limit - 1, self.centre + limit + 2
        span = np.full((len(model.satellites), len(BANDS), high - low), MISSING,
                       dtype=np.uint8)
        a, b = max(low, 0), min(high, model.slots)
        if a < b:
            span[:, :, a - low:b - low] = model.states[:, :, a:b]
        dark = (span != GOOD).all(axis=0)                  # (bands, window)
        middle = self.centre - low
        longest = 0
        for band in np.flatnonzero(dark[:, middle]):
            row = dark[band]
            left = middle
            while left > 0 and row[left - 1]:
                left -= 1
            right = middle
            while right < row.size - 1 and row[right + 1]:
                right += 1
            longest = max(longest, min(right - left + 1, limit + 1))
        return longest


def unfillable_times(times, usable_times, limit=MAX_DUAL_GAP):
    """The times inside a dual-unusable run longer than `limit` slots.

    `times` is a window's full cadence-spaced slot timeline; `usable_times` maps
    each satellite to the set of times where it has a GOOD frame of the band in
    question -- the caller applies this once per band.  A slot where no satellite is usable belongs to a dual
    gap; the returned set holds every time in a run of more than `limit` such
    slots -- the ones the fill policy (:data:`MAX_DUAL_GAP`) declines to touch.
    """
    dark = np.fromiter(
        (all(t not in usable for usable in usable_times.values()) for t in times),
        dtype=bool, count=len(times),
    )
    if not dark.any():
        return set()
    edges = np.diff(dark.astype(np.int8))
    starts = list(np.flatnonzero(edges == 1) + 1)
    ends = list(np.flatnonzero(edges == -1) + 1)
    if dark[0]:
        starts.insert(0, 0)
    if dark[-1]:
        ends.append(dark.size)
    beyond = set()
    for start, end in zip(starts, ends):
        if end - start > limit:
            beyond.update(times[start:end])
    return beyond


def run_lengths(states, value):
    """Lengths of every contiguous run of `value` in a 1-D state array."""
    mask = np.asarray(states) == value
    if not mask.any():
        return np.array([], dtype=np.int64)
    edges = np.diff(mask.astype(np.int8))
    starts = np.flatnonzero(edges == 1) + 1
    ends = np.flatnonzero(edges == -1) + 1
    if mask[0]:
        starts = np.concatenate([[0], starts])
    if mask[-1]:
        ends = np.concatenate([ends, [mask.size]])
    return ends - starts