#!/usr/bin/env python3
"""First-Hit Interacting Optimizer: exact 1-D solvable toy verification.

Particles start distance `a` from an absorbing target.  A particle with
constant targetward drift v and diffusion D has the exact inverse-Gaussian
first-passage distribution.  The pair-kick model is represented by its exact
marginal effective diffusion D+kappa*(N-1), which is the mechanism in the
paper's solvable Example B.
"""
import json
import math
from pathlib import Path
import numpy as np
from scipy.stats import invgauss

SEED = 1729
A = 1.0
D = 1.0
ALPHA = 1.0
KAPPA = 0.20
NS = np.array([4, 8, 16, 32, 64, 128, 256])
TRIALS = 30000


def hit_samples(n, trials, drift=0.0, diffusivity=D, rng=None):
    """Return trials of min first-passage time for n independent labels."""
    rng = np.random.default_rng() if rng is None else rng
    # X_t = -drift*t + sqrt(2*diffusivity) W_t, distance is A.
    # For drift > 0, T ~ IG(mean=A/drift, shape=A^2/(2D)).
    if drift > 0:
        shape = A * A / (2.0 * diffusivity)
        mean = A / drift
        # scipy's invgauss(mu, scale) has mean mu*scale and shape=scale.
        mu = mean / shape
        t = invgauss.rvs(mu, scale=shape, size=(trials, n), random_state=rng)
    else:
        # Levy law at zero drift: T = A^2 / (2D Z^2).
        z = rng.standard_normal((trials, n))
        t = A * A / (2.0 * diffusivity * z * z)
    return np.min(t, axis=1)


def fit_loglog(ns, means):
    return float(np.polyfit(np.log(ns), np.log(means), 1)[0])


def fit_loginv(ns, means):
    # log mean = c - beta log(log N); beta=1 is the logarithmic law.
    return float(np.polyfit(np.log(np.log(ns)), np.log(means), 1)[0])


def summarize(name, means):
    return {
        "name": name,
        "means": [float(x) for x in means],
        "loglog_slope": fit_loglog(NS, means),
        "loglog_slope_highN": fit_loglog(NS[-4:], means[-4:]),
        "loglogN_slope": fit_loginv(NS, means),
    }


def main():
    rng = np.random.default_rng(SEED)
    regimes = {}
    for name, drift_fn, diff_fn in [
        ("independent", lambda n: 0.0, lambda n: D),
        ("normalized_bounded", lambda n: ALPHA, lambda n: D),
        ("unnormalized_coherent", lambda n: ALPHA * (n - 1), lambda n: D),
        ("pair_kicks", lambda n: 0.0, lambda n: D + KAPPA * (n - 1)),
    ]:
        means = []
        for n in NS:
            means.append(np.mean(hit_samples(int(n), TRIALS, drift_fn(int(n)),
                                              diff_fn(int(n)), rng)))
        regimes[name] = summarize(name, np.array(means))

    # Quantitative mechanism checks, with predictions from the paper.
    # 1) normalized/independent remain in the log-extreme class (slope -1
    # versus log log N), while coherent drift tends to -1 versus log N.
    # 2) pair kicks have effective D_N=D+kappa(N-1), so N*log(N)*E[T]
    # tends to A^2/(4*kappa), here 1.25.
    # 3) direct parameter sweep verifies linear coherent force and diffusion.
    coherent_alpha = [0.8, 1.6, 3.2, 6.4]
    ncheck = 128
    coh_alpha_means = []
    for alpha in coherent_alpha:
        coh_alpha_means.append(float(np.mean(hit_samples(ncheck, TRIALS,
                                                           alpha*(ncheck-1), D, rng))))
    # In the zero-noise coherent limit the prediction is exact, providing a
    # clean parameter-scaling check independent of Monte Carlo crossover.
    deterministic_alpha_times_mean = [float(alpha * (A / (alpha * (ncheck - 1))))
                                      for alpha in coherent_alpha]
    kick_kappas = [0.05, 0.10, 0.20, 0.40]
    kick_kappa_means = []
    for kappa in kick_kappas:
        kick_kappa_means.append(float(np.mean(hit_samples(ncheck, TRIALS, 0.0,
                                                           D+kappa*(ncheck-1), rng))))
    checks = {
        "log_extreme_prediction": {
            "predicted": "slope of log(E[T]) vs log(log N) approximately -1",
            "independent_observed": regimes["independent"]["loglogN_slope"],
            "normalized_observed": regimes["normalized_bounded"]["loglogN_slope"],
        },
        "coherent_prediction": {
            "predicted": "E[T] approximately A/(alpha*(N-1)); high-N log-log slope -1",
            "observed_highN_slope": regimes["unnormalized_coherent"]["loglog_slope_highN"],
            "alpha_values": coherent_alpha,
            "alpha_times_mean": [a*m for a,m in zip(coherent_alpha, coh_alpha_means)],
            "target_alpha_times_mean": A/(ncheck-1),
            "zero_noise_alpha_times_mean": deterministic_alpha_times_mean,
        },
        "pair_kick_prediction": {
            "predicted": "N*log(N)*E[T] tends to A^2/(4*kappa)=%.4f" % (A*A/(4*KAPPA)),
            "observed_highN_NlogN_mean": [float(NS[i]*math.log(NS[i])*regimes["pair_kicks"]["means"][i]) for i in range(len(NS)-3, len(NS))],
            "observed_highN_loglog_slope": regimes["pair_kicks"]["loglog_slope_highN"],
            "kappa_values": kick_kappas,
            "kappa_times_NlogN_mean": [float(k*ncheck*math.log(ncheck)*m) for k,m in zip(kick_kappas, kick_kappa_means)],
            "target_kappa_times_NlogN_mean": A*A/4,
        },
    }
    out = {"seed": SEED, "A": A, "D": D, "alpha": ALPHA,
           "kappa": KAPPA, "trials": TRIALS,
           "note": "pair_kicks uses the paper's exact effective marginal diffusivity; cross-label covariance is not time-discretely simulated",
           "regimes": regimes, "checks": checks}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()