import json, math
from pathlib import Path
import numpy as np
from scipy.optimize import nnls

SEED = 2184
rng = np.random.default_rng(SEED)


def lasso_cd(A, b, lam=0.04, iters=1200):
    # objective 0.5||Ac-b||^2 + lam||c||_1, deterministic coordinate descent
    A = np.asarray(A, float); b = np.asarray(b, float)
    L = A.shape[1]; c = np.zeros(L)
    col = (A*A).sum(0) + 1e-12
    for _ in range(iters):
        old = c.copy()
        for k in range(L):
            rho = A[:, k] @ (b - A @ c + A[:, k] * c[k])
            z = abs(rho) - lam
            c[k] = np.sign(rho) * max(z, 0.) / col[k]
        if np.max(abs(c-old)) < 1e-9: break
    return c


def weak_patch_integral(n, sigma=1.0, repeats=300):
    # Integral of a smooth function with iid observation noise; Monte Carlo weak integration.
    vals=[]
    for _ in range(repeats):
        x = rng.random(n)
        phi = 6*x*(1-x) # compact-like bump on patch, normalized below
        truth = np.sin(2*np.pi*x) + .35*x
        vals.append(np.mean(phi*(truth + sigma*rng.normal(size=n))) / np.mean(phi))
    return np.std(vals, ddof=1)


def consensus_error(p, m, trials=30000):
    # binary local support decisions, modal support with random tie breaking
    x = rng.random((trials, m)) < p
    votes = x.sum(1)
    return np.mean(votes < (m/2)) if m % 2 else np.mean(votes <= (m//2-1))


def patch_regression_demo(npatch=80, q=24, noise=.22, lam=.045):
    # Two candidate weak operators; true support is term 0. Each local patch gets q
    # quadrature samples and noisy b. LASSO support is then pooled by modal consensus.
    local = []
    for j in range(npatch):
        # normalized test-function integrations create mildly varying local designs
        A = rng.normal(size=(3,2)); A[:,0] += 1.4
        ctrue = np.array([1.0, 0.0])
        b = A @ ctrue + noise*rng.normal(size=3)/math.sqrt(q)
        c = lasso_cd(A, b, lam)
        local.append(c)
    local=np.array(local)
    supports=np.abs(local)>0.16
    true_support=np.array([True, False])
    local_exact=np.all(supports == true_support[None, :], axis=1)
    modal_support=np.sum(supports, axis=0) >= np.ceil(npatch/2)
    modal_exact=bool(np.all(modal_support == true_support))
    return float(np.mean(local_exact)), modal_exact, local


def main():
    global rng
    # Prediction 1: weak integration averages independent noise, std proportional n^-1/2.
    ns=np.array([8,16,32,64,128,256])
    stds=np.array([weak_patch_integral(int(n)) for n in ns])
    slope=float(np.polyfit(np.log(ns), np.log(stds), 1)[0])

    # Prediction 2: local decisions only improve consensus when p > 1/2.
    ps=[.40,.55,.70,.85]
    ms=[3,7,15,31]
    rows=[]
    for p in ps:
        for m in ms:
            obs=consensus_error(p,m)
            bound=math.exp(-2*m*(p-.5)**2)
            rows.append({'p':p,'m':m,'observed_error':float(obs),'hoeffding_bound':float(bound)})

    # Prediction 3: with a fixed local accuracy p=.7, more patches reduce errors;
    # check monotonicity and compare the observed trend to the theoretical bound.
    fixed=[r for r in rows if r['p']==.70]
    monotonic=all(fixed[i]['observed_error'] >= fixed[i+1]['observed_error']-0.004 for i in range(len(fixed)-1))
    below_bound=all(r['observed_error'] <= r['hoeffding_bound']+0.015 for r in rows if r['p']>.5)

    # Secondary MVP comparison: local sparse fits versus modal support on noisy patches.
    rng=np.random.default_rng(SEED+9)
    reps=[]
    for _ in range(100):
        frac, modal, local=patch_regression_demo()
        reps.append((frac,modal))
    local_support=float(np.mean([x[0] for x in reps]))
    consensus_support_accuracy=float(np.mean([x[1] for x in reps]))

    result={
      'seed':SEED,
      'prediction_checks':{
        'weak_noise_scaling':{'predicted_loglog_slope':-0.5,'observed_loglog_slope':slope,'tolerance':0.12,'passed':abs(slope+0.5)<0.12,'n':ns.tolist(),'std':stds.tolist()},
        'consensus_transition':{'prediction':'majority improves for p>0.5 and worsens for p<0.5','rows':rows,'passed':(rows[0]['observed_error']>0.5 and rows[4]['observed_error']<0.5 and rows[8]['observed_error']<0.5)},
        'patch_count_scaling':{'prediction':'error decreases with m for p=.70 and stays below Hoeffding upper bound','rows':fixed,'monotonic_with_sampling_tolerance':monotonic,'below_bound_with_sampling_tolerance':below_bound,'passed':monotonic and below_bound}
      },
      'secondary_mvp':{'local_exact_support_accuracy_mean':local_support,'modal_consensus_exact_accuracy':consensus_support_accuracy,'interpretation':'consensus converts noisy local support votes into a region-level decision'},
      'all_mechanism_checks_passed':bool(abs(slope+0.5)<0.12 and rows[0]['observed_error']>0.5 and rows[4]['observed_error']<0.5 and monotonic and below_bound)
    }
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()