#!/usr/bin/env python3
"""Exact entropy-normalized stochastic-interaction MIPs and heuristic benchmarks.

All main results are generated, not imported from the manuscript. No network,
cloud, QPU, secret, or remote filesystem access. See protocol.json and README.md.
"""
from __future__ import annotations
import os
for _key in ('OPENBLAS_NUM_THREADS', 'OMP_NUM_THREADS', 'MKL_NUM_THREADS', 'NUMBA_NUM_THREADS'):
    os.environ[_key] = '1'
import argparse, hashlib, json, math, platform, sys, time
from datetime import datetime, timezone
from pathlib import Path
import numpy as np
import scipy
from scipy import linalg, optimize
from numba import njit

ROOT = Path(__file__).resolve().parents[1]
PROTOCOL = {
    'version': 'QN-MIP-20261012-v1',
    'root_seed': 20261012,
    'sizes': [8, 10, 12, 14],
    'families': ['dense', 'modular', 'sparse_ring'],
    'network_replicates': 5,
    'spectral_radius': 0.9,
    'innovation_covariance': 'identity',
    'objective': 'J(S)=Phi_SI(S)/min(H(X_S),H(X_complement))',
    'units': 'natural-log entropy; fixed unit-variance innovations',
    'exact_search': 'all 2^(n-1)-1 nonempty proper bipartitions, float64',
    'tie_atol': 1e-10,
    'tie_rtol': 1e-8,
    'sa': {'independent_runs':10, 'pilot_draws':32, 'restarts':4,
           'steps_per_restart':512, 'temperature_floor':1e-3,
           'T0':'max(std(pilot objective values),0.001)',
           'schedule':'geometric T0 to 0.001*T0'},
    'qaoa': {'depths':[1,2], 'optimizer':'COBYLA', 'optimizer_restarts':3,
             'max_evals_per_start':{'1':160,'2':240},
             'optimizer_rhobeg':0.7,'optimizer_tol':1e-5,
             'shot_replicates':10,'shots_per_replicate':256,
             'cost':'valid J/max_valid_J; invalid all-in-one state cost=2',
             'state':'uniform over all 2^(n-1) basis states',
             'mixer':'exp(-i beta sum_q X_q)',
             'cost_oracle':'full exact diagonal truth table; not a QUBO or efficient circuit',
             'selection':'lowest penalized expectation, never optimal-state probability'},
    'uniform': {'independent_runs':10,'draws':256,'support':'valid bipartitions'},
    'runtime': 'serial, one BLAS thread; exclude JIT warm-up and shared covariance setup',
    'analysis': 'report every fixed instance; success within atol+rtol*abs(Jstar)'
}


def jwrite(path: Path, obj: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(obj, indent=2, sort_keys=True, allow_nan=False)+'\n')


def rng_for(*keys: int) -> np.random.Generator:
    return np.random.default_rng(np.random.SeedSequence([PROTOCOL['root_seed'], *keys]))


def make_network(n: int, family: str, rep: int):
    fi=PROTOCOL['families'].index(family)
    rng=rng_for(n,fi,rep)
    W=rng.normal(size=(n,n))/math.sqrt(n)
    if family == 'modular':
        group=np.arange(n)//(n//2)
        W *= np.where(group[:,None] == group[None,:],1.0,0.15)
    elif family == 'sparse_ring':
        keep=rng.random((n,n))<0.2
        np.fill_diagonal(keep,True)
        for i in range(n):
            keep[i,(i+1)%n]=True
            keep[i,(i-1)%n]=True
        W *= keep
    perm=rng.permutation(n)
    W=W[np.ix_(perm,perm)]
    rho=float(np.max(np.abs(np.linalg.eigvals(W))))
    if rho<=0: raise ValueError('Degenerate matrix')
    W *= PROTOCOL['spectral_radius']/rho
    t=time.perf_counter()
    Sigma=linalg.solve_discrete_lyapunov(W,np.eye(n),method='bilinear')
    Sigma=(Sigma+Sigma.T)/2
    setup_sec=time.perf_counter()-t
    residual=float(np.linalg.norm(Sigma-W@[author/contact email redacted](n),'fro')/np.linalg.norm(Sigma,'fro'))
    if residual>1e-10 or np.min(np.linalg.eigvalsh(Sigma))<=0:
        raise ArithmeticError(f'Invalid covariance {residual=}')
    return W,Sigma,perm,setup_sec,residual


