from fastai import *

from fastai.basic_data import *

from fastai.basic_train import *

from fastai.tabular import *

from torch import nn

import torch.nn.functional as F

from torch.utils.data import Dataset

class KlsDataset(Dataset):

    """Kmer latent representation dataset"""

    def __init__(self, data,noise=0.):

        super().__init__()

        self.items = data.values if isinstance(data, pd.DataFrame) else data

        self.noise = noise

    def __len__(self):

        return len(self.items)

    def __getitem__(self, idx):

        item = self.items[idx,:]

        return (item if self.noise == 0. else self.mix_noise(item), item)

def wing(dims):

    layer_dims = list(zip(dims[:-1],dims[1:]))

    fcl = [nn.Linear(*x, bias=False) for x in layer_dims]

    relu = [nn.ReLU() for _ in range(len(fcl))]

    layers = np.asarray(list(zip(fcl, relu))).ravel()[:-1]

    return nn.Sequential(*layers)

def init_weights(m):

    if type(m) == nn.Linear:

        nn.init.xavier_uniform_(m.weight)

def print_weights(nlayer):

    print(list(net.encoder.modules())[nlayer].weight)

class KlsAutoEncoder (nn.Module):

    """Generic autoencoder"""

    def __init__(self, encoder_dims, decoder_dims):

        super().__init__(self)

        self.encoder = wing(encoder_dims)

        self.decoder = wing(decoder_dims)

    def forward(self, x):

        x = self.encoder(x)

        return self.decoder(x)

    def save_encoder(self,file:PathOrStr):

        torch.save(self.encoder.state_dict(), path)

class Encoder():

    """Encoder part of KlsAutoeEncoder ready for inference"""

    def __init__(self,file:PathOrStr,dims:Collection=[100,50,3]):

        e = wing(dims).double()

        e.load_state_dict(torch.load(file))

        e.eval()

        self.e = e

    def transform(self,data:Collection):

        """transform ```data``` to latent representaion"""

        return self.e.forward(tensor(data).double()).cpu().detach().numpy()