"""Training shards: a compact, portable copy of the archive for the learned filler.

The archive is 9.2 TB of tile-compressed FITS reached over three hops -- a virtiofs
re-export of an SMB share -- and the machine that trains the model is a different machine
on the other end of a gigabit link.  So training does not read the archive.  It reads
*shards*:
one file per (day, satellite) holding that day's frames at half resolution, which is
small enough to ship and to re-read every epoch.

Two decisions here are load-bearing.

**Shards store clean frames only.**  Corruption is applied on the fly during training
(:mod:`suvi.corruptions`), which costs nothing in storage, gives unlimited variety rather
than one frozen draw, and yields an exact ground-truth class label for the auxiliary
head.  Baking damage into the shards would fix the training distribution at extraction
time and make every experiment downstream a re-extraction.

**Pixels are stored as quantised ``asinh`` radiance.**  Solar radiance is heavy-tailed --
a sample across bands and epochs runs from -3.7 to 1231 -- so linear uint16 would spend
its whole range on the disc and quantise the corona to nothing.  ``asinh(x/s)`` is
linear below the knee `s` and logarithmic above it, which is what a signal that is
mostly noise-floor with occasional flares needs.

``log1p``, the space ``fillers._for_flow`` uses, was the obvious first choice and is
wrong here on two counts.  It is effectively linear below 1, and almost every pixel is:
median radiance is 0.004 to 0.19 depending on band.  Quantising log1p over [0, 8] gives
a 0.002 radiance pixel **6% error**, not the 0.05% the bright disc gets.  And it cannot
represent negative values at all, while 2-18% of pixels are negative after background
subtraction -- clipping those to zero would bias the noise floor the model has to learn
to reproduce.  ``asinh`` holds ~0.05% relative across the whole range and is signed.
"""

import hashlib
import json
import os
import struct

import numpy as np

from . import fitsio, paths

#: Half of the archive's native 1280.  Flow, gain and blend fields are smooth, so the
#: model predicts them here and they upsample cleanly to drive full-resolution frames.
SHARD_SIZE = 640
#: Knee of the asinh transform: linear below, logarithmic above.  Relative precision is
#: constant above the knee and decays below it, so the knee sits well under the faintest
#: band's median radiance (0.004 at 131A) -- far enough that the whole corona is in the
#: constant-precision regime, not just the disc.
ASINH_SCALE = 1e-4
#: Half-range of the stored asinh values, symmetric about zero.  +/-18 covers
#: +/-sinh(18)*scale = +/-3284 in radiance, against an observed span of -3.7 to 1231 --
#: room for a flare well past anything in the archive.  Anything beyond that is a
#: corruption, not an observation, and saturates rather than wrapping.
ASINH_RANGE = 18.0
#: Marker written for a slot whose frame would not read.  Distinguishes "we looked and
#: there was nothing usable" from "we never looked", which a plain absence cannot.
UNREADABLE = b"\x00"

MAGIC = b"SUVISHRD"
VERSION = 2
SLOTS_PER_DAY = 86400 // paths.CADENCE  # 360


def _zstd():
    """zstd compress/decompress, from the stdlib on 3.14+ or the pip package below it.

    The extraction host runs Python 3.14, where zstd is in the standard library; the
    training container ships whatever the ROCm image was built against.  Both produce
    ordinary zstd frames, so shards written by one are read by the other.
    """
    try:
        from compression import zstd

        return zstd.compress, zstd.decompress
    except ImportError:  # pragma: no cover - exercised only on the training host
        import zstandard

        return (
            lambda data, level=3: zstandard.ZstdCompressor(level=level).compress(data),
            lambda data: zstandard.ZstdDecompressor().decompress(data),
        )


# ------------------------------------------------------------------- pixel encoding


