import json
import math
from pathlib import Path

import numpy as np
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import accuracy_score


def symmetric_power_matrix(A, n=4):
    """R_n(A), coefficient order X^n, X^(n-1)Y, ..., Y^n."""
    a, b, c, d = np.asarray(A, dtype=float).reshape(2, 2).ravel()
    R = np.zeros((n + 1, n + 1), dtype=float)
    for k in range(n + 1):
        for u in range(n - k + 1):
            left = math.comb(n - k, u) * a ** (n - k - u) * c ** u
            for v in range(k + 1):
                R[u + v, k] += left * math.comb(k, v) * b ** (k - v) * d ** v
    return R


def polynomial_value(p, x, y):
    n = len(p) - 1
    return sum(p[k] * x ** (n - k) * y ** k for k in range(n + 1))


def random_sl2(rng):
    theta = rng.uniform(-math.pi, math.pi)
    rot = np.array([[math.cos(theta), -math.sin(theta)],
                    [math.sin(theta), math.cos(theta)]])
    t = rng.uniform(-1.1, 1.1)
    diag = np.diag([math.exp(t), math.exp(-t)])
    s = rng.uniform(-0.8, 0.8)
    return rot @ diag @ np.array([[1., s], [0., 1.]])


def monomials(points, n=4):
    x, y = points[:, 0], points[:, 1]
    return np.stack([x ** (n-k) * y ** k for k in range(n+1)], axis=1)


def rotation_points(rng, radii, angles):
    return np.stack([radii * np.cos(angles), radii * np.sin(angles)], axis=1)


def mini_classification(rng, n=4):
    # Binary radius task. The proposed representation has an exact SO(2)-invariant
    # quadratic norm in the binomially normalized coefficient basis.
    n_train = 500
    radii = np.where(rng.integers(0, 2, n_train) == 0, .8, 1.2)
    labels = (radii > 1).astype(int)
    angles = rng.uniform(-math.pi, math.pi, n_train)
    train_xy = rotation_points(rng, radii, angles)
    train_p = monomials(train_xy, n)
    weights = np.array([math.comb(n, k) for k in range(n+1)], dtype=float)
    train_inv = np.sum(weights[None, :] * train_p**2, axis=1)[:, None]

    # Same small MLP, trained only on the available rotated examples.
    baseline = MLPClassifier(hidden_layer_sizes=(16, 16), max_iter=500,
                             random_state=123, solver='lbfgs')
    baseline.fit(train_xy, labels)

    # Calibrate the equivariant scalar threshold on training data; no learned
    # transformation-specific parameters are used by this feature.
    threshold = (train_inv[labels == 0].mean() + train_inv[labels == 1].mean()) / 2

    # Held-out transformations are rotations in a disjoint angular sector.
    test_r = np.repeat([.8, 1.2], 400)
    test_y = np.repeat([0, 1], 400)
    test_angles = np.concatenate([rng.uniform(0.7, math.pi, 400),
                                  rng.uniform(0.7, math.pi, 400)])
    test_xy = rotation_points(rng, test_r, test_angles)
    test_p = monomials(test_xy, n)
    test_inv = np.sum(weights[None, :] * test_p**2, axis=1)
    equiv_pred = (test_inv > threshold).astype(int)
    base_pred = baseline.predict(test_xy)

    # Directly verify invariance of the proposed scalar on random rotations.
    invariant_err = []
    for _ in range(100):
        pxy = rng.normal(size=(1, 2))
        phi = rng.uniform(-math.pi, math.pi)
        A = np.array([[math.cos(phi), -math.sin(phi)],
                      [math.sin(phi), math.cos(phi)]])
        p = monomials(pxy, n)[0]
        pa = monomials(pxy @ A, n)[0]
        invariant_err.append(abs(np.dot(weights, p*p)-np.dot(weights, pa*pa)) /
                            (1 + abs(np.dot(weights, p*p))))
    return {
        'baseline_mlp_test_accuracy': float(accuracy_score(test_y, base_pred)),
        'symmetric_power_invariant_test_accuracy': float(accuracy_score(test_y, equiv_pred)),
        'invariant_scalar_max_relative_error': float(max(invariant_err)),
        'train_examples': n_train,
    }


def run(seed=17, n=4, copies=3, trials=200):
    rng = np.random.default_rng(seed)
    poly_err, comp_err = [], []
    for _ in range(trials):
        A, B = random_sl2(rng), random_sl2(rng)
        p = rng.normal(size=n + 1)
        q = symmetric_power_matrix(A, n) @ p
        for _ in range(3):
            x, y = rng.normal(size=2)
            lhs = polynomial_value(q, x, y)
            z = np.array([x, y]) @ A
            rhs = polynomial_value(p, z[0], z[1])
            poly_err.append(abs(lhs-rhs) / (1 + abs(rhs)))
        comp_err.append(np.linalg.norm(symmetric_power_matrix(A @ B, n) -
                                      symmetric_power_matrix(A, n) @ symmetric_power_matrix(B, n)))

    Wcopy = rng.normal(size=(copies, copies))
    Wgood = np.kron(Wcopy, np.eye(n + 1))
    Rmany = np.kron(np.eye(copies), symmetric_power_matrix(random_sl2(rng), n))
    p = rng.normal(size=copies * (n + 1))
    good_err = np.linalg.norm(Wgood @ Rmany @ p - Rmany @ Wgood @ p) / (1e-12 + np.linalg.norm(Rmany @ Wgood @ p))
    Wbad = rng.normal(size=Wgood.shape)
    bad_err = np.linalg.norm(Wbad @ Rmany @ p - Rmany @ Wbad @ p) / (1e-12 + np.linalg.norm(Rmany @ Wbad @ p))
    good, bad = [], []
    for _ in range(trials):
        Rm = np.kron(np.eye(copies), symmetric_power_matrix(random_sl2(rng), n))
        x = rng.normal(size=copies*(n+1))
        good.append(np.linalg.norm(Wgood @ Rm @ x - Rm @ Wgood @ x) / (1e-12 + np.linalg.norm(Rm @ Wgood @ x)))
        bad.append(np.linalg.norm(Wbad @ Rm @ x - Rm @ Wbad @ x) / (1e-12 + np.linalg.norm(Rm @ Wbad @ x)))
    result = {
        'seed': seed, 'degree': n, 'copies': copies, 'trials': trials,
        'max_polynomial_relative_error': float(max(poly_err)),
        'max_composition_absolute_error': float(max(comp_err)),
        'single_good_layer_relative_error': float(good_err),
        'single_bad_layer_relative_error': float(bad_err),
        'median_good_layer_error': float(np.median(good)),
        'median_bad_layer_error': float(np.median(bad)),
        'mean_good_layer_error': float(np.mean(good)),
        'mean_bad_layer_error': float(np.mean(bad)),
        'mini_experiment': mini_classification(rng, n),
    }
    return result


if __name__ == '__main__':
    out = run()
    print(json.dumps(out, indent=2))
    Path('results.json').write_text(json.dumps(out, indent=2) + '\n')