import torch

import torch.nn as nn

import torch.nn.functional as F

from equivariant_diffusion.egnn_new import EGNN, GNN

from equivariant_diffusion.en_diffusion import EnVariationalDiffusion

remove_mean_batch = EnVariationalDiffusion.remove_mean_batch

import numpy as np

class EGNNDynamics(nn.Module):

    def __init__(self, atom_nf, residue_nf,

                 n_dims, joint_nf=16, hidden_nf=64, device='cpu',

                 act_fn=torch.nn.SiLU(), n_layers=4, attention=False,

                 condition_time=True, tanh=False, mode='egnn_dynamics',

                 norm_constant=0, inv_sublayers=2, sin_embedding=False,

                 normalization_factor=100, aggregation_method='sum',

                 update_pocket_coords=True, edge_cutoff_ligand=None,

                 edge_cutoff_pocket=None, edge_cutoff_interaction=None,

                 reflection_equivariant=True, edge_embedding_dim=None):

        super().__init__()

        self.mode = mode

        self.edge_cutoff_l = edge_cutoff_ligand

        self.edge_cutoff_p = edge_cutoff_pocket

        self.edge_cutoff_i = edge_cutoff_interaction

        self.edge_nf = edge_embedding_dim

        self.atom_encoder = nn.Sequential(

            nn.Linear(atom_nf, 2 * atom_nf),

            act_fn,

            nn.Linear(2 * atom_nf, joint_nf)

        )

        self.atom_decoder = nn.Sequential(

            nn.Linear(joint_nf, 2 * atom_nf),

            act_fn,

            nn.Linear(2 * atom_nf, atom_nf)

        )

        self.residue_encoder = nn.Sequential(

            nn.Linear(residue_nf, 2 * residue_nf),

            act_fn,

            nn.Linear(2 * residue_nf, joint_nf)

        )

        self.residue_decoder = nn.Sequential(

            nn.Linear(joint_nf, 2 * residue_nf),

            act_fn,

            nn.Linear(2 * residue_nf, residue_nf)

        )

        self.edge_embedding = nn.Embedding(3, self.edge_nf) \

            if self.edge_nf is not None else None

        self.edge_nf = 0 if self.edge_nf is None else self.edge_nf

        if condition_time:

            dynamics_node_nf = joint_nf + 1

        else:

            print('Warning: dynamics model is _not_ conditioned on time.')

            dynamics_node_nf = joint_nf

        if mode == 'egnn_dynamics':

            self.egnn = EGNN(

                in_node_nf=dynamics_node_nf, in_edge_nf=self.edge_nf,

                hidden_nf=hidden_nf, device=device, act_fn=act_fn,

                n_layers=n_layers, attention=attention, tanh=tanh,

                norm_constant=norm_constant,

                inv_sublayers=inv_sublayers, sin_embedding=sin_embedding,

                normalization_factor=normalization_factor,

                aggregation_method=aggregation_method,

                reflection_equiv=reflection_equivariant

            )

            self.node_nf = dynamics_node_nf

            self.update_pocket_coords = update_pocket_coords

        elif mode == 'gnn_dynamics':

            self.gnn = GNN(

                in_node_nf=dynamics_node_nf + n_dims, in_edge_nf=self.edge_nf,

                hidden_nf=hidden_nf, out_node_nf=n_dims + dynamics_node_nf,

                device=device, act_fn=act_fn, n_layers=n_layers,

                attention=attention, normalization_factor=normalization_factor,

                aggregation_method=aggregation_method)

        self.device = device

        self.n_dims = n_dims

        self.condition_time = condition_time

    def forward(self, xh_atoms, xh_residues, t, mask_atoms, mask_residues):

        x_atoms = xh_atoms[:, :self.n_dims].clone()

        h_atoms = xh_atoms[:, self.n_dims:].clone()

        x_residues = xh_residues[:, :self.n_dims].clone()

        h_residues = xh_residues[:, self.n_dims:].clone()

        # embed atom features and residue features in a shared space

        h_atoms = self.atom_encoder(h_atoms)

        h_residues = self.residue_encoder(h_residues)

        # combine the two node types

        x = torch.cat((x_atoms, x_residues), dim=0)

        h = torch.cat((h_atoms, h_residues), dim=0)

        mask = torch.cat([mask_atoms, mask_residues])

        if self.condition_time:

            if np.prod(t.size()) == 1:

                # t is the same for all elements in batch.

                h_time = torch.empty_like(h[:, 0:1]).fill_(t.item())

            else:

                # t is different over the batch dimension.

                h_time = t[mask]

            h = torch.cat([h, h_time], dim=1)

        # get edges of a complete graph

        edges = self.get_edges(mask_atoms, mask_residues, x_atoms, x_residues)

        assert torch.all(mask[edges[0]] == mask[edges[1]])

        # Get edge types

        if self.edge_nf > 0:

            # 0: ligand-pocket, 1: ligand-ligand, 2: pocket-pocket

            edge_types = torch.zeros(edges.size(1), dtype=int, device=edges.device)

            edge_types[(edges[0] < len(mask_atoms)) & (edges[1] < len(mask_atoms))] = 1

            edge_types[(edges[0] >= len(mask_atoms)) & (edges[1] >= len(mask_atoms))] = 2

            # Learnable embedding

            edge_types = self.edge_embedding(edge_types)

        else:

            edge_types = None

        if self.mode == 'egnn_dynamics':

            update_coords_mask = None if self.update_pocket_coords \

                else torch.cat((torch.ones_like(mask_atoms),

                                torch.zeros_like(mask_residues))).unsqueeze(1)

            h_final, x_final = self.egnn(h, x, edges,

                                         update_coords_mask=update_coords_mask,

                                         batch_mask=mask, edge_attr=edge_types)

            vel = (x_final - x)

        elif self.mode == 'gnn_dynamics':

            xh = torch.cat([x, h], dim=1)

            output = self.gnn(xh, edges, node_mask=None, edge_attr=edge_types)

            vel = output[:, :3]

            h_final = output[:, 3:]

        else:

            raise Exception("Wrong mode %s" % self.mode)

        if self.condition_time:

            # Slice off last dimension which represented time.

            h_final = h_final[:, :-1]

        # decode atom and residue features

        h_final_atoms = self.atom_decoder(h_final[:len(mask_atoms)])

        h_final_residues = self.residue_decoder(h_final[len(mask_atoms):])

        if torch.any(torch.isnan(vel)):

            if self.training:

                vel[torch.isnan(vel)] = 0.0

            else:

                raise ValueError("NaN detected in EGNN output")

        if self.update_pocket_coords:

            # in case of unconditional joint distribution, include this as in

            # the original code

            vel = remove_mean_batch(vel, mask)

        return torch.cat([vel[:len(mask_atoms)], h_final_atoms], dim=-1), \

               torch.cat([vel[len(mask_atoms):], h_final_residues], dim=-1)

    def get_edges(self, batch_mask_ligand, batch_mask_pocket, x_ligand, x_pocket):

        adj_ligand = batch_mask_ligand[:, None] == batch_mask_ligand[None, :]

        adj_pocket = batch_mask_pocket[:, None] == batch_mask_pocket[None, :]

        adj_cross = batch_mask_ligand[:, None] == batch_mask_pocket[None, :]

        if self.edge_cutoff_l is not None:

            adj_ligand = adj_ligand & (torch.cdist(x_ligand, x_ligand) <= self.edge_cutoff_l)

        if self.edge_cutoff_p is not None:

            adj_pocket = adj_pocket & (torch.cdist(x_pocket, x_pocket) <= self.edge_cutoff_p)

        if self.edge_cutoff_i is not None:

            adj_cross = adj_cross & (torch.cdist(x_ligand, x_pocket) <= self.edge_cutoff_i)

        adj = torch.cat((torch.cat((adj_ligand, adj_cross), dim=1),

                         torch.cat((adj_cross.T, adj_pocket), dim=1)), dim=0)

        edges = torch.stack(torch.where(adj), dim=0)

        return edges