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

EPS = 1e-3


def bts_score(p_hat, y, eps=EPS):
    return math.log((p_hat + eps) / (y + eps))


def prediction_frequency_scaling():
    # At fixed y, changing p changes the score by exactly
    # log((p1+eps)/(p0+eps)); without epsilon this is log(p1/p0).
    y, p0 = 0.20, 0.25
    rows = []
    for rho in [0.5, 0.8, 1.0, 1.5, 2.0]:
        p1 = min(0.95, p0 * rho)
        observed = bts_score(p1, y) - bts_score(p0, y)
        predicted = math.log((p1 + EPS) / (p0 + EPS))
        rows.append(dict(rho=rho, observed=observed, predicted=predicted,
                         abs_error=abs(observed - predicted)))
    return rows


def monotonicity_sweep():
    # The derivative with respect to p_hat is 1/(p_hat+eps)>0.
    y = 0.35
    ps = np.linspace(0.02, 0.98, 25)
    rs = np.array([bts_score(float(p), y) for p in ps])
    return dict(min_adjacent_delta=float(np.min(np.diff(rs))),
                p_at_max=float(ps[np.argmax(rs)]), r_at_min=float(rs[0]),
                r_at_max=float(rs[-1]))


def variance_scaling(rng):
    # For K~Binomial(G,p), delta method predicts
    # Var[log((K/G+eps)/(y+eps))] ~= p(1-p)/(G*(p+eps)^2).
    p, y, trials = 0.30, 0.20, 50000
    rows = []
    for G in [8, 16, 32, 64, 128, 256]:
        k = rng.binomial(G, p, size=trials)
        scores = np.log((k / G + EPS) / (y + EPS))
        observed = float(np.var(scores, ddof=1))
        predicted = p * (1-p) / (G * (p + EPS)**2)
        rows.append(dict(G=G, observed=observed, predicted=predicted,
                         observed_times_G=observed*G,
                         predicted_times_G=predicted*G,
                         ratio=observed/predicted))
    return rows


def group_reward(answer, predicted, answers, eps=EPS):
    counts = np.bincount(answers, minlength=2) / len(answers)
    return np.log((counts[answer] + eps) / (predicted + eps))


def policy_simulation(rng, rounds=4000, G=16):
    # Two answer policies: honest follows latent truth; sycophantic follows
    # the user pressure. Each completion predicts the group's answer rate.
    # BTS is applied as a group-relative reward; baseline is majority/agreement.
    # This is a mechanism simulation, not neural-network training.
    results = {}
    for method in ["bts", "agreement"]:
        for condition in ["neutral", "pressure"]:
            correct = flips = 0
            reward_var = []
            for _ in range(rounds):
                truth = int(rng.random() < 0.5)
                pressure = truth ^ 1 if condition == "pressure" else truth
                # A mixed population: honest samples have 0.78 truth fidelity,
                # sycophantic samples have 0.84 pressure fidelity.
                honest = rng.random(G) < 0.55
                answers = np.where(honest,
                                  truth ^ (rng.random(G) > 0.78),
                                  pressure ^ (rng.random(G) > 0.84)).astype(int)
                # Prediction reports are calibrated to each respondent's rule.
                pred = np.where(honest, 0.55 + 0.23 * truth,
                                0.55 + 0.29 * pressure).astype(float)
                # pred is probability of answer 1; convert for each answer.
                ys = np.where(answers == 1, pred, 1-pred)
                if method == "bts":
                    scores = np.array([group_reward(int(a), float(y), answers)
                                       for a, y in zip(answers, ys)])
                    reward_var.append(float(np.var(scores)))
                    # Reward-weighted soft selection of a policy for the next
                    # prompt: BTS favors reports surprising relative to their
                    # own prediction, rather than raw consensus.
                    w = np.exp(np.clip(scores - scores.max(), -10, 10))
                    chosen = int(answers[rng.choice(G, p=w/w.sum())])
                else:
                    reward = np.bincount(answers, minlength=2) / G
                    reward_var.append(float(np.var(reward[answers])))
                    chosen = int(np.argmax(np.bincount(answers, minlength=2)))
                correct += chosen == truth
                flips += condition == "pressure" and chosen != truth
            results[f"{method}_{condition}"] = dict(accuracy=correct/rounds,
                                                       pressure_flip_rate=flips/rounds,
                                                       mean_reward_variance=float(np.mean(reward_var)))
    return results


def main():
    rng = np.random.default_rng(2735)
    out = {
        "prediction_1_log_scaling": prediction_frequency_scaling(),
        "prediction_2_monotonicity": monotonicity_sweep(),
        "prediction_3_variance_1_over_G": variance_scaling(rng),
        "policy_simulation": policy_simulation(rng),
    }
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()