"""
automl_lib.data_loader
~~~~~~~~~~~~~~~~~~~~~~
Universal data ingestion – returns train/val splits as PyTorch DataLoaders
or raw arrays depending on the modality.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset, random_split

from .config import AutoMLConfig
from .utils import get_logger, list_files, file_extension

logger = get_logger(__name__)


# ═══════════════════════════════════════════════════════════════════════════
#  Tabular
# ═══════════════════════════════════════════════════════════════════════════
class TabularDataset(Dataset):
    """Simple in-memory tabular dataset."""

    def __init__(self, X: np.ndarray, y: np.ndarray | None = None):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = (
            torch.tensor(y, dtype=torch.long)
            if y is not None
            else None
        )

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        if self.y is not None:
            return self.X[idx], self.y[idx]
        return self.X[idx]


# ═══════════════════════════════════════════════════════════════════════════
#  Image
# ═══════════════════════════════════════════════════════════════════════════
class ImageFolderDataset(Dataset):
    """Loads images from class-subfolder structure: root/<class>/<img>."""

    def __init__(self, root: str, transform=None):
        from PIL import Image

        self.samples: list[tuple[str, int]] = []
        self.classes: list[str] = sorted(
            d for d in os.listdir(root)
            if os.path.isdir(os.path.join(root, d))
        )
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}
        for cls_name in self.classes:
            cls_dir = os.path.join(root, cls_name)
            for fname in os.listdir(cls_dir):
                ext = file_extension(fname)
                if ext in {"jpg", "jpeg", "png", "bmp", "gif", "tiff", "webp"}:
                    self.samples.append(
                        (os.path.join(cls_dir, fname), self.class_to_idx[cls_name])
                    )
        self.transform = transform
        self._Image = Image

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        path, label = self.samples[idx]
        img = self._Image.open(path).convert("RGB")
        if self.transform:
            img = self.transform(img)
        return img, label


# ═══════════════════════════════════════════════════════════════════════════
#  Text
# ═══════════════════════════════════════════════════════════════════════════
class TextDataset(Dataset):
    """Simple text dataset from a CSV/TSV with text + label columns."""

    def __init__(self, texts: list[str], labels: list[int] | None = None,
                 max_length: int = 256):
        self.texts = texts
        self.labels = labels
        self.max_length = max_length
        self._tokenizer = None

    @property
    def tokenizer(self):
        if self._tokenizer is None:
            from transformers import AutoTokenizer
            self._tokenizer = AutoTokenizer.from_pretrained(
                "distilbert-base-uncased"
            )
        return self._tokenizer

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        enc = self.tokenizer(
            self.texts[idx],
            padding="max_length",
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt",
        )
        item = {k: v.squeeze(0) for k, v in enc.items()}
        if self.labels is not None:
            item["labels"] = torch.tensor(self.labels[idx], dtype=torch.long)
        return item


# ═══════════════════════════════════════════════════════════════════════════
#  Time-Series
# ═══════════════════════════════════════════════════════════════════════════
class TimeSeriesDataset(Dataset):
    """Sliding-window time-series dataset."""

    def __init__(self, data: np.ndarray, window: int = 30,
                 horizon: int = 1):
        self.data = torch.tensor(data, dtype=torch.float32)
        self.window = window
        self.horizon = horizon

    def __len__(self):
        return len(self.data) - self.window - self.horizon + 1

    def __getitem__(self, idx):
        x = self.data[idx: idx + self.window]
        y = self.data[idx + self.window: idx + self.window + self.horizon]
        return x, y.squeeze()


# ═══════════════════════════════════════════════════════════════════════════
#  Master loader
# ═══════════════════════════════════════════════════════════════════════════
def load_data(
    config: AutoMLConfig,
    data_type: str,
) -> dict[str, Any]:
    """Load data according to *data_type* and return a dict with keys:

    ``"train_loader"``, ``"val_loader"``, ``"num_classes"``,
    ``"input_shape"``, ``"dataset_info"``.
    """
    loaders = {
        "tabular": _load_tabular,
        "image": _load_image,
        "text": _load_text,
        "timeseries": _load_timeseries,
        "audio": _load_audio,
    }
    fn = loaders.get(data_type)
    if fn is None:
        raise ValueError(f"Unsupported data_type: {data_type}")
    return fn(config)


# ── Tabular ───────────────────────────────────────────────────────────────
def _load_tabular(config: AutoMLConfig) -> dict[str, Any]:
    ext = file_extension(config.data_path)
    if ext in ("csv", "tsv"):
        sep = "\t" if ext == "tsv" else ","
        df = pd.read_csv(config.data_path, sep=sep)
    elif ext in ("xlsx", "xls"):
        df = pd.read_excel(config.data_path)
    elif ext == "parquet":
        df = pd.read_parquet(config.data_path)
    elif ext == "json":
        df = pd.read_json(config.data_path)
    else:
        df = pd.read_csv(config.data_path)

    if config.max_samples:
        df = df.head(config.max_samples)

    target_col = config.target_column or df.columns[-1]
    y = df[target_col].values
    X = df.drop(columns=[target_col])

    # Basic auto-encoding
    from .preprocessing import preprocess_tabular
    X_processed, encoders = preprocess_tabular(X)

    # Encode target
    task = config.task
    if task is None:
        from .detector import detect_task
        task = detect_task(data=df, data_type="tabular", target_column=target_col)
        
    from sklearn.preprocessing import LabelEncoder
    class_weights = None
    if task == "classification":
        le = LabelEncoder()
        y = le.fit_transform(y)
        num_classes = len(le.classes_)
        
        # Compute class weights if requested
        if config.class_weights == "auto":
            from sklearn.utils.class_weight import compute_class_weight
            cw = compute_class_weight("balanced", classes=np.unique(y), y=y)
            class_weights = cw.tolist()
    else:
        y = y.astype(np.float32)
        num_classes = 1

    dataset = TabularDataset(X_processed, y)
    val_size = int(len(dataset) * config.val_split)
    train_size = len(dataset) - val_size
    train_ds, val_ds = random_split(dataset, [train_size, val_size])

    return {
        "train_loader": DataLoader(
            train_ds, batch_size=config.batch_size, shuffle=True,
            num_workers=0, pin_memory=config.pin_memory,
        ),
        "val_loader": DataLoader(
            val_ds, batch_size=config.batch_size, shuffle=False,
            num_workers=0, pin_memory=config.pin_memory,
        ),
        "num_classes": num_classes,
        "input_shape": (X_processed.shape[1],),
        "dataset_info": {
            "total_samples": len(dataset),
            "train_samples": train_size,
            "val_samples": val_size,
            "features": X_processed.shape[1],
            "target_column": target_col,
            "detected_task": task,
            "class_weights": class_weights,
        },
    }


# ── Image ─────────────────────────────────────────────────────────────────
def _load_image(config: AutoMLConfig) -> dict[str, Any]:
    from torchvision import transforms

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225]),
    ])
    val_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225]),
    ])

    full_ds = ImageFolderDataset(config.data_path, transform=transform)
    num_classes = len(full_ds.classes)
    val_size = int(len(full_ds) * config.val_split)
    train_size = len(full_ds) - val_size
    train_ds, val_ds = random_split(full_ds, [train_size, val_size])

    return {
        "train_loader": DataLoader(
            train_ds, batch_size=config.batch_size, shuffle=True,
            num_workers=config.num_workers, pin_memory=config.pin_memory,
        ),
        "val_loader": DataLoader(
            val_ds, batch_size=config.batch_size, shuffle=False,
            num_workers=config.num_workers, pin_memory=config.pin_memory,
        ),
        "num_classes": num_classes,
        "input_shape": (3, 224, 224),
        "dataset_info": {
            "total_samples": len(full_ds),
            "train_samples": train_size,
            "val_samples": val_size,
            "classes": full_ds.classes,
        },
    }


# ── Text ──────────────────────────────────────────────────────────────────
def _load_text(config: AutoMLConfig) -> dict[str, Any]:
    ext = file_extension(config.data_path)
    if ext in ("csv", "tsv"):
        sep = "\t" if ext == "tsv" else ","
        df = pd.read_csv(config.data_path, sep=sep)
    else:
        df = pd.read_csv(config.data_path)

    if config.max_samples:
        df = df.head(config.max_samples)

    # Heuristic: first text-like column = text, last column = label
    text_col = None
    for c in df.columns:
        if df[c].dtype == object:
            try:
                avg_len = df[c].str.len().mean()
                if avg_len > 20:
                    text_col = c
                    break
            except (AttributeError, TypeError):
                continue
    if text_col is None:
        text_col = df.columns[0]

    target_col = config.target_column or df.columns[-1]
    texts = df[text_col].astype(str).tolist()

    from sklearn.preprocessing import LabelEncoder
    le = LabelEncoder()
    labels = le.fit_transform(df[target_col].values)
    num_classes = len(le.classes_)

    full_ds = TextDataset(texts, labels.tolist())
    val_size = int(len(full_ds) * config.val_split)
    train_size = len(full_ds) - val_size
    train_ds, val_ds = random_split(full_ds, [train_size, val_size])

    return {
        "train_loader": DataLoader(
            train_ds, batch_size=config.batch_size, shuffle=True,
            num_workers=0,
        ),
        "val_loader": DataLoader(
            val_ds, batch_size=config.batch_size, shuffle=False,
            num_workers=0,
        ),
        "num_classes": num_classes,
        "input_shape": (256,),
        "dataset_info": {
            "total_samples": len(full_ds),
            "train_samples": train_size,
            "val_samples": val_size,
            "text_column": text_col,
            "target_column": target_col,
            "classes": le.classes_.tolist(),
        },
    }


# ── Time-Series ───────────────────────────────────────────────────────────
def _load_timeseries(config: AutoMLConfig) -> dict[str, Any]:
    df = pd.read_csv(config.data_path)
    if config.max_samples:
        df = df.head(config.max_samples)

    # Drop datetime-like columns and use remaining numeric columns
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    if not numeric_cols:
        raise ValueError("No numeric columns found for time-series data.")

    data = df[numeric_cols].values.astype(np.float32)
    # Normalize
    mean = data.mean(axis=0)
    std = data.std(axis=0) + 1e-8
    data = (data - mean) / std

    window = config.extra.get("window", 30)
    horizon = config.extra.get("horizon", 1)
    full_ds = TimeSeriesDataset(data, window=window, horizon=horizon)

    val_size = int(len(full_ds) * config.val_split)
    train_size = len(full_ds) - val_size
    train_ds, val_ds = random_split(full_ds, [train_size, val_size])

    return {
        "train_loader": DataLoader(
            train_ds, batch_size=config.batch_size, shuffle=False,
            num_workers=0,
        ),
        "val_loader": DataLoader(
            val_ds, batch_size=config.batch_size, shuffle=False,
            num_workers=0,
        ),
        "num_classes": data.shape[1],
        "input_shape": (window, data.shape[1]),
        "dataset_info": {
            "total_samples": len(full_ds),
            "train_samples": train_size,
            "val_samples": val_size,
            "features": data.shape[1],
            "window": window,
            "horizon": horizon,
            "norm_mean": mean.tolist(),
            "norm_std": std.tolist(),
        },
    }


# ── Audio (stub – extensible) ────────────────────────────────────────────
def _load_audio(config: AutoMLConfig) -> dict[str, Any]:
    """Placeholder for audio loading. Converts WAV files to mel-spectrograms
    and treats them like images (2D inputs)."""
    raise NotImplementedError(
        "Audio data loading is planned for a future release. "
        "Please convert audio to spectrograms and use the 'image' pipeline, "
        "or contribute an implementation!"
    )
