"""
=============================================================================
FUMD-AI Handover Forecaster -- self-contained preprocessing
=============================================================================

Everything needed to turn a `dataset_labeled_w<W>.csv` (the AI-ready output of
the FUMD-AI preprocessing workflow's Step 6) into the exact `(n, 6, 23)` input
this model was trained on -- with no dependency on the training package, on
`joblib`, or on a specific scikit-learn version.

All fitted parameters live in `preprocessing.json`, a plain-JSON export of the
`scaler.joblib` / `label_encoders.joblib` shipped alongside it.

Author(s):
  - Cristina Bernad (ORCID: 0000-0001-9537-415X)
  - Sonja Filiposka <sonja.filiposka@finki.ukim.mk> (ORCID: 0000-0003-0034-2855)
  - Katja Gilly (ORCID: 0000-0002-8985-0639)

Funding: This work has been funded by the FUMD-AI project, an EOSC GRAVITY -
Inter Project with Grant Number 25-EOSC-GRV-INTER-013.

SPDX-License-Identifier: MIT
"""

from __future__ import annotations

import json
import os

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))

TIME_COLUMN = "t"
GROUP_COLUMN = "veh_id"
TARGET_COLUMN = "servingCell"
CATEGORICAL_COLUMNS = ["lane"]
LABEL_ONLY_COLUMNS = ["migration", "destination"]

SEQUENCE_LENGTH = 6
FUTURE_STEPS = [1, 2, 3, 4, 5, 6, 7]
NUM_BASE_STATIONS = 10

#: The 23 input features, in the order the scaler and the encoder LSTM expect.
FEATURE_COLUMNS = [
    "angle", "speed", "pos", "lane", "signals",
    "averageCqiDl", "distance", "measuredSinrDl", "measuredSinrUl", "rcvdSinrDl",
    "rlcDelayDl", "rlcPduDelayDl", "rlcThroughputDl",
    "servingCell",
    "servingCell-1", "servingCell-2", "servingCell-3", "servingCell-4",
    "servingCell-5", "servingCell-6", "servingCell-7",
    "x", "y",
]


# ---------------------------------------------------------------------------
# Fitted parameters
# ---------------------------------------------------------------------------

class Preprocessor:
    """
    Applies the model's fitted lane encoding and standardisation.

    Transform-only by design: there is deliberately no `fit`. Re-fitting either
    the encoder or the scaler on new data would make its outputs incomparable
    with what the model was trained on (different lane->int codes, different
    z-score statistics) while still appearing to work.
    """

    def __init__(self, params: dict):
        scaler = params["scaler"]
        self.mean = np.asarray(scaler["mean"], dtype="float64")
        self.scale = np.asarray(scaler["scale"], dtype="float64")
        self.n_features = int(scaler["n_features_in"])
        if len(self.mean) != self.n_features or len(self.scale) != self.n_features:
            raise ValueError("preprocessing.json: scaler mean/scale length disagrees with n_features_in")

        self.encoders = {}
        for col, enc in params["label_encoders"].items():
            classes = [str(c) for c in enc["classes"]]
            self.encoders[col] = {
                "index": {c: i for i, c in enumerate(classes)},
                "unknown_code": int(enc["unknown_code"]),
                "n_classes": len(classes),
            }

    @classmethod
    def load(cls, path: str | None = None) -> "Preprocessor":
        path = path or os.path.join(HERE, "preprocessing.json")
        with open(path) as fh:
            return cls(json.load(fh))

    # -- categorical ------------------------------------------------------
    def encode_categoricals(self, df: pd.DataFrame) -> pd.DataFrame:
        """Map each categorical column's string values to their fitted codes.

        Values never seen at fit time map to that column's reserved unknown
        code instead of raising -- a road segment the training traffic simply
        never drove down is a legitimate input, not an error.
        """
        out = df.copy()
        for col, enc in self.encoders.items():
            if col not in out.columns:
                raise KeyError(f"expected categorical column {col!r} in the dataframe")
            codes = out[col].astype(str).map(enc["index"])
            n_unknown = int(codes.isna().sum())
            if n_unknown:
                print(f"encode_categoricals: {col}: {n_unknown} value(s) unseen in training "
                      f"-> unknown code {enc['unknown_code']}")
            out[col] = codes.fillna(enc["unknown_code"]).astype(int)
        return out

    # -- scaling ----------------------------------------------------------
    def scale_windows(self, X: np.ndarray) -> np.ndarray:
        """Z-score an (n, seq_len, n_features) array with the fitted statistics."""
        X = np.asarray(X, dtype="float64")
        if X.ndim != 3 or X.shape[-1] != self.n_features:
            raise ValueError(f"expected (n, seq_len, {self.n_features}), got {X.shape}")
        flat = X.reshape(-1, self.n_features)
        flat = (flat - self.mean) / self.scale
        return flat.reshape(X.shape).astype("float32")


