import json
import numpy as np
from db_transport import energy, explicit_step, positivity_dt_bound, dissipation


def make_graph(n=18, seed=7):
    rng = np.random.default_rng(seed)
    a = rng.uniform(.15, 1.0, (n, n)); c = (a + a.T) / 2
    c *= (rng.random((n, n)) < .28)
    c = np.triu(c, 1); c += c.T
    pi = rng.uniform(.4, 1.6, n); pi /= pi.sum()
    rho = rng.uniform(.15, 2.0, n); rho *= 1.0 / rho.sum()
    return rho, pi, c


def stability_sweep(rho, pi, c):
    bound = positivity_dt_bound(rho, pi, c)
    rows = []
    for factor in [.25, .5, .99, 1.01, 3., 10., 30., 100., 300., 1000.]:
        x = explicit_step(rho, pi, c, factor * bound)
        row = {'factor': factor, 'min_mass': float(x.min()),
               'mass_error': float(abs(x.sum() - rho.sum()))}
        row['energy_change'] = (energy(x, pi) - energy(rho, pi)) if np.all(x > 0) else None
        rows.append(row)
    first_negative = next((r['factor'] for r in rows if r['min_mass'] < 0), None)
    return bound, rows, first_negative


def dissipation_sweep(rho, pi, c):
    d = dissipation(rho, pi, c); rows = []
    for dt in [1e-5, 3e-5, 1e-4, 3e-4, 1e-3]:
        x = explicit_step(rho, pi, c, dt)
        measured = (energy(rho, pi) - energy(x, pi)) / dt
        rows.append({'dt': dt, 'measured': measured, 'predicted': d,
                     'relative_error': abs(measured-d)/d})
    return rows


def conductance_sweep(rho, pi, c):
    base = dissipation(rho, pi, c); rows = []
    for scale in [.25, .5, 1., 2., 4.]:
        actual = dissipation(rho, pi, c * scale)
        rows.append({'scale': scale, 'dissipation': actual, 'predicted': base * scale,
                     'relative_error': abs(actual-base*scale)/actual})
    return rows


def classification_utility(seed=11):
    rng = np.random.default_rng(seed); n = 80
    y = np.repeat(np.arange(2), n//2)
    c = np.full((n, n), .12/(n-1)); same = y[:, None] == y[None, :]
    c[same] = .88/(same.sum(1)[0]-1); np.fill_diagonal(c, 0); c = (c+c.T)/2
    scores = np.full((n, 2), .15); scores[np.arange(n), y] = .85
    pi = np.full((n, 2), 1/n); rho = scores / scores.sum(0, keepdims=True)
    # Baseline is an unconstrained residual with a deliberately large multiplier.
    x = scores.copy(); baseline_min = 1e9
    # Transport uses the explicit positivity bound, as required by the idea.
    dt_transport = .9 * positivity_dt_bound(rho, pi, c)
    tr = rho.copy(); transport_min = 1e9
    for _ in range(20):
        x = x + 2.0 * (c @ x - x)
        tr = explicit_step(tr, pi, c, dt_transport)
        baseline_min = min(baseline_min, float(x.min())); transport_min = min(transport_min, float(tr.min()))
    return {'baseline_accuracy': float((np.argmax(x, 1) == y).mean()),
            'transport_accuracy': float((np.argmax(tr / pi, 1) == y).mean()),
            'baseline_min_activation': baseline_min, 'transport_min_mass': transport_min,
            'transport_dt_over_bound': .9}


def main():
    rho, pi, c = make_graph(); bound, stability, first_negative = stability_sweep(rho, pi, c)
    diss = dissipation_sweep(rho, pi, c); conduct = conductance_sweep(rho, pi, c)
    result = {'seed': 7, 'n': len(rho), 'dt_bound': bound,
      'initial_energy': energy(rho, pi),
      'stability_prediction': 'nonnegative for dt <= bound; sufficiently larger dt can fail',
      'stability': stability, 'first_negative_factor_tested': first_negative,
      'dissipation_prediction': '-dF/dt equals edge dissipation with O(dt) Euler error',
      'dissipation': diss,
      'conductance_prediction': 'dissipation scales linearly with global conductance',
      'conductance': conduct, 'classification': classification_utility(),
      'max_mass_error': max(r['mass_error'] for r in stability),
      'max_conductance_relative_error': max(r['relative_error'] for r in conduct),
      'smallest_dt_dissipation_relative_error': diss[0]['relative_error']}
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()