def encode_frames(arrays):
    """Six 1280x1280 radiance arrays -> one 6 x 640 x 640 uint16 block.

    Downsampling is a 2x2 mean rather than decimation: SUVI frames carry read noise and
    cosmic-ray hits, and taking every other pixel would keep the hits at full amplitude
    while throwing away the averaging that suppresses them.
    """
    planes = []
    for array in arrays:
        image = np.nan_to_num(np.asarray(array, dtype=np.float64), nan=0.0,
                              posinf=0.0, neginf=0.0)
        if image.shape[0] % SHARD_SIZE or image.shape[1] % SHARD_SIZE:
            raise ValueError(f"cannot halve {image.shape} to {SHARD_SIZE}")
        factor = image.shape[0] // SHARD_SIZE
        small = image.reshape(SHARD_SIZE, factor, SHARD_SIZE, factor).mean(axis=(1, 3))
        unit = np.arcsinh(small / ASINH_SCALE) / ASINH_RANGE  # -> roughly [-1, 1]
        planes.append(np.clip(unit, -1.0, 1.0) * 32767.5 + 32767.5)
    return np.rint(np.stack(planes)).astype(np.uint16)


def decode_frames(block):
    """Inverse of :func:`encode_frames`, back to radiance."""
    return (np.sinh(coded_frames(block).astype(np.float64) * ASINH_RANGE)
            * ASINH_SCALE).astype(np.float32)


def coded_frames(block):
    """The stored block as the roughly-[-1, 1] values the model consumes.

    The uint16 in a shard *is* the model's input space -- ``encode_frames`` already
    applied the asinh -- so getting from one to the other is an affine rescale and
    nothing more.  Going via radiance instead costs a ``sinh`` and an ``arcsinh`` over
    47 million elements per training sample, which measured as a data loader pinning two
    cores at 200% while the GPU sat at 0% busy.  Only frames that are about to be
    corrupted need real radiance, because that is the space corruptions are defined in.
    """
    return ((np.asarray(block, dtype=np.float32) - 32767.5) / 32767.5).astype(np.float32)


# ------------------------------------------------------------------- shard container
#
# Layout:
#   offset 0        _PREFIX   MAGIC, VERSION                  (10 bytes)
#   offset 10       records back to back, sorted by t_start
#                   JSON header, carrying (offset, length) per slot
#   EOF - 12        _TRAILER  header length, MAGIC            (12 bytes)
#
# The index gives one record per slot, so a single record can be read without
# decompressing the rest -- which is what makes random sampling across the corpus
# cheap.  It lives at the *end* because its contents are not known until every
# payload has been compressed: with the index in front, writing a shard means
# buffering the whole ~1 GB of it, and shards are produced straight into a network
# stream that should never hold more than one record at a time.
#
# Record offsets are relative to the first record, which is therefore at a fixed
# position rather than one that depends on the header's length.

_PREFIX = struct.Struct("<8sH")
_TRAILER = struct.Struct("<I8s")


def _header_bytes(day, satellite, wavelengths, index):
    return json.dumps({
        "day": day,
        "satellite": int(satellite),
        "wavelengths": list(wavelengths),
        "size": SHARD_SIZE,
        "asinh_scale": ASINH_SCALE,
        "asinh_range": ASINH_RANGE,
        "index": index,
    }, separators=(",", ":")).encode()


def write_shard_stream(handle, day, satellite, wavelengths, records, level=3):
    """Write one shard to a binary stream.  Returns the bytes written.

    `records` is an *iterable* of (t_start, block) in ascending t_start order --
    an iterable, not a mapping, because the point of this function is that a
    shard can be produced without ever existing in memory or on disk.  Each
    payload is compressed and written as it arrives, so the caller may hand over
    a generator that reads and encodes one slot at a time and nothing larger than
    a single record is ever live.

    Ascending order is the caller's responsibility and is checked: offsets are
    cumulative, so an out-of-order record would silently produce an index that
    does not describe the file.
    """
    compress, _ = _zstd()
    index = {}
    offset = 0
    previous = None
    handle.write(_PREFIX.pack(MAGIC, VERSION))
    for t_start, block in records:
        if previous is not None and t_start <= previous:
            raise ValueError(
                f"records must ascend by t_start: {t_start} after {previous}")
        previous = t_start
        payload = UNREADABLE if block is None else compress(
            np.ascontiguousarray(block, dtype="<u2").tobytes(), level
        )
        handle.write(payload)
        index[str(t_start)] = [offset, len(payload)]
        offset += len(payload)

    header = _header_bytes(day, satellite, wavelengths, index)
    handle.write(header)
    handle.write(_TRAILER.pack(len(header), MAGIC))
    return _PREFIX.size + offset + len(header) + _TRAILER.size


