"""Optimise the measure weights of brf_model.py.

For fixed kap the objective 2^wF + 2^wX (2^kap - 1) is convex and every
branching constraint log2(sum 2^{-delta}) <= 0 is convex, so each inner
problem is solved by SLSQP from several starts; kap is scanned on a grid
and refined.
"""
import sys

import numpy as np
from scipy.optimize import minimize, minimize_scalar

from brf_model import branchings, full_weights, linear_constraints, objective, worst_branching


def solve_fixed_kap(kap, options, starts=8, seed=0):
    rng = np.random.default_rng(seed)
    names = [name for name, _ in branchings((1, .5, .3, 0, kap), options)]
    split = 'input-class' in options
    nz = 6 if split else 4

    def full(z):
        w = (z[0], z[1], z[2], z[3], kap)
        return w + (z[4], z[5]) if split else w

    keys = list(linear_constraints(full(np.zeros(nz) + 0.5)))

    def branch_g(z):
        w = full(z)
        return np.array([-np.log2(sum(2.0 ** -t for t in vec)) for _, vec in branchings(w, options)])

    def lin_g(z):
        lc = linear_constraints(full(z))
        return np.array([lc[k] for k in keys])

    cons = [{'type': 'ineq', 'fun': branch_g}, {'type': 'ineq', 'fun': lin_g}]
    best = None
    for s in range(starts):
        z0 = np.array([0.95, 0.6, 0.4, 0.0, 0.4, 0.0][:nz]) + (rng.normal(0, 0.1, nz) if s else 0)
        bounds = [(0, 2)] + [(-4, 2)] * (nz - 1)
        res = minimize(lambda z: objective(full(z)), z0, method='SLSQP', bounds=bounds,
                       constraints=cons, options={'maxiter': 2000, 'ftol': 1e-14})
        if not res.success:
            continue
        w = full(res.x)
        if worst_branching(w, options)[0] > 1 + 1e-9:
            continue
        if min(linear_constraints(w).values()) < -1e-9:
            continue
        if best is None or res.fun < best[0]:
            best = (res.fun, w)
    return best


def report(label, options):
    grid = [0.0, 0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0, 1.5]
    results = [(solve_fixed_kap(k, options), k) for k in grid]
    for r, k in results:
        if r:
            print(f"  kap={k:5.2f}: {r[0]:.8f}")
    (val, w), kap = min(((r, k) for r, k in results if r), key=lambda t: t[0][0])
    # refine around the best grid point
    for step in (0.02, 0.004, 0.0008, 0.00016):
        for k in (kap - 2 * step, kap - step, kap + step, kap + 2 * step):
            if k < 0:
                continue
            r = solve_fixed_kap(k, options)
            if r and r[0] < val:
                val, w = r
                kap = k
    wF, wB, wX, wL, kap, wY, wM = full_weights(w)
    worst = worst_branching(w, options)
    print(f"{label}: (b+d-d/c) = {val:.8f}  b={2**wF:.6f} c={2**kap:.6f} d={2**(kap+wY):.6f}")
    print(f"  weights wF={wF:.6f} wB={wB:.6f} wX={wX:.6f} wL={wL:.6f} kap={kap:.6f} wY={wY:.6f} wM={wM:.6f}; "
          f"worst branching {worst[1]} = {worst[0]:.9f}")
    return val, w


if __name__ == '__main__':
    opts = tuple(sys.argv[1:])
    report('options=' + (','.join(opts) or 'none'), opts)