"""Deterministic numerical checks for the claims in THEORY.md."""
import itertools
import math

import numpy as np


def simultaneous_mse(rng, depth, width, directions, trials=12000):
    gradients = rng.normal(size=(depth, width))
    gradients /= np.linalg.norm(gradients, axis=1, keepdims=True)
    estimate = np.zeros((trials, width), dtype=np.float64)
    for _ in range(directions):
        xi = rng.integers(0, 2, size=(trials, depth, width), dtype=np.int8)
        xi = 2.0 * xi - 1.0
        directional = np.einsum("tdw,dw->t", xi, gradients)
        estimate += xi[:, 0, :] * directional[:, None] / directions
    empirical = np.square(estimate - gradients[0]).sum(axis=1).mean()
    total_energy = np.square(gradients).sum()
    theoretical = (width * total_energy - np.square(gradients[0]).sum()) / directions
    return empirical, theoretical


def check_simultaneous_variance():
    print("SIMULTANEOUS RADEMACHER MSE")
    print("depth width K empirical theory ratio")
    cases = ((2, 8, 1), (4, 8, 1), (8, 8, 1),
             (4, 32, 1), (4, 32, 4), (4, 32, 16))
    for index, (depth, width, directions) in enumerate(cases):
        empirical, theoretical = simultaneous_mse(
            np.random.default_rng(100 + index), depth, width, directions)
        ratio = empirical / theoretical
        print(f"{depth:5d} {width:5d} {directions:2d} "
              f"{empirical:9.4f} {theoretical:9.4f} {ratio:6.3f}")
        assert abs(ratio - 1.0) < 0.06


def cubic_loss(z, linear, cubic):
    return linear @ z + cubic * np.power(z, 3).sum() / 6.0


def check_sigma_bias():
    # Exhaustive directions remove Monte Carlo error, exposing only the
    # centered-difference bias. For this separable cubic it is exactly O(sigma^2).
    dim = 8
    z = np.linspace(-0.4, 0.5, dim)
    linear = np.linspace(0.2, 0.9, dim)
    cubic = 1.7
    true_gradient = linear + 0.5 * cubic * np.square(z)
    directions = np.asarray(list(itertools.product((-1.0, 1.0), repeat=dim)))
    sigmas = np.asarray((0.005, 0.01, 0.02, 0.04))
    biases = []
    print("\nCENTERED FINITE-DIFFERENCE BIAS")
    print("sigma bias_norm")
    for sigma in sigmas:
        estimates = []
        for xi in directions:
            coefficient = (cubic_loss(z + sigma * xi, linear, cubic)
                           - cubic_loss(z - sigma * xi, linear, cubic)) / (2.0 * sigma)
            estimates.append(xi * coefficient)
        bias = np.linalg.norm(np.mean(estimates, axis=0) - true_gradient)
        biases.append(bias)
        print(f"{sigma:5.3f} {bias:.9e}")
    slope = np.polyfit(np.log(sigmas), np.log(biases), 1)[0]
    print(f"log-log slope={slope:.6f}")
    assert abs(slope - 2.0) < 1e-6


def check_descent_threshold():
    rng = np.random.default_rng(301)
    dim = 20
    gradient = rng.normal(size=dim)
    orthogonal = rng.normal(size=dim)
    orthogonal -= orthogonal.dot(gradient) * gradient / gradient.dot(gradient)
    orthogonal /= np.linalg.norm(orthogonal)
    neg_gradient = -gradient / np.linalg.norm(gradient)
    cosine = 0.35
    update = cosine * neg_gradient + math.sqrt(1.0 - cosine ** 2) * orthogonal
    beta = 2.3
    threshold = 2.0 * cosine * np.linalg.norm(gradient) / (beta * np.linalg.norm(update))

    def quadratic_change(step):
        return step * gradient.dot(update) + 0.5 * beta * step ** 2 * update.dot(update)

    below = quadratic_change(0.9 * threshold)
    above = quadratic_change(1.1 * threshold)
    print("\nSMOOTH DESCENT THRESHOLD")
    print(f"cos={cosine:.3f} eta_max={threshold:.6f} "
          f"delta_below={below:+.6f} delta_above={above:+.6f}")
    assert below < 0 < above


def check_predictor_timescale():
    rng = np.random.default_rng(401)
    examples = 50000
    cells = 12
    h = rng.normal(size=(examples, cells))
    slopes = np.exp(rng.normal(scale=0.25, size=cells))
    traffic = h * slopes
    print("\nNEUTRAL PREDICTOR TIMESCALE")
    print("eta steps residual_power empirical theory")
    for eta in (0.002, 0.01, 0.05):
        for steps in (0, 20, 100, 500):
            predictor = slopes * (1.0 - (1.0 - eta) ** steps)
            residual = traffic - h * predictor
            empirical = np.square(residual).mean() / np.square(traffic).mean()
            theoretical = (1.0 - eta) ** (2 * steps)
            print(f"{eta:5.3f} {steps:5d} {empirical:14.8f} {theoretical:14.8f}")
            assert abs(empirical - theoretical) < 1e-12


def squared_cosine(x, y):
    return float((x.ravel() @ y.ravel()) ** 2
                 / ((x.ravel() @ x.ravel()) * (y.ravel() @ y.ravel())))


def check_innovation_identification():
    rng = np.random.default_rng(501)
    examples = 100000
    cells = 8
    h = rng.normal(size=(examples, cells))
    traffic_slopes = np.linspace(0.8, 1.5, cells)
    teaching_slopes = np.linspace(0.3, 0.7, cells)
    innovation = rng.normal(scale=0.6, size=(examples, cells))
    noise = rng.normal(scale=0.15, size=(examples, cells))
    teaching = h * teaching_slopes + innovation
    traffic = h * traffic_slopes
    apical = teaching + traffic + noise

    neutral_residual = apical - traffic
    task_coeff = (apical * h).mean(axis=0) / np.square(h).mean(axis=0)
    task_residual = apical - h * task_coeff
    raw_alignment = squared_cosine(apical, teaching)
    neutral_alignment = squared_cosine(neutral_residual, teaching)
    task_alignment = squared_cosine(task_residual, teaching)
    predictable_fraction = np.square(h * teaching_slopes).mean() / np.square(teaching).mean()
    retained_fraction = np.square(teaching - h * teaching_slopes).mean() / np.square(teaching).mean()
    print("\nINNOVATION IDENTIFICATION")
    print(f"squared cosine raw={raw_alignment:.4f} neutral={neutral_alignment:.4f} "
          f"task_fit={task_alignment:.4f}")
    print(f"teaching predictable={predictable_fraction:.4f} "
          f"retained_after_task_fit={retained_fraction:.4f}")
    assert neutral_alignment > raw_alignment + 0.15
    assert task_alignment < neutral_alignment - 0.15
    assert abs(retained_fraction - (1.0 - predictable_fraction)) < 0.01


def main():
    check_simultaneous_variance()
    check_sigma_bias()
    check_descent_threshold()
    check_predictor_timescale()
    check_innovation_identification()
    print("\nALL THEORY CHECKS PASSED")


if __name__ == "__main__":
    main()