"""Dissipative softmax latent layer: finite-state CTMC verification."""
from __future__ import annotations
import numpy as np


def softmax(x):
    x = np.asarray(x, dtype=float)
    y = x - np.max(x)
    e = np.exp(y)
    return e / e.sum()


def conditional_rates(X, k=1.0):
    """Reversible conditional rates with log(W_ij/W_ji)=X_j-X_i."""
    X = np.asarray(X, dtype=float)
    d = X[None, :] - X[:, None]
    W = k * np.exp(d / 2)
    np.fill_diagonal(W, 0.0)
    return W


def generator(X, r, k=1.0, drive=0.15):
    """Row-generator on state 0 plus K conditional states.

    Reset exchange is  i <-> 0 with rates r and r*p_i.  Thus, without
    the driven cycle, the stationary law is P0=1/2 and P_i=p_i/2 for every r.
    The one-way cycle carries dissipation, but its occupation perturbation is
    O(drive/r) as r/k grows, which is the sector-separation prediction.
    """
    X = np.asarray(X, dtype=float)
    p = softmax(X); K = len(X)
    Q = np.zeros((K + 1, K + 1))
    W = conditional_rates(X, k)
    for i in range(K):
        for j in range(K):
            if i != j: Q[i+1, j+1] += W[i, j]
        Q[i+1, 0] += r
        Q[0, i+1] += r * p[i]
    # Directed reset cycle 0 -> 1 -> ... -> K -> 0.
    Q[0, 1] += drive
    for i in range(1, K): Q[i, i+1] += drive
    Q[K, 0] += drive
    np.fill_diagonal(Q, -Q.sum(axis=1))
    return Q, p


def stationary(Q):
    A = Q.T.copy(); A[-1] = 1.0
    b = np.zeros(Q.shape[0]); b[-1] = 1.0
    pi = np.linalg.solve(A, b)
    return np.maximum(pi, 0) / np.maximum(pi, 0).sum()


def cycle_current(pi, Q):
    edges = [(0, 1)] + [(i, i+1) for i in range(1, Q.shape[0]-1)] + [(Q.shape[0]-1, 0)]
    return float(np.mean([pi[i]*Q[i,j] - pi[j]*Q[j,i] for i,j in edges]))


def rate_ratio_error(X, W):
    X = np.asarray(X); mask = ~np.eye(len(X), dtype=bool)
    return float(np.max(np.abs(np.log(W[mask]/W.T[mask]) - (X[None,:]-X[:,None])[mask])))


def run_sweep():
    X = np.array([1.0, 0.2, -0.5, -1.0]); k = 1.0; drive = 0.15
    rows = []
    for ratio in [0.1, 0.3, 1, 3, 10, 30, 100, 300]:
        Q, p = generator(X, ratio*k, k, drive); pi = stationary(Q)
        cond = pi[1:] / pi[1:].sum()
        kl = float(np.sum(cond*np.log(np.maximum(cond,1e-300)/p)))
        rows.append((ratio, kl, float(np.max(np.abs(cond-p))), pi[0], cycle_current(pi,Q)))
    return X, rows, rate_ratio_error(X, conditional_rates(X,k))

if __name__ == '__main__':
    X, rows, ratio_err = run_sweep()
    print('X=', X)
    print('max conditional log-rate-ratio error=', ratio_err)
    print('r_over_k, KL(cond||p), max_abs_error, P0, cycle_current')
    for row in rows: print('%.3g %.8g %.8g %.8g %.8g' % row)
    tail = np.array(rows[-4:])
    print('asymptotic slope KL vs r/k (predicted -2):', np.polyfit(np.log(tail[:,0]), np.log(tail[:,1]), 1)[0])
    print('asymptotic slope error vs r/k (predicted -1):', np.polyfit(np.log(tail[:,0]), np.log(tail[:,2]), 1)[0])