"""

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 trainiqConfig

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: trainiqConfig,

    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: trainiqConfig) -> 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: trainiqConfig) -> 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: trainiqConfig) -> 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: trainiqConfig) -> 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: trainiqConfig) -> 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!"
    )