def write_shard(path, day, satellite, wavelengths, records, level=3):
    """Write one (day, satellite) shard to a file.

    `records` maps t_start -> 6x640x640 uint16 block (or None for a slot that was
    looked at and found unusable).  A thin wrapper over
    :func:`write_shard_stream` so there is only one implementation of the format.
    """
    temporary = f"{path}.partial"
    with open(temporary, "wb") as handle:
        write_shard_stream(handle, day, satellite, wavelengths,
                           ((t, records[t]) for t in sorted(records)), level)
    os.replace(temporary, path)   # never leave a half-written shard at the real name
    return os.path.getsize(path)


class Shard:
    """Random access to one shard.  Holds a single file handle, opened eagerly.

    Eagerly, because the pool rotator's consistency protocol depends on it: a
    rotated-out shard is *deleted* on the training host, and POSIX only keeps the
    bytes readable for handles that already exist.  A lazily-opened shard whose
    first read came after the delete once killed a run mid-epoch."""

    def __init__(self, path):
        self.path = path
        self._handle = open(path, "rb")
        size = os.fstat(self._handle.fileno()).st_size
        if size < _PREFIX.size + _TRAILER.size:
            raise ValueError(f"{path}: too short to be a shard ({size} bytes)")

        magic, version = _PREFIX.unpack(self._handle.read(_PREFIX.size))
        if magic != MAGIC:
            raise ValueError(f"{path}: not a shard")
        if version != VERSION:
            raise ValueError(f"{path}: shard version {version}, expected {VERSION}")

        length, tail_magic = _TRAILER.unpack(
            os.pread(self._handle.fileno(), _TRAILER.size, size - _TRAILER.size))
        # A shard is now streamed to its destination rather than written whole and
        # renamed, so a truncated one is a real possibility rather than a
        # theoretical one.  Every check below exists to make that fail here, at
        # open, instead of returning plausible-looking pixels from the wrong offset.
        if tail_magic != MAGIC:
            raise ValueError(f"{path}: truncated or not a shard (bad trailer)")
        start = size - _TRAILER.size - length
        if not _PREFIX.size <= start <= size - _TRAILER.size:
            raise ValueError(f"{path}: header length {length} does not fit the file")
        self.header = json.loads(os.pread(self._handle.fileno(), length, start))

        self._base = _PREFIX.size
        self.satellite = self.header["satellite"]
        self.day = self.header["day"]
        self.wavelengths = tuple(self.header["wavelengths"])
        self._index = {int(k): tuple(v) for k, v in self.header["index"].items()}
        end = max((o + n for o, n in self._index.values()), default=0)
        if self._base + end != start:
            raise ValueError(
                f"{path}: records end at {self._base + end}, header starts at {start}")

    def times(self):
        return sorted(self._index)

    def __contains__(self, t_start):
        return t_start in self._index

    def usable(self, t_start):
        """Whether the slot holds real frames, from the index alone.

        A vetting-dropped slot is stored as the one-byte UNREADABLE sentinel, so
        the index's length column answers without decompressing anything -- the
        segment verifier asks this for every slot of the corpus.
        """
        entry = self._index.get(t_start)
        return entry is not None and entry[1] > len(UNREADABLE)

    def raw(self, t_start):
        """The stored uint16 block for a slot, or None if it holds no usable frame."""
        entry = self._index.get(t_start)
        if entry is None:
            return None
        offset, length = entry
        if self._handle is None:            # reopened after an explicit close()
            self._handle = open(self.path, "rb")
        # pread, not seek+read: DataLoader workers fork *after* the shard opens
        # (they must -- the handle has to predate any pool-rotation delete), so
        # they share one file description.  A seek from one worker lands in the
        # middle of another's read; positional reads share nothing.
        payload = os.pread(self._handle.fileno(), length, self._base + offset)
        if payload == UNREADABLE:
            return None
        _, decompress = _zstd()
        flat = np.frombuffer(decompress(payload), dtype="<u2")
        return flat.reshape(len(self.wavelengths), SHARD_SIZE, SHARD_SIZE)

    def frames(self, t_start):
        """Radiance for a slot, 6 x 640 x 640 float32, or None."""
        block = self.raw(t_start)
        return None if block is None else decode_frames(block)

    def coded(self, t_start):
        """Model-space values for a slot, 6 x 640 x 640 float32, or None.

        The cheap path, and the one training uses.  See :func:`coded_frames`.
        """
        block = self.raw(t_start)
        return None if block is None else coded_frames(block)

    def close(self):
        if self._handle is not None:
            self._handle.close()
            self._handle = None

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()


