"""Turning shards into training samples: stacks, masks, and damage applied on the fly.

One sample is a target slot -- one satellite, one instant, all six bands -- plus the
stack of frames the model may draw on and a per-frame conditioning vector.  This module
is the only place that knows how a stack is laid out, so :mod:`suvi.model` and
:mod:`bench` cannot disagree about it.

Three things here decide what the model can learn.

**Offsets are exponentially spaced.**  Gaps in this archive run from one slot to three
hundred, and a fixed +/-K window either misses the far end of a long outage or wastes
most of its inputs on a short one.  Sampling at ``+/-{1, 4, 16, 64, 256}`` slots covers
four minutes to seventeen hours with ten frames per satellite, and an *anchor* -- the
nearest usable frame in each direction, however far -- guarantees the stack is never
empty even beyond that.

**Damage is applied here, not baked into the shards.**  Every epoch draws fresh
corruptions from :mod:`suvi.corruptions`, so the model sees far more variety than a
frozen dataset could hold, and the true mode is known exactly for the auxiliary head.

**Outages are episodic as well as scattered.**  Half of all samples carry a contiguous
simulated outage -- up to 400 slots, on one satellite or both -- because that is how
the real archive fails: measured over the two-satellite era, 11% of slots have both
spacecraft dark at once, in long runs.  Independent per-frame drops alone would never
construct that case, and a model cannot learn a regime it has never seen.

**Context frames get damaged too, not just the target.**  This is the requirement that
any supplied frame may itself be invalid.  A model trained on clean context and deployed
on an archive where 31% of slots are bad on both satellites would meet, at inference, a
distribution it had never seen.
"""

import math

import numpy as np

from . import align, corruptions, dataset, paths

#: Slot offsets sampled on each satellite, in both directions.  Exponentially spaced,
#: so one fixed-size stack serves every gap length: 4 minutes to 17 hours in five
#: steps, without needing hundreds of input frames to cover a 300-slot outage.  The
#: anchors (nearest usable frame either way, at any distance) still guarantee the
#: stack is never empty beyond the ladder's reach.  Cost is linear in stack depth in
#: the current model (no attention), so the two extra rungs are affordable where they
#: were not for the attention-based design this replaced.
OFFSETS = (1, 4, 16, 64, 256)
#: How often a sample carries a simulated *episodic* outage -- a contiguous run of
#: missing slots -- on top of the scattered per-frame drops.  Real outages are runs,
#: not confetti: measured against the archive, 11% of the two-satellite era has BOTH
#: spacecraft dark simultaneously, in 1,096 contiguous runs of up to 26 days.  With
#: independent per-frame drops the probability of reproducing that configuration in
#: a training stack is roughly p^14 -- the model would face at evaluation a regime
#: it had never once seen.
EPISODIC_PROBABILITY = 0.5
#: Given an episodic outage, how often it takes down *both* satellites at once (a
#: ground-segment or space-weather event) rather than one.
DUAL_OUTAGE_PROBABILITY = 0.4
#: Longest simulated outage, in slots; matches the anchor search limit, so a stack
#: always retains something real beyond the outage's edge.
MAX_OUTAGE_SLOTS = 400
#: Modes that leave a frame with no usable signal at all.  A frame damaged this way is
#: presented as `missing`; everything else is presented as `suspect` with its pixels.
NO_SIGNAL = frozenset({"all_zero", "drop_image_hdu", "nan_fill", "zblank_fill"})
#: Auxiliary-head classes: clean plus every catalogued mode, in a fixed order so a
#: checkpoint's class indices stay meaningful.
CLASSES = ("clean",) + tuple(sorted(corruptions.CATALOG))
CLASS_INDEX = {name: index for index, name in enumerate(CLASSES)}


def stack_layout(satellites, target_satellite):
    """The (satellite, offset) slots a stack is built from, in a fixed order.

    Deterministic and independent of what happens to be available, so a frame's position
    in the stack always means the same thing and the model can rely on it.

    The target satellite's own frame at offset 0 is the frame being reconstructed and is
    never an input.  The *counterpart's* frame at offset 0 is the single most valuable
    entry in the stack -- a real observation of the right Sun at the right instant, worth
    ~46 dB on its own once its instrument gain is known.
    """
    others = [s for s in satellites if s != target_satellite]
    layout = []
    for satellite in [target_satellite] + others:
        for offset in OFFSETS:
            layout.append((satellite, -offset))
            layout.append((satellite, offset))
    return layout + [(other, 0) for other in others]


