#!/usr/bin/env python3
"""Stage-2 native checks and untrusted Lean-certificate generation; run through vq."""
from dataclasses import asdict, replace
from fractions import Fraction as Q
import hashlib
import inspect
import json
import math
from pathlib import Path
import sys
import numpy as np
import vibeqc as v
from vibeqc import _vibeqc_core as core
from vibeqc.density_fitting import DensityFitting
from vibeqc.periodic.ccm.wigner_seitz import minimum_image
from vibeqc.neb import _minimum_image_diff
sys.path.insert(0, str(Path(__file__).resolve().parent))
from stage2_reference import (U, ETA, rational, dot_trace, dot_budget, df_bounds,
    certify_image, verify_image, cost, matvec, SearchBudgetExceeded,
    lean_image_case, lean_trace_case, lean_error_case)


def json_value(x):
    if isinstance(x, Q): return str(x)
    if isinstance(x, np.generic): return x.item()
    raise TypeError(type(x).__name__)


def main():
    failures_rejected = []
    controls = []
    def reject(label, action):
        try: action()
        except (ValueError, OverflowError): failures_rejected.append(label)
        else: raise AssertionError(label+': invalid input accepted')
    trace_cases = {
        'cancellation': [(-1.,1.), (2.**-53,1.), (1.,1.)],
        'subnormal': [(2.**-1074,.5)],
        'mixed': [(1e100,1e-100), (-3.1,.7), (1e-100,1e100)],
        'zero': [],
    }
    traces = {}
    for name,pairs in trace_cases.items():
        result, exact, bound, ops = dot_trace(pairs)
        traces[name] = {'inputs':pairs,'value':rational(result),'exact':exact,
                        'bound':bound,'error':abs(rational(result)-exact),'steps':len(ops)}
    assert traces['cancellation']['error'] > U*abs(traces['cancellation']['exact'])
    controls.append('relative_output_only_bound_fails_under_cancellation')
    reject('underflow_allowance_required', lambda: dot_trace(trace_cases['subnormal'], eta=Q(0)))
    reject('overflow', lambda: dot_trace([(sys.float_info.max,2.)]))
    reject('nan', lambda: dot_trace([(float('nan'),1.)]))

    rng = np.random.default_rng(20261102)
    synthetic = DensityFitting.__new__(DensityFitting)
    synthetic.n_orb = 3
    synthetic.B = rng.normal(size=(5,3,3))
    df_results = {}
    for name,D in [('synthetic',rng.normal(size=(6,6))[::2,::2]),
                   ('cancellation',np.array([[1e12,1e-10,-1e12],[-1.,2.,-3.],[4.,-5.,6.]]))]:
        rho = np.einsum('Pmn,mn->P',synthetic.B,D,optimize=True)
        J = synthetic.build_J(D)
        df_results[name] = df_bounds(synthetic.B,D,rho,J)
        if name == 'synthetic':
            wrong = J.copy(); wrong[0,1] += .01
            reject('perturbed_df_output', lambda: df_bounds(synthetic.B,D,rho,wrong))
            reject('transposed_density', lambda: df_bounds(synthetic.B,D.T,rho,J))
    # Full molecular calculation, and native resident DF with its own observed rho.
    coords = [[0.,0.,0.],[0.,1.4,1.1],[0.,-1.4,1.1]]
    mol = v.Molecule([v.Atom(z,r) for z,r in zip([8,1,1],coords)],0,1)
    basis = v.BasisSet(mol,'sto-3g'); aux = v.BasisSet(mol,'def2-universal-jkfit')
    opts = v.RHFOptions();opts.max_iter=150;opts.conv_tol_energy=1e-12;opts.conv_tol_grad=1e-10
    opts.density_fit=False;opts.scf_mode=v.SCFMode.CONVENTIONAL
    opts.initial_guess=v.InitialGuess.HCORE;opts.stability_check=False
    rhf = v.run_rhf(mol,basis,opts)
    assert rhf.converged
    storage = core._DensityFittingStorage(basis,aux)
    n,p = int(basis.nbasis),int(aux.nbasis)
    B = np.asarray(storage.mo_transform(np.eye(n),np.eye(n))).reshape(p,n,n)
    native_D = rng.normal(size=(2*n,2*n))[::2,::2]
    native_rho = np.asarray(storage.contract_B_with(native_D))
    native_J = np.asarray(storage.build_J(native_D))
    df_results['resident'] = df_bounds(B,native_D,native_rho,native_J)
    # Preserve exact inputs/outputs for replay, including the native first stage.
    np.savez('df-native-inputs.npz',B=B,D=native_D,rho=native_rho,J=native_J)

    cases = [
        ('skew',[[2,.5,.25],[0,1.5,.5],[0,0,1]],[.75,.25,.125]),
        ('tie',[[1,2,0],[0,1,0],[0,0,1]],[.5,.5,.5]),
        ('shear',[[1,8,0],[0,1,0],[0,0,1]],[.375,.25,.125]),
        ('origin',[[2,.5,.25],[0,1.5,.5],[0,0,1]],[0.,0.,0.]),
    ]
    for k in range(8):
        A=np.array([[2.,rng.integers(-6,7)/4,rng.integers(-6,7)/4],
                    [0.,1.5,rng.integers(-6,7)/4],[0.,0.,1.]])
        x=rng.integers(-24,25,size=3)/8
        cases.append(('random'+str(k),A.tolist(),x.tolist()))
    geometry = {}; kernel_cases = []; production_tie_gaps = {}
    for name,A,x in cases:
        cells,weights=minimum_image(np.array(x),np.array(A).T,tol_bohr=0.)
        proposal=tuple(int(v) for v in cells[0][0])
        c=certify_image(A,x,candidate=proposal)
        got={tuple(int(v) for v in row) for row in cells[0]}
        expected=set(c.ties)
        # Exact rational verification of the returned integer labels, not just close distances.
        assert got and got.issubset(expected), (name,got,expected)
        if got != expected:
            production_tie_gaps[name] = {'returned': sorted(got), 'certified_complete': sorted(expected)}
        assert all(cost(c.A,c.x,n)==cost(c.A,c.x,c.best) for n in got)
        delta=_minimum_image_diff(np.array(x)[None,:],np.array(A),3)[0]
        assert abs(float(delta@delta)-float(cost(c.A,c.x,c.best))) < 1e-12
        assert np.max(np.abs(weights[0]-1/len(got))) < 1e-15
        shift=(5,-3,2); translated=tuple(a+b for a,b in zip(c.x,matvec(c.A,shift)))
        ct=certify_image(c.A,translated,candidate=tuple(a+b for a,b in zip(c.best,shift)))
        assert set(ct.ties)=={tuple(a+b for a,b in zip(t,shift)) for t in c.ties}
        cr=certify_image(c.A,tuple(-v for v in c.x),candidate=tuple(-v for v in c.best))
        assert set(cr.ties)=={tuple(-v for v in t) for t in c.ties}
        geometry[name]={'certificate':asdict(c),'native_images':sorted(got), 'native_weights':weights[0].tolist(),
                        'certified_weights':[str(Q(1,len(expected)))]*len(expected),
                        'box_count':math.prod(h-l+1 for l,h in zip(c.lo,c.hi)),
                        'minimum_square':cost(c.A,c.x,c.best)}
        if name in ('skew','tie') or got != expected: kernel_cases.append((name,c))
    assert set(production_tie_gaps) == {'random1'}, production_tie_gaps
    c=kernel_cases[0][1]
    reject('singular_lattice',lambda:certify_image([[1,0,0],[0,0,0],[0,0,1]],[1,2,3]))
    reject('nonfinite_lattice',lambda:certify_image([[float('inf'),0,0],[0,1,0],[0,0,1]],[0,0,0]))
    reject('search_budget',lambda:certify_image(c.A,c.x,max_candidates=1))
    reject('bad_inverse',lambda:verify_image(replace(c,H=((Q(0),)*3,)*3)))
    reject('truncated_box',lambda:verify_image(replace(c,lo=c.best,hi=c.best)))
    reject('bad_incumbent',lambda:verify_image(replace(c,best=(100,0,0))))
    tied=kernel_cases[1][1]
    reject('missing_tie',lambda:verify_image(replace(tied,ties=tied.ties[:1])))
    rounded=tuple(round(v) for v in matvec(tied.H,tied.x))
    assert {rounded} != set(tied.ties);controls.append('single_rounding_drops_exact_ties')
    # A strongly sheared cell needs nonlocal integer labels even for a nearby Cartesian point.
    shear=geometry['shear']['certificate'];f=matvec(shear['H'],shear['x'])
    rounded=tuple(round(v) for v in f)
    assert cost(shear['A'],shear['x'],rounded)>geometry['shear']['minimum_square']
    controls.append('fractional_rounding_not_triclinic_nearest')
    # Same model input, but production's documented tolerance includes a nonminimal near tie.
    near=[1.+2.**-30,0.,0.];A=[[2.,.5,.25],[0.,1.5,.5],[0.,0.,1.]]
    near_cert=certify_image(A,near)
    near_cells,_=minimum_image(near,np.array(A).T)
    near_got={tuple(map(int,n)) for n in near_cells[0]}
    assert set(near_cert.ties).issubset(near_got) and near_got!=set(near_cert.ties)

    lean='import VibeQCStage2\nset_option maxRecDepth 100000\nset_option maxHeartbeats 4000000\n'
    lean+='namespace VibeQCStage2Cases\nopen VibeQCStage2\n'
    for name,c in kernel_cases: lean+=lean_image_case(name,c)
    lean+=lean_trace_case('cancel',trace_cases['cancellation'])
    lean+=lean_trace_case('underflow',trace_cases['subnormal'])
    # A native output gets a kernel-checked absolute enclosure for the exact supplied factors.
    row=df_results['resident'][1];i,j=row['i'],row['j']
    rho_exact=[sum((rational(B[a,k,l])*rational(native_D[k,l]) for k in range(n) for l in range(n)),Q(0)) for a in range(p)]
    exact=sum((rational(B[a,i,j])*rho_exact[a] for a in range(p)),Q(0))
    lean+=lean_error_case('native_df',rational(native_J[i,j]),exact,row['bound'])
    lean+='end VibeQCStage2Cases\n'
    Path('stage2-certificates.lean').write_text(lean)
    functions=[DensityFitting.build_J,minimum_image,_minimum_image_diff]
    report={'verdict':'PASS_WITH_LIMITATION','certificate_checks':'PASS',
      'production_exact_tie_completeness':'FAIL', 'production_tie_gaps':production_tie_gaps,'version':v.__version__,'native_module':core.__file__,
      'native_sha256':hashlib.sha256(Path(core.__file__).read_bytes()).hexdigest(),
      'sources':{f.__module__+'.'+f.__qualname__:{'path':inspect.getsourcefile(f),
                  'sha256':hashlib.sha256(inspect.getsource(f).encode()).hexdigest()} for f in functions},
      'rounding_allowances':{'u':U,'eta':ETA},'traces':traces,'df':df_results,'geometry':geometry,
      'rejected_controls':failures_rejected,'negative_controls':controls,'rhf_energy':float(rhf.energy),
      'near_tie':{'exact':near_cert.ties,'production_default':sorted(near_got)},
      'certificate_sha256':hashlib.sha256(lean.encode()).hexdigest(),
      'scope':'local rounding-model proofs; exact rational certificates; native per-input checks; no compiler/BLAS refinement'}
    Path('stage2-checks.json').write_text(json.dumps(report,default=json_value,indent=2)+'\n')
    print(json.dumps({'verdict':'PASS_WITH_LIMITATION','production_tie_gaps':production_tie_gaps,'native_version':v.__version__,'geometry_cases':len(geometry),
          'rejections':len(failures_rejected),'controls':len(controls),'rhf_energy':float(rhf.energy),
          'df_max_error':{k:max(float(x['error']) for x in rows) for k,rows in df_results.items()},
          'df_max_error_to_bound':{k:max(float(x['error']/x['bound']) for x in rows) for k,rows in df_results.items()}},indent=2))


if __name__=='__main__': main()