"""
Why the oracle-level guarantee of Theorem 2 does or does not show up in the
finite-sample Figure-6 experiment.  Sweeps the sample size N and the dimension
d, and separates two candidate causes:
   (i) the selected distribution is not linear-Gaussian, so Fisher-z is
       misspecified   -> tested by generating data DIRECTLY from G^+ (a clean
       linear-Gaussian SCM with no selection at all);
   (ii) PC simply has too little power at d=20, N=5,000 on the dense
        clique-augmented graph.
Run:  python3 exp6_sensitivity.py
"""
import json, os, warnings, itertools, time
import numpy as np

warnings.filterwarnings('ignore')
os.environ.setdefault('OMP_NUM_THREADS', '1')
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'outputs')
os.makedirs(OUT, exist_ok=True)


def one(task):
    d, N, seed, mode = task
    import evosel as E
    from evosel import (random_static_dag, clique_augmented, cpdag,
                        simulate_evolution, sem_params, SEL)
    from exp6_synthetic import _edges_from_cl, _score
    from causallearn.search.ConstraintBased.PC import pc
    rng = np.random.default_rng(1000 * d + seed)
    G = random_static_dag(d, rng, avg_deg=2.0)
    true = set((i, j) for j in range(d) for i in G.pa[j] if i != SEL)
    anS = G.ancestors([SEL]) - {SEL}
    gp = clique_augmented(G, d)
    dd, uu = cpdag(gp)
    orc = _score(dd, uu, true, anS)
    if mode == 'evolution':
        X, _ = simulate_evolution(G, d, 3, N, rng)
    else:                     # 'gplus': linear-Gaussian SCM on G^+ directly
        B = np.zeros((d, d))
        for (u, v) in gp.edges():
            B[u, v] = rng.uniform(0.5, 2.0) * (1 if rng.random() < .5 else -1)
        eps = rng.normal(0, np.sqrt(rng.uniform(1, 4, size=d)), size=(N, d))
        X = E._traits(B, eps)
    cg = pc(X, alpha=0.05, indep_test='fisherz', show_progress=False)
    a, b = _edges_from_cl(cg.G, d)
    fin = _score(a, b, true, anS)
    skel_true = set(frozenset(e) for e in gp.edges())
    skel = set(frozenset(e) for e in a) | b
    return {'d': d, 'N': N, 'seed': seed, 'mode': mode,
            'anS': len(anS), 'n_edges_Gplus': gp.n_edges(),
            'oracle_standard': orc['precision_standard'],
            'oracle_ours': orc['precision_ours'],
            'oracle_n_oriented': orc['n_oriented'],
            'standard': fin['precision_standard'], 'ours': fin['precision_ours'],
            'n_oriented': fin['n_oriented'], 'n_adjacencies': fin['n_adjacencies'],
            'skeleton_recall_vs_Gplus': len(skel & skel_true) / max(len(skel_true), 1),
            'skeleton_precision_vs_Gplus': len(skel & skel_true) / max(len(skel), 1)}


def agg(rows, keys):
    out = {}
    for r in rows:
        k = '_'.join(str(r[x]) for x in keys)
        out.setdefault(k, []).append(r)
    res = {}
    for k, v in out.items():
        res[k] = {'n': len(v)}
        for f in ('standard', 'ours', 'oracle_standard', 'oracle_ours',
                  'skeleton_recall_vs_Gplus', 'skeleton_precision_vs_Gplus',
                  'n_oriented', 'n_adjacencies'):
            a = np.array([x[f] for x in v], dtype=float)
            a = a[~np.isnan(a)]
            res[k][f] = float(a.mean())
        res[k]['ours_minus_standard'] = res[k]['ours'] - res[k]['standard']
        res[k]['frac_runs_ours_gt_standard'] = float(np.mean(
            [x['ours'] > x['standard'] for x in v if not np.isnan(x['ours'])]))
    return res


if __name__ == '__main__':
    import multiprocessing as mp
    tasks = []
    for d in (10, 15, 20):
        for N in (5000, 20000, 100000, 400000):
            for s in range(20):
                tasks.append((d, N, s, 'evolution'))
    for d in (10, 20):
        for N in (5000, 100000, 400000):
            for s in range(20):
                tasks.append((d, N, s, 'gplus'))
    t0 = time.time()
    with mp.Pool(processes=min(mp.cpu_count(), 12)) as pool:
        rows = pool.map(one, tasks, chunksize=1)
    res = {'by_d_N_mode': agg(rows, ['mode', 'd', 'N']),
           'runtime_sec': time.time() - t0, 'runs': rows}
    json.dump(res, open(os.path.join(OUT, 'sensitivity.json'), 'w'), indent=1)
    for k in sorted(res['by_d_N_mode']):
        e = res['by_d_N_mode'][k]
        print('%-22s std %.3f ours %.3f  diff %+.3f  win%% %.2f  skelrec %.2f  n_or %.1f/%.1f' %
              (k, e['standard'], e['ours'], e['ours_minus_standard'],
               e['frac_runs_ours_gt_standard'], e['skeleton_recall_vs_Gplus'],
               e['n_oriented'], e['n_adjacencies']))