"""
Does the Figure-6 finite-sample gap appear under ANY plausible reading of the
paper's simulation description?  Sweeps
   * generation T = 1,3,5,10 (the paper's x-axis runs to 10)
   * the noise s.d. of the S equation (selection strength; the paper only says
     noise variances are drawn from [1,4], so 0.05 = near-deterministic
     selection is a strictly more favourable setting than the text implies)
   * average degree 2 read as |E| = d (default) and as |E| = 2d
   * PC and GES, both with the paper's hyper-parameters
Run:  python3 exp6_variants.py
"""
import json, os, warnings, 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)
N = 5000


def one(task):
    d, T, snoise, deg, seed = task
    from evosel import random_static_dag, simulate_evolution, clique_augmented, cpdag, SEL
    from exp6_synthetic import _edges_from_cl, _score
    from causallearn.search.ConstraintBased.PC import pc
    from causallearn.search.ScoreBased.GES import ges
    rng = np.random.default_rng(97 * seed + 7 * d + T)
    G = random_static_dag(d, rng, avg_deg=deg)
    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)
    X, _ = simulate_evolution(G, d, T, N, rng, s_noise=snoise)
    out = {'d': d, 'T': T, 's_noise': snoise, 'avg_deg': deg, 'seed': seed,
           'anS': len(anS), 'oracle_standard': orc['precision_standard'],
           'oracle_ours': orc['precision_ours']}
    cg = pc(X, alpha=0.05, indep_test='fisherz', show_progress=False)
    a, b = _edges_from_cl(cg.G, d)
    s = _score(a, b, true, anS)
    out['PC_standard'], out['PC_ours'] = s['precision_standard'], s['precision_ours']
    out['PC_n_oriented'], out['PC_n_adj'] = s['n_oriented'], s['n_adjacencies']
    rec = ges(X, score_func='local_score_BIC', maxP=None,
              parameters={'lambda_value': 2.0})
    a, b = _edges_from_cl(rec['G'], d)
    s = _score(a, b, true, anS)
    out['GES_standard'], out['GES_ours'] = s['precision_standard'], s['precision_ours']
    out['GES_n_oriented'], out['GES_n_adj'] = s['n_oriented'], s['n_adjacencies']
    return out


if __name__ == '__main__':
    import multiprocessing as mp
    tasks = []
    for T in (1, 3, 5, 10):
        for sn in (0.05, 1.0, 2.0):
            for seed in range(20):
                tasks.append((20, T, sn, 2.0, seed))
    for T in (1, 3, 5, 10):
        for seed in range(20):
            tasks.append((20, T, 1.0, 4.0, seed))   # avg degree read as |E| = 2d
    t0 = time.time()
    with mp.Pool(processes=min(mp.cpu_count(), 12)) as pool:
        rows = pool.map(one, tasks, chunksize=1)
    agg = {}
    for r in rows:
        k = 'T%d_snoise%s_deg%s' % (r['T'], r['s_noise'], r['avg_deg'])
        agg.setdefault(k, []).append(r)
    summ = {}
    for k, v in agg.items():
        e = {'n': len(v)}
        for f in ('oracle_standard', 'oracle_ours', 'PC_standard', 'PC_ours',
                  'GES_standard', 'GES_ours', 'PC_n_oriented', 'PC_n_adj',
                  'GES_n_oriented', 'GES_n_adj'):
            a = np.array([x[f] for x in v], dtype=float)
            a = a[~np.isnan(a)]
            e[f] = float(a.mean())
        e['PC_gap'] = e['PC_ours'] - e['PC_standard']
        e['GES_gap'] = e['GES_ours'] - e['GES_standard']
        e['PC_win_rate'] = float(np.mean([x['PC_ours'] > x['PC_standard'] for x in v]))
        e['GES_win_rate'] = float(np.mean([x['GES_ours'] > x['GES_standard'] for x in v]))
        summ[k] = e
    json.dump({'summary': summ, 'runs': rows, 'runtime_sec': time.time() - t0},
              open(os.path.join(OUT, 'variants.json'), 'w'), indent=1)
    for k in sorted(summ):
        e = summ[k]
        print('%-24s PC %.3f/%.3f gap %+.3f win %.2f | GES %.3f/%.3f gap %+.3f win %.2f' %
              (k, e['PC_standard'], e['PC_ours'], e['PC_gap'], e['PC_win_rate'],
               e['GES_standard'], e['GES_ours'], e['GES_gap'], e['GES_win_rate']))