import json
import numpy as np
from scipy.optimize import brentq
import interval_equilibrium as ie


def roots(a, lo=-2.0, hi=2.0):
    grid = np.linspace(lo, hi, 40001)
    vals = a * np.tanh(grid) - grid
    ans = []
    for x, y, fx, fy in zip(grid[:-1], grid[1:], vals[:-1], vals[1:]):
        if fx == 0 or fx * fy < 0:
            z = x if fx == 0 else brentq(lambda t: a*np.tanh(t)-t, x, y)
            if not ans or abs(z - ans[-1]) > 1e-6:
                ans.append(float(z))
    return ans


def inclusion_check(a, lo, hi, p0, prad, n=20000, seed=7):
    rng = np.random.default_rng(seed)
    lo, hi = np.asarray(lo, float), np.asarray(hi, float)
    out = ie.krawczyk(lo, hi, a, p0, prad)
    jl, jh = ie.jac_interval(lo, hi, a)
    flo = a*np.tanh(lo) + (p0-prad) - hi
    fhi = a*np.tanh(hi) + (p0+prad) - lo
    ok = True
    for _ in range(n):
        z = rng.uniform(lo, hi)
        p = rng.uniform(p0-prad, p0+prad)
        f = a*np.tanh(z) + p - z
        j = a/np.cosh(z)**2 - 1
        ok = ok and np.all(f >= flo-1e-12) and np.all(f <= fhi+1e-12)
        ok = ok and np.all(j >= jl-1e-12) and np.all(j <= jh+1e-12)
    return {'samples': n, 'included': bool(ok),
            'q': None if out is None else float(out[2])}


def main():
    contraction = []
    for a in [.7, .9, .99, 1.0, 1.01, 1.2]:
        out = ie.krawczyk(np.array([-.2, -.2]), np.array([.2, .2]), a)
        contraction.append({'gain': a, 'q': None if out is None else float(out[2])})

    domains = []
    for a in [.7, 1.2, 1.6]:
        r = ie.classify(a, lo=(-2., -2.), hi=(2., 2.), budget=10000)
        domains.append({'gain': a, 'scalar_roots': len(roots(a)),
                        'status': r.status, 'certified': r.certified,
                        'excluded': r.excluded, 'inconclusive': r.inconclusive,
                        'visited': r.visited})

    uncertainty = []
    for prad in [0., .1, .2, .3, .5]:
        r = ie.classify(.7, prad=prad, lo=(-1., -1.), hi=(1., 1.), budget=10000)
        uncertainty.append({'prad': prad, 'status': r.status,
                            'visited': r.visited})

    baseline = []
    for a in [.7, .9, 1.2]:
        z, steps = ie.fixed_point(a, np.array([.05, .05]), steps=300)
        baseline.append({'gain': a, 'steps': steps,
                         'residual': float(np.max(np.abs(a*np.tanh(z)+.05-z)))})

    result = {'contraction_q': contraction, 'global_domains': domains,
              'parameter_uncertainty': uncertainty, 'fixed_point_baseline': baseline,
              'interval_inclusion': inclusion_check(.7, [-.4, -.3], [.6, .5], .1, .2)}
    with open('mvp_results.json', 'w') as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()