"""The exact optimum: the grader as a linear program, solved with HiGHS (scipy). Runs only when packs are built.

Everything in check.py is linear in the releases: storage is a running sum, gauge flow a weighted sum of earlier
releases, and both costs are convex piecewise-linear (rising weights), so an LP with one excess variable per tier
reaches the same cost. `optimal_plan` re-grades the LP's plan with check.evaluate and raises on a mismatch.
"""

from __future__ import annotations

import numpy as np
from scipy.optimize import linprog
from scipy.sparse import lil_matrix

from .basin import AF_PER_CFS_STEP, routing_weights
from .check import DAY_PER_STEP, FLOOD_W, POOL_TIERS, evaluate


class SolveError(RuntimeError):
    pass


def optimal_plan(task: dict) -> tuple[dict[str, list[float]], float]:
    H, T = task["steps"], task["tail_steps"]
    dams, points = task["dams"], task["points"]
    K = len(dams)
    # variable layout: R[k,t] | e[c,t,j] | u[k,j]
    nR = K * H
    nE = len(points) * (H + T) * len(FLOOD_W)
    nU = K * len(POOL_TIERS)
    N = nR + nE + nU
    iR = lambda k, t: k * H + t                                              # noqa: E731
    iE = lambda c, t, j: nR + (c * (H + T) + t) * len(FLOOD_W) + j           # noqa: E731
    iU = lambda k, j: nR + nE + k * len(POOL_TIERS) + j                      # noqa: E731

    cost = np.zeros(N)
    lb, ub = np.zeros(N), np.full(N, np.inf)
    for k, d in enumerate(dams):
        for t in range(H):
            lb[iR(k, t)], ub[iR(k, t)] = d["min_release_cfs"], d["max_release_cfs"]
        size = d["top_of_flood_af"] - d["bottom_of_flood_af"]
        lo = 0.0
        for j, (hi, w) in enumerate(POOL_TIERS):
            cost[iU(k, j)] = w / 1000
            ub[iU(k, j)] = (hi - lo) * size
            lo = hi
    for c in range(len(points)):
        for t in range(H + T):
            for j, (_, w) in enumerate(FLOOD_W):
                cost[iE(c, t, j)] = w / 1000 * DAY_PER_STEP

    rows, rhs = [], []
    A = lil_matrix((4 * K * H + 2 * K * H + len(points) * (H + T) * len(FLOOD_W) + K, N))
    r = 0

    def add(coefs: dict[int, float], b: float):
        nonlocal r
        for i, v in coefs.items():
            A[r, i] = v
        rhs.append(b)
        r += 1

    for k, d in enumerate(dams):
        # ramp
        for t in range(H):
            if t == 0:
                add({iR(k, 0): 1}, d["release_prev_cfs"] + d["ramp_cfs_per_step"])
                add({iR(k, 0): -1}, -d["release_prev_cfs"] + d["ramp_cfs_per_step"])
            else:
                add({iR(k, t): 1, iR(k, t - 1): -1}, d["ramp_cfs_per_step"])
                add({iR(k, t): -1, iR(k, t - 1): 1}, d["ramp_cfs_per_step"])
        # storage after step t: S0 + a*sum(I - R) within [inactive, top]
        cum_in = np.cumsum(d["inflow_cfs"]) * AF_PER_CFS_STEP
        for t in range(H):
            co = {iR(k, s): AF_PER_CFS_STEP for s in range(t + 1)}
            # S_t <= top  <=>  -a*sum R <= top - S0 - a*sum I  <=>  a*sum R >= S0 + a*sumI - top
            add({i: -v for i, v in co.items()}, d["top_of_flood_af"] - d["storage0_af"] - cum_in[t])
            add(co, d["storage0_af"] + cum_in[t] - d["top_of_inactive_af"])
        # end-of-window flood pool: S_H - bottom <= sum_j u_kj
        co = {iR(k, s): -AF_PER_CFS_STEP for s in range(H)}
        co.update({iU(k, j): -1 for j in range(len(POOL_TIERS))})
        add(co, d["bottom_of_flood_af"] - d["storage0_af"] - cum_in[-1])

    for c, p in enumerate(points):
        contrib: list[dict[int, float]] = [dict() for _ in range(H + T)]
        for k, d in enumerate(dams):
            lag = task["lags_h"].get(d["id"], {}).get(p["id"])
            if lag is None:
                continue
            w = routing_weights(lag)
            for s in range(H):
                for i, wi in enumerate(w):
                    if s + i < H + T:
                        contrib[s + i][iR(k, s)] = contrib[s + i].get(iR(k, s), 0) + wi
        for t in range(H + T):
            for j, (name, _) in enumerate(FLOOD_W):
                # Q_t - thr <= e  <=>  sum w R - e <= thr - base
                co = dict(contrib[t])
                co[iE(c, t, j)] = -1
                add(co, p["thresholds"][name] - p["base_cfs"][t])

    A = A[:r].tocsr()
    res = linprog(cost, A_ub=A, b_ub=np.array(rhs), bounds=list(zip(lb, ub)), method="highs")
    if res.status != 0:
        raise SolveError(f"LP status {res.status}: {res.message}")
    rel = {d["id"]: [float(res.x[iR(k, t)]) for t in range(H)] for k, d in enumerate(dams)}
    got = evaluate(task, rel)
    if not got.feasible or abs(got.cost - res.fun) > 1e-6 * max(1.0, abs(res.fun)):
        raise SolveError(f"LP and grader disagree: LP {res.fun:.6f}, grader {got.cost:.6f}, "
                         f"violations {got.violations[:3]}")
    # what the agent would type: whole cfs; re-check (rounding can cost a hair, never a rule)
    rounded = {k: [float(round(v)) for v in vs] for k, vs in rel.items()}
    g2 = evaluate(task, rounded)
    if g2.feasible:
        return rounded, g2.cost
    return rel, got.cost