from torch.utils.data import DataLoader
import yaml
import pandas as pd
import os 
from torch.utils.data import Dataset

from vox_bench.data_proc.loading_functions import get_loader_func


class AudioDataset(Dataset):
    def __init__(
        self,
        dataframe: pd.DataFrame,
        loader_name: str,
        **loader_kwargs
    ):
        self.df = dataframe.reset_index(drop=True)
        self.loader = get_loader_func(loader_name)
        self.loader_kwargs = loader_kwargs

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

    def __getitem__(self, idx):
      
      row = self.df.iloc[idx]
      rel_path = row["path"]
      root_dir = row["root_dir"]
      cls = row["class"]
    
      audio_path = os.path.join(root_dir, rel_path)
        
      if not os.path.exists(video_path):
        raise FileNotFoundError(f"Video not found: {video_path}")
      return self.loader(
        audio_path = audio_path,
        cls = cls,
        **self.loader_kwargs
        )

class DataModule:
    def __init__(self, config_path: str):

        with open(config_path) as f:
            self.cfg = yaml.safe_load(f)["data"]

        self.root_dirs = self.cfg.get("root_dir", [])
        self.train_val_csvs = self.cfg.get("train_val_csvs", [])
        self.test_csvs = self.cfg.get("test_csvs", [])

        self.data_params = self.cfg.get("data_params", {}).copy()
        self.batch_size = self.data_params.pop("batch_size", 16)
        self.num_workers = self.data_params.pop("num_workers", 4)

        self.train_ds = None
        self.val_ds = None
        self.test_ds = None

        self.train_loader = None
        self.val_loader = None
        self.test_loader = None

        self._validate_config()
        self._build_datasets()
        self._build_loaders()

    def _validate_config(self):

        if self.train_val_csvs:
            if not self.root_dirs:
                raise ValueError("root_dir must be provided if train_val_csvs are used")

            if len(self.root_dirs) != len(self.train_val_csvs):
                raise ValueError("root_dirs and train_val_csvs must match in length")

        if self.test_csvs:
            if not self.root_dirs:
                raise ValueError("root_dir must be provided if test_csvs are used")

            if len(self.root_dirs) != len(self.test_csvs):
                raise ValueError("root_dirs and test_csvs must match in length")

    def _build_datasets(self):

        train_dfs = []
        val_dfs = []
        test_dfs = []
        for root, train_val_csv in zip(self.root_dirs, self.train_val_csvs):

            df = pd.read_csv(train_val_csv)

            if "split" not in df.columns:
                raise ValueError(
                    f"{train_val_csv} must contain a 'split' column"
                )

            df["root_dir"] = root

            train_dfs.append(df[df["split"] == "train"])
            val_dfs.append(df[df["split"] == "val"])

        for root, test_csv in zip(self.root_dirs, self.test_csvs):

            df = pd.read_csv(test_csv)
            df["root_dir"] = root
            test_dfs.append(df)

        if train_dfs:
            train_df = pd.concat(train_dfs, ignore_index=True)
            self.train_ds = AudioDataset(train_df, **self.data_params)

        if val_dfs:
            val_df = pd.concat(val_dfs, ignore_index=True)
            self.val_ds = AudioDataset(val_df, **self.data_params)

        if test_dfs:
            test_df = pd.concat(test_dfs, ignore_index=True)
            self.test_ds = AudioDataset(test_df, **self.data_params)

    def _build_loaders(self):

        if self.train_ds is not None:
            self.train_loader = DataLoader(
                self.train_ds,
                batch_size=self.batch_size,
                shuffle=True,
                num_workers=self.num_workers,
                pin_memory=True
            )

        if self.val_ds is not None:
            self.val_loader = DataLoader(
                self.val_ds,
                batch_size=self.batch_size,
                shuffle=False,
                num_workers=self.num_workers,
                pin_memory=True
            )

        if self.test_ds is not None:
            self.test_loader = DataLoader(
                self.test_ds,
                batch_size=self.batch_size,
                shuffle=False,
                num_workers=self.num_workers,
                pin_memory=True
            )