class Oracle:
    """Fresh on-demand Gaussian Schur-complement oracle; no reference-table access."""
    def __init__(self, W: np.ndarray, Sigma: np.ndarray, cache=True):
        self.n=len(W); self.full=(1<<self.n)-1
        self.Sigma=Sigma; self.G=W@Sigma
        self.cache={} if cache else None
        self.calls=0; self.computed=0
        self.worst_negative=0.0

    def canonical(self, mask:int) -> int:
        return mask if mask&1 else self.full^mask

    def parts(self, mask:int):
        ids=np.array([i for i in range(self.n) if (mask>>i)&1],dtype=int)
        ix=np.ix_(ids,ids)
        P=self.Sigma[ix]; B=self.G[ix]
        factor=linalg.cho_factor(P,lower=True,check_finite=False)
        logdetP=2*np.log(np.diag(factor[0])).sum()
        [author/contact email redacted]_solve(factor,B.T,check_finite=False)
        R=(R+R.T)/2
        ch=linalg.cholesky(R,lower=True,check_finite=False)
        logdetR=2*np.log(np.diag(ch)).sum()
        H=0.5*(len(ids)*math.log(2*math.pi*math.e)+logdetP)
        return float(logdetR),float(H)

    def components(self, mask:int):
        mask=self.canonical(mask)
        if mask==self.full: return 0.0,0.0,0.0
        r1,h1=self.parts(mask);r2,h2=self.parts(self.full^mask)
        phi=0.5*(r1+r2)  # logdet(Q)=0 because Q=I
        self.worst_negative=min(self.worst_negative,phi)
        if phi < -1e-10: raise ArithmeticError(f'Negative Phi_SI: {phi}')
        phi=max(phi,0.0)
        k=min(h1,h2)
        if k<=0: raise ArithmeticError('Nonpositive entropy denominator')
        return phi/k,phi,k

    def __call__(self, mask:int)->float:
        self.calls+=1
        mask=self.canonical(int(mask))
        if mask==self.full: return 0.0
        if self.cache is not None and mask in self.cache: return self.cache[mask]
        val=self.components(mask)[0];self.computed+=1
        if self.cache is not None: self.cache[mask]=val
        return val


def exact_table(W,Sigma):
    oracle=Oracle(W,Sigma,cache=False);n=len(W);D=1<<(n-1)
    js=np.empty(D-1);phis=np.empty(D-1);ks=np.empty(D-1)
    t=time.perf_counter()
    for z in range(D-1):
        js[z],phis[z],ks[z]=oracle.components((z<<1)|1)
    sec=time.perf_counter()-t
    best=int(np.argmin(js));star=float(js[best])
    tolerance=PROTOCOL['tie_atol']+PROTOCOL['tie_rtol']*abs(star)
    opt=np.flatnonzero(js<=star+tolerance)
    return js,phis,ks,best,opt,sec


def queyranne(n:int, f):
    """Pendant-pair contraction; guarantee applies only to symmetric submodular f."""
    groups=[1<<i for i in range(n)]
    bestmask=groups[0];best=float('inf')
    while len(groups)>1:
        order=[groups[0]]; rest=groups[1:].copy(); union=groups[0]
        # Singleton terms can be reused in a pendant-pair pass.
        single={v:f(v) for v in rest}
        while rest:
            choice=min(rest,key=lambda v:(f(union|v)-single[v],v))
            order.append(choice);union |= choice;rest.remove(choice)
        s,t=order[-2:]
        ft=f(t)
        if ft<best:best=ft;bestmask=t
        groups=[v for v in groups if v!=s and v!=t]+[s|t]
    full=(1<<n)-1
    if not bestmask&1:bestmask=full^bestmask
    return bestmask,float(best)


