"""Test-only arrival-f GPU owner; CUDA compilation/execution requires admission."""
from dataclasses import replace
import hashlib
from numbers import Integral
from pathlib import Path

import numpy as np
import spacetime_reference as ref

PINS = {"spacetime_reference.py": "1fe13242b9a41062a1227db4a6f9b695165aefab7f4c0b3b2adff0c9023552bf",
        "cpu_reference.py": "bba2a06a9f6517b32079ffff15e4060d8c3b7ed35f16ca134e1877226c9e4e1e",
        "../taylor_green/oracle.py": "83e4d7425d45538278e9fb181af450a36358bd9a5ed6be57ff611dd9ae5eb35b"}


def check_sources():
    for name, expected in PINS.items():
        if hashlib.sha256((Path(__file__).parent/name).read_bytes()).hexdigest() != expected:
            raise RuntimeError(f"reference source changed: {name}")


def validate_config(config):
    cfg = replace(config)
    cfg.validate()
    if (cfg.nx != cfg.ny or not 8 <= cfg.nx <= 64 or cfg.nz != 1
            or cfg.collision != "SRT" or not cfg.is_d2q9
            or cfg.streaming != "double_buffer" or cfg.store_precision != "FP32"
            or cfg.data_type != "FP32" or cfg.ddf_shifting or cfg.les_enabled
            or cfg.velocity != 0 or cfg.wall_slot_width or cfg.wall_normal_velocity):
        raise ValueError("experiment requires periodic unforced square N8..64, unshifted FP32 D2Q9/SRT")
    fine = replace(cfg, nx=2*cfg.nx, ny=2*cfg.ny, viscosity=2*cfg.viscosity)
    fine.validate()
    return cfg, fine


def pack_plan(plan):
    """Pack reviewed coordinate events; no production transfer builder imports."""
    n, nf = plan.n, plan.nf
    groups, receivers = plan.groups, plan.receivers
    data = dict(samples=np.zeros((len(groups), 10), np.uint32),
                origins=np.empty(len(groups), np.uint32), sizes=np.empty(len(groups), np.int32),
                weights=np.zeros((len(groups), 4, 10), np.float64),
                targets=(receivers[:, 0]*n*n+receivers[:, 1]).astype(np.uint32))
    lookup = {tuple(row): i for i, row in enumerate(receivers)}
    packets = [[] for _ in receivers]; incoming = [[], []]
    for i, g in enumerate(groups):
        size = len(g.samples)
        if size not in (6, 10):
            raise ValueError("unexpected interpolation stencil")
        data['samples'][i, :size] = g.q*n*n+g.samples
        data['origins'][i] = g.q*n*n+g.source
        data['sizes'][i] = size
        data['weights'][i, :, :size] = g.weights
        for k, event in enumerate(g.events):
            if event is None:
                packets[lookup[(g.q, g.receiver)]].append((-1, 4*i+k))
            else:
                m, slot = event
                incoming[m].append((g.q*nf*nf+slot, 4*i+k))
    for m, outgoing in enumerate(plan.outgoing):
        for q, source, target in outgoing:
            packets[lookup[(int(q), int(target))]].append((m, int(q*nf*nf+source)))
    if any(len(row) != 4 or len(set(row)) != 4 for row in packets):
        raise ValueError("each receiver needs four distinct quarter-packets")
    if len({packet for row in packets for packet in row}) != 4*len(receivers):
        raise ValueError("packet has multiple receiver owners")
    data['times'] = np.array([[t for t, _ in row] for row in packets], np.int32)
    data['sources'] = np.array([[s for _, s in row] for row in packets], np.uint32)
    for m, rows in enumerate(incoming):
        data[f'incoming{m}'] = np.array(rows, np.uint32)
        if len({slot for slot, _ in rows}) != len(rows):
            raise ValueError("duplicate incoming writer")
    return {key: np.ascontiguousarray(value) for key, value in data.items()}


def table_hash(data):
    digest = hashlib.sha256()
    for key, value in sorted(data.items()):
        digest.update(f"{key}:{value.dtype}:{value.shape}".encode()+value.tobytes())
    return digest.hexdigest()


