import json
import numpy as np
from dissipative_softmax import generator, stationary, softmax


def entropy_production(pi, Q):
    s = 0.0
    for i in range(len(pi)):
        for j in range(i + 1, len(pi)):
            f, b = pi[i] * Q[i, j], pi[j] * Q[j, i]
            if f > 0 and b > 0:
                s += (f - b) * np.log(f / b)
    return float(s)


def baseline_and_layer():
    rng = np.random.default_rng(2750)
    K, n = 8, 500
    logits = rng.normal(size=(n, K))
    labels = np.array([rng.choice(K, p=softmax(x)) for x in logits])
    base_nll = -np.mean(np.log([softmax(x)[y] for x, y in zip(logits, labels)]))
    out = []
    for ratio in [1, 10, 100]:
        vals = []
        for x, y in zip(logits, labels):
            Q, _ = generator(x, ratio, 1.0, 0.15)
            pi = stationary(Q); q = pi[1:] / pi[1:].sum()
            vals.append(-np.log(max(q[y], 1e-300)))
        out.append((ratio, float(np.mean(vals))))
    return base_nll, out


def weak_affinity_check():
    """Build a cycle which is detailed-balanced at A=0, then tilt it by A."""
    X = np.array([1.0, .2, -.5, -1.0]); K = len(X); r = 30.; k = 1.; c = .2
    Q0, p = generator(X, r, k, drive=0.)
    pi0 = stationary(Q0)
    edges = [(0, 1)] + [(i, i+1) for i in range(1, K)] + [(K, 0)]
    results = []
    L = len(edges)
    for A in [0.005, .01, .02, .04, .08, .16]:
        Q = Q0.copy()
        for i, j in edges:
            # At A=0, added rates obey pi0_i w_ij = pi0_j w_ji.
            base_f = c * np.exp((np.log(pi0[j])-np.log(pi0[i]))/2)
            base_b = c * np.exp((np.log(pi0[i])-np.log(pi0[j]))/2)
            Q[i, j] += base_f * np.exp(A/(2*L))
            Q[j, i] += base_b * np.exp(-A/(2*L))
        np.fill_diagonal(Q, 0.)
        np.fill_diagonal(Q, -Q.sum(axis=1))
        pi = stationary(Q)
        Js = [pi[i]*Q[i,j] - pi[j]*Q[j,i] for i,j in edges]
        J = float(np.mean(Js)); sigma = entropy_production(pi, Q)
        results.append((A, J, sigma))
    slope_j = np.polyfit(np.log([x[0] for x in results]), np.log(np.abs([x[1] for x in results])), 1)[0]
    slope_s = np.polyfit(np.log([x[0] for x in results]), np.log([x[2] for x in results]), 1)[0]
    return results, slope_j, slope_s

if __name__ == '__main__':
    base, layer = baseline_and_layer()
    aff, sj, ss = weak_affinity_check()
    print(json.dumps({'baseline_nll': base, 'dissipative_nll': layer,
                      'weak_affinity': aff, 'loglog_slope_J_vs_A': sj,
                      'loglog_slope_sigma_vs_A': ss}, indent=2))