__author__ = "Xinqiang Ding <xqding@umich.edu>"
__date__ = "2017/10/16 02:50:08"

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from torch.distributions.categorical import Categorical

class MSA_Dataset(Dataset):
    '''
    Dataset class for multiple sequence alignment.
    '''

    def __init__(self, seq_msa_binary, seq_weight, seq_keys):
        '''
        seq_msa_binary: a two dimensional np.array.
                        size: [num_of_sequences, length_of_msa*num_amino_acid_types]
        seq_weight: one dimensional array.
                    size: [num_sequences].
                    Weights for sequences in a MSA.
                    The sum of seq_weight has to be equal to 1 when training latent space models using VAE
        seq_keys: name of sequences in MSA
        '''
        super(MSA_Dataset).__init__()
        self.seq_msa_binary = seq_msa_binary
        self.seq_weight = seq_weight
        self.seq_keys = seq_keys

    def __len__(self):
        assert(self.seq_msa_binary.shape[0] == len(self.seq_weight))
        assert(self.seq_msa_binary.shape[0] == len(self.seq_keys))
        return self.seq_msa_binary.shape[0]

    def __getitem__(self, idx):
        return self.seq_msa_binary[idx, :], self.seq_weight[idx], self.seq_keys[idx]

class VAE(nn.Module):
    def __init__(self, num_aa_type, dim_latent_vars, dim_msa_vars, num_hidden_units):
        super(VAE, self).__init__()

        # num of amino acid types
        self.num_aa_type = num_aa_type

        # dimension of latent space
        self.dim_latent_vars = dim_latent_vars

        # dimension of binary representation of sequences
        self.dim_msa_vars = dim_msa_vars

        # num of hidden neurons in encoder and decoder networks
        self.num_hidden_units = num_hidden_units

        # encoder
        self.encoder_linears = nn.ModuleList()
        self.encoder_linears.append(nn.Linear(dim_msa_vars, num_hidden_units[0]))
        for i in range(1, len(num_hidden_units)):
            self.encoder_linears.append(nn.Linear(num_hidden_units[i - 1], num_hidden_units[i]))
        self.encoder_mu = nn.Linear(num_hidden_units[-1], dim_latent_vars, bias=True)
        self.encoder_logsigma = nn.Linear(num_hidden_units[-1], dim_latent_vars, bias=True)

        # decoder
        self.decoder_linears = nn.ModuleList()
        self.decoder_linears.append(nn.Linear(dim_latent_vars, num_hidden_units[0]))
        for i in range(1, len(num_hidden_units)):
            self.decoder_linears.append(nn.Linear(num_hidden_units[i - 1], num_hidden_units[i]))
        self.decoder_linears.append(nn.Linear(num_hidden_units[-1], dim_msa_vars))

    def encoder(self, x):
        '''
        encoder transforms x into latent space z
        '''

        h = x
        for T in self.encoder_linears:
            h = T(h)
            h = torch.tanh(h)
        mu = self.encoder_mu(h)
        sigma = torch.exp(self.encoder_logsigma(h))
        return mu, sigma

    def decoder(self, z):
        '''
        decoder transforms latent space z into p, which is the log probability  of x being 1.
        '''

        h = z
        for i in range(len(self.decoder_linears) - 1):
            h = self.decoder_linears[i](h)
            h = torch.tanh(h)
        h = self.decoder_linears[-1](h)

        fixed_shape = tuple(h.shape[0:-1])
        h = torch.unsqueeze(h, -1)
        h = h.view(fixed_shape + (-1, self.num_aa_type))

        # h = torch.reshape(h, fixed_shape + (-1, self.num_aa_type))
        log_p = F.log_softmax(h, dim=-1)
        log_p = log_p.view(fixed_shape + (-1,))

        # h = h.view(h.size(0), -1, self.num_aa_type)
        # log_p = F.log_softmax(h, dim = 2)
        # log_p = log_p.view(log_p.size(0), -1)

        return log_p

    def compute_weighted_elbo(self, x, weight):
        # sample z from q(z|x)
        mu, sigma = self.encoder(x)
        eps = torch.randn_like(sigma)
        z = mu + sigma * eps

        # compute log p(x|z)
        log_p = self.decoder(z)
        log_PxGz = torch.sum(x * log_p, -1)

        # compute elbo
        elbo = log_PxGz - torch.sum(0.5 * (sigma**2 + mu**2 - 2 * torch.log(sigma) - 1), -1)
        weight = weight / torch.sum(weight)
        elbo = torch.sum(elbo * weight)

        return elbo

    def compute_p_importance_sampling(self, x, nsamples):

        with torch.no_grad():
            x = x.expand(nsamples, x.shape[0], x.shape[1])
            mu, sigma = self.encoder(x)
            eps = torch.randn_like(mu)
            z = mu + sigma * eps
            log_Pz = torch.sum(-0.5 * z**2 - 0.5 * torch.log(2 * z.new_tensor(np.pi)), -1)
            log_p = self.decoder(z)
            log_PxGz = torch.sum(x * log_p, -1)
            log_Pxz = log_Pz + log_PxGz

            log_QzGx = torch.sum(-0.5 * (eps)**2 -
                                 0.5 * torch.log(2 * z.new_tensor(np.pi))
                                 - torch.log(sigma), -1)

            log_Px = torch.logsumexp(log_Pxz - log_QzGx, 0) - torch.log(torch.tensor(nsamples))

        return log_Px

    def compute_p(self, x, nsamples):

        with torch.no_grad():
            # sample z from prior
            z = torch.randn((x.shape[0], nsamples, self.dim_latent_vars), device=x.device)

            # compute log p(x|z)
            log_p = self.decoder(z)
            log_PxGz = torch.sum(x.unsqueeze(1).expand(-1, nsamples, -1) * log_p, -1)

        return torch.mean(log_PxGz, -1)

    def compute_elbo_no_grad(self, x):
        with torch.no_grad():
            # sample z from q(z|x)
            mu, sigma = self.encoder(x)
            eps = torch.randn_like(sigma)
            z = mu + sigma * eps

            # compute log p(x|z)
            log_p = self.decoder(z)
            log_PxGz = torch.sum(x * log_p, -1)

            # compute elbo
            elbo = log_PxGz - torch.sum(0.5 * (sigma**2 + mu**2 - 2 * torch.log(sigma) - 1), -1)

        return elbo

    def sample(self, nsamples):

        # sample z from prior
        device = next(self.parameters()).device
        z = torch.randn((nsamples, self.dim_latent_vars), device=device)
        log_p = self.decoder(z)

        dist = Categorical(logits=log_p.reshape(nsamples, -1, self.num_aa_type))

        data = dist.sample()

        return data

    def compute_elbo(self, x):
        # sample z from q(z|x)
        mu, sigma = self.encoder(x)
        eps = torch.randn_like(sigma)
        z = mu + sigma * eps

        # compute log p(x|z)
        log_p = self.decoder(z)
        log_PxGz = torch.sum(x * log_p, -1)

        # compute elbo
        elbo = log_PxGz - torch.sum(0.5 * (sigma**2 + mu**2 - 2 * torch.log(sigma) - 1), -1)

        return elbo

    def compute_elbo_with_multiple_samples(self, x, num_samples):
        with torch.no_grad():
            x = x.expand(num_samples, x.shape[0], x.shape[1])
            mu, sigma = self.encoder(x)
            eps = torch.randn_like(mu)
            z = mu + sigma * eps
            log_Pz = torch.sum(-0.5 * z**2 - 0.5 * torch.log(2 * z.new_tensor(np.pi)), -1)
            log_p = self.decoder(z)
            log_PxGz = torch.sum(x * log_p, -1)
            log_Pxz = log_Pz + log_PxGz

            log_QzGx = torch.sum(-0.5 * (eps)**2 -
                                 0.5 * torch.log(2 * z.new_tensor(np.pi))
                                 - torch.log(sigma), -1)
            log_weight = (log_Pxz - log_QzGx).detach().data
            log_weight_max = torch.max(log_weight, 0)[0]
            log_weight = log_weight - log_weight_max
            weight = torch.exp(log_weight)
            elbo = torch.log(torch.mean(weight, 0)) + log_weight_max
            return elbo

    def sample_latent_var(self, mu, sigma):
        eps = torch.ones_like(sigma).normal_()
        z = mu + sigma * eps
        return z

    def forward(self, x):
        mu, sigma = self.encoder(x)
        z = self.sample_latent_var(mu, sigma)
        p = self.decoder(z)
        return mu, sigma, p


