"""Generated data: seeded order / van generator and instance (de)serialisation.

Real vs generated: travel times come from real taxi trips; everything here (orders,
time windows, demands, service times, van capacities and shifts) is synthetic and fully
determined by the seed. See docs/data_dictionary.md for every field.
"""
from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass, field
from pathlib import Path

import numpy as np

DEFAULT_SHIFTS = [(8 * 60, 17 * 60), (9 * 60, 18 * 60), (10 * 60, 19 * 60)]
WINDOW_WIDTHS = (30, 45, 60, 90, 120)


@dataclass
class Order:
    id: int
    zone: int            # zone index 0..Z-1
    earliest: float      # minutes after midnight
    latest: float
    demand: int          # load units
    service: float       # minutes spent at the stop


@dataclass
class Van:
    id: int
    capacity: int
    shift_start: float
    shift_end: float
    max_stops: int = 10


@dataclass
class Instance:
    orders: list
    vans: list
    depot: int
    weekday: int
    name: str = "instance"
    seed: int = 0
    return_to_depot: bool = True
    meta: dict = field(default_factory=dict)

    def order(self, oid: int) -> Order:
        return self._by_id[oid]

    def __post_init__(self):
        self._by_id = {o.id: o for o in self.orders}

    # ---- io
    def to_dict(self):
        return {"name": self.name, "seed": self.seed, "depot": self.depot, "weekday": self.weekday,
                "return_to_depot": self.return_to_depot, "meta": self.meta,
                "orders": [asdict(o) for o in self.orders], "vans": [asdict(v) for v in self.vans]}

    @classmethod
    def from_dict(cls, d):
        return cls([Order(**o) for o in d["orders"]], [Van(**v) for v in d["vans"]],
                   d["depot"], d["weekday"], d.get("name", "instance"), d.get("seed", 0),
                   d.get("return_to_depot", True), d.get("meta", {}))

    def save(self, path: Path):
        Path(path).write_text(json.dumps(self.to_dict(), indent=1))

    @classmethod
    def load(cls, path: Path):
        return cls.from_dict(json.loads(Path(path).read_text()))

    def widened(self, minutes: float, only=None) -> "Instance":
        """Copy with `latest` pushed back by `minutes` (all orders, or only ids in `only`)."""
        d = self.to_dict()
        for o in d["orders"]:
            if only is None or o["id"] in only:
                o["latest"] += minutes
        d["name"] = f"{self.name}+{int(minutes)}min"
        return Instance.from_dict(d)


def _r15(x):
    return 15 * round(x / 15)


def generate_instance(n_orders: int, seed: int, n_zones: int = 20, n_vans: int = 3,
                      zone_weights=None, depot: int | None = None, weekday: int | None = None,
                      cap_slack: float = 1.25, max_stops: int = 10,
                      day=(9 * 60, 17 * 60), widths=WINDOW_WIDTHS,
                      shifts=DEFAULT_SHIFTS, return_to_depot: bool = True) -> Instance:
    """Dispatcher-style instance: n orders with windows inside `day`, n_vans vans."""
    rng = np.random.default_rng(seed)
    w = None if zone_weights is None else np.asarray(zone_weights, float) / np.sum(zone_weights)
    depot = int(rng.integers(n_zones)) if depot is None else depot
    weekday = int(rng.integers(5)) if weekday is None else weekday
    orders = []
    for k in range(n_orders):
        width = float(rng.choice(widths))
        earliest = float(_r15(rng.uniform(day[0], day[1] - width)))
        orders.append(Order(
            id=k, zone=int(rng.choice(n_zones, p=w)),
            earliest=earliest, latest=earliest + width,
            demand=int(rng.choice([1, 2, 3], p=[0.5, 0.3, 0.2])),
            service=float(rng.choice([5, 7, 10]))))
    total = sum(o.demand for o in orders)
    cap = max(3, math.ceil(total * cap_slack / n_vans))
    vans = [Van(v, cap, *shifts[v % len(shifts)], max_stops=max_stops) for v in range(n_vans)]
    return Instance(orders, vans, depot, weekday, f"n{n_orders}_s{seed}", seed, return_to_depot,
                    {"generator": "generate_instance", "cap_slack": cap_slack})


def generate_route_instance(n_stops: int, seed: int, tm, n_zones: int = 20, start: float = 8 * 60,
                            return_to_depot: bool = True) -> Instance:
    """Single-van instance that is feasible by construction (for the search experiment).

    A hidden random tour is simulated with the time model; each window is placed so that
    the hidden tour's arrival lies inside it. The optimal tour is usually a different one.
    """
    rng = np.random.default_rng(seed)
    depot = int(rng.integers(n_zones))
    zones = rng.choice(n_zones, size=n_stops, replace=True)
    t, cur, orders = start, depot, []
    for k in range(n_stops):
        z = int(zones[k])
        t += tm.tt(cur, z, t) + float(rng.uniform(0, 25))   # idle gap spreads windows out
        width = float(rng.choice([60, 90, 120]))
        earliest = t - float(rng.uniform(0, width))
        service = float(rng.choice([5, 7, 10]))
        orders.append(Order(k, z, round(earliest, 1), round(earliest + width, 1), 1, service))
        t, cur = t + service, z
    order_perm = rng.permutation(n_stops)            # hide the construction order
    orders = [orders[i] for i in order_perm]
    for new_id, o in enumerate(orders):
        o.id = new_id
    van = Van(0, n_stops, start, 23 * 60, max_stops=n_stops)
    return Instance(orders, [van], depot, 1, f"route{n_stops}_s{seed}", seed, return_to_depot,
                    {"generator": "generate_route_instance"})