"""Experimental FP64 BVP coupling: one periodic rectangle, D2Q9, ratio two.

Guzik et al., https://www.osti.gov/servlets/purl/1226215, sections 3.2--3.5.
Coordinate-generated stream events; CPUREF collision only, no production imports.
This is a mathematical prototype, not a certified accuracy implementation.
"""
from dataclasses import dataclass
from functools import lru_cache
import argparse
import hashlib
import json
from pathlib import Path
import resource
import sys
import subprocess
import time

import numpy as np
import cpu_reference as cpu


class Spec(cpu.Spec):
    @property
    def n_fine(self):
        return 2 * self.n_coarse


def geometry(n, bounds):
    x0, y0, x1, y1 = bounds
    if not (2 <= x0 < x1 <= n - 2 and 2 <= y0 < y1 <= n - 2):
        raise ValueError("one interior rectangle with two coarse cells of clearance required")
    y, x = np.indices((n, n))
    covered = (x >= x0) & (x < x1) & (y >= y0) & (y < y1)
    y, x = np.indices((2*n, 2*n))
    active = (x >= 2*x0) & (x < 2*x1) & (y >= 2*y0) & (y < 2*y1)
    inner = (x >= 2*x0-2) & (x < 2*x1+2) & (y >= 2*y0-2) & (y < 2*y1+2) & ~active
    rows, links = [], []
    for ky, kx in zip(*np.nonzero(~covered)):
        for q, (cy, cx) in enumerate(zip(cpu.CY, cpu.CX)):
            if any(active[2*ky+dy-cy, 2*kx+dx-cx]
                   for dy in (0, 1) for dx in (0, 1)
                   if 0 <= 2*ky+dy-cy < 2*n and 0 <= 2*kx+dx-cx < 2*n):
                rows.append((q, ky, kx))
                sy, sx = (ky-cy) % n, (kx-cx) % n
                if not covered[sy, sx]:
                    links.append(cpu.CornerLink(f"{kx}_{ky}_{q}", q, (sx, sy), (kx, ky)))
    kind = np.where(active, 0, np.where(inner, 1, 2))
    return cpu.Maps(~covered, covered, active, inner, kind == 2, kind,
                    np.arange(2*n)//2, (np.arange(2*n)+.5)/2,
                    np.array(rows, dtype=int), tuple(links))


@lru_cache(None)
def interpolation(kind, points):
    """Return data coordinates, evaluation matrix, and constrained coefficient map."""
    if kind == "face":
        samples = np.array([(s, t) for t in (0., 1.) for s in (-1., 0., 1.)])
        a = np.array([[1, s, t, s*t] for s, t in samples])
        e = np.array([[1, s, t, s*t] for s, t in points])
        b = e.mean(axis=0, keepdims=True)
        d = np.zeros((1, 6)); d[0, 1] = 1
    else:
        cross = ((0, 0), (-1, 0), (1, 0), (0, -1), (0, 1))
        samples = np.array([(x, y, t) for t in (0., 1.) for x, y in cross])
        a = np.array([[1, x, y, t] for x, y, t in samples])
        e = np.array([[1, x, y, t] for x, y, t in points])
        b = np.stack((e.mean(axis=0), [0, 0, 0, 1]))
        d = np.zeros((2, 10)); d[0, 0] = 1; d[1, 0] = -1; d[1, 5] = 1
    kkt = np.block([[a.T @ a, b.T], [b, np.zeros((len(b), len(b)))]])
    coeff = np.linalg.solve(kkt, np.vstack((a.T, d)))[:4]
    if not np.allclose(b @ coeff, d, rtol=0, atol=2e-14):
        raise AssertionError("invalid constrained interpolation matrix")
    return samples, a, e, coeff


@dataclass
class Group:
    q: int
    source: int
    receiver: int
    kind: str
    points: tuple
    samples: np.ndarray
    weights: np.ndarray
    events: tuple  # (substep, fine ghost slot), or None for a leftover