@njit(cache=True)
def qaoa_state(cost,theta):
    D=cost.size;psi=np.ones(D,dtype=np.complex128)/np.sqrt(D)
    for layer in range(theta.size//2):
        gamma=theta[2*layer];beta=theta[2*layer+1]
        for j in range(D):psi[j] *= np.exp(-1j*gamma*cost[j])
        c=np.cos(beta);s=-1j*np.sin(beta)
        step=1
        while step<D:
            for base in range(0,D,2*step):
                for offset in range(step):
                    a=base+offset;b=a+step;pa=psi[a];pb=psi[b]
                    psi[a]=c*pa+s*pb;psi[b]=s*pa+c*pb
            step*=2
    return psi


@njit(cache=True)
def qaoa_energy(cost,theta):
    psi=qaoa_state(cost,theta)
    energy=0.0
    for j in range(cost.size):energy += cost[j]*(psi[j].real**2+psi[j].imag**2)
    return energy


def train_qaoa(cost:np.ndarray,p:int,seedkeys:tuple):
    rng=rng_for(*seedkeys,p,420)
    best_e=float('inf');best_theta=None;log=[];statuses=[]
    started=time.perf_counter()
    for restart in range(PROTOCOL['qaoa']['optimizer_restarts']):
        theta0=np.empty(2*p)
        theta0[0::2]=rng.uniform(0,2*math.pi,p)
        theta0[1::2]=rng.uniform(0,math.pi,p)
        def objective(theta):
            nonlocal best_e,best_theta
            value=float(qaoa_energy(cost,np.asarray(theta)))
            log.append([float(restart),value,*map(float,theta)])
            if value<best_e:
                best_e=value;best_theta=np.asarray(theta).copy()
            return value
        result=optimize.minimize(objective,theta0,method='COBYLA',options={
            'maxiter':PROTOCOL['qaoa']['max_evals_per_start'][str(p)],
            'rhobeg':PROTOCOL['qaoa']['optimizer_rhobeg'],
            'tol':PROTOCOL['qaoa']['optimizer_tol']})
        statuses.append({'restart':restart,'nfev':int(result.nfev),
                         'success':bool(result.success),'message':str(result.message)})
    psi=qaoa_state(cost,best_theta)
    probs=(psi.real**2+psi.imag**2)
    norm=float(probs.sum())
    if abs(norm-1)>1e-10:raise ArithmeticError(f'Nonunitary QAOA: {norm}')
    probs=probs/norm
    return probs,best_theta,float(best_e),time.perf_counter()-started,np.asarray(log),statuses


def run_sa(W,Sigma,seedkeys):
    cfg=PROTOCOL['sa'];rng=rng_for(*seedkeys,811)
    n=len(W);D=1<<(n-1);oracle=Oracle(W,Sigma)
    bestz=-1;best=float('inf')
    def evaluate(z):
        nonlocal best,bestz
        val=oracle((int(z)<<1)|1)
        if val<best:best=val;bestz=int(z)
        return val
    started=time.perf_counter()
    pilot=np.array([evaluate(int(rng.integers(D-1))) for _ in range(cfg['pilot_draws'])])
    T0=max(float(np.std(pilot)),0.001)
    temps=np.geomspace(T0,T0*cfg['temperature_floor'],cfg['steps_per_restart'])
    rejected_invalid=0
    for _ in range(cfg['restarts']):
        z=int(rng.integers(D-1));val=evaluate(z)
        for temperature in temps:
            new=z^(1<<int(rng.integers(n-1)))
            if new==D-1:
                rejected_invalid+=1;continue
            newval=evaluate(new);delta=newval-val
            if delta<=0 or rng.random()<math.exp(-delta/temperature):z=new;val=newval
    return bestz,float(best),time.perf_counter()-started,oracle.calls,oracle.computed,T0,rejected_invalid


def sample_output(probs,js,opt,seedkeys,shots=256):
    rng=rng_for(*seedkeys)
    sample=rng.choice(len(probs),size=shots,p=probs)
    valid=sample[sample<len(js)]
    if len(valid):
        z=int(valid[np.argmin(js[valid])]);value=float(js[z])
    else:z=-1;value=None
    hit=bool(np.isin(sample,opt).any())
    return sample,z,value,hit,int(len(sample)-len(valid))


def validate():
    """Independent correctness checks without looking at benchmark outcomes."""
    checks={};rng=np.random.default_rng(4928501)
    # Queyranne on known symmetric submodular weighted cuts.
    cut_tests=0
    for n in range(3,10):
        for _ in range(10):
            M=rng.uniform(size=(n,n));M=(M+M.T)/2;np.fill_diagonal(M,0)
            def f(mask):
                ids=[i for i in range(n) if (mask>>i)&1]
                rest=[i for i in range(n) if not ((mask>>i)&1)]
                return float(M[np.ix_(ids,rest)].sum())
            bm,bv=queyranne(n,f)
            ev=min(f((z<<1)|1) for z in range((1<<(n-1))-1))
            if abs(bv-ev)>1e-10:raise AssertionError('Queyranne weighted-cut failure')
            cut_tests+=1
    checks['queyranne_weighted_cut_checks_passed']=cut_tests
    # Queyranne on unnormalized Gaussian mutual information.
    mi_tests=0
    for n in range(3,10):
        for _ in range(5):
            A=rng.normal(size=(n,n));P=A@A.T+np.eye(n)
            full=(1<<n)-1;ldfull=np.linalg.slogdet(P)[1]
            def f(mask):
                if mask in (0,full):return 0.0
                a=[i for i in range(n) if (mask>>i)&1];b=[i for i in range(n) if not ((mask>>i)&1)]
                return float(0.5*(np.linalg.slogdet(P[np.ix_(a,a)])[1]+np.linalg.slogdet(P[np.ix_(b,b)])[1]-ldfull))
            bm,bv=queyranne(n,f)
            ev=min(f((z<<1)|1) for z in range((1<<(n-1))-1))
            if abs(bv-ev)>1e-10:raise AssertionError('Queyranne MI failure')
            mi_tests+=1
    checks['queyranne_gaussian_MI_checks_passed']=mi_tests
    # QAOA comparison to independent dense matrix exponentials.
    maxerr=0.0
    for m in (2,3,4):
        D=1<<m;cost=rng.uniform(size=D);theta=rng.normal(size=4)
        psi=np.ones(D,dtype=complex)/math.sqrt(D)
        mixer=np.zeros((D,D))
        for z in range(D):
            for q in range(m):mixer[z,z^(1<<q)]+=1
        for p in range(2):
            psi=np.exp(-1j*theta[2*p]*cost)*psi
            psi=linalg.expm(-1j*theta[2*p+1]*mixer)@psi
        got=qaoa_state(cost,theta)
        maxerr=max(maxerr,float(np.max(np.abs(psi-got))))
    if maxerr>1e-12:raise AssertionError('QAOA dense-unitary mismatch')
    checks['qaoa_dense_expm_max_amplitude_error']=maxerr
    # Independent SI formula via logdet of joint covariance, all n=6 cuts.
    W,P,*_=make_network(6,'dense',99);G=W@P;o=Oracle(W,P);err=0
    joint=np.block([[P,G.T],[G,P]])
    for z in range((1<<5)-1):
        mask=(z<<1)|1;direct=0
        for mask_part in (mask,((1<<6)-1)^mask):
            s=[i for i in range(6) if (mask_part>>i)&1];both=s+[i+6 for i in s]
            direct+=np.linalg.slogdet(joint[np.ix_(both,both)])[1]-np.linalg.slogdet(P[np.ix_(s,s)])[1]
        err=max(err,abs(o.components(mask)[1]-0.5*direct))
    if err>1e-11:raise AssertionError('Conditional determinant identity failure')
    checks['SI_joint_vs_conditional_logdet_max_error']=float(err)
    # Independent blocks must have zero loss across the known disconnection.
    W=np.diag([0.3,0.4,0.5,0.6]);P=linalg.solve_discrete_lyapunov(W,np.eye(4))
    zero=max(Oracle(W,P).components(mask)[0] for mask in range(1,15))
    if zero>1e-12:raise AssertionError('Independent-node zero check failed')
    checks['independent_node_max_normalized_loss']=float(zero)
    checks['passed']=True
    jwrite(ROOT/'validation/unit_tests.json',checks)
    return checks


def run_instance(n,family,rep):
    fi=PROTOCOL['families'].index(family);id_=f'n{n:02d}_{family}_r{rep:02d}'
    meta_path=ROOT/'results'/f'{id_}.json'
    if meta_path.exists():return json.loads(meta_path.read_text())
    start=time.perf_counter();W,P,perm,setup_sec,residual=make_network(n,family,rep)
    js,phis,ks,best,opt,table_sec=exact_table(W,P)
    star=float(js[best]);tol=PROTOCOL['tie_atol']+PROTOCOL['tie_rtol']*abs(star)
    arrays={'W':W,'Sigma':P,'node_permutation':perm,'J':js,'Phi_SI':phis,'normalizer':ks,
            'optimal_z':opt}
    rec={'instance':id_,'n':n,'family':family,'rep':rep,
         'covariance_sec':setup_sec,'covariance_relative_residual':residual,
         'rho':float(np.max(np.abs(np.linalg.eigvals(W)))),
         'exact_sec':table_sec,'partition_count':int(len(js)),
         'exact_z':best,'exact_mask':(best<<1)|1,'Jstar':star,
         'Phi_at_MIP':float(phis[best]),'normalizer_at_MIP':float(ks[best]),
         'n_tied_optima':int(len(opt)), 'uniform_pstar':float(len(opt)/len(js)),
         'exact_smaller_part':min(((best<<1)|1).bit_count(),n-((best<<1)|1).bit_count())}
    # Verify determinant identity on predetermined random cuts in every instance.
    G=W@P;joint=np.block([[P,G.T],[G,P]]);check_rng=rng_for(n,fi,rep,12)
    err=0
    for z in check_rng.integers(len(js),size=8):
        mask=(int(z)<<1)|1;d=0
        for part in (mask,((1<<n)-1)^mask):
            s=[i for i in range(n) if (part>>i)&1];b=s+[i+n for i in s]
            d+=np.linalg.slogdet(joint[np.ix_(b,b)])[1]-np.linalg.slogdet(P[np.ix_(s,s)])[1]
        err=max(err,abs(float(phis[z])-0.5*d))
    if err>1e-10:raise AssertionError('Instance determinant check failed')
    rec['SI_crosscheck_max_abs_error']=float(err)
    oracle=Oracle(W,P);t=time.perf_counter();qm,qv=queyranne(n,oracle);qsec=time.perf_counter()-t
    rec['queyranne']={'z':qm>>1,'J':qv,'hit':bool(qv<=star+tol),'sec':qsec,
                      'calls':oracle.calls,'unique_computations':oracle.computed}
    trials=[]
    for r in range(PROTOCOL['sa']['independent_runs']):
        sz,sv,sec,calls,unique,T0,invalid=run_sa(W,P,(n,fi,rep,r))
        trials.append({'algorithm':'SA','trial':r,'z':sz,'J':sv,'hit':bool(sv<=star+tol),
                       'sec':sec,'calls':calls,'unique_computations':unique,
                       'T0':T0,'invalid_proposals':invalid})
    # A valid-uniform baseline with the same readout count as QAOA.
    uprob=np.full(len(js),1/len(js));us=[]
    for r in range(PROTOCOL['uniform']['independent_runs']):
        st=time.perf_counter();sample,z,value,hit,invalid=sample_output(uprob,js,opt,(n,fi,rep,r,315),256)
        sec=time.perf_counter()-st;us.append(sample)
        trials.append({'algorithm':'Uniform','trial':r,'z':z,'J':value,'hit':hit,'sec':sec,
                       'invalid_shots':invalid})
    arrays['uniform_shots']=np.asarray(us)
    # The oracle construction is exponential and separately accounted for.
    cost=np.concatenate([js/float(js.max()),[2.0]])
    arrays['qaoa_cost']=cost
    for p in PROTOCOL['qaoa']['depths']:
        probs,theta,energy,train_sec,history,statuses=train_qaoa(cost,p,(n,fi,rep))
        pstar=float(probs[opt].sum());samples=[]
        rec[f'qaoa_p{p}']={'training_sec':train_sec,'nfev':len(history),
                         'pstar':pstar,'p_invalid':float(probs[-1]),'penalized_energy':energy,
                         'theoretical_hit256':float(-np.expm1(256*np.log1p(-min(pstar,1-1e-16)))),
                         'optimizer_status':statuses,'theta':theta.tolist()}
        arrays[f'qaoa_p{p}_probabilities']=probs
        arrays[f'qaoa_p{p}_theta']=theta;arrays[f'qaoa_p{p}_optimizer_history']=history
        for r in range(PROTOCOL['qaoa']['shot_replicates']):
            st=time.perf_counter()
            sample,z,value,hit,invalid=sample_output(probs,js,opt,(n,fi,rep,p,r,621),256)
            sec=time.perf_counter()-st;samples.append(sample)
            trials.append({'algorithm':f'QAOA-p{p}','trial':r,'z':z,'J':value,'hit':hit,'sec':sec,
                          'invalid_shots':invalid})
        arrays[f'qaoa_p{p}_shots']=np.asarray(samples)
    rec['trials']=trials;rec['total_instance_sec']=time.perf_counter()-start
    np.savez_compressed(ROOT/'data'/f'{id_}.npz',**arrays)
    jwrite(meta_path,rec)
    print(f'{id_}: J*={star:.6g} Q={int(rec["queyranne"]["hit"])} '
          f'SA={sum(x["hit"] for x in trials if x["algorithm"]=="SA")}/10 '
          f'QAOA2={sum(x["hit"] for x in trials if x["algorithm"]=="QAOA-p2")}/10 '
          f'[{rec["total_instance_sec"]:.2f}s]',flush=True)
    return rec


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--validate-only',action='store_true')
    parser.add_argument('--limit',type=int,default=0)
    args=parser.parse_args()
    for folder in ('data','results','validation'): (ROOT/folder).mkdir(exist_ok=True,parents=True)
    config=ROOT/'protocol.json'
    if config.exists() and json.loads(config.read_text())!=PROTOCOL:
        raise RuntimeError('Frozen protocol differs: use a new output directory')
    if not config.exists():jwrite(config,PROTOCOL)
    if not (ROOT/'environment.json').exists():
        cpu='unknown'
        if Path('/proc/cpuinfo').exists():
            for line in Path('/proc/cpuinfo').read_text().splitlines():
                if line.startswith('model name'):cpu=line.split(':',1)[1].strip();break
        jwrite(ROOT/'environment.json',{
            'started_utc':datetime.now(timezone.utc).isoformat(),'python':sys.version,
            'numpy':np.__version__,'scipy':scipy.__version__,'platform':platform.platform(),
            'cpu':cpu,'thread_count':1,
            'benchmark_code_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            'protocol_sha256':hashlib.sha256(config.read_bytes()).hexdigest()})
    print('Validating and warming up ...',flush=True);print(validate(),flush=True)
    if args.validate_only:return
    started=time.perf_counter();recs=[]
    for n in PROTOCOL['sizes']:
        for family in PROTOCOL['families']:
            for rep in range(PROTOCOL['network_replicates']):
                recs.append(run_instance(n,family,rep))
                if args.limit and len(recs)>=args.limit:
                    print('Limit reached',flush=True);return
    jwrite(ROOT/'execution_summary.json',{'instance_count':len(recs),
        'last_execution_elapsed_sec':time.perf_counter()-started,
        'sum_instance_sec':sum(r['total_instance_sec'] for r in recs),
        'partition_count_total':sum(r['partition_count'] for r in recs),
        'finished_utc':datetime.now(timezone.utc).isoformat()})
    print('FINISHED',flush=True)

if __name__=='__main__':main()