"""Critical-line zero ordinates, one cached front door.

Values come from ``mpmath.zetazero`` at the requested working precision.
That is a computation, not a literature pin: a load-bearing claim that
depends on verified-zero ranges must cite the pinned Platt / Platt–Trudgian
sources recorded in canon witnesses, not this module.

The Platt-table reader handles the compressed binary format published by
LMFDB.  Each stored midpoint lies on a grid of spacing ``2^-101`` and has
certified absolute radius ``2^-102``.  The caller remains responsible for
pinning source hashes and the theorem establishing completeness.
"""

from __future__ import annotations

import math
import struct
from dataclasses import dataclass
from functools import lru_cache
from fractions import Fraction
from pathlib import Path
from typing import Iterable, Iterator

import mpmath as mp
from flint import arb

__all__ = [
    "PLATT_ZERO_RADIUS",
    "PlattZeroRecord",
    "iter_platt_zero_enclosures",
    "iter_platt_zero_midpoints",
    "iter_platt_zero_records",
    "zero_ordinate",
    "zero_ordinates",
]

DEFAULT_DPS = 50
PLATT_ZERO_RADIUS = Fraction(1, 1 << 102)
_PLATT_GRID_DENOMINATOR = 1 << 101
_FILE_COUNT = struct.Struct("<Q")
_BLOCK_HEADER = struct.Struct("<ddQQ")
_DELTA = struct.Struct("<QIB")


@dataclass(frozen=True)
class PlattZeroRecord:
    """One zero from a Platt binary table.

    ``block_t0`` is the exact binary64 block origin and ``offset_units`` is
    the cumulative integer offset in units of ``2^-101``.
    """

    index: int
    block_t0: float
    offset_units: int

    def midpoint_float(self) -> float:
        """Return the midpoint as a binary64 exploration value."""
        return self.block_t0 + math.ldexp(self.offset_units, -101)

    def midpoint_fraction(self) -> Fraction:
        """Return the stored midpoint exactly."""
        return Fraction.from_float(self.block_t0) + Fraction(
            self.offset_units, _PLATT_GRID_DENOMINATOR
        )

    def enclosure(self) -> arb:
        """Return the certified ``2^-102``-radius Arb enclosure."""
        midpoint = self.midpoint_fraction()
        midpoint_arb = arb(midpoint.numerator) / midpoint.denominator
        radius_arb = (
            arb(PLATT_ZERO_RADIUS.numerator) / PLATT_ZERO_RADIUS.denominator
        )
        return arb(midpoint_arb, radius_arb)


def _ordered_platt_paths(paths: Iterable[str | Path]) -> list[Path]:
    entries: list[tuple[float, Path]] = []
    for raw_path in paths:
        path = Path(raw_path)
        with path.open("rb") as stream:
            header = stream.read(_FILE_COUNT.size + _BLOCK_HEADER.size)
        if len(header) != _FILE_COUNT.size + _BLOCK_HEADER.size:
            raise ValueError(f"truncated Platt zero file header: {path}")
        block_count = _FILE_COUNT.unpack_from(header, 0)[0]
        if block_count == 0:
            raise ValueError(f"empty Platt zero file: {path}")
        t0 = _BLOCK_HEADER.unpack_from(header, _FILE_COUNT.size)[0]
        entries.append((t0, path))
    entries.sort()
    return [path for _, path in entries]


def iter_platt_zero_records(
    paths: Iterable[str | Path],
    *,
    count: int | None = None,
) -> Iterator[PlattZeroRecord]:
    """Yield contiguous records from one or more Platt binary files.

    Files are ordered by their first block origin, so callers may pass paths
    in arbitrary order.  Headers, block counts, index continuity, and positive
    encoded increments are validated while streaming.
    """
    if count is not None and count < 0:
        raise ValueError(f"count must be nonnegative, got {count}")
    if count == 0:
        return

    yielded = 0
    expected_index = 0
    expected_t0: float | None = None
    for path in _ordered_platt_paths(paths):
        with path.open("rb") as stream:
            raw_count = stream.read(_FILE_COUNT.size)
            block_count = _FILE_COUNT.unpack(raw_count)[0]
            for block_number in range(block_count):
                raw_header = stream.read(_BLOCK_HEADER.size)
                if len(raw_header) != _BLOCK_HEADER.size:
                    raise ValueError(
                        f"truncated block {block_number} header in {path}"
                    )
                t0, t1, nt0, nt1 = _BLOCK_HEADER.unpack(raw_header)
                if not t0 < t1 or nt0 > nt1:
                    raise ValueError(
                        f"invalid block {block_number} header in {path}"
                    )
                if nt0 != expected_index:
                    raise ValueError(
                        f"noncontiguous zero index in {path}: "
                        f"expected {expected_index}, found {nt0}"
                    )
                if expected_t0 is not None and t0 != expected_t0:
                    raise ValueError(
                        f"noncontiguous block origin in {path}: "
                        f"expected {expected_t0}, found {t0}"
                    )

                offset_units = 0
                for index in range(nt0 + 1, nt1 + 1):
                    raw_delta = stream.read(_DELTA.size)
                    if len(raw_delta) != _DELTA.size:
                        raise ValueError(
                            f"truncated zero {index} in {path}"
                        )
                    low, middle, high = _DELTA.unpack(raw_delta)
                    delta = low + (middle << 64) + (high << 96)
                    if delta <= 0:
                        raise ValueError(
                            f"nonpositive encoded gap at zero {index} in {path}"
                        )
                    offset_units += delta
                    yield PlattZeroRecord(index, t0, offset_units)
                    yielded += 1
                    if count is not None and yielded >= count:
                        return

                expected_index = nt1
                expected_t0 = t1


def iter_platt_zero_midpoints(
    paths: Iterable[str | Path],
    *,
    count: int | None = None,
) -> Iterator[float]:
    """Yield binary64 midpoint approximations for exploration."""
    for record in iter_platt_zero_records(paths, count=count):
        yield record.midpoint_float()


def iter_platt_zero_enclosures(
    paths: Iterable[str | Path],
    *,
    count: int | None = None,
) -> Iterator[arb]:
    """Yield certified Arb enclosures from pinned Platt tables."""
    for record in iter_platt_zero_records(paths, count=count):
        yield record.enclosure()


@lru_cache(maxsize=4096)
def _cached_ordinate(index: int, dps: int) -> mp.mpf:
    with mp.workdps(dps + 10):
        value = mp.im(mp.zetazero(index))
    with mp.workdps(dps):
        return +value


def zero_ordinate(index: int, *, dps: int = DEFAULT_DPS) -> mp.mpf:
    """The imaginary part of the ``index``-th nontrivial zero (1-indexed)."""
    if index < 1:
        raise ValueError(f"zero index must be positive, got {index}")
    return _cached_ordinate(index, dps)


def zero_ordinates(count: int, *, dps: int = DEFAULT_DPS) -> tuple[mp.mpf, ...]:
    """The first ``count`` zero ordinates in increasing order."""
    return tuple(zero_ordinate(index, dps=dps) for index in range(1, count + 1))
