import json
import random
from pathlib import Path

import numpy as np
from sklearn.metrics import mutual_info_score


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)


def entropy(x: np.ndarray) -> float:
    _, counts = np.unique(x, return_counts=True)
    p = counts.astype(float) / len(x)
    return float(-(p * np.log(p + 1e-15)).sum())


def mi(x: np.ndarray, y: np.ndarray) -> float:
    return float(mutual_info_score(x, y))


def quantize(x: np.ndarray, delta: float) -> np.ndarray:
    return np.floor(x / delta + 0.5).astype(np.int64)


def sample(seed: int, n: int = 120_000, noise: float = 0.35):
    set_seed(seed)
    theta = np.random.normal(size=n)
    g = theta + np.random.normal(scale=noise, size=n)
    # A fixed-magnitude update is useful on the quadratic iff it moves toward 0.
    useful_action = (theta * g > 0).astype(np.int64)
    return theta, g, useful_action


def channel_sweep(seed=7):
    theta, g, utility = sample(seed)
    deltas = [2.0, 1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125]
    out = []
    for delta in deltas:
        z = quantize(g, delta)
        acquired = entropy(z)
        useful = mi(z, utility)
        unavailable = acquired - useful
        # Controller: sign observation, with zero-bin rejection.
        eta = 0.5
        update = -eta * np.sign(z)
        theta_new = theta + update
        progress = float(np.mean(0.5 * theta**2 - 0.5 * theta_new**2))
        action_correct = float(np.mean((z != 0) & (np.sign(z) == np.sign(theta))))
        out.append({"delta": delta, "I_acq": acquired, "I_use": useful,
                    "I_un": unavailable, "progress": progress,
                    "action_correct": action_correct,
                    "zero_fraction": float(np.mean(z == 0))})
    return out


def learning_sweep(seed=7):
    # Compare raw high-resolution SGD with the actionable finite-resolution
    # controller on L(theta)=theta^2/2, averaged over fixed initial points.
    set_seed(seed)
    initials = np.random.normal(size=2000)
    rows = []
    for delta in [None, 2.0, 1.0, 0.5, 0.25, 0.125]:
        theta = initials.copy()
        for _ in range(25):
            g = theta + np.random.normal(scale=0.35, size=len(theta))
            if delta is None:
                z = g
            else:
                z = quantize(g, delta) * delta
            theta -= 0.35 * np.sign(z) * np.minimum(np.abs(z), 2.0)
        rows.append({"method": "raw" if delta is None else f"quantized_{delta}",
                     "final_loss": float(np.mean(0.5 * theta**2))})
    return rows


def predictions(rows):
    # Prediction 1: acquired information grows approximately linearly in
    # -log(delta), with slope equal to the effective one-dimensional channel.
    x = np.array([-np.log(r["delta"]) for r in rows[-5:]])
    y = np.array([r["I_acq"] for r in rows[-5:]])
    slope, intercept = np.polyfit(x, y, 1)
    r2 = float(np.corrcoef(x, y)[0, 1] ** 2)
    # Prediction 2: useful information saturates while acquired information
    # keeps rising; compare last two increments.
    du_acq = rows[-1]["I_acq"] - rows[-2]["I_acq"]
    du_use = rows[-1]["I_use"] - rows[-2]["I_use"]
    # Prediction 3: progress saturates at fine resolution.
    dp = rows[-1]["progress"] - rows[-2]["progress"]
    return {"acq_log_slope": float(slope), "acq_log_r2": r2,
            "fine_delta_I_acq": float(du_acq),
            "fine_delta_I_use": float(du_use),
            "fine_delta_progress": float(dp),
            "predicted_acq_slope": 1.0,
            "predicted_useful_trend": "plateau relative to I_acq",
            "predicted_progress_trend": "plateau relative to I_acq"}


def main():
    channel = channel_sweep()
    result = {"channel": channel, "learning": learning_sweep(),
              "predictions": predictions(channel)}
    Path("toy_results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()