# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_dataloaders.ipynb.

# %% auto 0
__all__ = ['get_dls', 'CrossSpectraDataset', 'SpectralDataset', 'SpectralEmbeddingDataset']

# %% ../nbs/03_dataloaders.ipynb 2
from pathlib import Path
from tqdm import tqdm
from collections import namedtuple
import fastcore.all as fc

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

# %% ../nbs/03_dataloaders.ipynb 5
def get_dls(
    train_ds:Dataset, # Train dataset
    valid_ds:Dataset, # Valid dataset
    bs:int, # Batch size 
    **kwargs
    ) -> DataLoader:  # NamedTuple `DataLoader`with `.train` and `.valid` field names
    "Train and valid dataloaders."
    Dataloaders = namedtuple('Dataloader', ['train', 'valid'])
    return Dataloaders(
        DataLoader(train_ds, batch_size=bs, shuffle=True, **kwargs),
        DataLoader(valid_ds, batch_size=bs*2, **kwargs))


# %% ../nbs/03_dataloaders.ipynb 6
class CrossSpectraDataset(Dataset):
    "Custom Pytorch dataset for IR instrument cross-calibration."
    def __init__(self, 
                 df:pd.DataFrame, # Infrared spectra with associated sample id
                 pair_idxs:list, # List of tuples (pairs) of replicate indices to sample from the dataframe 
                 spectra_at:int=2 # Index of the column where the spectroscopy data starts
                 ):
        self.data = tensor(df.iloc[:, spectra_at:].to_numpy(),
                           dtype=torch.float32)
        self.pair_idxs = [list(p) for p in pair_idxs]

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

    def __getitem__(self, idx):
        X, y = self.data[self.pair_idxs[idx]]
        return X[None, :], y[None, :]

# %% ../nbs/03_dataloaders.ipynb 11
class SpectralDataset(Dataset):
    "(Infrared Spectra, soil property) custom PyTorch Dataset."

    def __init__(self,
                 X:np.ndarray, # Spectra
                 y:np.ndarray, # Analyte
                 metadata:np.ndarray=None
                 ):
        self.X = X
        self.y = y
        self.metadata = metadata

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

    def __getitem__(self, idx):
        return (tensor(self.X[idx, :][None, :], dtype=torch.float32),
                tensor(self.y[idx, :], dtype=torch.float32))

# %% ../nbs/03_dataloaders.ipynb 14
class SpectralEmbeddingDataset(Dataset):
    "(Infrared Spectra, Infrared Spectra) custom PyTorch Dataset."

    def __init__(self,
                 X:np.ndarray, # Spectra
                 ):
        self.X = X

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

    def __getitem__(self, idx):
        return (tensor(self.X[idx, :][None, :], dtype=torch.float32),
                tensor(self.X[idx, :][None, :], dtype=torch.float32))
