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

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

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

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


# %% ../nbs/03_dataloaders.ipynb 6
class CrossSpectraDataset(Dataset):
    "Custom Pytorch Dataset accessing pairs of spectra replicates."
    def __init__(self, 
                 df, 
                 pair_idxs, 
                 spectra_at=2):
        self.df = df
        self.spectra_at = spectra_at
        self.pair_idxs = pair_idxs

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

    def __getitem__(self, idx):
        X, y = [self.df.iloc[id, self.spectra_at:].to_numpy(dtype=float)
                for id in self.pair_idxs[idx]]
        return (tensor(X[None, :], dtype=torch.float32),
                tensor(y[None, :], dtype=torch.float32))
        # return (tensor(X, dtype=torch.float32),
        #         tensor(y, dtype=torch.float32))


# %% ../nbs/03_dataloaders.ipynb 7
def get_dls(train_ds, valid_ds, bs, **kwargs):
    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 10
class SpectralDataset(Dataset):
    "Infrared Spectra 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))