import torch
import torch.nn as nn
import numpy as np

# Constants
N_PAIR = 'n-pair'
ANGULAR = 'angular'
N_PAIR_ANGULAR = 'n-pair-angular'
MAIN_LOSS_CHOICES = (N_PAIR, ANGULAR, N_PAIR_ANGULAR)

CROSS_ENTROPY = 'cross-entropy'

class BlendedLoss(object):
    def __init__(self, main_loss_type, cross_entropy_flag):
        super(BlendedLoss, self).__init__()
        self.main_loss_type = main_loss_type
        assert main_loss_type in MAIN_LOSS_CHOICES, "invalid main loss: %s" % main_loss_type

        self.metrics = []
        if self.main_loss_type == N_PAIR:
            self.main_loss_fn = NPairLoss()
        elif self.main_loss_type == ANGULAR:
            self.main_loss_fn = AngularLoss()
        elif self.main_loss_type == N_PAIR_ANGULAR:
            self.main_loss_fn = NPairAngularLoss()
        else:
            raise ValueError

        self.cross_entropy_flag = cross_entropy_flag
        self.lambda_blending = 0
        if cross_entropy_flag:
            self.cross_entropy_loss_fn = nn.CrossEntropyLoss()
            self.lambda_blending = 0.3

    def calculate_loss(self, target, output_embedding, output_cross_entropy=None):
        if target is not None:
            target = (target, )

        loss_dict = {}
        blended_loss = 0
        if self.cross_entropy_flag:
            assert output_cross_entropy is not None, "Outputs for cross entropy loss is needed"

            loss_inputs = self._gen_loss_inputs(target, output_cross_entropy)
            cross_entropy_loss = self.cross_entropy_loss_fn(*loss_inputs)
            blended_loss += self.lambda_blending * cross_entropy_loss
            loss_dict[CROSS_ENTROPY + '-loss'] = [cross_entropy_loss.item()]

        loss_inputs = self._gen_loss_inputs(target, output_embedding)
        main_loss_outputs = self.main_loss_fn(*loss_inputs)
        main_loss = main_loss_outputs[0] if type(main_loss_outputs) in (tuple, list) else main_loss_outputs
        blended_loss += (1 - self.lambda_blending) * main_loss
        loss_dict[self.main_loss_type + '-loss'] = [main_loss.item()]

        for metric in self.metrics:
            metric(output_embedding, target, main_loss_outputs)

        return blended_loss, loss_dict

    @staticmethod
    def _gen_loss_inputs(target, embedding):
        if type(embedding) not in (tuple, list):
            embedding = (embedding, )
        loss_inputs = embedding
        if target is not None:
            if type(target) not in (tuple, list):
                target = (target, )
            loss_inputs += target
        return loss_inputs

