"""Independent coordinate oracle, exact kernel host shim and closed pair budgets."""
import ctypes
from contextlib import nullcontext
from copy import deepcopy
from dataclasses import replace
import hashlib
import json
from pathlib import Path
import subprocess
import tempfile
import unittest
from unittest.mock import patch

import numpy as np
import run_case as rc
import shear_axis_square as axis
import shear_centered as candidate
import test_shear_wave as legacy
from test_axis_square import masks
from CelerisLab.lbm.hierarchy import build_hierarchy
from CelerisLab.lbm.blocks import BlockLayout2D
from CelerisLab.lbm.refinement_tables import projected_explosion_tables
from CelerisLab.lbm.refinement import ACTIVE, OUTER_GHOST
from tests.unit import test_projected_explosion as old

REPO = Path(__file__).resolve().parents[3]
SOURCE = REPO/'src/CelerisLab/lbm/kernels/blocks.cu'


def oracle(c, pair, centered=True):
    """Uses only dense coordinates/legal donors, never production neighbor/mask arrays."""
    out = np.empty((9, len(pair.ghost_fine)))
    legal = set(map(int, pair.ghost_coarse))
    for r, (donor, target) in enumerate(zip(pair.ghost_coarse, pair.ghost_fine)):
        y, x = divmod(int(donor), 128); fy, fx = divmod(int(target), 256)
        delta = np.array([(fx+.5)/2-x-.5, (fy+.5)/2-y-.5])
        usable = [(128*(y-dy)+x-dx in legal and 128*(y+dy)+x+dx in legal)
                  for dx,dy in [(1,0),(0,1)]]
        straight = ((x in (30,31,96,97) and 36 <= y < 92 and usable == [False,True])
                    or (y in (30,31,96,97) and 36 <= x < 92 and usable == [True,False]))
        for q, e in enumerate(old.E):
            grad = np.array([(c[q,y+dy,x+dx]-c[q,y-dy,x-dx])/2 if ok else 0.
                             for (dx,dy),ok in zip([(1,0),(0,1)],usable)])
            if q: grad -= e*(grad@e)/(e@e)
            out[q,r] = c[q,y,x] + (2 if centered and straight and np.all(e) else 1)*(grad@delta)
    return out


class CenteredTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.h = build_hierarchy(rc.config(128,100), axis.REFINEMENTS)
        cls.p = cls.h.pairs[0]
        cls.mask = candidate.routing(cls.h)
        cls.layout = BlockLayout2D.build(128,128,8,8,periodic=(True,True))
        cls.ids, cls.offsets = projected_explosion_tables(cls.p,cls.layout,cls.h.kinds[0])
        cls.neighbors = old.decode(cls.layout,cls.ids)

    def test_mask_constants_quartets_and_transpose(self):
        self.assertEqual(candidate.metadata(self.mask), dict(variant=candidate.VARIANT,
            mask_sha256='112de63ce9bafd2e205a9e7029ef6ce5683fe14407b6ac04444579c84960b7ea',
            ghost_rows=2112, selected_rows=1792,
            rule='diagonal projected correction x2; audited straight faces only'))
        c=np.random.default_rng(47).normal(0,.001,(9,128,128))
        actual=oracle(c,self.p); baseline=oracle(c,self.p,False)
        base,g=old.reconstruct(c,self.p,self.neighbors,self.offsets)
        g[np.ix_(np.flatnonzero(np.all(old.E,axis=1)),np.flatnonzero(self.mask),[0,1])]*=2
        expected=c.reshape(9,-1)[:,self.p.ghost_coarse]+np.einsum('qra,ra->qr',g,self.offsets)
        np.testing.assert_allclose(actual,expected,atol=1e-18,rtol=0)
        np.testing.assert_array_equal(actual[:,~self.mask.astype(bool)],baseline[:,~self.mask.astype(bool)])
        for d in np.unique(self.p.ghost_coarse):
            rows=self.p.ghost_coarse==d
            np.testing.assert_allclose(actual[:,rows].mean(axis=1),c.reshape(9,-1)[:,d],atol=1e-18)
        const=np.broadcast_to(old.W[:,None,None],c.shape)
        np.testing.assert_array_equal(oracle(const,self.p),np.broadcast_to(old.W[:,None],actual.shape))
        qswap=[list(map(tuple,old.E)).index(tuple(e[::-1])) for e in old.E]
        index={int(t):r for r,t in enumerate(self.p.ghost_fine)}
        swap=[index[int(t%256)*256+int(t//256)] for t in self.p.ghost_fine]
        np.testing.assert_array_equal(self.mask,self.mask[swap])
        np.testing.assert_allclose(oracle(c[qswap].transpose(0,2,1),self.p),actual[qswap][:,swap],atol=1e-18)
        qreflect=[list(map(tuple,old.E)).index(tuple((-e[0],e[1]))) for e in old.E]
        reflect=[index[int(t//256)*256+255-int(t%256)] for t in self.p.ghost_fine]
        np.testing.assert_array_equal(self.mask,self.mask[reflect])
        np.testing.assert_allclose(oracle(c[qreflect,:,::-1],self.p),actual[qreflect][:,reflect],atol=1e-18)
        # Central dependency corridors remain in a constant mask regime beyond radius8*2.5.
        for name,(x,y) in axis.GROUPS.items():
            tangent = y if name.endswith(('left','right')) else x
            self.assertGreater(min(tangent+.25-36,92-(tangent+.75)),20)

    def test_actual_generated_kernel_shifted_unshifted(self):
        child=self.p.fine_layout((8,4))  # unequal owner strides, unused slots poisoned
        patched=candidate.patch_source(SOURCE.read_text(),self.mask)
        kernel=patched.split('__device__ __constant__ unsigned char centered_face_mask',1)[1]
        kernel='unsigned char centered_face_mask'+kernel.split('// Targets are receiver-local SOURCE slots',1)[0]
        kernel=kernel.replace('extern "C" __global__','extern "C"')
        code='#include <cstdint>\nusing fpxx=float;\n'
        code+=f'constexpr unsigned NQ=9,TOTAL_CELLS={np.prod(child.local_shape)};\n'
        code+='struct Id{unsigned x;}; Id blockIdx{0},blockDim{1},threadIdx{0};\n'
        code+='unsigned index_f(unsigned k,unsigned q){return q*TOTAL_CELLS+k;}\n'
        code+='int d_cx[]={0,1,-1,0,0,1,-1,1,-1},d_cy[]={0,0,0,1,-1,1,-1,-1,1};\n'+kernel
        original=SOURCE.read_text().split('extern "C" __global__ void BlockExplodeProjected(',1)[1].split('// Targets are receiver-local SOURCE slots',1)[0]
        code+='void Original('+original
        code+='extern "C" void invoke(float*c,float*f,uint64_t*t,uint64_t*d,unsigned n,unsigned long l,uint64_t*b,float*o){for(unsigned r=0;r<n;r++){threadIdx.x=r;BlockExplodeProjected(c,f,t,d,n,l,b,o);}}'
        code+=code[code.index('extern "C" void invoke'):].replace('void invoke(', 'void invoke_old(').replace('BlockExplodeProjected(c,f,','Original(c,f,')
        with tempfile.TemporaryDirectory() as tmp:
            src=Path(tmp)/'kernel.cpp'; lib=Path(tmp)/'kernel.so'; src.write_text(code)
            subprocess.run(['g++','-std=c++17','-O2','-shared','-fPIC',str(src),'-o',str(lib)],check=True,capture_output=True)
            library=ctypes.CDLL(str(lib));fn=library.invoke;fn_old=library.invoke_old
            ptr=lambda a:a.ctypes.data_as(ctypes.c_void_p)
            targets=np.ascontiguousarray(child.owner.ravel()[self.p.ghost_fine],np.uint64)
            donors=np.ascontiguousarray(self.layout.owner.ravel()[self.p.ghost_coarse],np.uint64)
            for shifted in (True,False):
                c=np.random.default_rng(8).normal(0,.001,(9,128,128))
                if not shifted:c+=old.W[:,None,None]
                c=c.astype(np.float32); c[:,self.p.covered]=np.nan
                packed=np.full(self.layout.nblocks*9*np.prod(self.layout.local_shape),np.nan,np.float32)
                for q in range(9):packed[self.layout.slot(self.layout.owner.ravel(),q)]=c[q].ravel()
                fine=np.full(child.nblocks*9*np.prod(child.local_shape),np.nan,np.float32)
                fn(ptr(packed),ptr(fine),ptr(targets),ptr(donors),ctypes.c_uint(len(targets)),ctypes.c_ulong(np.prod(self.layout.local_shape)),ptr(self.ids),ptr(self.offsets))
                original=np.full_like(fine,np.nan)
                fn_old(ptr(packed),ptr(original),ptr(targets),ptr(donors),ctypes.c_uint(len(targets)),ctypes.c_ulong(np.prod(self.layout.local_shape)),ptr(self.ids),ptr(self.offsets))
                output=np.array([fine[child.slot(targets,q)] for q in range(9)])
                unchanged=np.array([original[child.slot(targets,q)] for q in range(9)])
                np.testing.assert_array_equal(output[:,~self.mask.astype(bool)],unchanged[:,~self.mask.astype(bool)])
                np.testing.assert_array_equal(output[:5],unchanged[:5])
                np.testing.assert_allclose(output,oracle(c.astype(float),self.p),atol=6e-8,rtol=2e-7)

    def test_mixed_mask_full_pair_mass_momentum(self):
        rng=np.random.default_rng(12); p=self.p
        c=old.W[:,None,None]*(1+rng.normal(0,.002,(9,128,128)))
        f=old.W[:,None,None]*(1+rng.normal(0,.002,(9,256,256)))
        start=old.totals(c,f,p); c[:,p.covered]=np.nan
        for _ in range(2):
            f.reshape(9,-1)[:,p.ghost_fine]=oracle(c,p)
            for _ in range(2):
                f=np.array([np.roll(f[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(old.E)])
                f[:,p.fine_kind==OUTER_GHOST]=np.nan
                f[:,p.fine_kind==ACTIVE]=old.collide(f[:,p.fine_kind==ACTIVE],p.tau_fine)
            out=np.array([np.roll(c[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(old.E)])
            out.reshape(-1)[p.coalesce_coarse]=f.reshape(-1)[p.coalesce_fine].mean(axis=1)
            out[:,~p.covered]=old.collide(out[:,~p.covered],p.tau_coarse)
            out[:,p.covered]=np.nan; c=out
            np.testing.assert_allclose(old.totals(c,f,p),start,rtol=0,atol=4e-10)

    def test_private_build_identity_cleanup_and_old_path(self):
        from CelerisLab.cuda import compiler_v2 as compiler
        from assess_projected_pairs import BUILD_SOURCES
        original=compiler.generate_config; report={}
        with candidate.private_build(self.h,report,candidate.VARIANT):
            for level,cfg in enumerate(self.h.configs):
                cfg=replace(cfg,nx=10,ny=10);cfg.validate(); build=compiler.generate_config(cfg)
                try:
                    meta=json.loads((Path(build)/'.celeris-build').read_text())
                    self.assertEqual(meta['source_sha256'],candidate.source_digest(build))
                    self.assertEqual(meta['source_sha256'],candidate.BUILD_SHA256 if level else BUILD_SOURCES[128])
                    if level:self.assertEqual(meta['experimental_reconstruction'],candidate.metadata(self.mask))
                finally:compiler.cleanup_build_dir(build)
        self.assertIs(compiler.generate_config,original)
        with candidate.private_build(None,{},'current'):
            self.assertIs(compiler.generate_config,original)
        self.assertEqual(hashlib.sha256(SOURCE.read_bytes()).hexdigest(),'668e5cf14e88bdcfa281c2fa6c75c0f8f8cdbbfb8ba2ca03064ae02ef7ad0057')

    def mock_records(self):
        import argparse
        helper=legacy.ProbeTraceTests();helper.setUp()
        parse=argparse.ArgumentParser.parse_args
        records=[]
        for n,depth,variant in candidate.CASES:
            scheme='constant' if depth==1 else 'projected_linear'
            def parse_extra(parser,*a,**k):
                result=parse(parser,*a,**k);result.test_reconstruction=variant;return result
            def arrays(m,n,t):return [axis.initial(v.shape,l,n,.03*n/100) for l,v in enumerate(m)]
            with patch.object(argparse.ArgumentParser,'parse_args',parse_extra), \
                 patch.object(helper,'exact_arrays',side_effect=arrays), \
                 patch.object(candidate,'private_build',return_value=nullcontext()) as private:
                r,counts,_=helper.run_mocked(n,depth,candidate.DIAGNOSTIC,scheme,masks(n,depth),initializer=axis.initial)
            self.assertEqual(r['test_reconstruction'],variant)
            self.assertEqual(r['probe_trace']['steps'],list(range(9)))
            self.assertEqual([v['step'] for v in r['rows']],[0,8*(n//128)])
            self.assertLess(r['observed_shifted_max'],candidate.POP_ENVELOPE)
            self.assertEqual(counts[-1],8*(n//128));self.assertEqual(private.call_args.args[2],variant)
            records.append(r)
        return records

    def test_four_mock_runs_admission_and_old_diagnostic_unchanged(self):
        self.assertEqual(len(self.mock_records()),4)
        self.assertEqual(hashlib.sha256(Path(candidate.__file__).read_bytes()).hexdigest(),rc.SHEAR_CENTERED_SHA256)
        from tests.validation.multigrid import check_projected_pairs as gate
        with patch.object(gate,'preflight',return_value=rc.NUMERICAL_REVISION), \
             patch.object(rc,'provenance',return_value=dict(source_tree_sha256=rc.SOURCE_SHA256)):
            rc.verify_source(REPO,'r','shear_centered',probes=True)
            with patch.object(candidate,'__file__','/stale/shear_centered.py'), self.assertRaises(RuntimeError):
                rc.verify_source(REPO,'r','shear_centered',probes=True)
        import argparse, io, sys
        from contextlib import redirect_stdout
        with tempfile.TemporaryDirectory() as tmp:
            base=['run_case.py','--repo',str(REPO),'--expected-revision','f'*40,
                  '--output',str(Path(tmp)/'r.json'),'--re','100','--n','128','--depth','2',
                  '--flow','shear_wave','--probe-trace','--explosion','projected_linear','--dry-run']
            cases=[(['--diagnostic',candidate.DIAGNOSTIC],False),
                   (['--diagnostic',rc.AXIS_SQUARE,'--test-reconstruction',candidate.VARIANT],False),
                   (['--diagnostic',candidate.DIAGNOSTIC,'--test-reconstruction',candidate.VARIANT,'--depth','1'],False),
                   (['--diagnostic',candidate.DIAGNOSTIC,'--test-reconstruction',candidate.VARIANT],True)]
            for extra,valid in cases:
                with patch.object(sys,'argv',base+extra), patch.object(rc,'verify_source',return_value={}) as verify, \
                     patch.object(rc,'runtime_source',return_value={}), redirect_stdout(io.StringIO()), patch('sys.stderr',io.StringIO()):
                    if valid:self.assertEqual(rc.main(),0)
                    else:
                        with self.assertRaises(SystemExit):rc.main()
                        verify.assert_not_called()

    def test_frozen_comparison_gates_and_corruptions(self):
        import test_projected_pairs as fixture
        records=self.mock_records(); source=fixture.source()
        for r in records:
            r.update(source=deepcopy(source),runtime_identity=deepcopy(fixture.RUNTIME),
                     runtime_source=rc.runtime_source(REPO),cuda_visible_devices='0',
                     actual_build_sidecars=[fixture.build_metadata(r['n']*2**l) for l in range(r['depth'])])
        records[3]['reconstruction']=candidate.metadata(self.mask)
        records[3]['actual_build_sidecars'][1].update(source_sha256=candidate.BUILD_SHA256,
                                                    experimental_reconstruction=candidate.metadata(self.mask))
        def change(r,step,group,component,value):
            for i,p in enumerate(r['probe_trace']['probes']):
                if p['group']==group:r['probe_trace']['values'][step][3*i+component]=value
        for r,bottom,top in [(records[2],-1e-5,1e-6),(records[3],-1e-6,8e-7)]:
            change(r,1,'inner_bottom',1,bottom);change(r,1,'inner_top',1,top)
        for r in records:r["probe_trace"]["bytes"]=axis.validate(r["probe_trace"])
        result=candidate.compare(records,source)
        self.assertTrue(result['advance_to_review']);self.assertFalse(result['qualified'])
        np.testing.assert_allclose(result['moment_deadband'],[1.1925e-6,7.95e-7,7.95e-7],rtol=2e-16,atol=0)
        for group,step,value,check in [('inner_bottom',1,-9.8e-6,'bottom_step1_clear'),
                ('inner_top',1,4e-6,'top_step1_nonworse'),('outer_top',8,4e-6,'top_step8_nonworse'),
                ('inner_left',1,4e-6,'vertical_null')]:
            bad=deepcopy(records);change(bad[3],step,group,1,value)
            bad[3]['probe_trace']['bytes']=axis.validate(bad[3]['probe_trace'])
            self.assertFalse(candidate.compare(bad,source)['checks'][check])
        for mutation in [lambda r:r['source'].update(revision='wrong'),
                         lambda r:r['runtime_identity'].update(device='wrong'),
                         lambda r:r.update(viscosity=.1),lambda r:r.update(refinements=[]),
                         lambda r:r.update(observed_shifted_max=.04),
                         lambda r:r['actual_build_sidecars'][1].update(source_sha256='wrong'),
                         lambda r:r['probe_trace']['values'][0].__setitem__(0,1.01),
                         lambda r:r['probe_trace']['steps'].pop(),
                         lambda r:r['probe_trace']['probes'][0].update(i=1)]:
            bad=deepcopy(records);mutation(bad[3])
            with self.assertRaises((ValueError,KeyError)):candidate.compare(bad,source)