#######################
## import modules 
#######################

from sklearn.preprocessing import MinMaxScaler, StandardScaler
import torch
import torch.nn as nn
import torch.nn.functional as FDataset
from torch.utils.data import DataLoader
from torch.utils.data import Subset
from torch.utils.data import Dataset
import torch.nn.functional as F
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd

class vae_hi(object):
    def __init__(self, train_data, scale_type=2, fill_dict={np.nan:0}, used_vars=[],
                learning_rate = 0.001, batch_size = 128, num_epochs = 10, num_hidden_1=12, num_latent=4):
        '''函数功能
        初始化AE类
        
        Args:
        train_data (pandas.DataFrame) : 建模数据集
        scale_type (int): 默认值为2，1--StandardScaler, 2--MinMaxScaler 0--将不进行归一化
        fill_dict (dict): 默认为{np.nan:0}，缺失值替换列表 
        used_vars (list) : 入模变量列表
        learning_rate (float) : 学习率
        batch_size (int) : 默认为128，批次数据
        num_epochs (int) : 默认为10，训练轮数
        num_hidden_1 (int) : 默认为12， 隐藏层1节点数
        num_latent (int) : 默认为4
        '''
        self.train_data = train_data
        self.scale_type = scale_type
        self.fill_dict = fill_dict 
        self.st = None
        self.learning_rate = learning_rate 
        self.batch_size = batch_size 
        self.num_epochs = num_epochs 
        self.used_vars = used_vars
        self.num_hidden_1 = num_hidden_1
        self.num_latent = num_latent

    def _preprocess(self, data, batch_size, shuffle, is_train):
        idx = data.index
        data_x = data[self.used_vars].replace(self.fill_dict)
        if is_train:
            if self.scale_type==1:
                self.st = StandardScaler()
            elif self.scale_type ==2:
                self.st = MinMaxScaler()
            if self.scale_type !=1 and self.scale_type!=2  and self.scale_type!=0:
                raise ValueError('scale_type must be 1--StandardScaler, 2--MinMaxScaler')
            if self.scale_type==1 or self.scale_type==2 :
                self.st.fit_transform(data_x) 
        if self.st is not None:
            data_x = self.st.transform(data_x)
        else:
            data_x = data_x.values

        data_x = pd.DataFrame(data_x, 
                       columns=self.used_vars,
                       index=idx)

        dataset = UDFDataset(data_x.values)

        loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
        for features in loader:
            print(features.shape)
            break
        return data_x, dataset, loader

    def train(self):
        device = torch.device('cpu')
        data_x, dataset, loader = self._preprocess(self.train_data, self.batch_size, True, True)
        model = VariationalAutoencoder(data_x.shape[1],
                                                 self.num_hidden_1, self.num_latent, device)
        self.model = model.to(device)
        ##########################
        ### COST AND OPTIMIZER
        ##########################

        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)  
        for epoch in range(self.num_epochs):
            self.model.train()
            
            for batch_idx, features in enumerate(loader):
                features = features.to(device)
                z_mean, z_log_var, encoded, decoded = self.model(features)
                kl_divergence = (.5 * (z_mean ** 2 + torch.exp(z_log_var) - z_log_var - 1)).sum()
                pixelwise_bce = F.binary_cross_entropy(decoded, features, reduction='sum')
                loss = kl_divergence + pixelwise_bce
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                if not batch_idx % 100:
                    print('Epoch %03d/%03d |loss %.2f' % (epoch + 1, self.num_epochs, loss))
            print('Epoch %03d/%03d |loss %.2f' % (epoch + 1, self.num_epochs, loss))

        self.train_data = None
        return self

    def predict(self, oot_data, output_method=1):
        device = torch.device('cpu')
        data_x, dataset, loader = self._preprocess(oot_data, oot_data.shape[0], False, False)
        self.model.eval()
        with torch.set_grad_enabled(False):
            for features in loader:
                features = features.to(device)
                z_mean, z_log_var, encoded, decoded = self.model.forward_predict(features)
                loss0 = torch.mean(torch.abs(features - decoded), dim=1)
                loss_kl = torch.mean((.5 * (z_mean ** 2 + torch.exp(z_log_var) - z_log_var - 1)), dim=1)
        if output_method == 1:
                
            loss = loss0
        elif output_method == 2:
            loss = loss0 + loss_kl
        scored = pd.DataFrame(index=oot_data.index)

        scored['Loss_mae'] = loss.detach().numpy()
        result = pd.concat([scored, oot_data], axis=1)
        return result


