"""UX-77: where the cheap run rolls on. Display only; control is unchanged.

A cheap run is sized for the batteries charging together, at their full
rating, and starts at the first of the cheapest hours. The peak limit leaves
less power than that, so a battery is still short when the run ends, and
control plans the next run then (`Controller._price_runs()`). On a flat night
that goes on until the batteries are at their level or the cheap hours end.

`roll_end()` asks what control asks at each run end, of the same functions
(`modetrade.energy()`, `priceruns.plan()`), until nothing is short or the next
run does not follow on. `for_today()` reads the batteries at the level they
will have by then: the charge power is the peak limit's room (peak aim minus
the house, as `dayplan` counts it), shared as control shares it: by room
left to full (`packshare.weights()`), read again as the levels move, so a
battery that fills drops out and its share goes to the others. Where that
room is not known (no house profile or no aim), the rating is used.
DECISIONS/UX-54-forecast-tab.md has the record.
"""

import dataclasses
import math
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from . import arbitrage, dayplan, modetrade, packshare, priceruns
from .dayplan import EFF
from .price import required_spread, uncertainty_of

SLOT_S = 900                    # a price slot
STEP = timedelta(minutes=5)     # how finely the level is followed
LIMIT = 24                      # rolls at most (a day of one-hour runs)


def roll_end(run, points, packs_at, *, required, floor_pct, ceiling_pct,
             targets=None, limit=LIMIT):
    """Where the cheap run `run` ends once control has rolled it on.

    `packs_at(end)` gives the `modetrade.Pack`s as they stand when a run
    ends at `end`. At each end the charge still short is asked of
    `modetrade.energy()` and the next run of `priceruns.plan()`, as control
    asks them. It stops when nothing is short or the next run does not follow
    on, and is `run.end` when control would not roll it."""
    end_t = run.end
    for _ in range(limit):
        e = modetrade.energy(packs_at(end_t), floor_pct, ceiling_pct,
                             targets=targets)
        if e[0] <= 1e-6:
            break
        new = priceruns.plan(points, end_t, required_eur_kwh=required,
                             charge_kwh=e[0], charge_kw=e[1], sell_kwh=0.0,
                             sell_kw=0.0)
        nxt = new.cheap
        if nxt is None or not new.charge_useful() or nxt.start > end_t \
                or nxt.end <= end_t:
            break
        end_t = nxt.end
    return end_t


@dataclass(frozen=True)
class Roll:
    """`end`: where control stops charging, `stops`: name -> when that battery
    reaches its level (rounded up to a price slot), None while it still
    charges at `end`. A battery not in `stops` does not charge in the run."""
    end: datetime
    stops: dict

    def end_for(self, name, base_end):
        """The end of `name`'s block: its stop, no later than `end`."""
        if name not in self.stops:
            return base_end
        stop = self.stops[name]
        return self.end if stop is None else min(self.end, stop)

    def last(self, base_end):
        """The last end of any battery's block (the chart's cheap bar)."""
        return max([base_end] + [self.end_for(n, base_end)
                                 for n in self.stops])


def _ceil_slot(t):
    return datetime.fromtimestamp(
        math.ceil(t.timestamp() / SLOT_S) * SLOT_S, tz=t.tzinfo
        or timezone.utc)