class NPairLoss(nn.Module):
    """
    N-Pair loss
    Sohn, Kihyuk. "Improved Deep Metric Learning with Multi-class N-pair Loss Objective," Advances in Neural Information
    Processing Systems. 2016.
    http://papers.nips.cc/paper/6199-improved-deep-metric-learning-with-multi-class-n-pair-loss-objective
    """
    def __init__(self, l2_reg=0.02, **kwargs):
        super(NPairLoss, self).__init__()
        self.l2_reg = l2_reg

    def forward(self, embeddings, target):
        n_pairs, n_negatives = self.get_n_pairs(target)

        if embeddings.is_cuda:
            n_pairs = n_pairs.cuda()
            n_negatives = n_negatives.cuda()

        anchors = embeddings[n_pairs[:, 0]]  # (n, embedding_size)
        positives = embeddings[n_pairs[:, 1]]  # (n, embedding_size)
        negatives = embeddings[n_negatives]  # (n, n-1, embedding_size)

        losses = self.n_pair_loss(anchors, positives, negatives) \
            + self.l2_reg * self.l2_loss(anchors, positives)

        return losses

    @staticmethod
    def get_n_pairs(labels):
        """
        Get index of n-pairs and n-negatives
        :param labels: label vector of mini-batch
        :return: A tuple of n_pairs (n, 2)
                        and n_negatives (n, n-1)
        """
        labels = labels.cpu().data.numpy()
        n_pairs = []
        for label in set(labels):
            label_mask = (labels == label)
            label_indices = np.where(label_mask)[0]
            if len(label_indices) < 2:
                continue
            anchor, positive = np.random.choice(label_indices, 2, replace=False)
            n_pairs.append([anchor, positive])

        n_pairs = np.array(n_pairs)
        n_negatives = []
        for i in range(len(n_pairs)):
            negative = np.concatenate([n_pairs[:i, 1], n_pairs[i + 1:, 1]])
            n_negatives.append(negative)

        n_negatives = np.array(n_negatives)
        return torch.LongTensor(n_pairs), torch.LongTensor(n_negatives)

    @staticmethod
    def n_pair_loss(anchors, positives, negatives):
        """
        Calculates N-Pair loss
        :param anchors: A torch.Tensor, (n, embedding_size)
        :param positives: A torch.Tensor, (n, embedding_size)
        :param negatives: A torch.Tensor, (n, n-1, embedding_size)
        :return: A scalar
        """
        anchors = torch.unsqueeze(anchors, dim=1)  # (n, 1, embedding_size)
        positives = torch.unsqueeze(positives, dim=1)  # (n, 1, embedding_size)

        x = torch.matmul(anchors, (negatives - positives).transpose(1, 2))  # (n, 1, n-1)
        x = torch.sum(torch.exp(x), 2)  # (n, 1)
        loss = torch.mean(torch.log(1 + x))
        return loss

    @staticmethod
    def l2_loss(anchors, positives):
        """
        Calculates L2 norm regularization loss
        :param anchors: A torch.Tensor, (n, embedding_size)
        :param positives: A torch.Tensor, (n, embedding_size)
        :return: A scalar
        """
        return torch.sum(anchors**2 + positives**2) / anchors.shape[0]

