import json
from pathlib import Path
import numpy as np

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


def phi(u, a, b, c, order="12"):
    """Path energy for F1=-a*u1+b*u2 and F2=-a*u2+c*u1."""
    u = np.asarray(u)
    cross = c if order == "12" else b
    return -0.5 * a * np.sum(u * u, axis=-1) + cross * u[..., 0] * u[..., 1]


def field(u, a, b, c):
    u = np.asarray(u)
    return np.stack((-a*u[..., 0] + b*u[..., 1],
                     -a*u[..., 1] + c*u[..., 0]), axis=-1)


def grid_distribution(a, b, c, gamma, order, n=501, lo=-1.0, hi=1.0):
    x = np.linspace(lo, hi, n)
    X, Y = np.meshgrid(x, x, indexing="ij")
    U = np.stack((X, Y), axis=-1)
    logits = phi(U, a, b, c, order) / gamma
    logits -= logits.max()
    weights = np.exp(logits)
    weights /= weights.sum()
    return x, X, Y, weights


def moments(a, b, c, gamma, order, n=501):
    _, X, Y, W = grid_distribution(a, b, c, gamma, order, n=n)
    return {
        "cross": float(np.sum(X * Y * W)),
        "var1": float(np.sum(X * X * W)),
        "var2": float(np.sum(Y * Y * W)),
    }


def sample_joint(a, b, c, gamma, order, count=50000, n=301):
    _, X, Y, W = grid_distribution(a, b, c, gamma, order, n=n)
    indices = rng.choice(W.size, size=count, p=W.ravel())
    return np.stack((X.ravel()[indices], Y.ravel()[indices]), axis=1)


def sample_independent(a, gamma, count=50000, n=10001):
    # Independent Gibbs baseline from each self-quadratic q_i=-a*u_i^2/2.
    x = np.linspace(-1., 1., n)
    p = np.exp(-0.5*a*x*x/gamma)
    p /= p.sum()
    i = rng.choice(n, size=(count, 2), p=p)
    return x[i]


def main():
    a = 2.0
    records = {}

    # Prediction 1: Delta Phi=(c-b)u1u2, whose uniform-box RMS is |c-b|/3.
    disagreement = []
    grid = np.linspace(-1., 1., 401)
    X, Y = np.meshgrid(grid, grid, indexing="ij")
    U = np.stack((X, Y), axis=-1)
    for d in [0., .2, .5, 1.0, 1.5]:
        b, c = .4 - d/2, .4 + d/2
        observed = float(np.sqrt(np.mean((phi(U,a,b,c,"12") - phi(U,a,b,c,"21"))**2)))
        predicted = abs(c-b)/3.
        disagreement.append({"asymmetry_abs_c_minus_b": d,
                             "predicted_rms": predicted,
                             "observed_rms": observed})
    records["path_disagreement_scaling"] = disagreement

    # Prediction 2: if b=c, the field is conservative and the two paths coincide.
    conservative = []
    for k in [-1., -.2, 0., .4, 1.0]:
        d = float(np.max(np.abs(phi(U,a,k,k,"12") - phi(U,a,k,k,"21"))))
        conservative.append({"b_equals_c": k, "max_path_difference": d})
    records["conservative_field_check"] = conservative

    # Prediction 3: on an unbounded quadratic Gibbs model, Cov(u1,u2)=gamma*k/(a^2-k^2).
    # Compare this scaling against exact bounded-box quadrature for a stable k.
    b = 0.0
    c = 0.5
    temperature = []
    for gamma in [.20, .35, .50, .75, 1.0, 1.5]:
        observed = moments(a,b,c,gamma,"12",n=501)["cross"]
        predicted = gamma*c/(a*a-c*c)
        temperature.append({"gamma": gamma, "predicted_unbounded_cov": predicted,
                           "observed_bounded_cov": observed,
                           "ratio_observed_to_predicted": observed/predicted})
    records["temperature_covariance_scaling"] = temperature

    # Direct numerical gradient verification of the path construction.
    u = np.array([.31, -.47])
    eps = 1e-5
    grad_num = []
    for j in range(2):
        plus, minus = u.copy(), u.copy()
        plus[j] += eps; minus[j] -= eps
        grad_num.append((phi(plus,a,.1,.7,"12") - phi(minus,a,.1,.7,"12"))/(2*eps))
    expected = np.array([-a*u[0] + .7*u[1], -a*u[1] + .7*u[0]])
    records["gradient_check"] = {"numerical": grad_num, "expected_path_gradient": expected.tolist(),
                                  "max_abs_error": float(np.max(np.abs(np.array(grad_num)-expected)))}

    # Secondary coordination comparison: coherent path Gibbs versus independent Gibbs.
    joint = sample_joint(a, .0, .8, .5, "12")
    independent = sample_independent(a, .5, count=len(joint))
    records["coordination_comparison"] = {
        "path_joint_mean_product": float(np.mean(joint[:,0]*joint[:,1])),
        "independent_mean_product": float(np.mean(independent[:,0]*independent[:,1])),
        "path_joint_abs_product": float(np.mean(np.abs(joint[:,0]*joint[:,1]))),
        "independent_abs_product": float(np.mean(np.abs(independent[:,0]*independent[:,1])))
    }

    out = Path("results.json")
    out.write_text(json.dumps(records, indent=2))
    print(json.dumps(records, indent=2))


if __name__ == "__main__":
    main()