## this is main file for generating synthetic data for training
## it works as a dataset loader for the neural network
## it generates random synthetic beta values sample with label
## it uses distribution_generator.py to generate beta values
## all functions implemented in distribution_generator.py are used here automatically

import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn


from numpy.random import default_rng

import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from torchvision.transforms import ToTensor
from . import distribution_generator as dsg


class SynMethDatasetLoader(Dataset):
    def __init__(self, dataset_size, *args, **kwargs):
        """
        Generates random synthetic beta values sample with label
        """

        self.rnd1 = default_rng()
        self.rnd2 = default_rng()
        self.size_ = 100
        self.dataset_size = dataset_size
        self.item_no_start = 1

    def random_dsg(self):
        # ths function randomly and automatically selects a distribution generator from distribution_generator.py
        dsgs = sorted(
            name
            for name in dsg.__dict__
            if name.islower()
            and not name.startswith("_")
            and callable(dsg.__dict__[name])
            and name != "default_rng"
            and name != "main"
            and name != "minmax_scale"
        )

        n_dsg = self.rnd1.integers(len(dsgs))
        n_methylation_status = self.rnd2.integers(3)
        if n_methylation_status == 0:
            methylation_status = True
        elif n_methylation_status == 1:
            methylation_status = False
        else:
            methylation_status = None

        beta = dsg.__dict__[dsgs[n_dsg]](
            size=self.size_, methylation_status=methylation_status
        )
        return beta, n_methylation_status

    def item_no(self):
        self.item_no_start = self.item_no_start + 1
        return self.item_no_start

    def convert_random_beta_to_tensor(self, beta_sample):
        self.size_ = len(beta_sample)
        beta_sample = torch.tensor(beta_sample)
        beta_sample = torch.reshape(beta_sample, [1, self.size_]).float()
        return beta_sample

    # This are modified but required functions for Dataset class
    def __len__(self):

        return self.dataset_size

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        beta_sample, label = self.random_dsg()

        if label == 0:
            labelname = "hyper"
        elif label == 1:
            labelname = "hypo"
        else:
            labelname = "intedermined"

        beta_sample = self.convert_random_beta_to_tensor(beta_sample)
        label = torch.tensor(label)

        sample = [beta_sample, label]

        return sample