def shard_name(day, satellite):
    return f"{day}_g{satellite:d}.shard"


# ------------------------------------------------------------------------ extraction


def read_slot(root, rows):
    """Decode one slot's six bands from the archive.

    `rows` maps wavelength -> archive-relative path.  Returns the encoded block, or
    None if any band is missing or unreadable -- a partial slot is not a training
    sample, and admitting one would teach the model that bands go missing
    independently when in practice a satellite drops out whole.
    """
    arrays = []
    for wavelength in paths.WAVELENGTHS:
        relpath = rows.get(wavelength)
        if relpath is None:
            return None
        image, _ = fitsio.read_image(paths.abspath(relpath, root))
        if image is None or image.shape != (1280, 1280):
            return None
        arrays.append(image)
    return encode_frames(arrays)


# ----------------------------------------------------------------------- day choice


#: Smallest plausible real frame.  A healthy tile-compressed SUVI frame is 1.4-1.8 MB;
#: a failed download leaves a 5,760-byte header-only stub that is indexed like any other
#: frame and reads back as None.  There are 125,732 of them in the archive -- 2.3% -- and
#: one chosen day turned out to be 100% stubs, producing an empty shard after a full
#: extraction pass over it.  Counting rows is not the same as counting frames.
MIN_FRAME_BYTES = 100_000


def choose_days(conn, count, t_from, t_to, exclude=(), satellites=(16, 18),
                wavelengths=paths.WAVELENGTHS, min_bytes=MIN_FRAME_BYTES):
    """Days where both satellites have near-complete coverage, spread evenly.

    Spread rather than sampled at random: the point of the training set is to span the
    solar cycle, and independent draws over a 2.5-year range clump.  Complete rather
    than best-effort: a day missing half its slots would silently bias the gap-length
    distribution the model trains against.  And *real* rather than merely present -- see
    :data:`MIN_FRAME_BYTES`.
    """
    wanted = len(wavelengths) * len(satellites) * SLOTS_PER_DAY
    placeholders = ",".join("?" * len(wavelengths))
    satellite_places = ",".join("?" * len(satellites))
    rows = conn.execute(
        f"""
        SELECT strftime('%Y-%m-%d', t_start, 'unixepoch') AS day, COUNT(*) AS n
        FROM frame
        WHERE t_start >= ? AND t_start < ?
          AND wavelength IN ({placeholders}) AND satellite IN ({satellite_places})
          AND size_bytes >= ?
        GROUP BY day HAVING n >= ?
        ORDER BY day
        """,
        (t_from, t_to, *wavelengths, *satellites, int(min_bytes), int(wanted * 0.98)),
    ).fetchall()

    excluded = set(exclude)
    days = [row[0] for row in rows if row[0] not in excluded]
    if not days:
        return []
    if len(days) <= count:
        return days
    step = len(days) / count
    return [days[int(i * step)] for i in range(count)]


def check_split(days_by_split):
    """Assert the splits are disjoint and no two blocks abut in time.

    Cheap, and it fails loudly.  A leak here would not crash anything -- it would just
    produce validation numbers that look excellent and mean nothing.
    """
    import datetime

    seen = {}
    for name, days in days_by_split.items():
        for day in days:
            if day in seen:
                raise ValueError(f"day {day} is in both {seen[day]!r} and {name!r}")
            seen[day] = name

    ordered = sorted(seen)
    adjacent = []
    for first, second in zip(ordered, ordered[1:]):
        a = datetime.date.fromisoformat(first)
        b = datetime.date.fromisoformat(second)
        if (b - a).days == 1 and seen[first] != seen[second]:
            adjacent.append((first, second))
    return adjacent