class AngularLoss(NPairLoss):
    """
    Angular loss
    Wang, Jian. "Deep Metric Learning with Angular Loss," CVPR, 2017
    https://arxiv.org/pdf/1708.01682.pdf
    """
    def __init__(self, l2_reg=0.02, angle_bound=1., lambda_ang=2, **kwargs):
        super(AngularLoss, self).__init__()
        self.l2_reg = l2_reg
        self.angle_bound = angle_bound
        self.lambda_ang = lambda_ang
        self.softplus = nn.Softplus()

    def forward(self, embeddings, target):
        n_pairs, n_negatives = self.get_n_pairs(target)
        if embeddings.is_cuda:
            n_pairs = n_pairs.cuda()
            n_negatives = n_negatives.cuda()

        anchors = embeddings[n_pairs[:, 0]]  # (n, embedding_size)
        positives = embeddings[n_pairs[:, 1]]  # (n, embedding_size)
        negatives = embeddings[n_negatives]  # (n, n-1, embedding_size)

        losses = self.angular_loss(anchors, positives, negatives, self.angle_bound) \
                 + self.l2_reg * self.l2_loss(anchors, positives)

        return losses, 0, 0, 0

    @staticmethod
    def angular_loss(anchors, positives, negatives, angle_bound=1.):
        """
        Calculates angular loss
        :param anchors: A torch.Tensor, (n, embedding_size)
        :param positives: A torch.Tensor, (n, embedding_size)
        :param negatives: A torch.Tensor, (n, n-1, embedding_size)
        :param angle_bound: tan^2 angle
        :return: A scalar
        """
        print(anchors, positives, negatives);exit(1)
        anchors = torch.unsqueeze(anchors, dim=1)  # (n, 1, embedding_size)
        positives = torch.unsqueeze(positives, dim=1)  # (n, 1, embedding_size)

        x = 4. * angle_bound * torch.matmul((anchors + positives), negatives.transpose(1, 2)) \
            - 2. * (1. + angle_bound) * torch.matmul(anchors, positives.transpose(1, 2))  # (n, 1, n-1)

        # Preventing overflow
        with torch.no_grad():
            t = torch.max(x, dim=2)[0]

        x = torch.exp(x - t.unsqueeze(dim=1))
        x = torch.log(torch.exp(-t) + torch.sum(x, 2))
        loss = torch.mean(t + x)

        return loss

class NPairAngularLoss(AngularLoss):
    """
    Angular loss
    Wang, Jian. "Deep Metric Learning with Angular Loss," CVPR, 2017
    https://arxiv.org/pdf/1708.01682.pdf
    """
    def __init__(self, l2_reg=0.02, angle_bound=1., lambda_ang=2, **kwargs):
        super(NPairAngularLoss, self).__init__()
        self.l2_reg = l2_reg
        self.angle_bound = angle_bound
        self.lambda_ang = lambda_ang

    def forward(self, embeddings, target):
        n_pairs, n_negatives = self.get_n_pairs(target)

        if embeddings.is_cuda:
            n_pairs = n_pairs.cuda()
            n_negatives = n_negatives.cuda()

        anchors = embeddings[n_pairs[:, 0]]  # (n, embedding_size)
        positives = embeddings[n_pairs[:, 1]]  # (n, embedding_size)
        negatives = embeddings[n_negatives]  # (n, n-1, embedding_size)

        losses = self.n_pair_angular_loss(anchors, positives, negatives, self.angle_bound) \
            + self.l2_reg * self.l2_loss(anchors, positives)

        return losses, 0, 0, 0

    def n_pair_angular_loss(self, anchors, positives, negatives, angle_bound=1.):
        """
        Calculates N-Pair angular loss
        :param anchors: A torch.Tensor, (n, embedding_size)
        :param positives: A torch.Tensor, (n, embedding_size)
        :param negatives: A torch.Tensor, (n, n-1, embedding_size)
        :param angle_bound: tan^2 angle
        :return: A scalar, n-pair_loss + lambda * angular_loss
        """
        n_pair = self.n_pair_loss(anchors, positives, negatives)
        angular = self.angular_loss(anchors, positives, negatives, angle_bound)

        return (n_pair + self.lambda_ang * angular) / (1 + self.lambda_ang)


# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
'''neighbor list update'''
from mindspore import numpy as np
from mindspore import ops, Tensor
from .common import get_periodic_displacement, get_range_tensor, get_neighbour_index
from .crd_to_uint_crd import crd_to_uint_crd

step = Tensor(1, np.int32)

