import json
import math
import numpy as np
import torch
from orientation_pooling import doubled_angle_pool, orientation_from_pool


def energies_for_angles(angles, weights, B=256):
    # Direct impulses at arbitrary angles represented exactly by continuous moments;
    # bins are used only for the neural pooling API, with nearest-bin assignment.
    e = torch.zeros((B,), dtype=torch.float64)
    for a, w in zip(angles, weights):
        b = int(round((a % math.pi) / math.pi * B)) % B
        e[b] += w
    return e


def circular_axis_error(a, b):
    d = abs(float(a-b)) % math.pi
    return min(d, math.pi-d)


def main():
    torch.manual_seed(7)
    np.random.seed(7)
    B = 256
    # Prediction 1: orientation reversal is exactly identical because bins live on [0, pi).
    rng = np.random.default_rng(7)
    reversal_err = []
    for _ in range(100):
        angles = rng.uniform(0, math.pi, 8)
        weights = rng.uniform(.1, 2., 8)
        e1 = energies_for_angles(angles, weights, B)
        e2 = energies_for_angles(np.r_[angles, angles + math.pi], np.r_[weights, np.zeros(8)], B)
        p1, _ = doubled_angle_pool(e1)
        p2, _ = doubled_angle_pool(e2)
        reversal_err.append(float((p1-p2).abs().max()))

    # Prediction 2: equal two-orientation mixture has q=|cos(delta)|.
    deltas = np.linspace(0, math.pi/2, 19)
    q_obs, q_pred, phi_err = [], [], []
    base = .37
    for d in deltas:
        e = energies_for_angles([base, base+d], [1., 1.], B)
        p, _ = doubled_angle_pool(e)
        q = float(p[-1]); q_obs.append(q); q_pred.append(abs(math.cos(d)))
        expected = (base + d/2) % math.pi if d < math.pi/2 else 0.
        # At exact cancellation angle is undefined; gated output is explicitly zero.
        ph = float(orientation_from_pool(p))
        phi_err.append(0. if abs(d-math.pi/2)<1e-12 else circular_axis_error(ph, expected))

    # Prediction 3: unequal weighted pair follows the exact resultant formula.
    deltas2 = np.linspace(.0, math.pi/2, 13)
    ratios = [.25, .5, 1., 2., 4.]
    weighted_errors = []
    for ratio in ratios:
        for d in deltas2:
            a, b = 1., ratio
            e = energies_for_angles([base, base+d], [a,b], B)
            p, _ = doubled_angle_pool(e)
            pred = math.sqrt(a*a+b*b+2*a*b*math.cos(2*d))/(a+b)
            weighted_errors.append(abs(float(p[-1])-pred))

    # Tiny downstream check: classify a randomly chosen unoriented bin. A model
    # trained on raw directed 2B bins sees reversal as a new class; pooled channels do not.
    # We instead test nearest-prototype generalization after training on angles in [0,pi).
    train_angles = np.linspace(0, math.pi, 32, endpoint=False)
    prototypes = []
    for a in train_angles:
        e = energies_for_angles([a], [1.], B)
        prototypes.append(doubled_angle_pool(e)[0].numpy())
    prototypes = np.asarray(prototypes)
    test = [(a+math.pi) % (2*math.pi) for a in train_angles]
    pooled_correct = 0
    for a in test:
        e = energies_for_angles([a], [1.], B)
        x = doubled_angle_pool(e)[0].numpy()
        pred = np.argmin(((prototypes-x)**2).sum(1))
        pooled_correct += (pred == int(round((a % math.pi)/math.pi*32)) % 32)
    # Baselines on the same prototype task. Directed bins use [0, 2pi),
    # so reversal shifts the representation by pi; scalar averaging has no
    # orientation information and must guess among classes.
    def directed_hist(angle, n=64):
        h = np.zeros(n)
        h[int(round((angle % (2*math.pi)) / (2*math.pi) * n)) % n] = 1.
        return h
    directed_prototypes = np.asarray([directed_hist(a) for a in train_angles])
    directed_correct = 0
    scalar_correct = 0
    for a in test:
        h = directed_hist(a)
        pred = np.argmin(((directed_prototypes-h)**2).sum(1))
        directed_correct += (pred == int(round((a % math.pi)/math.pi*32)) % 32)
        # Equal-energy scalar representations tie; deterministic argmin guesses class 0.
        scalar_correct += (0 == int(round((a % math.pi)/math.pi*32)) % 32)

    results = {
        "config": {"B": B, "seed": 7},
        "prediction_1_reversal_max_abs_error": max(reversal_err),
        "prediction_1_expected": 0.0,
        "prediction_2_deltas_deg": [round(float(x*180/math.pi), 3) for x in deltas],
        "prediction_2_q_observed": q_obs,
        "prediction_2_q_predicted_abs_cos_delta": q_pred,
        "prediction_2_max_q_error": max(abs(a-b) for a,b in zip(q_obs,q_pred)),
        "prediction_2_max_orientation_error_non_cancel_deg": max(phi_err)*180/math.pi,
        "prediction_3_max_weighted_q_error": max(weighted_errors),
        "prediction_3_expected": 0.0,
        "reversal_prototype_accuracy": pooled_correct/len(test),
        "directed_bin_baseline_reversal_accuracy": directed_correct/len(test),
        "scalar_average_baseline_reversal_accuracy": scalar_correct/len(test),
        "cancellation_q_at_90deg": q_obs[-1],
    }
    print(json.dumps(results, indent=2))

if __name__ == '__main__':
    main()