# ---------------------------------------------------------------------------
# Windowing
# ---------------------------------------------------------------------------

def build_sequences(
    df: pd.DataFrame,
    feature_columns: list[str] = FEATURE_COLUMNS,
    target_column: str = TARGET_COLUMN,
    sequence_length: int = SEQUENCE_LENGTH,
    future_steps: list[int] = FUTURE_STEPS,
    group_col: str = GROUP_COLUMN,
    time_col: str = TIME_COLUMN,
    label_columns: list[str] = LABEL_ONLY_COLUMNS,
):
    """
    Build per-vehicle sliding windows, identical to how the model was trained.

    Returns
    -------
    X : (n_windows, sequence_length, n_features) -- unscaled
    y : (n_windows, len(future_steps))           -- true serving cell per horizon
    extra : {label_col: (n_windows,)}            -- labels at the first forecast
            step, for evaluation only; never model input.

    Windows never cross a vehicle boundary, and a vehicle with fewer than
    `sequence_length + max(future_steps)` rows contributes none.
    """
    future_steps = list(future_steps)
    needed = sequence_length + max(future_steps)

    X_list, y_list = [], []
    extra_lists = {c: [] for c in label_columns if c in df.columns}
    n_skipped = 0
    has_target = target_column in df.columns

    for _, group in df.groupby(group_col, sort=False):
        group = group.sort_values(time_col).reset_index(drop=True)
        n = len(group)
        if n < needed:
            n_skipped += 1
            continue

        feats = group[feature_columns].to_numpy(dtype=float)
        target = group[target_column].to_numpy() if has_target else None
        label_arrays = {c: group[c].to_numpy() for c in extra_lists}

        for i in range(n - needed):
            X_list.append(feats[i:i + sequence_length])
            if has_target:
                y_list.append([target[i + sequence_length + s] for s in future_steps])
            for c, arr in label_arrays.items():
                extra_lists[c].append(arr[i + sequence_length + future_steps[0]])

    if n_skipped:
        print(f"build_sequences: skipped {n_skipped} vehicle(s) with < {needed} rows")

    X = np.asarray(X_list, dtype=float)
    y = np.asarray(y_list, dtype=float) if has_target else None
    extra = {c: np.asarray(v) for c, v in extra_lists.items()}
    print(f"build_sequences: {X.shape[0]} windows, X {X.shape}")
    return X, y, extra


def prepare(
    df: pd.DataFrame,
    preprocessor: Preprocessor | None = None,
    **kwargs,
):
    """Encode -> window -> scale, in one call. Returns `(X_scaled, y, extra)`."""
    preprocessor = preprocessor or Preprocessor.load()
    df = df.sort_values([GROUP_COLUMN, TIME_COLUMN]).reset_index(drop=True)
    df = preprocessor.encode_categoricals(df)
    X, y, extra = build_sequences(df, **kwargs)
    return preprocessor.scale_windows(X), y, extra