def not_excluded_mask(atom_numbers, excluded_list_start, excluded_list, excluded_numbers):
    not_excluded = np.full((atom_numbers, atom_numbers), True, np.bool_)
    for i, v in enumerate(excluded_list_start):
        if excluded_numbers[i] > 0:
            excluded_serial = excluded_list[ops.tensor_range(v, v + excluded_numbers[i], step)]
            not_excluded[i, excluded_serial] = False
    return not_excluded

def find_atom_neighbors(atom_numbers, uint_crd, uint_dr_to_dr_cof, cutoff_skin_square):
    dr = get_periodic_displacement(uint_crd, np.expand_dims(uint_crd, -2), uint_dr_to_dr_cof)
    dr2 = np.sum(dr ** 2, -1)
    atom_idx = get_range_tensor(atom_numbers)
    nl_mask = np.logical_and(atom_idx.reshape(-1, 1) < atom_idx, dr2 < cutoff_skin_square)
    return nl_mask

def delete_excluded_atoms_serial_in_neighbor_list(
        atom_numbers, max_neighbor_numbers, nl_mask, not_excluded):
    mask = np.logical_and(nl_mask, not_excluded)
    serial_idx = get_neighbour_index(atom_numbers, atom_numbers)
    nl_serial = np.where(mask, serial_idx, atom_numbers)
    nl_serial = np.sort(nl_serial, -1)[:, : max_neighbor_numbers]
    nl_numbers = np.sum(mask, -1)
    return nl_numbers, nl_serial

def crd_periodic_map(crd, box_length):
    crd = np.where(crd < 0, crd + box_length, crd)
    crd = np.where(crd > box_length, crd - box_length, crd)
    return crd

def find_atom_in_grid_serial(grid_length_inverse, crd, grid_N, Nxy, atom_in_grid_serial):
    grid_idx = (crd * grid_length_inverse).astype(np.int32)
    grid_idx = np.where(grid_idx < grid_N, grid_idx, 0)
    atom_in_grid_serial = grid_idx[..., 2] * Nxy + grid_idx[..., 1] * grid_N[0] + grid_idx[..., 0]
    return atom_in_grid_serial

