"""Pristine contiguous segments: the unit the training corpus is built from.

A segment is a maximal run of slots in which *both* satellites are GOOD -- all
six wavelengths present and detector-clean -- taken from the archive profile
(:mod:`suvi.availability`) and then verified against the extracted shards.  The
model trains only on targets whose entire input span lies inside one segment, so
every frame it is fed is pre-verified pristine and every hole or damage it sees
is simulated, against complete ground truth.  Day boundaries mean nothing here;
the minimum length is the model's own input requirement, ``2*span + 1`` frames.

Segments are stored as ``[t0, t1)`` unix-time pairs in ``segments.json`` beside
the shard manifest, and the train/val/test split assigns whole segments, so the
two sides of a split can never share a contiguous stretch of Sun.
"""

import datetime
import json
import os

import numpy as np

from . import paths


def from_availability(model, minimum_slots):
    """Maximal jointly-pristine runs of at least `minimum_slots`, as (t0, t1)."""
    pristine = model.pristine()
    edges = np.diff(pristine.astype(np.int8))
    starts = list(np.flatnonzero(edges == 1) + 1)
    ends = list(np.flatnonzero(edges == -1) + 1)
    if pristine.size and pristine[0]:
        starts.insert(0, 0)
    if pristine.size and pristine[-1]:
        ends.append(pristine.size)
    # int(), not numpy's: these bounds are written to JSON and compared against
    # shard keys, and a stray int64 breaks json.dump at the very end of a
    # multi-hour build.
    return [
        (int(model.t0 + a * model.cadence), int(model.t0 + b * model.cadence))
        for a, b in zip(starts, ends)
        if (b - a) >= minimum_slots
    ]


def decorrelation_slots(cadence=paths.CADENCE):
    """Frames after which a segment is showing an essentially different Sun.

    Half a synodic rotation: by then the hemisphere that was facing us has gone
    round the back and the one that was hidden is in view, so two frames that far
    apart share no surface features.  This is the honest scale for "how much
    diversity does one contiguous segment hold" -- a segment shorter than this is
    one evolving view of one hemisphere, and truncating it below this length
    discards genuinely new Sun rather than redundancy.

    Equatorial synodic rotation is the fastest and so the conservative choice:
    26.1 days, giving ~4,700 frames at the 4-minute cadence.  Higher latitudes
    take half again as long.
    """
    from . import fillers

    period_days = 360.0 / fillers.rotation_rate(0.0, synodic=True)
    return int(period_days / 2 * 86400 / cadence)


def days_touched(segments, cadence=paths.CADENCE):
    """Every calendar day any segment's slots fall on -- the days to extract."""
    days = set()
    for t0, t1 in segments:
        for t in range(t0, t1, cadence):
            days.add(datetime.datetime.fromtimestamp(
                t, datetime.timezone.utc).strftime("%Y-%m-%d"))
    return sorted(days)


def verify(segments, usable, cadence=paths.CADENCE, minimum_slots=1):
    """Split segments wherever extraction disagrees with the profile.

    The profile is header-only; ``disc_v1`` runs at extraction and can reject a
    slot inside a nominally pristine run (measured ~0.3%).  `usable(t)` says
    whether every band on every satellite actually made it into the shards at
    `t`.  Any slot that did not splits its segment; fragments shorter than
    `minimum_slots` are dropped.  The corpus is only final after this pass.
    """
    verified = []
    for t0, t1 in segments:
        start = t0
        for t in range(t0, t1, cadence):
            if not usable(t):
                if (t - start) // cadence >= minimum_slots:
                    verified.append((start, t))
                start = t + cadence
        if (t1 - start) // cadence >= minimum_slots:
            verified.append((start, t1))
    return verified


def targets(segments, span_slots, cadence=paths.CADENCE):
    """Every slot time at least `span_slots` from both ends of its segment."""
    margin = span_slots * cadence
    out = []
    for t0, t1 in segments:
        out.extend(range(t0 + margin, t1 - margin, cadence))
    return out


def split(segments, seed, val_targets, test_targets, span_slots,
          cadence=paths.CADENCE):
    """Assign whole segments to val/test until each holds enough targets.

    Deterministic in `seed`.  Splitting by segment, not by day or slot, is the
    leakage rule: two targets in one contiguous run share most of their context,
    so a run must land entirely on one side.
    """
    order = list(segments)
    np.random.default_rng(seed).shuffle(order)

    def mass(segment):
        t0, t1 = segment
        return max((t1 - t0) // cadence - 2 * span_slots, 0)

    splits = {"val": [], "test": [], "train": []}
    need = {"val": val_targets, "test": test_targets}
    for segment in order:
        for name in ("val", "test"):
            if need[name] > 0:
                splits[name].append(tuple(segment))
                need[name] -= mass(segment)
                break
        else:
            splits["train"].append(tuple(segment))
    return {name: sorted(chosen) for name, chosen in splits.items()}


def containing(segments, t):
    """The segment holding time `t`, or None."""
    for t0, t1 in segments:
        if t0 <= t < t1:
            return (t0, t1)
    return None


def save(path, segments):
    payload = {"segments": [[int(a), int(b)] for a, b in sorted(segments)]}
    # Whole-file write, then rename: a half-written segments.json is worse than
    # none, because the next run would load it and train on a truncated corpus.
    with open(path + ".tmp", "w") as handle:
        json.dump(payload, handle)
    os.replace(path + ".tmp", path)


def load(path):
    with open(path) as handle:
        return [tuple(pair) for pair in json.load(handle)["segments"]]