import json
import math
from pathlib import Path
import numpy as np


def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -50.0, 50.0)))


def q_values(p_m, eps, beta, p_a):
    q_h = (1 - p_m) * eps + p_m * (1 - eps)
    return q_h, (q_h + beta * p_a) / (1 + beta)


def update_pm(p_m, rho, kappa, delta_pi):
    return (1 - rho) * p_m + rho * sigmoid(kappa * delta_pi)


def map_decision(reports, qs):
    reports = np.asarray(reports)
    qs = np.clip(np.asarray(qs), 1e-3, 1 - 1e-3)
    score = np.sum(np.where(reports == 1, np.log((1-qs)/qs),
                            np.log(qs/(1-qs))))
    return 1 if score >= 0 else -1


def implied_pa_threshold(p_m, eps, beta):
    """P_a solving q=1/2 under the stated q formula."""
    qh, _ = q_values(p_m, eps, beta, 0.0)
    return (0.5 * (1 + beta) - qh) / beta


def math_checks():
    eps, beta = .12, .35
    p = 0.0
    for _ in range(300):
        p = update_pm(p, .2, 3., 1.)
    qh0, q0 = q_values(0, eps, beta, .7)
    qh1, q1 = q_values(1, eps, beta, .7)
    return {
        "bounded_dynamics": bool(0 <= p <= 1),
        "q_formula_endpoints": bool(abs(qh0-eps) < 1e-12 and abs(qh1-(1-eps)) < 1e-12),
        "q_at_pa_half_p0": q_values(0, eps, beta, .5)[1],
        "q_at_pa_half_p1": q_values(1, eps, beta, .5)[1],
        "implied_pa_threshold_p0": implied_pa_threshold(0, eps, beta),
        "implied_pa_threshold_p1": implied_pa_threshold(1, eps, beta),
        "claimed_universal_threshold_holds": bool(
            abs(q_values(0, eps, beta, .5)[1]-.5) < 1e-10 and
            abs(q_values(1, eps, beta, .5)[1]-.5) < 1e-10),
    }


def one_round(rng, n_honest, beta, eps, pa, pm, method, adaptive):
    n_bad = int(round(beta*n_honest))
    truth = 1 if rng.random() < .5 else -1
    sensed = np.where(rng.random(n_honest) < eps, -truth, truth)
    honest = np.where(rng.random(n_honest) < pm, -sensed, sensed)
    bad = np.where(rng.random(n_bad) < pa, -truth, truth)
    reports = np.r_[honest, bad]
    q = q_values(pm, eps, beta, pa)[1]
    if method == 'majority':
        pred = 1 if reports.sum() >= 0 else -1
    elif method == 'map':
        pred = map_decision(reports, np.full(reports.size, q))
    else:
        raise ValueError(method)
    observed_error = np.mean(reports != truth)
    # Simple bounded-rational conformity proxy, included only to exercise the
    # proposed state update; it is not claimed to reproduce the paper's payoff.
    next_pm = update_pm(pm, .25, 4., 2*(observed_error-.5)) if adaptive else pm
    return int(pred == truth), next_pm, q, observed_error


def simulate(rng, pa, beta, eps, method, adaptive, rounds=20):
    pm, acc, qs, observed = 0., [], [], []
    for _ in range(rounds):
        a, pm, q, e = one_round(rng, 20, beta, eps, pa, pm, method, adaptive)
        acc.append(a); qs.append(q); observed.append(e)
    return np.mean(acc), np.mean(qs), pm, np.mean(observed)


def run_experiment(seed=7):
    rng = np.random.default_rng(seed)
    rows = []
    # Two beta values make dependence on attacker population explicit.
    for beta in (.25, 1.0, 2.0):
        for eps in (.05, .20):
            for pa in np.linspace(.1, .9, 9):
                for method, adaptive in [('majority', False), ('map', False),
                                         ('map', True)]:
                    vals = [simulate(rng, float(pa), beta, eps, method, adaptive)
                            for _ in range(60)]
                    rows.append({
                        'beta': beta, 'eps': eps, 'p_a': float(pa),
                        'method': method, 'adaptive': adaptive,
                        'accuracy': float(np.mean([v[0] for v in vals])),
                        'mean_q': float(np.mean([v[1] for v in vals])),
                        'final_pm': float(np.mean([v[2] for v in vals])),
                        'observed_error': float(np.mean([v[3] for v in vals]))})
    return rows


def main():
    out = {'math_checks': math_checks(), 'results': run_experiment()}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out['math_checks'], indent=2))
    for r in out['results']:
        if r['eps'] == .05 and r['method'] in ('majority', 'map') and not r['adaptive']:
            print(r)


if __name__ == '__main__':
    main()