def neighbor_list_update(
        grid_numbers, atom_numbers, not_first_time, Nxy, excluded_atom_numbers,
        cutoff_square, half_skin_square, cutoff_with_skin, half_cutoff_with_skin, cutoff_with_skin_square,
        refresh_interval, cutoff, skin, max_atom_in_grid_numbers, max_neighbor_numbers,
        atom_numbers_in_grid_bucket, bucket, crd, box_length, grid_N, grid_length_inverse, atom_in_grid_serial,
        old_crd, crd_to_uint_crd_cof, uint_crd, gpointer, nl_atom_numbers, nl_atom_serial, uint_dr_to_dr_cof,
        not_excluded, need_refresh_flag, refresh_count):
    """
    Update (or construct if first time) the Verlet neighbor list for the
    calculation of short-ranged force. Assume the number of atoms is n,
    the number of grids divided is G, the maximum number of atoms in one
    grid is m, the maximum number of atoms in single atom's neighbor list
    is L, and the number of total atom in excluded list is E.

    Args:
        grid_numbers (int32): the total number of grids divided.
        not_first_time (int32): whether to construct the neighbor
            list first time or not.
        Nxy (int32): the total number of grids divided in xy plane.
        excluded_atom_numbers (int32): the total atom numbers in the excluded list.
        cutoff (float32): the cutoff distance for short-range force calculation. Default: 10.0.
        skin (float32): the overflow value of cutoff to maintain a neighbor list. Default: 2.0.
        cutoff_square (float32): the suqare value of cutoff.
        half_skin_square (float32): skin*skin/4, indicates the maximum
            square value of the distance atom allowed to move between two updates.
        cutoff_with_skin (float32): cutoff + skin, indicates the
            radius of the neighbor list for each atom.
        half_cutoff_with_skin (float32): cutoff_with_skin/2.
        cutoff_with_skin_square (float32): the square value of cutoff_with_skin.
        refresh_interval (int32): the number of iteration steps between two updates of neighbor
            list. Default: 20.
        max_atom_in_grid_numbers (int32): the maximum number of atoms in one grid. Default: 64.
        max_neighbor_numbers (int32): The maximum number of neighbors. Default: 800.
        atom_numbers_in_grid_bucket (Tensor, int32) - [G,], the number of atoms in each grid bucket.
        bucket (Tensor, int32) - (Tensor,int32) - [G, m], the atom indices in each grid bucket.
        crd (Tensor, float32) - [n,], the coordinates of each atom.
        box_length (Tensor, float32) - [3,], the length of 3 dimensions of the simulation box.
        grid_n (Tensor, int32) - [3,], the number of grids divided of 3 dimensions of the
            simulation box.
        grid_length_inverse (float32) - the inverse value of grid length.
        atom_in_grid_serial (Tensor, int32) - [n,], the grid index for each atom.
        old_crd (Tensor, float32) - [n, 3], the coordinates before update of each atom.
        crd_to_uint_crd_cof (Tensor, float32) - [3,], the scale factor
            between the unsigned int value and the real space coordinates.
        uint_crd (Tensor, uint32) - [n, 3], the unsigned int coordinates value fo each atom.
        gpointer (Tensor, int32) - [G, 125], the 125 nearest neighbor grids (including self) of each
            grid. G is the number of nearest neighbor grids.
        nl_atom_numbers (Tensor, int32) - [n,], the number of atoms in neighbor list of each atom.
        nl_atom_serial (Tensor, int32) - [n, L], the indices of atoms in neighbor list of each atom.
        uint_dr_to_dr_cof (Tensor, float32) - [3,], the scale factor between
            the real space coordinates and the unsigned int value.
        excluded_list_start (Tensor, int32) - [n,], the start excluded index in excluded list for
            each atom.
        excluded_numbers (Tensor, int32) - [n,], the number of atom excluded in excluded list for
            each atom.
        not_excluded (Tensor, bool) - [n, n], marking the excluded atoms for each atom, where each
            element ij indicates whether atom j is not excluded for atom i.
        need_refresh_flag (Tensor, int32) - [n,], whether the neighbor list of each atom need update
            or not.
        refresh_count (Tensor, int32) - [1,], count how many iteration steps have passed since last
            update.

    Outputs:
        nl_atom_numbers (Tensor, int32) - [n,], the number of atoms in neighbor list of each atom.
        nl_atom_serial (Tensor, int32) - [n, L], the indices of atoms in neighbor list of each atom.
        crd (Tensor, float32) - [n,], the coordinates of each atom.
        old_crd (Tensor, float32) - [n, 3], the coordinates before update of each atom.
        need_refresh_flag (Tensor, int32) - [n,], whether the neighbor list of each atom need update
            or not.
        refresh_count (Tensor, int32) - [1,], count how many iteration steps have passed since last
            update.

    Supported Platforms:
        ``GPU``
    """
    half_crd_to_uint_crd_cof = 0.5 * crd_to_uint_crd_cof
    if not_first_time:
        if refresh_interval > 0:
            refresh_cond = (refresh_count % refresh_interval) == 0
            trans_vec = np.full(3, -skin, np.float32)
            crd = np.where(refresh_cond, crd + trans_vec, crd)
            crd = np.where(refresh_cond, crd_periodic_map(crd, box_length), crd)
            crd = np.where(refresh_cond, crd - trans_vec, crd)
            old_crd = np.where(refresh_cond, crd, old_crd)

            uint_crd = np.where(refresh_cond,
                                crd_to_uint_crd(half_crd_to_uint_crd_cof, crd).astype(np.int32),
                                uint_crd.astype(np.int32)).astype(np.uint32)

            nl_mask = find_atom_neighbors(
                atom_numbers, uint_crd, uint_dr_to_dr_cof, cutoff_square)

            nl_atom_numbers_updated, nl_atom_serial_updated = delete_excluded_atoms_serial_in_neighbor_list(
                atom_numbers, max_neighbor_numbers, nl_mask, not_excluded)
            nl_atom_numbers = np.where(refresh_cond, nl_atom_numbers_updated, nl_atom_numbers)
            nl_atom_serial = np.where(refresh_cond, nl_atom_serial_updated, nl_atom_serial)

            refresh_count += 1
        else:
            r1 = crd - old_crd
            r1_2 = np.sum(r1, -1)
            if (r1_2 > half_skin_square).any():
                trans_vec = np.full(3, skin, np.float32)
                crd += trans_vec
                crd = crd_periodic_map(crd, box_length)
                crd -= trans_vec
                old_crd[...] = crd

                uint_crd = crd_to_uint_crd(half_crd_to_uint_crd_cof, crd)

                nl_mask = find_atom_neighbors(
                    atom_numbers, uint_crd, uint_dr_to_dr_cof, cutoff_with_skin_square)

                nl_atom_numbers, nl_atom_serial = delete_excluded_atoms_serial_in_neighbor_list(
                    atom_numbers, max_neighbor_numbers, nl_mask, not_excluded)

                need_refresh_flag[0] = 0
    else:
        trans_vec = np.full(3, skin, np.float32)
        crd = crd_periodic_map(crd, box_length)
        crd += trans_vec
        old_crd[...] = crd

        uint_crd = crd_to_uint_crd(half_crd_to_uint_crd_cof, crd)

        nl_mask = find_atom_neighbors(
            atom_numbers, uint_crd, uint_dr_to_dr_cof, cutoff_with_skin_square)

        nl_atom_numbers, nl_atom_serial = delete_excluded_atoms_serial_in_neighbor_list(
            atom_numbers, max_neighbor_numbers, nl_mask, not_excluded)

    return nl_atom_numbers, nl_atom_serial, crd, old_crd, need_refresh_flag, refresh_count


