import json
import numpy as np
from numpy.linalg import eigvalsh, norm, cholesky, inv

SEED = 2240

def sinkhorn(R, pi, rho, n_iter=10000, tol=1e-12):
    a = np.ones(len(pi)); b = np.ones(len(rho))
    for it in range(n_iter):
        a = pi / np.maximum(R @ b, 1e-300)
        b = rho / np.maximum(R.T @ a, 1e-300)
        if it % 20 == 0:
            P = a[:, None] * R * b[None, :]
            err = max(np.max(abs(P.sum(1)-pi)), np.max(abs(P.sum(0)-rho)))
            if err < tol:
                break
    P = a[:, None] * R * b[None, :]
    err = max(np.max(abs(P.sum(1)-pi)), np.max(abs(P.sum(0)-rho)))
    return P, it + 1, float(err)

def gaussian_sample(means, covs, weights, n, rng):
    labels = rng.choice(len(weights), n, p=weights)
    x = np.empty((n, means.shape[1]))
    for k in range(len(weights)):
        q = labels == k
        if q.any():
            x[q] = rng.multivariate_normal(means[k], covs[k], q.sum())
    return x, labels

def mmd2_rbf(x, y, bandwidth=1.0):
    def kernel(a, b):
        d = ((a[:, None, :] - b[None, :, :]) ** 2).sum(-1)
        return np.exp(-d / (2 * bandwidth * bandwidth))
    return float(kernel(x,x).mean() + kernel(y,y).mean() - 2*kernel(x,y).mean())

def component_bridge_sample(ms, Cs, mt, Ct, pi, rho, P, n, rng):
    flat = P.ravel(); pair = rng.choice(flat.size, n, p=flat)
    x = np.empty((n, ms.shape[1]))
    for q in range(flat.size):
        ij = np.flatnonzero(pair == q)
        if len(ij) == 0: continue
        i, j = divmod(q, mt.shape[0])
        # Simple Gaussian endpoint-preserving bridge kernel: independent target draw.
        x[ij] = rng.multivariate_normal(mt[j], Ct[j], len(ij))
    return x

def global_gaussian_sample(mt, Ct, rho, n, rng):
    mean = (rho[:,None] * mt).sum(0)
    cov = sum(rho[j] * (Ct[j] + np.outer(mt[j]-mean, mt[j]-mean)) for j in range(len(rho)))
    return rng.multivariate_normal(mean, cov, n)

def main():
    rng = np.random.default_rng(SEED)
    pi = np.array([.45, .55]); rho = np.array([.5, .5])
    ms = np.array([[-4., 0.], [4., 0.]])
    mt = np.array([[-4., 2.], [4., -2.]])
    Cs = np.array([[[.35,0],[0,.35]], [[.35,0],[0,.35]]])
    Ct = np.array([[[.35,.05],[.05,.45]], [[.45,-.05],[-.05,.35]]])
    dist2 = ((ms[:,None,:] - mt[None,:,:])**2).sum(-1)
    tau = 3.0
    R = pi[:,None] * rho[None,:] * np.exp(-dist2/(2*tau*tau))

    # Prediction 1: entropic projection has exactly prescribed marginals; truncation error falls.
    marginal = []
    for steps in [1, 2, 5, 10, 30, 100, 300]:
        P, _, err = sinkhorn(R, pi, rho, n_iter=steps, tol=0)
        marginal.append([steps, err])

    # Prediction 2: epsilon^2 I raises the smallest eigenvalue by epsilon^2 exactly.
    raw = np.array([[1e-10,0],[0,2.]])
    inflation = []
    for eps in [0., 1e-4, 1e-3, 1e-2, 1e-1, .5, 1.]:
        vals = eigvalsh(raw + eps*eps*np.eye(2))
        inflation.append([eps, float(vals[0]), float(vals[0]-eigvalsh(raw)[0])])

    # Prediction 3: the stated perturbation bound is valid and scales quadratically in map error.
    d = 4; rg = np.random.default_rng(SEED+1)
    Sigma = rg.standard_normal((d,d)); Sigma = Sigma@Sigma.T + .2*np.eye(d)
    chi = rg.standard_normal((d,d))
    bound_rows = []
    for scale in [.01,.03,.1,.3,1.0]:
        D = scale * rg.standard_normal((d,d))
        lhs = np.trace(D @ chi @ Sigma @ chi.T @ D.T)
        rhs = eigvalsh(Sigma).max() * norm(chi,2)**2 * norm(D,'fro')**2
        bound_rows.append([scale, float(lhs), float(rhs), float(lhs/rhs)])

    # Stability proxy: Cholesky failure versus inflation for nearly indefinite estimates.
    failures = []
    for eps in [0., 1e-5, 1e-4, 1e-3, 1e-2, 1e-1]:
        fail = 0; conds = []
        for k in range(100):
            A = np.array([[1e-8, 0.0],[0.0, 1.]])
            noise = rng.normal(scale=2e-8, size=(2,2)); noise = (noise+noise.T)/2
            C = A + noise + eps*eps*np.eye(2)
            try:
                cholesky(C); conds.append(eigvalsh(C).max()/eigvalsh(C).min())
            except np.linalg.LinAlgError:
                fail += 1
        failures.append([eps, fail, float(np.median(conds)) if conds else None])

    # Mini comparison at target: label-coupled mixture versus one global Gaussian.
    N = 1200
    target, _ = gaussian_sample(mt, Ct, rho, N, np.random.default_rng(SEED+2))
    P, _, final_err = sinkhorn(R, pi, rho)
    idea = component_bridge_sample(ms, Cs, mt, Ct, pi, rho, P, N, np.random.default_rng(SEED+3))
    base = global_gaussian_sample(mt, Ct, rho, N, np.random.default_rng(SEED+4))
    # Mode recall: fraction within radius 1.8 of either target component mean, per mode.
    def recall(samples):
        vals=[]
        for j in range(2):
            vals.append(float(np.mean(norm(samples-mt[j],axis=1)<1.8)))
        return vals
    result = {
        'seed': SEED,
        'coupling': {'R': R.tolist(), 'P': P.tolist(), 'final_marginal_error': final_err, 'iterations': 10000},
        'prediction_1_marginal_error_by_iterations': marginal,
        'prediction_2_inflation_eigenvalue': inflation,
        'prediction_3_bound_by_map_scale': bound_rows,
        'stability_cholesky_failures_and_median_condition': failures,
        'comparison': {
            'target_mmd_self_reference': mmd2_rbf(target,target),
            'idea_mmd2': mmd2_rbf(idea,target),
            'global_gaussian_mmd2': mmd2_rbf(base,target),
            'idea_mode_recall': recall(idea),
            'global_gaussian_mode_recall': recall(base)
        }
    }
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()