import json
import math
import numpy as np
try:
    from scipy.special import sph_harm
except ImportError:
    from scipy.special import sph_harm_y
    def sph_harm(m, n, phi, theta):
        # New SciPy API takes (degree, order, polar, azimuth).
        return sph_harm_y(n, m, theta, phi)

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

SEED = 635
rng = np.random.default_rng(SEED)


def fibonacci_sphere(m):
    # Nearly uniform quadrature points on S^2.
    k = np.arange(m, dtype=float)
    z = 1.0 - 2.0 * (k + 0.5) / m
    phi = (math.pi * (3.0 - math.sqrt(5.0)) * k) % (2.0 * math.pi)
    r = np.sqrt(np.maximum(0.0, 1.0 - z * z))
    return np.column_stack((r * np.cos(phi), r * np.sin(phi), z))


def real_harmonics(points, lmax):
    # Real orthonormal spherical harmonics, evaluated in scipy's convention.
    x, y, z = points.T
    theta = np.arccos(np.clip(z, -1, 1))
    phi = np.arctan2(y, x) % (2 * np.pi)
    cols, degrees = [], []
    for ell in range(lmax + 1):
        for mm in range(-ell, ell + 1):
            a = sph_harm(abs(mm), ell, phi, theta)
            if mm < 0:
                h = math.sqrt(2.0) * (-1 if mm % 2 else 1) * a.imag
            elif mm == 0:
                h = a.real
            else:
                h = math.sqrt(2.0) * (-1 if mm % 2 else 1) * a.real
            cols.append(h)
            degrees.append(ell)
    return np.asarray(cols).T.real, np.asarray(degrees)


def make_funk_matrix(points, neighbors=120):
    # For each u, average samples closest to u-perpendicular (small |u.v|).
    dots = np.abs(points @ points.T)
    idx = np.argpartition(dots, neighbors, axis=1)[:, :neighbors]
    A = np.zeros_like(dots)
    rows = np.arange(len(points))[:, None]
    # A local strip is a simple positive approximation to the great-circle average.
    weights = np.exp(-(dots[rows, idx] / 0.055) ** 2)
    weights /= weights.sum(axis=1, keepdims=True)
    A[rows, idx] = weights
    return A


def projectors(H, degrees):
    # Weighted-by-sampling least-squares projectors; Fibonacci weights are uniform.
    p0 = degrees == 0
    p2 = degrees == 2
    B0, B2 = H[:, p0], H[:, p2]
    P0 = B0 @ np.linalg.pinv(B0)
    P2 = B2 @ np.linalg.pinv(B2)
    return P0, P2


def phi_squared(f, A):
    # In n=3, c_3/kappa_2 = 1 and Phi(f)=R(f^2).
    z = A @ (f ** 2)
    return A @ (z ** 2)


def rel_norm(x):
    return float(np.sqrt(np.mean(x * x)))


def diagnostic(points, A, H, degrees):
    P0, P2 = projectors(H, degrees)
    Pge4 = np.eye(len(points)) - P0 - P2
    results = {}
    # Small perturbations make the finite difference equal to the claimed derivative.
    for ell in (2, 4, 6):
        j = np.flatnonzero(degrees == ell)[0]
        h = H[:, j] / rel_norm(H[:, j])
        eps = 1e-3
        out = (phi_squared(1.0 + eps * h, A) - 1.0) / eps
        if ell == 2:
            ratio = rel_norm(P2 @ out) / rel_norm(h)
            results['degree2_multiplier'] = ratio
        else:
            ratio = rel_norm(Pge4 @ out) / rel_norm(h)
            results[f'degree{ell}_multiplier'] = ratio
    # Direct nonlinear test, with mixed ellipsoidal and high-frequency perturbation.
    h2 = H[:, np.flatnonzero(degrees == 2)[2]] / rel_norm(H[:, np.flatnonzero(degrees == 2)[2]])
    h4 = H[:, np.flatnonzero(degrees == 4)[1]] / rel_norm(H[:, np.flatnonzero(degrees == 4)[1]])
    eps = 1e-2
    before2, before4 = rel_norm(P2 @ (eps*h2)), rel_norm(Pge4 @ (eps*h4))
    after = phi_squared(1 + eps*h2 + eps*h4, A) - 1
    results['mixed_degree2_ratio'] = rel_norm(P2 @ after) / before2
    results['mixed_high_ratio'] = rel_norm(Pge4 @ after) / before4
    results['predicted_high_multiplier'] = 9.0 / 16.0
    results['predicted_gap'] = 7.0 / 16.0
    return results


def toy_classification(points, A, H, degrees, n_samples=1200):
    # Same tiny supervised task: infer the sign of an ellipsoidal (degree-2) coefficient
    # amid high-frequency angular contamination. Compare raw features with the proposed map.
    i2 = np.flatnonzero(degrees == 2)[2]
    hi = np.flatnonzero(degrees >= 4)
    X, Xmap, y = [], [], []
    for _ in range(n_samples):
        sign = 1 if rng.random() > 0.5 else -1
        f = np.ones(len(points))
        f += sign * 0.18 * H[:, i2] / rel_norm(H[:, i2])
        coeff = rng.normal(0, 0.055, size=len(hi))
        f += H[:, hi] @ coeff / math.sqrt(len(hi))
        f = np.maximum(f, 0.03)
        X.append(f - 1.0)
        Xmap.append(phi_squared(f, A) - 1.0)
        y.append(sign > 0)
    X, Xmap, y = np.asarray(X), np.asarray(Xmap), np.asarray(y)
    xa, xb, ya, yb = train_test_split(X, y, test_size=.3, random_state=SEED, stratify=y)
    ma = LogisticRegression(C=1.0, max_iter=200, random_state=SEED).fit(xa, ya)
    xma, xmb, _, _ = train_test_split(Xmap, y, test_size=.3, random_state=SEED, stratify=y)
    # split indices above are independently random but identically seeded; labels align.
    mm = LogisticRegression(C=1.0, max_iter=200, random_state=SEED).fit(xma, ya)
    return {'raw_accuracy': float(accuracy_score(yb, ma.predict(xb))),
            'idea_accuracy': float(accuracy_score(yb, mm.predict(xmb)))}


def main():
    m = 1400
    points = fibonacci_sphere(m)
    H, degrees = real_harmonics(points, 6)
    A = make_funk_matrix(points)
    d = diagnostic(points, A, H, degrees)
    c = toy_classification(points, A, H, degrees)
    out = {'seed': SEED, 'points': m, 'neighbor_count': 120,
           'diagnostic': d, 'classification': c}
    print(json.dumps(out, indent=2, sort_keys=True))


if __name__ == '__main__':
    main()