"""Baselines the NN is compared against.

HistoricalAverage : mean duration per (pickup zone, drop-off zone, hour) on the train split,
                    backing off to (pickup, drop-off) and then the global mean when a cell has
                    fewer than `min_count` trips (rare pairs at night).
DistanceLinear    : ordinary least squares duration = a + b * zone-pair mean distance
                    (numpy least squares; scikit-learn was dropped to save disk space).
                    No time-of-day information at all.
"""
from __future__ import annotations

import numpy as np


class HistoricalAverage:
    name = "hist-avg"

    def __init__(self, n_zones: int, min_count: int = 3):
        self.Z, self.min_count = n_zones, min_count

    def fit(self, trips):
        Z = self.Z
        y = trips.duration_min.values
        self.global_mean = float(y.mean())
        s3 = np.zeros((Z, Z, 24)); n3 = np.zeros((Z, Z, 24))
        np.add.at(s3, (trips.pu.values, trips["do"].values, trips.hour.values), y)
        np.add.at(n3, (trips.pu.values, trips["do"].values, trips.hour.values), 1)
        s2, n2 = s3.sum(2), n3.sum(2)
        pair = np.where(n2 >= self.min_count, s2 / np.maximum(n2, 1), self.global_mean)
        cell = np.where(n3 >= self.min_count, s3 / np.maximum(n3, 1), pair[:, :, None])
        self.table, self.counts = cell, n3
        return self

    def predict(self, pu, do, hour, weekday=None):
        return self.table[np.asarray(pu, int), np.asarray(do, int), np.asarray(hour, int)]


class DistanceLinear:
    name = "distance-linear"

    def __init__(self, pair_dist: np.ndarray):
        self.pair_dist = np.asarray(pair_dist, float)

    def fit(self, trips):
        x = self.pair_dist[trips.pu.values, trips["do"].values]
        X = np.column_stack([np.ones_like(x), x])
        (self.intercept_, self.slope_), *_ = np.linalg.lstsq(X, trips.duration_min.values, rcond=None)
        return self

    def predict(self, pu, do, hour=None, weekday=None):
        x = self.pair_dist[np.asarray(pu, int), np.asarray(do, int)]
        return np.clip(self.intercept_ + self.slope_ * x, 1.0, 120.0)


class NNPredictor:
    """Adapter so the NN has the same predict() signature as the baselines."""
    name = "nn"

    def __init__(self, model):
        self.model = model

    def predict(self, pu, do, hour, weekday):
        return self.model.predict(pu, do, hour, weekday)