class UDFDataset(Dataset):
    def __init__(self, x):
        self.x = x
    
    def __getitem__(self, index):
        features = self.x[index]
        features = torch.tensor(features, dtype=torch.float32)
        return features
        
    def __len__(self):
        return self.x.shape[0]


class VariationalAutoencoder(torch.nn.Module):

    def __init__(self, num_features, num_hidden_1, num_latent, device):
        super(VariationalAutoencoder, self).__init__()
        
        self.device = device
        
        ### ENCODER
        self.hidden_1 = torch.nn.Linear(num_features, num_hidden_1)
        self.z_mean = torch.nn.Linear(num_hidden_1, num_latent)
        # in the original paper (Kingma & Welling 2015, we use
        # have a z_mean and z_var, but the problem is that
        # the z_var can be negative, which would cause issues
        # in the log later. Hence we assume that latent vector
        # has a z_mean and z_log_var component, and when we need
        # the regular variance or std_dev, we simply use 
        # an exponential function
        self.z_log_var = torch.nn.Linear(num_hidden_1, num_latent)
        
        
        ### DECODER
        self.linear_3 = torch.nn.Linear(num_latent, num_hidden_1)
        self.linear_4 = torch.nn.Linear(num_hidden_1, num_features)

    def reparameterize(self, z_mu, z_log_var):
        # Sample epsilon from standard normal distribution
        eps = torch.randn(z_mu.size(0), z_mu.size(1)).to(self.device)
        # note that log(x^2) = 2*log(x); hence divide by 2 to get std_dev
        # i.e., std_dev = exp(log(std_dev^2)/2) = exp(log(var)/2)
        z = z_mu + eps * torch.exp(z_log_var/2.) 
        return z
        
    def encoder(self, features):

        ### ENCODER
        x = self.hidden_1(features)
        x = F.leaky_relu(x)
        z_mean = self.z_mean(x)
        z_log_var = self.z_log_var(x)
        encoded = self.reparameterize(z_mean, z_log_var)
        return z_mean, z_log_var, encoded
    
    def decoder(self, encoded):
        ### DECODER
        x = self.linear_3(encoded)
        x = F.leaky_relu(x)
        x = self.linear_4(x)
        decoded = torch.sigmoid(x)
        return decoded

    def forward(self, features):
        
        z_mean, z_log_var, encoded = self.encoder(features)
        decoded = self.decoder(encoded)
        
        return z_mean, z_log_var, encoded, decoded
    
    def reparameterize_predict(self, z_mu, z_log_var):
        # Sample epsilon from standard normal distribution
        eps = torch.randn(z_mu.size(0), z_mu.size(1)).to(self.device)
        # note that log(x^2) = 2*log(x); hence divide by 2 to get std_dev
        # i.e., std_dev = exp(log(std_dev^2)/2) = exp(log(var)/2)
        z = z_mu 
        return z
        
    def encoder_predict(self, features):
        ### ENCODER
        x = self.hidden_1(features)
        x = F.leaky_relu(x)
        z_mean = self.z_mean(x)
        z_log_var = self.z_log_var(x)
        encoded = self.reparameterize_predict(z_mean, z_log_var)
        return z_mean, z_log_var, encoded

    def forward_predict(self, features):
        
        z_mean, z_log_var, encoded = self.encoder_predict(features)
        decoded = self.decoder(encoded)
        
        return z_mean, z_log_var, encoded, decoded