class _Charge:
    """The batteries' levels followed from `t0` while the chargers charge at
    the room the peak limit leaves: `min(sum of ratings, aim - house)`, shared
    by room left to full as control does (`_shares()`), each stopping at `modetrade.charge_to()`."""

    def __init__(self, packs, t0, ceiling_pct, targets, aim_w, house):
        self.packs = {p.name: p for p in packs}
        self.t = t0
        self.ceil = ceiling_pct
        self.targets = targets
        self.aim = aim_w
        self.house = house
        self.stops = {}

    def at(self, end):
        """The packs when a run ends at `end`."""
        self._to(end)
        return list(self.packs.values())

    def _room_w(self, hour, total):
        if self.aim is None or self.house is None:
            return total
        house_w, base_w = self.house
        return min(total, max(0.0, float(self.aim)
                              - max(house_w(hour), base_w, 0.0)))

    def _to(self, end):
        while self.t < end:
            nxt = min(end, self.t + STEP)
            for s, e, hour in dayplan._hour_slices(self.t, nxt):
                self._charge(s, e, hour)
            self.t = nxt

    def _shares(self, ch, w):
        """Watts per charger: `w` shared by `packshare.weights()` over all
        packs (room left to full, as control passes it), no charger above its
        rating, the surplus going to the others. By rating where the room left
        is not known (no capacity)."""
        wt = packshare.weights(list(self.packs.values()), -1.0)
        if wt is None:
            wt = {p.name: abs(p.rating_w) for p in ch}
        wt = {p.name: wt.get(p.name, 0.0) for p in ch}
        cap = {p.name: abs(p.rating_w) for p in ch}
        out, open_, left = {}, set(wt), float(w)
        while open_:
            tot = sum(wt[n] for n in open_)
            if tot <= 0.0:
                for n in open_:
                    out[n] = 0.0
                break
            over = [n for n in open_ if left * wt[n] / tot > cap[n] + 1e-9]
            if not over:
                for n in open_:
                    out[n] = left * wt[n] / tot
                break
            for n in over:
                out[n] = cap[n]
                left -= cap[n]
                open_.discard(n)
        return out

    def _charge(self, s, e, hour):
        ch = modetrade.chargers(list(self.packs.values()), self.ceil,
                                targets=self.targets)
        total = sum(abs(p.rating_w) for p in ch)
        if not ch or total <= 0.0:
            return
        w = self._room_w(hour, total)
        hrs = (e - s).total_seconds() / 3600.0
        share = self._shares(ch, w)
        for p in ch:
            self.stops.setdefault(p.name, None)
            if not p.capacity_kwh:
                continue
            to = modetrade.charge_to(p, self.ceil, self.targets)
            kwh = share[p.name] * hrs / 1000.0 * EFF
            room = (to - p.soc) / 100.0 * float(p.capacity_kwh)
            if kwh <= 0.0:
                continue
            if kwh >= room - 1e-9:
                self.stops[p.name] = _ceil_slot(
                    s + (e - s) * (room / kwh))
                kwh = room
            self.packs[p.name] = dataclasses.replace(
                p, soc=p.soc + kwh / float(p.capacity_kwh) * 100.0)


def for_today(c, rp, now):
    """The `Roll` of the cheap run `rp.cheap` now under way or to come today,
    or None when control does not roll it on (no run left, charging not
    useful, the trade gate shut, nothing short at its end). Display only:
    nothing here is fed back to control."""
    if rp is None or not rp.on or rp.cheap is None or rp.cheap.end <= now \
            or not rp.charge_useful() \
            or getattr(c, "mode_gate", None) is not None:
        return None
    plan = c.fresh_tariff_plan(now)
    if plan is None:
        return None
    floor, ceil = arbitrage.soc_bounds(c.o)
    hp = getattr(c, "house_profile", None)
    sim = _Charge(c._mode_packs(), max(now, rp.cheap.start), ceil,
                  c.plan_targets(), c.peak_aim_w(),
                  None if hp is None else hp.estimator())
    required = required_spread(
        arbitrage.threshold_eur_kwh(c.o),
        uncertainty_of(plan.read.product)
        | arbitrage.weather_uncertainty(True, c.solar_fc,
                                        bool(c.weather_held(now))))
    end = roll_end(rp.cheap, plan.read.points, sim.at, required=required,
                   floor_pct=floor, ceiling_pct=ceil,
                   targets=c.plan_targets())
    if end <= rp.cheap.end:
        return None
    return Roll(end, dict(sim.stops))