def rounding_ledger(plan, tables, before, after, g0, predicted, fine_phases, pre, h, registers):
    """Audit actual stored FP32 phases; never call a missing packet 'rounding'."""
    maps, n = plan.maps, plan.n
    f64 = lambda a: np.asarray(a, dtype=np.float64)
    collision = (f64(g0[:, maps.coarse_active])-f64(before.coarse_f[:, maps.coarse_active])).sum(axis=1)
    for f, g in fine_phases:
        collision += .25*(f64(g[:, maps.fine_active])-f64(f[:, maps.fine_active])).sum(axis=1)
    np.testing.assert_array_equal(h, pre.astype(np.float32))
    origins = f64(g0).ravel()[tables['origins']]
    np.testing.assert_allclose(pre.mean(axis=1), origins, rtol=0, atol=2e-14)
    group = np.bincount(tables['origins']//(n*n), weights=f64(h).mean(axis=1)-origins, minlength=9)
    packet = np.empty(tables['times'].shape, np.float64)
    for when, source in ((-1, h), (0, fine_phases[0][1]), (1, fine_phases[1][1])):
        mask = tables['times'] == when
        packet[mask] = f64(source).ravel()[tables['sources'][mask]]
    p = f64(predicted).ravel()[tables['targets']]
    ideal = -p + .25*packet.sum(axis=1)
    reg_error = registers-ideal
    bound = 16*np.finfo(float).eps*(abs(p)+.25*abs(packet).sum(axis=1))
    if np.any(abs(reg_error) > bound):
        raise AssertionError("receiver packet register mismatch (not roundoff)")
    actual = after.coarse_f.ravel()[tables['targets']]
    np.testing.assert_array_equal(actual, (p+registers).astype(np.float32))
    q = tables['targets']//(n*n)
    store_error = np.bincount(q, weights=f64(actual)-(p+registers), minlength=9)
    register_error = np.bincount(q, weights=reg_error, minlength=9)
    # directional() must sum converted arrays, not accumulate in FP32.
    total = lambda s: ref.directional(ref.cpu.State(f64(s.coarse_f), f64(s.fine_f)), maps)
    raw = total(after)-total(before)-collision
    unexplained = raw-group-store_error-register_error
    limit = float(64*np.finfo(float).eps*total(before).sum())
    np.testing.assert_allclose(unexplained, 0, rtol=0, atol=limit)
    return dict(raw=raw.tolist(), physical_collision=collision.tolist(),
                child_rounding=group.tolist(), receiver_rounding=store_error.tolist(),
                register_rounding=register_error.tolist(), unexplained=unexplained.tolist(),
                unexplained_bound=limit, minimum_pre=float(pre.min()), minimum_stored=float(h.min()))


class SpaceTimeGPU:
    """Caller-owned context, one stream, explicit synchronized arrival state."""
    def __init__(self, config, bounds, *, context, mode='spacetime', identity=False):
        check_sources()
        self.configs = validate_config(config)
        if mode not in ('spacetime', 'frozen') or type(identity) is not bool:
            raise ValueError("unsupported interpolation or debug collision mode")
        self.bounds = tuple(bounds)
        if len(self.bounds) != 4 or any(isinstance(v, bool) or not isinstance(v, Integral) for v in self.bounds):
            raise ValueError("rectangle bounds must be four integers")
        self.plan = ref.Plan(config.nx, self.bounds)
        self.tables = pack_plan(self.plan)
        self.ctx, self.mode, self.identity = context, mode, identity
        self.allocations, self.modules, self.build_metadata = [], [], []
        self.status, self.step_count, self.device_bytes = 'new', 0, 0
        self.health_counts = np.zeros(6, np.uint64)
        import pycuda.driver as cuda
        from CelerisLab.cuda import compiler_v2 as compiler
        root = Path(__file__).resolve().parents[3]
        if Path(compiler._KERNEL_DIR).resolve() != root/'src/CelerisLab/lbm/kernels':
            raise RuntimeError("wrong production kernel checkout")
        self.cuda = cuda
        with context.activate():
            self.stream = cuda.Stream()
            try:
                source = Path(__file__).with_suffix('.cu')
                for cfg in self.configs:
                    build = compiler.generate_config(cfg)
                    try:
                        Path(build, source.name).write_bytes(source.read_bytes())
                        ptx = compiler.compile_kernel(context.sm_arch, source=source.name, build_dir=build)
                        self.modules.append(compiler.load_module(ptx))
                        import json
                        metadata = json.loads(Path(ptx+'.json').read_text())
                        metadata['test_entry_sha256'] = hashlib.sha256(source.read_bytes()).hexdigest()
                        self.build_metadata.append(metadata)
                    finally:
                        compiler.cleanup_build_dir(build)
                self.device_tables = {key: self.upload(value) for key, value in self.tables.items()}
                self.active = self.upload(self.plan.maps.fine_active.astype(np.uint8))
                self.coarse = [self.allocate(9*config.nx**2*4) for _ in range(4)]
                self.fine = [self.allocate(9*(2*config.nx)**2*4) for _ in range(3)]
                self.pre = self.allocate(len(self.plan.groups)*4*8)
                self.h = self.allocate(len(self.plan.groups)*4*4)
                self.registers = self.allocate(len(self.plan.receivers)*8)
                self.counts = self.allocate(6*4)
            except BaseException:
                try:
                    self.close()
                except Exception:
                    pass
                raise

    def allocate(self, size):
        pointer = self.cuda.mem_alloc(size)
        self.allocations.append(pointer); self.device_bytes += size
        return pointer

    def upload(self, array):
        pointer = self.allocate(array.nbytes)
        self.cuda.memcpy_htod(pointer, array)
        return pointer

    def read(self, pointer, shape, dtype=np.float32):
        self.stream.synchronize()
        out = np.empty(shape, dtype)
        self.cuda.memcpy_dtoh(out, pointer)
        return out

    def launch(self, level, name, *args, count=None):
        cfg = self.configs[level]
        nt = cfg.threads_per_block
        grid = ((cfg.nx+nt-1)//nt, cfg.ny) if count is None else ((count+nt-1)//nt, 1)
        self.modules[level].get_function(name)(*args, block=(nt,1,1), grid=grid, stream=self.stream)

    def initialize(self, state):
        if (self.status not in ('new', 'ready') or state.phase != 'arrival_f'
                or isinstance(state.tick, bool) or not isinstance(state.tick, Integral) or state.tick != 0):
            raise ValueError("initial synchronized arrival state with tick=0 required")
        arrays = [np.array(a, dtype=np.float32, copy=True, order='C') for a in (state.coarse_f,state.fine_f)]
        for a, mask in zip(arrays, (self.plan.maps.coarse_active,self.plan.maps.fine_active)):
            if a.shape != (9,*mask.shape) or not np.isfinite(a[:,mask]).all() or (a[:,mask] <= 0).any():
                raise ValueError("positive finite leaf populations and exact shapes required")
            a[:,~mask] = np.nan
        with self.ctx.activate():
            try:
                self.cuda.memcpy_htod(self.coarse[0], arrays[0])
                self.cuda.memcpy_htod(self.fine[0], arrays[1])
            except BaseException:
                self.status = 'failed'
                raise
        self.status, self.step_count = 'ready', 0
        self.health_counts.fill(0)

    def download(self):
        if self.status != 'ready':
            raise RuntimeError(f"GPU owner is {self.status}")
        try:
            with self.ctx.activate():
                arrays = [self.read(self.coarse[0], (9,self.plan.n,self.plan.n)),
                          self.read(self.fine[0], (9,self.plan.nf,self.plan.nf))]
        except BaseException:
            self.status = 'failed'
            raise
        for a, mask in zip(arrays, (self.plan.maps.coarse_active,self.plan.maps.fine_active)):
            a[:,~mask] = np.nan
        return ref.cpu.State(*arrays, tick=2*self.step_count)

    def step(self, *, mutation=0, audit=True):
        """One interval; audit=False skips only host ledger snapshots, never health checks."""
        if (self.status != 'ready' or type(mutation) is not int or mutation not in range(5)
                or type(audit) is not bool or (mutation and not audit)):
            raise RuntimeError("ready owner and known test mutation required")
        before = self.download() if audit else None
        n, nf, t = self.plan.n, self.plan.nf, self.device_tables
        ng, nr = len(self.plan.groups), len(self.plan.receivers)
        u32, null = np.uint32, np.uintp(0)
        omega = [np.float32(0 if self.identity else cfg.omega) for cfg in self.configs]
        with self.ctx.activate():
            try:
                self.status = 'advancing'
                self.cuda.memset_d32_async(self.counts, 0, 6, self.stream)
                fc, g0, predicted, gp = self.coarse
                self.launch(0, 'CoarseStart', fc,self.fine[0],g0,omega[0],
                            *(u32(v) for v in self.bounds),self.counts)
                self.launch(0, 'Pull', g0,predicted,null,self.counts)
                self.launch(0, 'Collide', predicted,gp,null,omega[0],self.counts)
                self.launch(0, 'Interpolate', g0,gp,t['samples'],t['origins'],t['sizes'],t['weights'],
                            self.pre,self.h,u32(ng),u32(self.mode=='frozen'),self.counts,count=ng)
                if audit:
                    coarse_g = self.read(g0,(9,n,n)); prediction = self.read(predicted,(9,n,n))
                fine_phases = []
                for m in (0,1):
                    f, out, g = self.fine
                    if audit:
                        old_f = self.read(f,(9,nf,nf))
                    self.launch(1, 'Collide', f,g,self.active,omega[1],self.counts)
                    if audit:
                        fine_phases.append((old_f,self.read(g,(9,nf,nf))))
                    self.launch(0, 'Register', predicted,g,self.h,t['targets'],t['sources'],t['times'],
                                self.registers,u32(nr),np.int32(m),np.int32(mutation),count=nr)
                    rows = self.tables[f'incoming{m}']
                    self.launch(1, 'Inject', g,self.h,t[f'incoming{m}'],u32(len(rows)),u32(mutation==4),count=len(rows))
                    self.launch(1, 'Pull', g,out,self.active,self.counts)
                    self.fine[0], self.fine[1] = out, f
                self.launch(0, 'ApplyRegisters', predicted,self.registers,t['targets'],u32(nr),self.counts,count=nr)
                counts = self.read(self.counts,(6,),np.uint32)
                self.health_counts += counts.astype(np.uint64)
                if counts.any():
                    raise AssertionError(f"boundary/leaf health failed: {counts.tolist()}")
                ledger = dict(boundary_budget_status='not_measured',unexplained=None)
                if audit:
                    coarse_after = self.read(predicted,(9,n,n))
                    coarse_after[:,~self.plan.maps.coarse_active] = np.nan
                    after = ref.cpu.State(coarse_after,self.read(self.fine[0],(9,nf,nf)),tick=2*(self.step_count+1))
                    ledger = rounding_ledger(self.plan,self.tables,before,after,coarse_g,prediction,fine_phases,
                                             self.read(self.pre,(ng,4),np.float64),self.read(self.h,(ng,4)),
                                             self.read(self.registers,(nr,),np.float64))
                    ledger['boundary_budget_status'] = 'measured'
                ledger['health_counts'] = counts.tolist()
                self.coarse[0], self.coarse[2] = predicted, fc
                self.step_count += 1
                self.status = 'ready'
                return ledger
            except BaseException:
                self.status = 'failed'
                raise

    def close(self):
        if self.status == 'closed':
            return
        with self.ctx.activate():
            error = None
            try:
                self.stream.synchronize()
            except Exception as exc:
                error = exc
            failed = []
            for pointer in self.allocations:
                try:
                    pointer.free()
                except Exception as exc:
                    error = error or exc
                    failed.append(pointer)
            self.allocations = failed
            self.status = 'cleanup_failed' if failed else 'closed'
            if not failed:
                self.modules.clear()
            if error:
                raise error

    def __enter__(self):
        return self

    def __exit__(self, kind, value, traceback):
        try:
            self.close()
        except Exception:
            if kind is None:
                raise