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

The archive is 9.2 TB of tile-compressed FITS on a filesystem that cannot survive being
walked repeatedly, 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 = 1
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: MAGIC | version | header length | JSON header | records back to back.
# The JSON header carries an (offset, length) per slot, so one record can be read
# without decompressing the rest -- which is what makes random sampling across a
# 60-day dataset cheap.

_PREFIX = struct.Struct("<8sHI")


def write_shard(path, day, satellite, wavelengths, records, level=3):
    """Write one (day, satellite) shard.  `records` maps t_start -> 6x640x640 uint16."""
    compress, _ = _zstd()
    index = {}
    blobs = []
    offset = 0
    for t_start in sorted(records):
        block = records[t_start]
        payload = UNREADABLE if block is None else compress(
            np.ascontiguousarray(block, dtype="<u2").tobytes(), level
        )
        index[str(t_start)] = [offset, len(payload)]
        blobs.append(payload)
        offset += len(payload)

    header = 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()

    temporary = f"{path}.partial"
    with open(temporary, "wb") as handle:
        handle.write(_PREFIX.pack(MAGIC, VERSION, len(header)))
        handle.write(header)
        for blob in blobs:
            handle.write(blob)
    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 lazily."""

    def __init__(self, path):
        self.path = path
        self._handle = None
        with open(path, "rb") as handle:
            magic, version, length = _PREFIX.unpack(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}")
            self.header = json.loads(handle.read(length))
            self._base = _PREFIX.size + length
        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()}

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

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

    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:
            self._handle = open(self.path, "rb")
        self._handle.seek(self._base + offset)
        payload = self._handle.read(length)
        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 block_size(days, when_consecutive=5):
    """How many days to keep together when splitting.

    Splitting per *sample* would leak almost perfectly -- consecutive frames are four
    minutes apart -- but splitting per *day* only leaks if two days in the set are
    themselves adjacent, and then only across the midnight seam.  When the chosen days
    are spread (60 days over two and a half years puts them ~15 days apart) each day is
    an independent view of the Sun and blocks buy nothing while costing resolution in
    the split.  So: group only when there is something to guard against.
    """
    import datetime

    ordered = sorted(days)
    adjacent = any(
        (datetime.date.fromisoformat(second) - datetime.date.fromisoformat(first)).days <= 1
        for first, second in zip(ordered, ordered[1:])
    )
    if not adjacent:
        return 1
    # Never group so coarsely that three splits become impossible.  With very few days
    # the guard cannot be had at all; `check_split` reports the seams that remain, which
    # is more useful than refusing to split.
    return max(1, min(when_consecutive, len(ordered) // 3))


def split_days(days, block=None, train=8, val=1, test=1):
    """Partition days into disjoint train/val/test sets.

    Assignment is round-robin over blocks so that val and test are *interleaved* through
    the date range rather than carved off one end.  That matters more than it sounds:
    this archive spans the rise of solar cycle 25, so a validation set drawn from one
    stretch measures the model on one level of activity and says nothing about the rest.

    `block` defaults to :func:`block_size`, which uses whole-day granularity unless the
    chosen days are actually adjacent.  Fixing it at 5 with only 12 blocks to hand out
    gave val and test one block each -- five consecutive chosen days, four months of
    calendar -- which is exactly the failure this is meant to avoid.
    """
    days = sorted(days)
    block = block_size(days) if block is None else block
    blocks = [days[i : i + block] for i in range(0, len(days), block)]
    if len(blocks) < 3:
        raise ValueError(
            f"{len(days)} days in blocks of {block} gives {len(blocks)} blocks; "
            "need at least three to make three splits"
        )

    # Take proportions, then place the held-out blocks at evenly spaced positions.
    # Walking a repeating ['train'...,'val','test'] cycle instead looks equivalent and
    # is not: with 8 blocks and a cycle of 10 it never reaches 'val' or 'test' at all,
    # and returns empty held-out sets without complaining.  The failure then surfaces
    # much later, as a training run that cannot find any validation shards.
    total = train + val + test
    count = len(blocks)
    wanted = {"val": max(1, round(count * val / total)),
              "test": max(1, round(count * test / total))}
    if wanted["val"] + wanted["test"] >= count:
        raise ValueError(f"{count} blocks cannot yield train, val and test")

    held = wanted["val"] + wanted["test"]
    stride = count / held
    assignment = {}
    for position in range(held):
        index = min(count - 1, int(position * stride + stride / 2))
        while index in assignment:            # collisions when stride is near 1
            index = (index + 1) % count
        assignment[index] = "val" if position % 2 == 0 else "test"

    out = {"train": [], "val": [], "test": []}
    for number, group in enumerate(blocks):
        out[assignment.get(number, "train")].extend(group)
    return out


def manifest(days_by_split, extra=None):
    """A split manifest plus its digest, so a checkpoint can name the data it saw."""
    payload = {
        "version": VERSION,
        "size": SHARD_SIZE,
        "asinh_scale": ASINH_SCALE,
        "asinh_range": ASINH_RANGE,
        "splits": {name: sorted(days) for name, days in days_by_split.items()},
    }
    if extra:
        payload.update(extra)
    body = json.dumps(payload, sort_keys=True, separators=(",", ":"))
    payload["digest"] = hashlib.sha256(body.encode()).hexdigest()[:16]
    return payload


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