import numpy as np
from scipy.special import logsumexp

SEED = 3114
L = 8
BETA = 0.65
J = 1.0
N = 4096


def energy(x):
    x = np.asarray(x)
    s = 2*x - 1
    return -J * np.sum(s * np.roll(s, -1, axis=-1), axis=-1)


def state_id(x):
    return np.asarray(x, dtype=np.int64) @ (1 << np.arange(L))


def exact_target():
    states = ((np.arange(1 << L)[:, None] >> np.arange(L)) & 1).astype(np.int64)
    e = energy(states)
    lp = -BETA * e
    p = np.exp(lp - logsumexp(lp))
    return states, e, p


def normalized_ess(logw):
    # Paper's ESS: (sum exp A)^2 / (M sum exp(2A)); computed stably.
    m = len(logw)
    return float(np.exp(2*logsumexp(logw) - np.log(m) - logsumexp(2*logw)))


def sample_proposal(rng, n):
    # Collapsed autoregressive special case: independent Bernoulli conditionals.
    return (rng.random((n, L)) < 0.90).astype(np.int64)


def log_q(x):
    x = np.asarray(x)
    return np.sum(x*np.log(.9) + (1-x)*np.log(.1), axis=1)


def rate(y, i):
    z = y.copy()
    z[i] = 1 - z[i]
    de = float(energy(z) - energy(y))
    return min(1.0, np.exp(-BETA * de))


def ctmc_one(y, horizon, rng):
    y = y.copy()
    logpath = 0.0
    t = 0.0
    jumps = 0
    while t < horizon:
        rates = np.array([rate(y, i) for i in range(L)])
        lam = rates.sum()
        dt = rng.exponential(1.0 / lam)
        if t + dt >= horizon:
            logpath -= lam * (horizon - t)
            break
        t += dt
        i = int(rng.choice(L, p=rates / lam))
        logpath += np.log(rates[i]) - lam * dt
        y[i] = 1-y[i]
        jumps += 1
    return y, logpath, jumps


def ctmc_batch(x, horizon, rng):
    ys, lps, js = [], [], []
    for y in x:
        z, lp, j = ctmc_one(y, horizon, rng)
        ys.append(z); lps.append(lp); js.append(j)
    return np.array(ys), np.array(lps), np.array(js)


def empirical_distribution(x):
    counts = np.bincount(state_id(x), minlength=1 << L).astype(float)
    return counts / len(x)


def kl(p, q):
    # q is exact target and strictly positive.
    mask = p > 0
    return float(np.sum(p[mask] * (np.log(p[mask]) - np.log(q[mask]))))


def generator_check():
    y = np.zeros(L, dtype=np.int64)
    Q = np.zeros((1 << L, 1 << L))
    for sid in range(1 << L):
        x = ((sid >> np.arange(L)) & 1).astype(np.int64)
        for i in range(L):
            z=x.copy(); z[i]=1-z[i]
            Q[state_id(z), sid] = rate(x, i)
        Q[sid, sid] = -np.sum(Q[:, sid])
    colsum = np.max(np.abs(Q.sum(axis=0)))
    # Small-time transition check from all-zero state, using expm-free first order.
    sid = 0
    eps = 1e-7
    first_order = np.eye(1 << L)[:, sid] + eps*Q[:, sid]
    valid_nonnegative = bool(first_order.min() >= -1e-12)
    return float(colsum), valid_nonnegative


def path_check(rng):
    y=np.zeros(L,dtype=np.int64)
    z, lp, jumps=ctmc_one(y, 0.5, rng)
    # Every accepted jump contributes log(rate)-lambda*waiting; survival contributes too.
    return bool(np.isfinite(lp) and jumps >= 0), float(lp), int(jumps)


def main():
    states, energies, target = exact_target()
    rng = np.random.default_rng(SEED)
    x0 = sample_proposal(rng, N)
    a0 = -BETA*energy(x0) - log_q(x0)
    proposal_ess = normalized_ess(a0)
    proposal_kl = kl(empirical_distribution(x0), target)

    # Uniform mutation baseline: same number of local updates as the CTMC's mean.
    # One mutation per site-independent random event, with no energy-aware rates.
    uniform = x0.copy()
    for _ in range(8):
        rows=rng.integers(N,size=N); cols=rng.integers(L,size=N)
        uniform[rows,cols] = 1-uniform[rows,cols]
    uniform_kl = kl(empirical_distribution(uniform), target)

    # CTMC refinement, then importance weights based on the proposal (the stable option
    # recommended by the idea for fixed augmentation). Compare raw sample distribution
    # and weighted ESS; path likelihood is also recorded as a diagnostic.
    refined, path_lp, jumps = ctmc_batch(x0, 1.0, rng)
    refined_kl = kl(empirical_distribution(refined), target)
    refined_a0 = -BETA*energy(refined) - log_q(refined)
    refined_ess = normalized_ess(refined_a0)
    corrected_ess = normalized_ess(refined_a0 - path_lp)

    # Mode recovery: count exact states represented among samples; target modes are
    # the two ground states of the ferromagnetic ring.
    ground = set(np.flatnonzero(energies == energies.min()).tolist())
    def ground_hits(x): return int(np.isin(state_id(x), list(ground)).sum())
    print('CONFIG', {'L':L,'beta':BETA,'N':N,'seed':SEED})
    print('MATH', {'generator_max_column_sum':generator_check()[0],
                   'generator_first_order_nonnegative':generator_check()[1],
                   'path_finite_and_valid':path_check(np.random.default_rng(SEED+1))})
    print('RESULT', {
        'proposal_ess':proposal_ess,
        'proposal_kl':proposal_kl,
        'uniform_mutation_kl':uniform_kl,
        'ctmc_kl':refined_kl,
        'ctmc_ess_A0':refined_ess,
        'ctmc_ess_path_corrected':corrected_ess,
        'mean_ctmc_jumps':float(jumps.mean()),
        'ground_hits_proposal':ground_hits(x0),
        'ground_hits_uniform':ground_hits(uniform),
        'ground_hits_ctmc':ground_hits(refined),
        'unique_proposal':int(np.unique(state_id(x0)).size),
        'unique_uniform':int(np.unique(state_id(uniform)).size),
        'unique_ctmc':int(np.unique(state_id(refined)).size),
    })

if __name__ == '__main__':
    main()