import json
import numpy as np
from sklearn.datasets import make_moons
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import accuracy_score, log_loss
from convex_transport import QuadraticConvexTransport, change_of_variables_kl, standard_normal_logpdf

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


def exact_kl(d, s):
    return 0.5 * d * (s*s - 1.0 - 2.0*np.log(s))


def verify_math():
    # Prediction 1: identity is the unique zero-KL member of this isotropic family.
    scales = np.array([0.75, 0.9, 1.0, 1.1, 1.3])
    identity_kls = np.array([exact_kl(3, s) for s in scales])
    # Prediction 2: KL ~= d*(s-1)^2 around identity (Taylor expansion).
    eps = np.array([0.005, 0.01, 0.02, 0.04])
    small_kl = np.array([exact_kl(3, 1+e) for e in eps])
    quad_ratio = small_kl / (3 * eps**2)
    # Prediction 3: at fixed scale, KL is exactly linear in dimension.
    dims = np.array([1, 2, 4, 8, 16])
    dim_kls = np.array([exact_kl(int(d), 1.2) for d in dims])
    dim_slope = np.polyfit(dims, dim_kls, 1)[0]
    expected_slope = exact_kl(1, 1.2)
    # Direct numerical change-of-variables check in 2D.
    x = rng.normal(size=(300000, 2))
    A = np.diag([1.15, 0.85])
    t = QuadraticConvexTransport(A)
    mc = change_of_variables_kl(x, t, standard_normal_logpdf)
    analytic = 0.5 * np.sum(np.diag(A)**2 - 1 - 2*np.log(np.diag(A)))
    return {
        'identity_scales': scales.tolist(),
        'identity_family_kl': identity_kls.tolist(),
        'identity_min_at_scale_1': bool(np.argmin(identity_kls) == 2 and identity_kls[2] < 1e-12),
        'small_eps': eps.tolist(),
        'small_kl_over_d_eps2': quad_ratio.tolist(),
        'quadratic_prediction_mean_ratio': 1.0,
        'quadratic_observed_mean_ratio': float(np.mean(quad_ratio)),
        'dimension_values': dims.tolist(),
        'dimension_kls': dim_kls.tolist(),
        'dimension_predicted_slope': float(expected_slope),
        'dimension_observed_slope': float(dim_slope),
        'cov_kl_analytic': float(analytic),
        'cov_kl_monte_carlo': float(mc),
        'cov_relative_error': float(abs(mc-analytic)/analytic),
    }


def train_moons():
    X, y = make_moons(n_samples=1200, noise=0.20, random_state=SEED)
    # fixed held-out shift: radial convex transport plus modest noise
    tr, te = np.arange(900), np.arange(900, 1200)
    # make_moons is ordered randomly by generator; use deterministic split
    Xtr, ytr = X[tr], y[tr]
    Xte, yte = X[te], y[te]
    shift = QuadraticConvexTransport(np.eye(2)*1.18)
    Xshift = shift.map(Xte) + rng.normal(0, .08, Xte.shape)
    rho = exact_kl(2, 1.18)
    # Use identical small MLP and data budget. Convex augmentation is a class-agnostic
    # strongly convex gradient map, selected on the KL boundary.
    configs = {
        'ERM': Xtr,
        'Gaussian': Xtr + rng.normal(0, .12, Xtr.shape),
        'ConvexTransport': shift.map(Xtr),
    }
    out = {'rho': float(rho), 'n_train': len(tr)}
    for name, Xa in configs.items():
        model = MLPClassifier(hidden_layer_sizes=(32, 32), activation='tanh', solver='adam',
                              alpha=1e-3, batch_size=64, learning_rate_init=2e-3,
                              max_iter=350, random_state=SEED, early_stopping=False)
        model.fit(Xa, ytr)
        out[name] = {
            'clean_accuracy': float(accuracy_score(yte, model.predict(Xte))),
            'shift_accuracy': float(accuracy_score(yte, model.predict(Xshift))),
            'clean_logloss': float(log_loss(yte, model.predict_proba(Xte))),
        }
    return out


if __name__ == '__main__':
    result = {'math_verification': verify_math(), 'mini_experiment': train_moons()}
    with open('results.json', 'w') as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))