import json, math, random, time
import numpy as np
from coherence_sampler import LaplacianColumns, coherent_sample, uniform_sample


def weighted_graph(n, edges):
    a=[[] for _ in range(n)]
    for i,j,w in edges:
        a[i].append((j,w)); a[j].append((i,w))
    return a

def two_cliques(k, alpha=0.0):
    n=2*k; e=[]
    for off in (0,k):
        for i in range(off,off+k):
            for j in range(i+1,off+k): e.append((i,j,1.0))
    if alpha:
        # One weak bridge keeps the block structure explicit.
        e.append((k-1,k,alpha))
    return weighted_graph(n,e)

def complete(n):
    return weighted_graph(n,[(i,j,1.0) for i in range(n) for j in range(i+1,n)])

def dense_laplacian(adj):
    n=len(adj); L=np.zeros((n,n))
    for i,nbrs in enumerate(adj):
        L[i,i]=sum(w for _,w in nbrs)
        for j,w in nbrs: L[i,j]-=w
    return L

def coverage(sel,k):
    return len(set(i//k for i in sel))

def mean_uniform_coverage(n,k,batch,reps=3000):
    vals=[]
    for r in range(reps): vals.append(coverage(uniform_sample(n,batch,random.Random(10000+r)),k))
    return float(np.mean(vals))

def verify_sparse_math():
    adj=two_cliques(5,0.23); c=LaplacianColumns(adj); L=dense_laplacian(adj)
    max_col=max(np.max(np.abs(np.array([c.column(i).get(j,0.) for j in range(10)])-L[:,i])) for i in range(10))
    max_dot=max(abs(c.dot(i,j)-float(L[:,i]@L[:,j])) for i in range(10) for j in range(10))
    return {"max_column_error":float(max_col),"max_inner_product_error":float(max_dot)}

def prediction_disconnected_sweep():
    # Prediction 1: at zero inter-block coupling, cross-block signatures have exactly zero coherence.
    out=[]
    for k in [3,5,8,12]:
        c=LaplacianColumns(two_cliques(k,0.0))
        cross=max(c.coherence(i,j) for i in range(k) for j in range(k,2*k))
        within=min(c.coherence(i,j) for i in range(k) for j in range(1,k))
        sel=coherent_sample(c,list(range(2*k)),2,random.Random(7))
        out.append({"k":k,"observed_cross_coherence":cross,"predicted":0.0,
                    "selected_blocks":coverage(sel,k),"predicted_blocks":2,
                    "within_coherence_min":within})
    return out

def prediction_pool_sweep():
    # Prediction 2: with two disconnected equal blocks and an all-node candidate pool,
    # a batch of two selects both blocks; smaller random pools can fail only when
    # the pool itself misses a block.
    k=10; c=LaplacianColumns(two_cliques(k,0.0)); n=2*k; b=2
    out=[]
    for m in [2,3,4,6,10,15,20]:
        vals=[]; selected_both=0
        for r in range(1000):
            rng=random.Random(200000+m*1000+r)
            cand=rng.sample(range(n),m)
            s=coherent_sample(c,cand,b,rng)
            vals.append(coverage(s,k)); selected_both += coverage(s,k)==2
        pool_has_both=1-(2*math.comb(k,m) / math.comb(2*k,m) if m<=k else 0)
        out.append({"candidate_pool":m,"observed_mean_blocks":float(np.mean(vals)),
                    "observed_both_fraction":selected_both/1000,
                    "predicted_pool_has_both":float(pool_has_both)})
    return out

def prediction_coupling_sweep():
    # Prediction 3: the cross-block column coherence is zero at alpha=0 and
    # increases smoothly with bridge weight; selection's block coverage falls
    # only once alpha is large enough to make signatures less distinct.
    out=[]; k=8
    for alpha in [0,1e-4,1e-3,1e-2,0.05,0.1,0.25,0.5,1.0,2.0]:
        c=LaplacianColumns(two_cliques(k,alpha))
        cross=max(c.coherence(i,j) for i in range(k) for j in range(k,2*k))
        # The maximum is attained by the two bridge endpoints.  For a k-clique
        # joined by one edge of weight alpha, direct substitution in c(i,S)
        # gives this exact prediction.
        d=(k-1)+alpha
        predicted=2*alpha*d/(d*d+(k-1)+alpha*alpha)
        vals=[]
        for r in range(200):
            vals.append(coverage(coherent_sample(c,list(range(2*k)),2,random.Random(3000+r)),k))
        uniform=[]
        for r in range(200):
            uniform.append(coverage(uniform_sample(2*k,2,random.Random(9000+r)),k))
        out.append({"alpha":alpha,"max_cross_coherence":float(cross),
                    "predicted_bridge_coherence":float(predicted),
                    "coherent_mean_blocks":float(np.mean(vals)),
                    "uniform_mean_blocks":float(np.mean(uniform))})
    return out

def estimator_variance():
    # A transparent small loss-vector experiment: empirical inclusion p is
    # estimated from repeated candidate pools, then HT weights are evaluated.
    n=20; k=10; c=LaplacianColumns(two_cliques(k,0.0)); B=4; M=10; R=3000
    counts=np.zeros(n)
    for r in range(R):
        rng=random.Random(4000+r); cand=rng.sample(range(n),M)
        for i in coherent_sample(c,cand,B,rng): counts[i]+=1
    p=counts/R; loss=np.linspace(0.2,1.2,n)
    estimates=[]; uni=[]
    for r in range(2000):
        rng=random.Random(8000+r); cand=rng.sample(range(n),M); s=coherent_sample(c,cand,B,rng)
        estimates.append(sum(loss[i]/max(p[i],1/R) for i in s)/n)
        u=uniform_sample(n,B,rng); uni.append(n*sum(loss[i] for i in u)/(B*n))
    truth=float(np.mean(loss))
    return {"true_mean_loss":truth,"coherent_HT_mean":float(np.mean(estimates)),
            "coherent_HT_std":float(np.std(estimates)),"uniform_mean":float(np.mean(uni)),
            "uniform_std":float(np.std(uni)),"min_estimated_p":float(np.min(p))}

def main():
    t=time.time()
    result={"sparse_math":verify_sparse_math(),
            "disconnected_sweep":prediction_disconnected_sweep(),
            "candidate_pool_sweep":prediction_pool_sweep(),
            "coupling_sweep":prediction_coupling_sweep(),
            "estimator":estimator_variance(),
            "runtime_sec":time.time()-t}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__=="__main__": main()