class Plan:
    def __init__(self, n, bounds):
        self.n, self.nf = n, 2*n
        self.maps = geometry(n, bounds)
        active = self.maps.fine_active
        lo, hi = np.array(bounds[1::-1]), np.array(bounds[3:1:-1])
        self.groups = []
        for iy, ix in zip(*np.nonzero(self.maps.coarse_active)):
            origin = np.array((iy, ix))
            for q in range(1, 9):
                c = np.array((cpu.CY[q], cpu.CX[q]))
                starts = [2*origin + (dy, dx) for dy in (0, 1) for dx in (0, 1)]
                events = []
                for p in starts:
                    event = None
                    for m in (0, 1):
                        old, new = p+m*c, p+(m+1)*c
                        inside = lambda z: bool(np.all(z >= 0) and np.all(z < 2*n) and active[tuple(z)])
                        if not inside(old) and inside(new):
                            event = (m, int(old[0]*2*n+old[1]))
                            break
                    events.append(event)
                if all(event is None for event in events):
                    continue
                axes = [j for j in range(2) if origin[j] < lo[j] or origin[j] >= hi[j]]
                kind = "corner" if len(axes) == 2 else "face"
                points = []
                for p in starts:
                    x = (p+.5)/2
                    times = [(float(lo[j] if c[j] > 0 else hi[j])-x[j])/c[j] for j in axes]
                    time = max(times)
                    offset = x+c*time-(origin+.5+c*.5)
                    points.append((float(offset[1]), float(offset[0]), float(time-.5))
                                  if kind == "corner" else (float(offset[1-axes[0]]), float(time-.5)))
                points = tuple(points)
                samples, _, e, coeff = interpolation(kind, points)
                indices = []
                for row in samples:
                    shift = np.array((row[1], row[0]), dtype=int) if kind == "corner" else np.eye(2, dtype=int)[1-axes[0]]*int(row[0])
                    pos = (origin+shift) % n
                    indices.append(int(pos[0]*n+pos[1]))
                target = (origin+c) % n
                self.groups.append(Group(q, int(iy*n+ix), int(target[0]*n+target[1]), kind,
                                         points, np.array(indices), e @ coeff, tuple(events)))
        self.outgoing = []
        # m is zero-based: a substep-0 exit travels two fine spacings by
        # synchronization; a substep-1 exit travels one.
        for m in (0, 1):
            records = []
            for q in range(1, 9):
                c = np.array((cpu.CY[q], cpu.CX[q]))
                for y, x in zip(*np.nonzero(active)):
                    p = np.array((y, x)); endpoint = p+c
                    if not active[tuple(endpoint)]:
                        target = (p+(2-m)*c)//2
                        assert self.maps.coarse_active[tuple(target)]
                        records.append((q, int(y*2*n+x), int(target[0]*n+target[1])))
            self.outgoing.append(np.array(records, dtype=int))
        incoming = [(m, g.q, slot) for g in self.groups for event in g.events
                    if event is not None for m, slot in [event]]
        assert len(incoming) == len(set(incoming)), "duplicate incoming writer"
        expected = {(m, int(q), int((y-cpu.CY[q])*2*n+x-cpu.CX[q]))
                    for m in (0, 1) for q in range(1, 9)
                    for y, x in zip(*np.nonzero(active)) if not active[y-cpu.CY[q], x-cpu.CX[q]]}
        assert set(incoming) == expected, "incoming coverage mismatch"
        receivers = {(int(q), int(target)) for out in self.outgoing for q, _, target in out}
        receivers |= {(g.q, g.receiver) for g in self.groups if None in g.events}
        assert receivers == {(int(q), int(y*n+x)) for q, y, x in self.maps.coalesce}
        self.receivers = np.array(sorted(receivers), dtype=int)


def directional(state, maps):
    return state.coarse_f[:, maps.coarse_active].sum(axis=1) + .25*state.fine_f[:, maps.fine_active].sum(axis=1)