class Sampler:
    """Builds training samples from a set of shards.

    `shards` maps (day, satellite) -> :class:`suvi.dataset.Shard`.
    """

    def __init__(self, shards, satellites=(16, 18), damage_probability=0.30,
                 drop_probability=0.40, seed=0):
        self.shards = shards
        self.satellites = tuple(satellites)
        #: Ceilings, not rates.  Each *sample* draws its own severity uniformly up to
        #: these, rather than every frame being damaged at one fixed probability.
        #:
        #: Fixed rates of 0.25 drop / 0.35 damage left barely half the stack clean in
        #: every single sample, so the model never saw the easy case and learned a
        #: hedging average instead of "trust the counterpart".  It showed up as a model
        #: that gained 5 dB on its own harsh validation set and ~1 dB on the bench,
        #: where neighbours are mostly clean.  Real outages are episodic -- long clean
        #: stretches broken by bad runs -- so drawing severity per sample covers both
        #: regimes, and roughly a fifth of samples come through almost untouched.
        self.damage_probability = damage_probability
        self.drop_probability = drop_probability
        self.rng = np.random.default_rng(seed)
        self._by_satellite = {}
        self._day_of = {}
        for (day, satellite), shard in shards.items():
            self._day_of[id(shard)] = day
            self._by_satellite.setdefault(satellite, {}).update(
                {time: shard for time in shard.times()}
            )
        #: (day, target satellite, other satellite, pair time) -> per-band (gains,
        #: offsets).  See :meth:`_calibration`.
        self._transfers = {}

    def targets(self):
        """Every (satellite, time) that could be reconstructed, in a stable order."""
        found = []
        for satellite in self.satellites:
            for time in sorted(self._by_satellite.get(satellite, {})):
                found.append((satellite, time))
        return found

    def _read(self, satellite, time):
        """One slot in *model* space -- the cheap path, used for every clean frame.

        Shards store the asinh representation already, so this is a rescale rather than
        the sinh/arcsinh round trip that going via radiance would cost.
        """
        shard = self._by_satellite.get(satellite, {}).get(time)
        if shard is None:
            return None
        return shard.coded(time) if hasattr(shard, "coded") else encode_for_model(
            shard.frames(time)
        )

    def _donor(self, satellite, time, rng, reach=60):
        """Another real frame, for the modes that splice one in.

        ``frozen`` repeats a neighbour, ``wrong_time`` files one under the wrong stamp
        and ``torn_frame`` merges two.  All three need a *plausible* second observation,
        so the donor is drawn from nearby on the same satellite rather than invented.
        """
        available = self._by_satellite.get(satellite, {})
        for _ in range(8):
            offset = int(rng.integers(-reach, reach + 1))
            if offset == 0:
                continue
            candidate = time + offset * paths.CADENCE
            if candidate in available:
                return self._read(satellite, candidate)
        return None

    def _anchor(self, satellite, time, direction, exclude=frozenset(), limit=400):
        """Nearest slot with a frame outside `exclude`, up to `limit` slots away.

        `exclude` holds the times inside a simulated outage.  The production stack
        builder picks anchors from *clean* times, so the sampler must too: an anchor
        that lands inside the outage and arrives `missing` would leave the stack with
        no real long-range frame at all, which is not what the pipeline would see.
        """
        available = self._by_satellite.get(satellite, {})
        for step in range(1, limit + 1):
            candidate = time + direction * step * paths.CADENCE
            if candidate in available and candidate not in exclude:
                return candidate
        return None

    def _outages(self, rng, time):
        """Per-satellite times inside this sample's simulated episodic outage.

        One contiguous run per affected satellite, log-uniform in length from one
        slot to :data:`MAX_OUTAGE_SLOTS`, positioned so the target usually sits
        inside it -- reproducing the archive's real failure mode of equipment and
        comms dropping out for hours to days, sometimes on both spacecraft at once.
        """
        blocked = {satellite: frozenset() for satellite in self.satellites}
        # Gated on the drop setting: drop_probability=0 means "no synthetic losses
        # of any kind", which diagnostics and tests rely on.
        if self.drop_probability <= 0 or rng.random() >= EPISODIC_PROBABILITY:
            return blocked
        length = int(round(math.exp(rng.uniform(0.0, math.log(MAX_OUTAGE_SLOTS)))))
        centre = time + int(rng.integers(-length, length + 1)) * paths.CADENCE
        start = centre - (length // 2) * paths.CADENCE
        interval = frozenset(start + i * paths.CADENCE for i in range(length))
        if rng.random() < DUAL_OUTAGE_PROBABILITY:
            hit = self.satellites
        else:
            hit = (self.satellites[int(rng.integers(len(self.satellites)))],)
        for satellite in hit:
            blocked[satellite] = interval
        return blocked

    def _calibration(self, day, satellite, other, exclude):
        """Per-band (gains, offsets) putting `other`'s radiance on `satellite`'s scale.

        Fitted from one simultaneous pair per day -- the transfer drifts ~12% per
        *week*, so a fit at most 24 hours stale is well inside its own noise, and one
        pair per day is what a production pipeline could equally afford.

        `exclude` is the time being reconstructed.  A pair at that instant would fit
        the transfer against the answer itself -- the oracle gain the whole exercise
        exists to estimate honestly -- so the fit steps to a neighbouring pair
        instead.  Returns identity when the day has no usable pair.
        """
        local = self.shards.get((day, satellite))
        counterpart = self.shards.get((day, other))
        if local is None or counterpart is None:
            return None
        candidates = sorted(set(local.times()) & set(counterpart.times()) - {exclude})
        if not candidates:
            return None
        pair_time = candidates[len(candidates) // 2]

        key = (day, satellite, other, pair_time)
        if key not in self._transfers:
            self._transfers[key] = align.fit_photometry(
                counterpart.frames(pair_time), local.frames(pair_time)
            )
        return self._transfers[key]

    def build(self, satellite, time, rng=None):
        """One sample: the stack, its per-frame metadata, and the withheld target.

        `frames` is (S, 6, H, W) in *model* space -- roughly [-1, 1] asinh radiance --
        with zeros where a frame is missing; `target` is the withheld frame in the same
        space.  `gains`/`offsets` are the per-frame radiance transfers that put a
        cross-satellite frame on the target instrument's scale (identity rows for
        same-satellite and missing frames), and `b0`/`radius` the solar geometry the
        rotation warp needs -- both consumed by :func:`suvi.align.align_stack` rather
        than here, so the CPU loader stays cheap.

        Returns None if the target itself is absent, or if not one frame in the stack
        carries pixels -- there is nothing to reconstruct *from*, and the model's
        contract is to decline rather than emit a fabricated frame.
        """
        rng = rng if rng is not None else self.rng
        target = self._read(satellite, time)
        if target is None:
            return None
        day = self._day_of.get(id(self._by_satellite[satellite][time]))

        # This sample's severity, drawn once and applied to every frame in its stack.
        drop_probability = float(rng.uniform(0.0, self.drop_probability))
        damage_probability = float(rng.uniform(0.0, self.damage_probability))
        outages = self._outages(rng, time)

        layout = list(stack_layout(self.satellites, satellite))
        # Anchors: the nearest usable frame either way on the target satellite --
        # outside any simulated outage, as the production stack builder would pick
        # them -- so a long gap still has something real to work from.
        for direction in (-1, 1):
            anchor = self._anchor(satellite, time, direction, exclude=outages[satellite])
            if anchor is not None:
                offset = (anchor - time) // paths.CADENCE
                if (satellite, offset) not in layout:
                    layout.append((satellite, offset))

        frames, states, dts, same, classes, transfers = [], [], [], [], [], []
        for source_satellite, offset in layout:
            when = time + offset * paths.CADENCE
            image = self._read(source_satellite, when)
            state, label = "available", "clean"

            if image is None:
                state = "missing"
            elif when in outages[source_satellite]:
                # Inside this sample's simulated episodic outage.
                image, state = None, "missing"
            elif rng.random() < drop_probability:
                # A scattered single-slot loss, independent of the episodic runs.
                image, state = None, "missing"
            elif rng.random() < damage_probability:
                mode = str(rng.choice(sorted(corruptions.CATALOG)))
                corruption = corruptions.CATALOG[mode]
                label = mode
                if corruption.kind != "array":
                    # A file-level corruption destroys the container, so by the time a
                    # frame would reach the model there is nothing left of it.
                    image, state = None, "missing"
                elif mode in NO_SIGNAL:
                    image, state = None, "missing"
                else:
                    donor = self._donor(source_satellite, when, rng) \
                        if corruption.needs_donor else None
                    if corruption.needs_donor and donor is None:
                        label, state = "clean", "available"  # nothing to draw from
                    else:
                        image = _damage(image, mode, int(rng.integers(0, 2**31 - 1)),
                                        float(rng.uniform(0.5, 1.0)), donor)
                        state = "suspect"

            shape = target.shape
            frames.append(np.zeros(shape, dtype=np.float32) if image is None else image)
            states.append(state)
            dts.append(float(offset * paths.CADENCE))
            same.append(source_satellite == satellite)
            classes.append(CLASS_INDEX[label])
            transfer = None
            if source_satellite != satellite and image is not None:
                transfer = self._calibration(day, satellite, source_satellite, time)
            transfers.append(transfer)

        if all(state == "missing" for state in states):
            return None

        identity = (np.ones(6, dtype=np.float32), np.zeros(6, dtype=np.float32))
        b0, radius = align.solar_ephemeris(float(time))
        return {
            "frames": np.stack(frames),
            "states": states,
            "dts": np.array(dts, dtype=np.float32),
            "same_satellite": np.array(same, dtype=bool),
            "classes": np.array(classes, dtype=np.int64),
            "gains": np.stack([(t or identity)[0] for t in transfers]),
            "offsets": np.stack([(t or identity)[1] for t in transfers]),
            "b0": float(b0),
            "radius": float(radius),
            "target": target,
            "slot": (satellite, time),
        }


def _damage(coded, mode, seed, severity, donor=None):
    """Apply one array corruption to all six bands of a frame, in model space.

    Corruptions are *defined* on radiance -- ``gain_shift`` multiplies a physical
    quantity, ``saturate`` clips against a physical ceiling -- so the frame is decoded,
    damaged, and re-encoded.  This is the expensive path, which is why only the frames
    actually being corrupted take it.

    Each band gets the same seed, so a fault hits the frame coherently rather than
    differently per band, which is how the real ones behave: one instrument produces all
    six.  ``apply_array`` returns (pixels, header_overrides); the overrides describe a
    FITS header this pipeline does not carry, so only the pixels are kept.
    """
    radiance = decode_from_model(coded)
    donor_radiance = None if donor is None else decode_from_model(donor)
    bands = []
    for index, band in enumerate(radiance):
        damaged, _ = corruptions.apply_array(
            mode, band, seed, severity,
            donor=None if donor_radiance is None else donor_radiance[index],
        )
        bands.append(damaged)
    return encode_for_model(np.stack(bands))


def encode_for_model(radiance):
    """Radiance -> the roughly-[-1, 1] asinh space the model works in.

    Corruptions are applied in *radiance*, because that is where they are defined -- a
    ``gain_shift`` multiplies a physical quantity and a ``saturate`` clips against a
    physical ceiling.  Applying them after the transform would model a different fault.
    So the encoding happens here, once, on the way into the network.
    """
    coded = np.arcsinh(
        np.nan_to_num(np.asarray(radiance, dtype=np.float32), nan=0.0,
                      posinf=0.0, neginf=0.0) / dataset.ASINH_SCALE
    ) / dataset.ASINH_RANGE
    return np.clip(coded, -1.0, 1.0).astype(np.float32)


def decode_from_model(coded):
    """Inverse of :func:`encode_for_model`, back to radiance."""
    scaled = np.asarray(coded, dtype=np.float64) * dataset.ASINH_RANGE
    return (np.sinh(np.clip(scaled, -dataset.ASINH_RANGE, dataset.ASINH_RANGE))
            * dataset.ASINH_SCALE).astype(np.float32)


def to_tensors(sample, torch):
    """Pack a sample into tensors for :func:`suvi.align.align_stack` and the model.

    The stack is *unaligned* here: alignment is elementwise math plus a warp, which
    the GPU does in milliseconds and the CPU loader should not spend its budget on.
    """
    from . import model as model_module

    condition = torch.stack([
        model_module.frame_conditioning(
            sample["states"][i], bool(sample["same_satellite"][i]), float(sample["dts"][i])
        )
        for i in range(len(sample["states"]))
    ])
    # Already in model space: the sampler works there throughout, so nothing to convert.
    return {
        "stack": torch.from_numpy(np.ascontiguousarray(sample["frames"])),
        "condition": condition,
        "dts": torch.from_numpy(np.ascontiguousarray(sample["dts"])),
        "gains": torch.from_numpy(np.ascontiguousarray(sample["gains"])),
        "offsets": torch.from_numpy(np.ascontiguousarray(sample["offsets"])),
        "b0": torch.tensor(sample["b0"], dtype=torch.float32),
        "radius": torch.tensor(sample["radius"], dtype=torch.float32),
        "target": torch.from_numpy(np.ascontiguousarray(sample["target"])),
    }