import json
import numpy as np
from scipy.optimize import minimize_scalar


def bdd_q(M, dims):
    starts = np.cumsum([0] + list(dims))
    qs = []
    for i in range(len(dims)):
        a, b = starts[i], starts[i + 1]
        D = M[a:b, a:b]
        off = np.concatenate([M[a:b, starts[j]:starts[j+1]]
                              for j in range(len(dims)) if j != i], axis=1)
        try:
            z = np.linalg.solve(D, off)
            qs.append(float(np.max(np.sum(np.abs(z), axis=1))))
        except np.linalg.LinAlgError:
            qs.append(float('inf'))
    return max(qs), qs


def fixed_point(r, b, max_steps=2000, tol=1e-8):
    x = np.zeros(2)
    for k in range(1, max_steps + 1):
        xn = b + r * np.array([x[1], x[0]])
        if np.max(np.abs(xn - x)) < tol:
            return True, k, xn
        x = xn
    return False, max_steps, x


def toy_sweep():
    rows = []
    b = np.array([1.0, -0.4])
    for r in np.linspace(0, 1.2, 49):
        M = np.array([[1., -r], [-r, 1.]])
        q, block_q = bdd_q(M, [1, 1])
        sv = np.linalg.svd(M, compute_uv=False)
        cond = float(sv[0] / sv[-1]) if sv[-1] > 1e-14 else float('inf')
        ok, steps, _ = fixed_point(r, b)
        rows.append(dict(r=float(r), q=float(q), block_q=block_q,
                         determinant=float(np.linalg.det(M)), cond2=cond,
                         converged=ok, iterations=steps))
    below = [z['r'] for z in rows if z['converged']]
    above = [z['r'] for z in rows if not z['converged']]
    fit = [(z['q'], z['cond2']) for z in rows if z['q'] <= .9]
    scaled = [c * (1-q) for q, c in fit if np.isfinite(c)]
    return {
        'rows': rows,
        'predictions': {
            'q_equals_coupling_max_abs_error': max(abs(z['q'] - z['r']) for z in rows),
            'predicted_singularity_q': 1.0,
            'observed_zero_determinant_r': next(z['r'] for z in rows if abs(z['determinant']) < 1e-12),
            'observed_fixed_point_boundary_midpoint': (max(below) + min(above)) / 2,
            'predicted_fixed_point_boundary_q': 1.0,
            'condition_scaled_c_times_1_minus_q_median_q_le_0.9': float(np.median(scaled)),
            'condition_scaled_range_q_le_0.9': [float(min(scaled)), float(max(scaled))]
        }
    }


def constrained_fit(seed=7, lam=0.0, delta=.2):
    rng = np.random.default_rng(seed)
    r_true = .95
    b = np.array([1.0, -.4])
    y = np.linalg.solve(np.array([[1., -r_true], [-r_true, 1.]]), b)
    y = y + .001 * rng.normal(size=2)

    def parts(r):
        M = np.array([[1., -r], [-r, 1.]])
        pred = np.linalg.solve(M, b)
        task = .5 * np.sum((pred-y)**2)
        q = abs(r)
        penalty = lam * max(0., q-(1-delta))**2
        return task + penalty, task, q

    opt = minimize_scalar(lambda r: parts(r)[0], bounds=(-.99, .99), method='bounded',
                          options={'xatol': 1e-12, 'maxiter': 1000})
    total, task, q = parts(opt.x)
    return {'r_learned': float(opt.x), 'q': float(q), 'task_loss': float(task),
            'total_loss': float(total), 'target_q': r_true,
            'certificate_limit': 1-delta, 'optimizer_success': bool(opt.success)}


def random_certificate_check(seed=19, n=1000):
    rng = np.random.default_rng(seed)
    worst_q = 0.0
    violations = 0
    for _ in range(n):
        # Two 2x2 diagonal blocks, with off-diagonal blocks scaled to q<0.9.
        D1 = np.eye(2) + 0.2 * rng.normal(size=(2, 2))
        D2 = np.eye(2) + 0.2 * rng.normal(size=(2, 2))
        C12 = rng.normal(size=(2, 2))
        C21 = rng.normal(size=(2, 2))
        # Scale both cross blocks until the exact BDD q is below 0.9.
        M = np.block([[D1, C12], [C21, D2]])
        q0, _ = bdd_q(M, [2, 2])
        scale = 0.85 / max(q0, 1e-12)
        M = np.block([[D1, scale*C12], [scale*C21, D2]])
        q, _ = bdd_q(M, [2, 2])
        worst_q = max(worst_q, q)
        if abs(np.linalg.det(M)) < 1e-8:
            violations += 1
    return {'samples': n, 'worst_q': float(worst_q),
            'near_singular_count': violations}


def main():
    out = {'toy': toy_sweep(),
           'random_block_certificate': random_certificate_check(),
           'fit_unconstrained': constrained_fit(lam=0.0),
           'fit_bdd_lam_10': constrained_fit(lam=10.0),
           'fit_bdd_lam_1000': constrained_fit(lam=1000.0),
           'fit_bdd_lam_100000': constrained_fit(lam=100000.0)}
    with open('results.json', 'w') as f:
        json.dump(out, f, indent=2)
    p = out['toy']['predictions']
    print('q/coupling max error:', p['q_equals_coupling_max_abs_error'])
    print('observed singularity r:', p['observed_zero_determinant_r'],
          'predicted:', p['predicted_singularity_q'])
    print('observed fixed-point boundary:', p['observed_fixed_point_boundary_midpoint'],
          'predicted:', p['predicted_fixed_point_boundary_q'])
    print('median cond*(1-q):', p['condition_scaled_c_times_1_minus_q_median_q_le_0.9'])
    for k in ('fit_unconstrained', 'fit_bdd_lam_10', 'fit_bdd_lam_1000'):
        print(k + ':', out[k])


if __name__ == '__main__':
    main()