def advance(spec, plan, state, mode="spacetime", mutation=None):
    """One arrival-to-arrival interval. 'frozen' disables space AND time variation."""
    if mode not in ("frozen", "spacetime"):
        raise ValueError(mode)
    n, nf, maps = plan.n, plan.nf, plan.maps
    if state.phase != "arrival_f" or state.tick % 2:
        raise ValueError("synchronized arrival state required")
    fc, ff = state.coarse_f.copy(), state.fine_f.copy()
    covered_mean = ff.reshape(9, n, 2, n, 2).mean(axis=(2, 4))
    fc[:, maps.coarse_covered] = covered_mean[:, maps.coarse_covered]
    all_cells = np.ones((n, n), dtype=bool)
    def collide(f, mask, tau):
        return cpu._collide(spec, f, np.isfinite(f), tau, mask, np.zeros(mask.shape, dtype=int))
    gc = collide(fc, all_cells, spec.tau_coarse)
    collision = (gc[:, maps.coarse_active]-fc[:, maps.coarse_active]).sum(axis=1)
    predicted = cpu._stream_periodic(gc)
    gp = collide(predicted, all_cells, spec.tau_coarse) if mode == "spacetime" else gc
    registers = np.zeros((9, n*n))
    q, target = plan.receivers.T
    registers[q, target] = -predicted.reshape(9, -1)[q, target]
    incoming = [[], []]
    group_error = 0.
    min_boundary = float("inf")
    negative_boundary = nonfinite_boundary = 0
    dropped = False
    for group in plan.groups:
        q = group.q
        count = len(group.samples)//2
        data = np.concatenate((gc.reshape(9, -1)[q, group.samples[:count]], gp.reshape(9, -1)[q, group.samples[count:]]))
        value = (group.weights @ data if mode == "spacetime"
                 else np.full(4, gc.reshape(9, -1)[q, group.source]))
        min_boundary = min(min_boundary, float(value.min()))
        negative_boundary += int(np.count_nonzero(value < 0))
        nonfinite_boundary += int(np.count_nonzero(~np.isfinite(value)))
        group_error = max(group_error, abs(value.mean()-gc.reshape(9, -1)[q, group.source]))
        for k, event in enumerate(group.events):
            if event is None:
                if mutation == "omit_leftover" and not dropped:
                    dropped = True
                else:
                    registers[q, group.receiver] += .25*value[k]
            else:
                m, slot = event
                incoming[m].append((q, slot, value[k]))
    for m in (0, 1):
        gf = collide(ff, maps.fine_active, spec.tau_fine)
        collision += .25*(gf[:, maps.fine_active]-ff[:, maps.fine_active]).sum(axis=1)
        out = plan.outgoing[m]
        if mutation != f"omit_outgoing_{m}":
            np.add.at(registers, (out[:, 0], out[:, 2]), .25*gf.reshape(9, -1)[out[:, 0], out[:, 1]])
        for q, slot, value in incoming[m]:
            gf.reshape(9, -1)[q, slot] = value
        fresh = np.full_like(ff, np.nan)
        for q in range(9):
            moved = cpu._shift(gf[q], int(cpu.CY[q]), int(cpu.CX[q]), np.nan)
            fresh[q, maps.fine_active] = moved[maps.fine_active]
        if not np.isfinite(fresh[:, maps.fine_active]).all():
            raise AssertionError("invalid fine incoming population")
        ff = fresh
    final_c = predicted + registers.reshape(9, n, n)
    final_c[:, maps.coarse_covered] = np.nan
    final = cpu.State(final_c, ff, tick=state.tick+2)
    residual = directional(final, maps)-directional(state, maps)-collision
    minimum = min(float(final_c[:, maps.coarse_active].min()), float(ff[:, maps.fine_active].min()))
    return final, dict(directional_residual=residual, collision=collision,
                       group_error=group_error, min_boundary=min_boundary, min_population=minimum,
                       negative_boundary_count=negative_boundary, nonfinite_boundary_count=nonfinite_boundary)


