#!/usr/bin/env python3
"""Production exact-image acceptance, including optional Python-only source overlay.

Run through the calculation scheduler with a qualified native installation.
The overlay loads only the two named geometry modules, without modifying the
installation. It is recorded as source-under-test, not as a deployed runtime.
Supply check_stage2.py and stage2_reference.py alongside this script.
"""
import argparse
import hashlib
import importlib.abc
import importlib.util
import json
import os
from pathlib import Path
import runpy
import sys
import time


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--source-overlay', type=Path)
    parser.add_argument('--pytest-file', type=Path)
    args = parser.parse_args()
    here = Path(__file__).resolve().parent
    sys.path.insert(0, str(here))
    overlay = {}
    if args.source_overlay:
        for name in ('wigner_seitz', '_exact_images'):
            path = (args.source_overlay / (name+'.py')).resolve()
            if not path.is_file(): raise ValueError('missing overlay source: '+str(path))
            overlay['vibeqc.periodic.ccm.'+name] = path
        class GeometrySource(importlib.abc.MetaPathFinder):
            def find_spec(self, fullname, path=None, target=None):
                if fullname in overlay:
                    return importlib.util.spec_from_file_location(fullname, overlay[fullname])
                return None
        sys.meta_path.insert(0, GeometrySource())

    import numpy as np
    import vibeqc as v
    from vibeqc import _vibeqc_core as core
    from vibeqc.periodic.ccm.wigner_seitz import minimum_image
    from vibeqc.periodic.ccm._exact_images import certify_image
    from vibeqc.periodic.ccm.system import CCMSystem
    from vibeqc.periodic.ccm.scf import run_ccm_rhf
    from vibeqc.periodic.ccm.padded import ccm_nuclear_repulsion
    from stage2_reference import lean_image_case, lean_q, cost

    if args.pytest_file:
        os.environ['PYTEST_DISABLE_PLUGIN_AUTOLOAD'] = '1'
        import pytest
        status = pytest.main(['-q', str(args.pytest_file.resolve()), '-o', 'addopts='])
        if status: raise AssertionError('geometry regression tests failed')

    saved_argv = sys.argv[:]
    sys.argv = ['check_stage2.py']
    try: runpy.run_path(str(here/'check_stage2.py'), run_name='__main__')
    finally: sys.argv = saved_argv
    stage2 = json.loads(Path('stage2-checks.json').read_text())
    assert stage2['verdict'] == 'PASS' and stage2['production_tie_gaps'] == {}

    # Export certificates produced by the production module, and compare wrapper outputs.
    A = [[2.,-.5,0.],[0.,1.5,-1.],[0.,0.,1.]]
    x = [1.125,-.625,.25]
    samples = [('production_gap',A,x),
               ('production_half',np.diag([2.,3.,4.]).tolist(),[1.,0.,0.]),
               ('production_near',np.diag([2.,3.,4.]).tolist(),[np.nextafter(1.,2.),0.,0.])]
    certificates = Path('stage2-certificates.lean').read_text().replace('end VibeQCStage2Cases\n','')
    for name,cell,disp in samples:
        cert=certify_image(cell,disp)
        cells,weights=minimum_image(disp,np.asarray(cell).T,tol_bohr=0.)
        assert [tuple(map(int,n)) for n in cells[0]] == list(cert.ties)
        assert np.array_equal(weights[0], np.full(len(cert.ties),1/len(cert.ties)))
        certificates += lean_image_case(name,cert)
        params=f'{name}A {name}H {name}x {lean_q(cert.R)} {name}lo {name}hi {name}best'
        certificates += f'theorem {name}_nonempty : (certifiedTies {name}A {name}x {name}lo {name}hi {name}best).Nonempty :=\n'
        certificates += f'  certified_ties_nonempty {params} {name}_valid\n'
        certificates += f'theorem {name}_weights : (∑ _n ∈ certifiedTies {name}A {name}x {name}lo {name}hi {name}best, (1 / ((certifiedTies {name}A {name}x {name}lo {name}hi {name}best).card : ℝ))) = 1 :=\n'
        certificates += f'  certified_weights_sum_one {params} {name}_valid\n'
    certificates = certificates.replace('open VibeQCStage2\n','open VibeQCStage2\nopen scoped BigOperators\n')
    certificates += 'end VibeQCStage2Cases\n'
    Path('production-image-certificates.lean').write_text(certificates)

    # A full experimental CCM calculation whose two-center geometry contains the tie.
    columns=4*np.asarray(A);position=4*np.asarray(x)
    unit=v.PeriodicSystem(3,columns,[v.Atom(1,[0.,0.,0.]),v.Atom(1,position.tolist())],0,1)
    strict=CCMSystem(unit,(1,1,1),'sto-3g',weight_tol_bohr=0.)
    tolerant=CCMSystem(unit,(1,1,1),'sto-3g')
    exact_weights=strict.cell_weight_matrices(); float_weights=tolerant.cell_weight_matrices()
    assert set(exact_weights)==set(float_weights)
    assert all(np.array_equal(exact_weights[g],float_weights[g]) for g in exact_weights)
    kw=dict(method='aiccm2026dev-a',initial_guess='HCORE',conv_tol=1e-12,conv_tol_grad=1e-10,max_iter=150)
    started=time.perf_counter(); result=run_ccm_rhf(strict,**kw); elapsed=time.perf_counter()-started
    reference=run_ccm_rhf(tolerant,**kw)
    assert result.converged and reference.converged
    assert abs(result.energy-reference.energy)<1e-10
    cert=certify_image(columns.tolist(),position.tolist())
    nuclear_reference=1/np.sqrt(float(cost(cert.A,cert.x,cert.best)))
    assert abs(ccm_nuclear_repulsion(strict)-nuclear_reference)<1e-14

    # Molecular limit of the same CCM driver, matched to the existing water ORCA reference.
    coords=[[0.,0.,0.],[0.,1.4,1.1],[0.,-1.4,1.1]]
    water=v.PeriodicSystem(3,np.diag([32.,32.,32.]),[v.Atom(z,r) for z,r in zip([8,1,1],coords)],0,1)
    molecular=CCMSystem(water,(1,1,1),'sto-3g',weight_tol_bohr=0.)
    water_result=run_ccm_rhf(molecular,**kw)
    assert water_result.converged
    orca_energy=-74.960337068945
    assert abs(water_result.energy-orca_energy)<2e-8

    loaded={name: {'path':str(Path(sys.modules[name].__file__).resolve()),
                  'sha256':hashlib.sha256(Path(sys.modules[name].__file__).read_bytes()).hexdigest()}
            for name in ['vibeqc.periodic.ccm.wigner_seitz','vibeqc.periodic.ccm._exact_images']}
    report={'verdict':'PASS','runtime_version':v.__version__,'native_module':core.__file__,
      'native_sha256':hashlib.sha256(Path(core.__file__).read_bytes()).hexdigest(),
      'source_overlay':bool(overlay),'geometry_sources':loaded,
      'stage2_status':stage2['verdict'],'stage2_geometry_cases':len(stage2['geometry']),
      'ccm_tie_energy':float(result.energy),'ccm_positive_tolerance_energy':float(reference.energy),
      'ccm_tie_seconds':elapsed,'ccm_nuclear_energy':float(ccm_nuclear_repulsion(strict)),
      'ccm_nuclear_reference':float(nuclear_reference),'ccm_water_energy':float(water_result.energy),
      'orca_water_reference':orca_energy,'water_absolute_tolerance':2e-8,
      'certificate_sha256':hashlib.sha256(certificates.encode()).hexdigest(),
      'limits':['Exact represented-input geometry only','No proof of Python interpreter or NumPy conversion',
                'Positive-tolerance selection is a distinct floating-point policy',
                'Finite triclinic CCM SCF has no matched ORCA/CRYSTAL counterpart; molecular-limit reference is independent']}
    Path('exact-image-checks.json').write_text(json.dumps(report,indent=2)+'\n')
    print(json.dumps(report,indent=2))


if __name__=='__main__': main()