"""Simulations for /guides/pearson-vs-spearman-correlation/. Seed 20261110, fixed before running.
Python 3.13, NumPy, SciPy 1.18.1. Run: python simulation.py  -> simulation_output.json

1. Width of the 95% Fisher-z interval for Pearson r by n (exact formula, no simulation).
2. How far a sample r strays from the true correlation: bivariate normal, true rho = 0.5.
3. Coverage of 95% intervals (bivariate normal, rho = 0.5):
   Pearson Fisher-z interval; Spearman Fisher-z interval with variance 1/(n-3) and with (1 + rs^2/2)/(n-3).
   Target for Spearman: the population Spearman correlation (6/pi) asin(rho/2).
4. False-positive rate of scipy.stats.spearmanr's p-value at small n (independent normal data).
"""
import json
import math
import pathlib

import numpy as np
from scipy import stats

here = pathlib.Path(__file__).parent
rng = np.random.default_rng(20261110)
out = {}
q = stats.norm.ppf(0.975)

# 1. CI width
widths = {}
for r in (0.3, 0.5, 0.8):
    for n in (5, 6, 8, 10, 12, 20, 30, 50, 100):
        z = math.atanh(r); se = 1 / math.sqrt(n - 3)
        widths[f"r{r}_n{n}"] = [round(math.tanh(z - q * se), 4), round(math.tanh(z + q * se), 4)]
out["ci_width"] = widths


def bvn(n, rho, reps):
    a = rng.standard_normal((reps, n))
    b = rho * a + math.sqrt(1 - rho ** 2) * rng.standard_normal((reps, n))
    return a, b


def rowcorr(a, b):
    a = a - a.mean(1, keepdims=True); b = b - b.mean(1, keepdims=True)
    return (a * b).sum(1) / np.sqrt((a * a).sum(1) * (b * b).sum(1))


def ranks(m):
    return np.argsort(np.argsort(m, axis=1), axis=1).astype(float)


REPS = 20_000
spread = {}
for n in (8, 12, 30, 100):
    a, b = bvn(n, 0.5, REPS)
    r = rowcorr(a, b)
    spread[str(n)] = [round(float(np.quantile(r, 0.025)), 3), round(float(np.quantile(r, 0.975)), 3)]
out["spread_rho0.5"] = spread

# 3. coverage
rho = 0.5
rho_s = 6 / math.pi * math.asin(rho / 2)
cov = {}
for n in (8, 12, 20, 40):
    a, b = bvn(n, rho, REPS)
    r = rowcorr(a, b)
    z = np.arctanh(np.clip(r, -0.999999, 0.999999)); se = 1 / math.sqrt(n - 3)
    cp = np.mean((np.tanh(z - q * se) <= rho) & (rho <= np.tanh(z + q * se)))
    rs = rowcorr(ranks(a), ranks(b))
    zs = np.arctanh(np.clip(rs, -0.999999, 0.999999))
    c1 = np.mean((np.tanh(zs - q * se) <= rho_s) & (rho_s <= np.tanh(zs + q * se)))
    se2 = np.sqrt((1 + rs ** 2 / 2) / (n - 3))
    c2 = np.mean((np.tanh(zs - q * se2) <= rho_s) & (rho_s <= np.tanh(zs + q * se2)))
    cov[str(n)] = {"pearson_fisher": round(float(cp), 4), "spearman_var_1_over_n3": round(float(c1), 4),
                   "spearman_var_1_plus_rs2_over2": round(float(c2), 4)}
out["coverage_rho0.5"] = {"population_spearman": rho_s, "by_n": cov}

# 4. scipy spearmanr p at small n under independence (no ties)
fp = {}
for n in (5, 6, 8, 10, 15):
    a, b = bvn(n, 0.0, REPS)
    rs = rowcorr(ranks(a), ranks(b))
    # SciPy's spearmanr p-value: t = rs*sqrt((n-2)/(1-rs^2)) on n-2 df (same formula as scipy uses)
    with np.errstate(divide="ignore"):
        t = rs * np.sqrt((n - 2) / np.maximum(1 - rs ** 2, 1e-300))
    p = 2 * stats.t.sf(np.abs(t), n - 2)
    fp[str(n)] = round(float(np.mean(p < 0.05)), 4)
# check that the vectorised formula matches scipy on a few samples
a, b = bvn(8, 0.0, 5)
check = [abs(float(stats.spearmanr(a[i], b[i]).pvalue) -
             float(2 * stats.t.sf(abs(rowcorr(ranks(a[i:i+1]), ranks(b[i:i+1]))[0]) *
                                  math.sqrt(6 / (1 - rowcorr(ranks(a[i:i+1]), ranks(b[i:i+1]))[0] ** 2)), 6)))
         for i in range(5)]
out["scipy_spearman_false_positive"] = {"by_n": fp, "max_abs_diff_vs_scipy": max(check)}
out["reps"] = REPS
out["seed"] = 20261110

json.dump(out, open(here / "simulation_output.json", "w"), indent=2)
print(json.dumps(out, indent=1))