def write_receipt(path, result):
    """Normalize NumPy scalars at the JSON boundary; reject other unknown types."""
    def scalar(value):
        if isinstance(value, np.generic):
            return value.item()
        raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")
    path.write_text(json.dumps(result, indent=2, allow_nan=False, default=scalar)+"\n")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--n64", action="store_true", help="explicitly run the frozen 1024-interval experiment")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--max-seconds",type=float,default=180.,help="wall-time cap, including plan generation")
    parser.add_argument("--max-rss-mib",type=float,default=256.,help="Linux peak-RSS screening cap")
    args = parser.parse_args()
    if not args.n64:
        parser.error("full trajectory requires --n64")
    if args.output.exists():
        parser.error("output exists; preserve prior evidence and choose a new path")
    if not (np.isfinite(args.max_seconds) and args.max_seconds>0 and np.isfinite(args.max_rss_mib) and args.max_rss_mib>0):
        parser.error("positive finite resource caps required")
    started = time.monotonic()
    sys.path.insert(0, str(Path(__file__).parents[1]/"taylor_green"))
    import oracle as tg
    plan = Plan(64, (8, 8, 56, 48))
    spec = Spec(n_coarse=64, patch_lo=8, halo=16, tau_coarse=.5576, tau_fine=.6152, collision="srt")
    arrays = [tg.initial((64,64), 0, 64, .0192), tg.initial((128,128), 1, 64, .0192)]
    for a, mask in zip(arrays, (plan.maps.coarse_active, plan.maps.fine_active)):
        a[:, ~mask] = np.nan
    initial = cpu.State(*arrays)
    mass0 = tg.metrics(arrays, (plan.maps.coarse_active, plan.maps.fine_active), 0, 64, .0192)["mass"]
    metrics = ("velocity_l2", "pressure_l2", "energy_relative")
    frozen_baseline = dict(zip(metrics, (.01131134574353476, .03592961319244149, .012514764591008386)))
    budget_bound = float(64*np.finfo(float).eps*mass0)
    root = Path(__file__).resolve().parents[3]
    source = dict(commit=subprocess.check_output(["git", "rev-parse", "HEAD"],cwd=root,text=True).strip(),
                  dirty=bool(subprocess.check_output(["git","status","--porcelain"],cwd=root,text=True)),
                  hashes={p.name:hashlib.sha256(p.read_bytes()).hexdigest()
                          for p in (Path(__file__),Path(cpu.__file__),Path(tg.__file__))},numpy=np.__version__)
    result = dict(scope="experimental CPU FP64; no GPU/accuracy certification", source=source,
                  screen=dict(velocity_factor=.8, pressure_factor=1., energy_factor=1.,
                              directional_step_bound=budget_bound, mass_relative=1e-11,
                              momentum_absolute=1e-9, baseline_metric_atol=1e-10),
                  limits=dict(seconds=args.max_seconds,rss_mib=args.max_rss_mib),runs=[],completed=False)
    def save():
        result["resources"] = dict(seconds=time.monotonic()-started,peak_rss_mib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024)
        write_receipt(args.output, result)
    save()
    try:
        for mode in ("constant_ivp", "spacetime"):
            measured = mode == "spacetime"
            state = initial; record = dict(mode=mode,rows=[],
                                          boundary_budget_status="measured" if measured else "not_measured",
                                          max_directional_residual=0. if measured else None,
                                          minimum_boundary=None,negative_boundary_count=0 if measured else None,
                                          nonfinite_boundary_count=0 if measured else None)
            result["runs"].append(record)
            for k in range(1, 1025):
                if time.monotonic()-started>args.max_seconds or resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024>args.max_rss_mib:
                    raise RuntimeError("experiment resource cap exceeded")
                if mode == "constant_ivp":
                    state,_ = cpu.advance_coarse_once(spec,plan.maps,state)
                else:
                    state,diagnostic = advance(spec,plan,state)
                    error = float(np.max(np.abs(diagnostic["directional_residual"])))
                    record["max_directional_residual"] = max(record["max_directional_residual"],error)
                    if np.isfinite(diagnostic["min_boundary"]):
                        record["minimum_boundary"] = (diagnostic["min_boundary"] if record["minimum_boundary"] is None else min(record["minimum_boundary"],diagnostic["min_boundary"]))
                    for name in ("negative_boundary_count","nonfinite_boundary_count"):
                        record[name] += diagnostic[name]
                    if (not np.isfinite(error) or error>budget_bound or diagnostic["group_error"]>2e-14
                            or diagnostic["negative_boundary_count"] or diagnostic["nonfinite_boundary_count"]
                            or not np.isfinite(diagnostic["min_boundary"]) or diagnostic["min_boundary"]<=0
                            or not np.isfinite(diagnostic["min_population"]) or diagnostic["min_population"]<=0):
                        record["failed_diagnostic"] = {name: repr(value) for name,value in diagnostic.items()}
                        raise AssertionError("candidate directional/health invariant failed")
                if k in (256,512,768,1024):
                    row = tg.metrics((state.coarse_f,state.fine_f),(plan.maps.coarse_active,plan.maps.fine_active),k,64,.0192,mass0)
                    row["momentum_change"] = (cpu.MOMENTS.T@(directional(state,plan.maps)-directional(initial,plan.maps)))[1:].tolist()
                    record["rows"].append(row)
                    save(); print(mode,k,json.dumps(row),flush=True)
            record["worst"] = {key:max(row[key] for row in record["rows"]) for key in metrics}
            if mode == "constant_ivp":
                for key,value in frozen_baseline.items():
                    if abs(record["worst"][key]-value)>1e-10:
                        raise AssertionError(f"frozen baseline mismatch: {key}")
        baseline,candidate = result["runs"]
        result["checks"] = {key:candidate["worst"][key]<=factor*baseline["worst"][key]
                            for key,factor in zip(metrics,(.8,1.,1.))}
        result["checks"]["mass_momentum"] = all(row["mass_relative"]<=1e-11 and max(map(abs,row["momentum_change"]))<=1e-9 for row in candidate["rows"])
        result["completed"] = True
        result["experimental_pass"] = all(result["checks"].values())
        save()
        return 0 if result["experimental_pass"] else 1
    except Exception as exc:
        result["failure"] = dict(mode=mode,step=k,error=repr(exc))
        save()
        raise


if __name__ == "__main__":
    raise SystemExit(main())