"""Aggregation/dispatch optimization skill v1 (docs/05 §2.2 聚合/调度优化).

Decompose a province-level target curve to aggregation units (docs/04 §1:
never below unit granularity). Per interval, an LP:

    minimise  Σ_u w_u · x_{u,t}  +  M · s_t
    s.t.      Σ_u x_{u,t} + s_t = target_t          (balance, s_t = shortfall)
              0 ≤ x_{u,t} ≤ available_{u,t}

Intervals are independent in v1 (no ramp/energy coupling); the whole day is
one sparse LP solved by HiGHS. Shortfall is reported, never hidden — a
non-zero shortfall is what the power-balance simulation escalates on.
"""

from __future__ import annotations

import time
from decimal import Decimal

import numpy as np
import scipy
from scipy.optimize import Bounds, LinearConstraint, milp
from scipy.sparse import csr_matrix, hstack, identity, kron

from vpp_contracts.dispatch_optimization_request import DispatchOptimizationRequest
from vpp_contracts.dispatch_optimization_result import DispatchOptimizationResult

from . import SKILL_VERSIONS
from .numeric import INTERVALS, SCALE_MW, SCALE_MWH, curve, quantize, to_array

SKILL = "dispatch-optimization"
SHORTFALL_PENALTY = 1e6


def optimize_dispatch(req: DispatchOptimizationRequest) -> DispatchOptimizationResult:
    if req.target_mw.date != req.market_date or any(u.available_mw.date != req.market_date for u in req.units):
        raise ValueError("all curves must be for market_date")
    ids = [u.unit_id for u in req.units]
    if len(set(ids)) != len(ids):
        raise ValueError("duplicate unit_id")

    target = to_array(v.root for v in req.target_mw.values)
    avail = np.stack([to_array(v.root for v in u.available_mw.values) for u in req.units])  # (U, T)
    weights = np.array([float(Decimal(u.cost_weight)) for u in req.units])
    if np.any(target < 0) or np.any(avail < 0) or np.any(weights < 0):
        raise ValueError("target, availability and cost weights must be non-negative")

    n_u, n_t = avail.shape
    # Variables: x (U·T, unit-major) then s (T).
    c = np.concatenate([np.repeat(weights, n_t), np.full(n_t, SHORTFALL_PENALTY)])
    ub = np.concatenate([avail.ravel(), target])
    bounds = Bounds(np.zeros(n_u * n_t + n_t), ub)
    # Balance: for each t, Σ_u x_{u,t} + s_t = target_t
    a_x = kron(csr_matrix(np.ones((1, n_u))), identity(n_t, format="csr"))
    a = hstack([a_x, identity(n_t, format="csr")])
    constraints = [LinearConstraint(a, target, target)]

    t0 = time.perf_counter()
    res = milp(c, constraints=constraints, bounds=bounds)
    wall_ms = int(round((time.perf_counter() - t0) * 1000))
    status = {0: "OPTIMAL", 2: "INFEASIBLE", 3: "INFEASIBLE"}.get(int(res.status), "ERROR")

    if res.x is None:
        x = np.zeros((n_u, n_t))
        shortfall = target.copy()
    else:
        x = np.clip(res.x[: n_u * n_t].reshape(n_u, n_t), 0.0, None)
        shortfall = np.clip(target - x.sum(axis=0), 0.0, None)

    allocations = [{"unit_id": uid, "target_mw": curve(x[i], req.market_date, SCALE_MW)} for i, uid in enumerate(ids)]
    # Recompute shortfall from the *quantized* allocations so the headline is consistent with the curves.
    alloc_sum = np.zeros(n_t)
    for a_ in allocations:
        alloc_sum += to_array(a_["target_mw"]["values"])
    short_mwh = float(np.clip(target - alloc_sum, 0.0, None).sum() * 0.25)
    if short_mwh < 1e-3:
        short_mwh = 0.0

    return DispatchOptimizationResult.model_validate(
        {
            "market_date": req.market_date,
            "total_target_mw": curve([v.root for v in req.target_mw.values], req.market_date, SCALE_MW),
            "allocations": allocations,
            "shortfall_mwh": quantize(short_mwh, SCALE_MWH),
            "solver": {"name": "highs", "version": f"scipy-{scipy.__version__}", "status": status, "wall_time_ms": wall_ms},
            "skill_version": SKILL_VERSIONS[SKILL],
        }
    )