#This class is exactly like neural transfer however within a class
"""
Created on Sun Jul 12 21:53:59 2020
@author: kun-je, Adanna Obibuaku
NST project in spyder
Thisis project is done using "A Neural Algorithm of Artistic Style
by. Leon A. Gatys,  Alexander S. Ecker, Matthias Bethge" as a reference
"""
import os
import numpy as np
import tensorflow as tf
import tensorflow.keras.preprocessing.image as img
from tensorflow.keras.applications.vgg19 import decode_predictions
from tensorflow.keras.applications.vgg19 import VGG19
from tensorflow.keras.applications.vgg19 import preprocess_input

PATH =  os.path.dirname(__file__)
MAIN_PATH =  os.path.join(PATH, "static/utils/")
MODEL = VGG19()
IMG_WIDTH = 224
IMG_HEIGHT = 224
LEARNING_RATE = 0.2
CHANNEL = 3

class Neural():

    def __init__(self, alpha, beta, content_layers, style_layers, style_path, content_path, **kwargs):
        self.alpha = alpha
        self.beta = beta
        self.content_layers = content_layers
        self.style_layers = style_layers
        self.content_path = content_path
        self.style_path = style_path

        self.style_path = os.path.join(MAIN_PATH, self.style_path)
        self.content_path = os.path.join(MAIN_PATH, 'input', self.content_path)
        self.c_image, self.g_image, self.s_image = self.tensor_inputs(self.content_path, self.content_path, self.style_path)

    def load_image(self, image_path):
        """
            Description:
                As we are using a pre-trained version VGG16 we have to resize and normalise
                the inputs.
            Args:
                image_path (str): This takes a given an image path
            Returns:
                <class 'numpy.ndarray'> : This would convert the given image into array
                <class 'PIL.Image.Image'>: This would convert the given image into PIL format
        """
        image_array = img.load_img(image_path, target_size=(IMG_HEIGHT, IMG_WIDTH))
        image = img.img_to_array(image_array)
        image = image.reshape((1, IMG_HEIGHT, IMG_WIDTH, CHANNEL))
        image = preprocess_input(image)
        return tf.convert_to_tensor(image), image_array

    def tensor_inputs(self, c_image_path, g_image_path, s_image_path):
        """
            Description:
                This is used to take return the tensor image of our content image,
                generate image and style image
            Args:
                c_image_path ():
                g_image_path ():
                s_image_path ():
            Returns:
        """
        c_image = self.load_image(c_image_path)[0]
        g_image = self.load_image(g_image_path)[0]
        s_image = self.load_image(s_image_path)[0]
        return c_image, g_image, s_image

    def deprocess_img(self, image):
        """
            Description:
                This is used to reverse the depressing of the image. This is used in order
                to get the image.
            Args:
                image (<class 'numpy.ndarray'>) : Take in the given image in a preprocess format
            Returns:
        """
        temp_image = image
        temp_image = temp_image[0] # Gets one image, from samples
        temp_image = temp_image.reshape((IMG_HEIGHT, IMG_WIDTH, CHANNEL)) # converts it into 3-dimentions
        temp_image[:,:,0] += 103.939 #This adds the mean rgb back to the image, which the preprocess to off
        temp_image[:,:,1] += 116.779
        temp_image[:,:,2] += 123.68
        temp_image = temp_image[:,:,::-1]
        temp_image = np.clip(temp_image, 0, 255)
        return temp_image.astype('uint8')

    def save_image(self, file_name, array_image):
        """
            Description:
                This saves a given tensor image and saves the generated file into
                an output folder
            Args:
                file_name (string): This takes in the given file name
                array_image (): This takes in the given array
        """
        file_name = os.path.join(MAIN_PATH, "output/", file_name)
        img.save_img(file_name, self.deprocess_img(array_image))
        return True

    def MSE(self, matrix_content, matrix_generated):
        """
            Args:
                matrix_content (<class 'numpy.ndarray'>):
                matrix_generated (<class 'numpy.ndarray'>):
            Returns:
                int: A number made by perform substraction operation from each matrix (tensor), followed by
                    squared operation with each substraction operation. The operation reduce mean is then applied.
        """
        return tf.reduce_mean(tf.square(matrix_content - matrix_generated))

    def get_layer(self, c_image, s_image, g_image, layer_name):
        """
            Description:
                This returns the activation of the input image.
            Args:
                image (<class 'numpy.ndarray'>): A given image array
                layer_name (str): A given layer name within the cnn model
            Returns:

                <class 'numpy.ndarray'> :
        """
        tensor_image = tf.concat([c_image, s_image, g_image], axis = 0) #put images within one array
        layer = tf.keras.Model(inputs=MODEL.inputs, outputs=MODEL.get_layer(layer_name).output)
        feature = layer(tensor_image) #This will return the activations of the function
        return feature

    def get_feature(self, c_image, s_image, g_image, layer_name):
        """
            Description:
                This function takes in the tensor repersentations c_image, s_image and g_image
                and returns their feauture activations.
            Args:
                c_image (): This is a tensor repersentation of the content image
                s_image (): This is a tensor repersentation of the style image
                g_image (): This is a tensor repersentation of the generated image
            Returns:
                : features of content image
                : features of style image
                : features of generated image

        """
        layer_feature = self.get_layer(c_image, s_image, g_image, layer_name)
        c_feature = layer_feature[0, :, :, :]
        s_feature = layer_feature[1, :, :, :]
        g_feature = layer_feature[2, :, :, :]
        return c_feature, s_feature, g_feature

    def content_loss_function(self, c_feature, g_feature):
        #todo need to change doc string as type was changed
        """
            Args:
                layer_name (str): To take in the layer name

            Returns:
                int: The loss content. A low integer denotes the content is similar
                to the generated image. A high integer denotes the content is not similar
                to the generated image
        """
        WEIGHT = 0.5
        loss = self.MSE(g_feature, c_feature)
        return WEIGHT*loss

    def gram_matrix(self, tensor):
        """
            Args:
                tensor (tensor): take 3D tensor
            Returns:
                gram (tensor) : gram matrix which is 2D array of the multiplication
                of the reshape matrix and its transpose
        """
        m_shape = []
        m_shape.append(tensor.shape[2])
        m_shape.append(tensor.shape[0]*tensor.shape[1])
        tensor = tf.reshape(tensor,m_shape)
        gram = tf.matmul(tensor,tf.transpose(tensor))
        return gram

    def style_loss_function(self, s_feature, g_feature):
        """
            Args:
                c_image_path (str): To take the style image path
                g_image_path (str): To take the generate image path
            Returns:
                int: The loss content. A low integer denotes the content is similar
                to the generated image. A high integer denotes the content is not similar
                to the generated image
        """

        #finding gram matrix of s and g image from perticular layer
        generated_gram = self.gram_matrix(g_feature)
        style_gram = self.gram_matrix(s_feature)

        img_size = IMG_HEIGHT * IMG_WIDTH

        loss = self.MSE(generated_gram, style_gram)/(4*(CHANNEL**2)*(img_size**2))
        return loss

    def total_loss_function(self, c_image,s_image,g_image,alpha,beta):
        """
            Args:
                c_image_path (str): To take the content image path
                s_image_path (str): To take the style image path
                g_image_path (str): To take the generate image path
            Returns:
                int: The totoal loss of style and content.
        """
        content_loss = 0
        c_feature, s_feature, g_feature = self.get_feature(c_image, s_image, g_image, self.content_layers[0])
        content_loss += self.content_loss_function(c_feature, g_feature)

        style_loss = 0
        for layer in self.style_layers:
            c_feature, s_feature, g_feature = self.get_feature(c_image, s_image, g_image, layer)
            style_loss += self.style_loss_function(s_feature, g_feature)

        content_loss *= alpha
        style_loss *= beta

        #total loss
        loss = style_loss + content_loss
        return loss

    def gradient_total_loss(self, c_image, s_image, g_image, alpha, beta):
        """
            Description:
                The purpose of this function is to find the current gradient at
                the generate image variable (g_image),
            Args:
                c_image ():
                s_image ():
                g_image ():
            Returns:
        """
        with tf.GradientTape() as tape:
            tape.watch(g_image)
            loss =  self.total_loss_function(c_image,s_image,g_image,alpha,beta)
        dy_dx = tape.gradient(loss, g_image)
        return loss, dy_dx

    def regression_total_loss(self):
            """
                Description:
                    The purpose of this is to use an optimization algorthium called gradient
                    descent to minimise our loss function (total loss function) in the direction
                    of the steepest descent. In other words where the graph is at its lowest point.
                    In this circumstances we want c_image, s_image to remain static while within
                    every loop we countisouly change g_image to minimise the loss of our total loss function.
                Args:
                    c_image ()
                    s_image ()
                    g_image ()
                Returns

            """
            opt = tf.keras.optimizers.Adam(learning_rate=0.2)
            self.g_image  = tf.Variable(self.g_image)
            iteration = 8000
            for _ in range(iteration+1):
                loss, dy_dx = self.gradient_total_loss(self.c_image, self.s_image, self.g_image, self.alpha, self.beta)
                print("\t Iteration: %d\t Loss: %f" % (_, loss))
                opt.apply_gradients([(dy_dx, self.g_image)]) # Apply gradient to varaiable
                if _ % 10 == 0:
                    fname = "img_%d.jpg" % (_)
                    self.save_image(fname, self.g_image.numpy())