"""
Dataset concreto do pipeline ``atualizado`` — dados tabulares + PyTorch.

A maior parte da lógica está em :class:`base_dataset.BaseDataset`
(``docker/base_dataset.py``). Aqui só fica o que é específico:

  - :py:meth:`build_features` — seleciona colunas numéricas do DataFrame,
    descartando colunas de identificação (``id``, ``image_path``, ...).
  - :py:meth:`get_data_loader` — devolve ``DataLoader`` do PyTorch
    montado em cima dos arrays preparados pela base.
"""

from __future__ import annotations

import os
import sys
from typing import Optional, Tuple

import numpy as np
import pandas as pd

try:
    import torch
    from torch.utils.data import DataLoader, TensorDataset
    _TORCH_AVAILABLE = True
except ImportError:  # pragma: no cover
    _TORCH_AVAILABLE = False

# Importa BaseDataset de ../base_dataset.py
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from base_dataset import BaseDataset  # noqa: E402


class TrainingDataset(BaseDataset):
    """Dataset tabular para o pipeline em ``docker/atualizado``."""

    # ------------------------------------------------------------------
    # Hooks da base
    # ------------------------------------------------------------------
    def build_features(self) -> None:
        if self.df is None:
            raise RuntimeError("self.df está vazio; load_raw_data() não rodou.")

        # Prioridade: lista explícita no metadata.json
        explicit = self.metadata.get("features")
        if explicit:
            missing = [c for c in explicit if c not in self.df.columns]
            if missing:
                raise ValueError(
                    f"Colunas declaradas em metadata['features'] não encontradas no DataFrame: {missing}"
                )
            self.feature_columns = list(explicit)
            self.X = self.df[self.feature_columns].to_numpy(dtype=np.float32)
            return

        # Fallback: inferência automática de colunas numéricas
        feats = []
        for col in self.df.columns:
            if col == self.target_column:
                continue
            if col.lower() in self.NON_FEATURE_COLS:
                continue
            if not pd.api.types.is_numeric_dtype(self.df[col]):
                continue
            feats.append(col)

        if not feats:
            raise ValueError(
                "Nenhuma coluna numérica encontrada para usar como feature."
            )

            self.feature_columns = feats
            self.X = self.df[feats].to_numpy(dtype=np.float32)

    def get_data_loader(
        self,
        batch_size: int = 32,
        train_ratio: Optional[float] = None,
        seed: Optional[int] = None,
    ) -> Tuple["DataLoader", "DataLoader"]:
        if not _TORCH_AVAILABLE:
            raise RuntimeError(
                "PyTorch não está instalado; instale-o ou use get_arrays()."
            )

        train_ratio, seed = self._resolve_split_params(train_ratio, seed)
        idx_train, idx_val = self._split_indices(train_ratio, seed)

        assert self.X is not None and self.y is not None  # pra mypy

        X_train = torch.from_numpy(self.X[idx_train])
        y_train = torch.from_numpy(self.y[idx_train]).long()
        X_val = torch.from_numpy(self.X[idx_val])
        y_val = torch.from_numpy(self.y[idx_val]).long()

        bs = self._resolve_batch_size(batch_size)
        train_loader = DataLoader(
            TensorDataset(X_train, y_train), batch_size=bs, shuffle=True
        )
        val_loader = DataLoader(
            TensorDataset(X_val, y_val), batch_size=bs, shuffle=False
        )

        print(
            f"  - Train samples: {len(X_train)} | Val samples: {len(X_val)}"
            f" | Batch size: {bs}"
        )
        return train_loader, val_loader