"""Read individual examples and align forecasts by their full sample key."""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Sequence

import pyarrow.parquet as pq

from src.config import SAMPLE_KEYS, Method, load_methods


def keyed_rows(table) -> dict[tuple, int]:
    columns = [table[name].to_pylist() for name in SAMPLE_KEYS]
    keys = list(zip(*columns))
    if any(any(value is None for value in key) for key in keys):
        raise ValueError("Null sample keys")
    result = {key: index for index, key in enumerate(keys)}
    if len(result) != len(keys):
        raise ValueError("Duplicate sample keys")
    return result


def clean_series(values: list) -> list:
    # Keep missing observations at their original time positions.
    return [float(value) if value is not None and math.isfinite(value) else None for value in values]


class SampleRepository:
    def __init__(self, directory: Path, methods: Sequence[Method] | None = None):
        self.directory = directory
        self.methods = load_methods(directory) if methods is None else tuple(methods)
        self._evaluations = {}
        self.samples = pq.read_table(directory / "data" / "structured_effects.parquet")
        self.index = keyed_rows(self.samples)
        if not self.index:
            raise ValueError("No examples available")
        micro = self.samples["micro"].to_pylist()
        self.entries = [
            {"id": row, **dict(zip(SAMPLE_KEYS, key)), "effects": sorted({
                "/".join(str(event[field]) for field in ("family", "subtype") if event.get(field))
                for event in json.loads(micro[row])
            })}
            for key, row in self.index.items()
        ]

    def catalog(self) -> dict:
        return {"samples": self.entries, "methods": [method.to_dict() for method in self.methods]}

    def _evaluation(self, method: Method):
        if method.id not in self._evaluations:
            table = pq.read_table(self.directory / method.relative_path,
                                  columns=SAMPLE_KEYS + ["forecast", "mase"])
            index = keyed_rows(table)
            if index.keys() != self.index.keys():
                raise ValueError(f"{method.filename}: sample keys do not match the examples")
            self._evaluations[method.id] = table, index
        return self._evaluations[method.id]

    def sample(self, sample_id: int) -> dict:
        if sample_id < 0 or sample_id >= len(self.samples):
            raise KeyError(sample_id)
        result = {name: self.samples[name][sample_id].as_py() for name in self.samples.column_names}
        key = tuple(result[name] for name in SAMPLE_KEYS)
        result["id"] = sample_id
        result["macro"] = json.loads(result["macro"])
        result["micro"] = json.loads(result["micro"])
        for name in ("history", "future", "tsfm_average_forecast"):
            result[name] = clean_series(result[name])
        if len(result["history"]) != result["n_history"] or len(result["future"]) != result["n_horizon"]:
            raise ValueError("Series lengths do not match example metadata")
        forecasts = []
        for method in self.methods:
            table, index = self._evaluation(method)
            row = index[key]
            values = clean_series(table["forecast"][row].as_py())
            mase = table["mase"][row].as_py()
            if len(values) != result["n_horizon"]:
                raise ValueError(f"{method.filename}: forecast length does not match the horizon")
            if mase is None or not math.isfinite(mase) or mase < 0:
                raise ValueError(f"{method.filename}: invalid MASE")
            forecasts.append({**method.to_dict(), "forecast": values, "mase": float(mase)})
        result["forecasts"] = forecasts
        return result