#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Jul 31 19:25:08 2021

@author: akshat
"""
import os, copy
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

import rdkit
from rdkit import Chem
from rdkit import RDLogger
from rdkit.Chem import Descriptors
RDLogger.DisableLog('rdApp.*')

import inspect
from collections import OrderedDict

import multiprocessing
from typing import List
from .features import get_mol_info

def get_mol_feature(smi: str) -> np.array:
    """
    Given a SMILES string (smi), a user needs to provide code for creating descriptors
    for the molecules. This will be used as features for the (1) classifier NN and (2)
    the discriminator (if active).

    This function will be parallelised by processes.

    Parameters
    ----------
    smi : str
        Valid SMILE string generated by JANUS.

    Returns
    -------
    descriptor : array
        Property value of SMILES string 'smi'.

    """
    # mol = Chem.MolFromSmiles(smi)
    # descriptor = AllChem.GetMorganFingerprintAsBitVect(mol, 3, 1024)
    descriptor = get_mol_info(smi)

    return np.array(descriptor)


def obtain_features(smi_list: List[str], num_workers: int = 1):
    assert num_workers <= multiprocessing.cpu_count(), 'num_workers exceed cpu count'
    with multiprocessing.Pool(num_workers) as pool:
        dataset_x = pool.map(
            get_mol_feature, smi_list
        )
        
    return np.array(dataset_x)


class MLP(nn.Module):
    def __init__(self, h_sizes: List[int], n_input: int, n_output: int):
        super(MLP, self).__init__()
        # Layers
        self.hidden = nn.ModuleList([nn.Linear(n_input, h_sizes[0])])
        for k in range(len(h_sizes)-1):
            self.hidden.append(nn.Linear(h_sizes[k], h_sizes[k+1]))
        self.predict = nn.Linear(h_sizes[-1], n_output)

    def forward(self, x):
        for layer in self.hidden:
            x = torch.sigmoid(layer(x))
        output= torch.sigmoid(self.predict(x))
        return output

class EarlyStopping():
    ''' Class that checks criteria for early stopping. Saves the best model weights.
    '''
    def __init__(self, patience, min_delta, mode='minimize'):
        self.patience = patience
        self.best_weights = None
        self.checkpoint = 0
        self.best_epoch = 0
        if mode == 'maximize':
            self.monitor_fn = lambda a, b: np.greater(a - min_delta, b)
            self.best_val = -np.inf
        elif mode == 'minimize':
            self.monitor_fn = lambda a, b: np.less(a + min_delta, b)
            self.best_val = np.inf
        else:
            raise ValueError(f'Mode should be either minimize or maximize.')

    def check_criteria(self, net, epoch, new_val):
        ''' Compare with value in memory. If there is an improvement, reset the checkpoint and
        save the model weights.
        Return True if stopping criteria is met (checkpoint is exceeds patience), otherwise, return False.
        '''
        if self.monitor_fn(new_val, self.best_val):
            self.best_val = new_val
            self.checkpoint = 0
            self.best_weights = copy.deepcopy(net.state_dict())
            self.best_epoch = epoch
        else:
            self.checkpoint += 1
        
        return self.checkpoint > self.patience

    def restore_best(self, net, verbose=True):
        if verbose: print(f'        Early stopping at epoch: {self.best_epoch}       loss: {self.best_val}')
        net.load_state_dict(self.best_weights)
        return net


def get_device(use_gpu: bool):
    if use_gpu:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'   
        if device == 'cpu':
            print('No GPU available, defaulting to CPU.')
    else:
        device = 'cpu'

    return device

############################################################################################

def create_network(n_hidden: List[int], n_input: int, n_output: int, device: str):
    ''' Obtain network
    
    Parameters:
    n_hidden             (list) : Intermediate discrm layers (e.g. [100, 10])
    device               (str)  : Device discrm. will be initialized 
    
    Returns:
    net : torch model
    optimizer   : Loss function optimized (Adam)
    loss_func   : Loss (Cross-Entropy )
    '''
    net = MLP(n_hidden, n_input, n_output).to(device)
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001, weight_decay=1e-4)
    loss_func = nn.BCELoss()
    
    return net, optimizer, loss_func


def train_valid_split(data_x: np.array, data_y: np.array, train_ratio: float = 0.8, 
        seed: int = 30624700):
    ''' Return a random split of training and validation data. 
    Ratio determines the size of training set. Avoids use of sklearn.
    '''
    n = data_x.shape[0]
    train_n = np.floor(n*train_ratio).astype(int)
    indices = np.random.RandomState(seed=seed).permutation(n)
    
    train_x, train_y = data_x[indices[:train_n]], data_y[indices[:train_n]]
    valid_x, valid_y = data_x[indices[train_n:]], data_y[indices[train_n:]]

    return train_x, train_y, valid_x, valid_y


def do_x_training_steps(data_x: np.array, data_y: np.array, net: nn.Module, 
        optimizer: torch.optim.Optimizer, loss_func: nn.Module, 
        steps: int, batch_size: int, device: str):
    ''' Do steps for training. Set batch_size to -1 for full batch training, 
    and 1 for SGD. 
    '''
    train_x, train_y, valid_x, valid_y = train_valid_split(data_x, data_y)
    train_x = torch.tensor(train_x, device=device, dtype=torch.float)
    train_y = torch.tensor(train_y, device=device, dtype=torch.float)
    valid_x = torch.tensor(valid_x, device=device, dtype=torch.float)
    valid_y = torch.tensor(valid_y, device=device, dtype=torch.float)

    if batch_size == -1:
        batch_size = train_x.shape[0]

    train_loader = DataLoader(TensorDataset(train_x, train_y), batch_size=batch_size, shuffle=True)
    valid_loader = DataLoader(TensorDataset(valid_x, valid_y), batch_size=batch_size)

    early_stop = EarlyStopping(patience=500, min_delta=1e-7, mode='minimize')
    net.train()
    for t in range(steps):
        # training steps
        for x, y in train_loader:
            pred_y = net(x)
            loss = loss_func(pred_y, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # validation steps
        val_loss = 0
        net.eval()
        with torch.no_grad():
            for x, y in valid_loader:
                pred_y = net(x)
                loss = loss_func(pred_y, y)
                val_loss += loss
        net.train()
        val_loss /= len(valid_loader)

        if t % 1000 == 0: 
            print('        Epoch:{} Loss:{}'.format(t, loss.item()))
            print(f'                Validation loss: {val_loss.item()}')

        # check early stopping criteria
        stop = early_stop.check_criteria(net, t, val_loss.item())
        if stop:
            net = early_stop.restore_best(net)
            break

    return net


def create_and_train_network(smi_list: List[str], targets: np.array, n_hidden: List[int] = [100, 10], 
        use_gpu: bool = True, num_workers: int = 1): 
    ''' Featurize smiles and train classifier network. Return trained network.
    '''
    # featurize
    dataset_x = obtain_features(smi_list, num_workers=num_workers)
    avg_val = np.percentile(targets, 80) # np.average(targets)
    dataset_y  = np.expand_dims([1 if x >= avg_val else 0 for x in targets ], -1)
    
    # create network
    device = get_device(use_gpu)
    net, optimizer, loss_func = create_network(n_hidden, dataset_x.shape[-1], dataset_y.shape[-1], device)

    # train network
    net = do_x_training_steps(
        data_x=dataset_x, data_y=dataset_y, net=net, 
        optimizer=optimizer,  loss_func=loss_func, steps=20000, 
        batch_size=1024, device=device
    )
    
    return net


def obtain_model_pred(smi_list: List[str], net: nn.Module, use_gpu: bool = True,
        num_workers: int = 1, batch_size: int = 1024): 
    predictions = []

    device = get_device(use_gpu)
    data_x = obtain_features(smi_list, num_workers=num_workers)
    data_x = torch.tensor(data_x, device=device, dtype=torch.float)
    loader = DataLoader(TensorDataset(data_x), batch_size=batch_size)

    num_batches = -(-data_x.shape[0] // batch_size)     # ceil division

    print('Number of batches: ', num_batches)
    with torch.no_grad():
        for i, x in enumerate(loader): 
            print('        Predicting Batch: {}/{}'.format(i, num_batches))
            outputs = net(x[0])
            out_    = outputs.detach().cpu().numpy()
            predictions.append(out_)
    
    predictions = np.concatenate(predictions, axis=0)   # concatenate in the batch axis

    return predictions

