import math

from argparse import Namespace

from typing import Optional

from time import time

from pathlib import Path

import numpy as np

import torch

import torch.nn.functional as F

from torch.utils.data import DataLoader

import pytorch_lightning as pl

import wandb

from torch_scatter import scatter_add, scatter_mean

from Bio.PDB import PDBParser

from Bio.PDB.Polypeptide import three_to_one

from constants import dataset_params, FLOAT_TYPE, INT_TYPE

from equivariant_diffusion.dynamics import EGNNDynamics

from equivariant_diffusion.en_diffusion import EnVariationalDiffusion

from equivariant_diffusion.conditional_model import ConditionalDDPM, \

    SimpleConditionalDDPM

from dataset import ProcessedLigandPocketDataset

import utils

from analysis.visualization import save_xyz_file, visualize, visualize_chain

from analysis.metrics import BasicMolecularMetrics, CategoricalDistribution, \

    MoleculeProperties

from analysis.molecule_builder import build_molecule, process_molecule

from analysis.docking import smina_score

class LigandPocketDDPM(pl.LightningModule):

    def __init__(

            self,

            outdir,

            dataset,

            datadir,

            batch_size,

            lr,

            egnn_params: Namespace,

            diffusion_params,

            num_workers,

            augment_noise,

            augment_rotation,

            clip_grad,

            eval_epochs,

            eval_params,

            visualize_sample_epoch,

            visualize_chain_epoch,

            auxiliary_loss,

            loss_params,

            mode,

            node_histogram,

            pocket_representation='CA',

            virtual_nodes=False

    ):

        super(LigandPocketDDPM, self).__init__()

        self.save_hyperparameters()

        ddpm_models = {'joint': EnVariationalDiffusion,

                       'pocket_conditioning': ConditionalDDPM,

                       'pocket_conditioning_simple': SimpleConditionalDDPM}

        assert mode in ddpm_models

        self.mode = mode

        assert pocket_representation in {'CA', 'full-atom'}

        self.pocket_representation = pocket_representation

        self.dataset_name = dataset

        self.datadir = datadir

        self.outdir = outdir

        self.batch_size = batch_size

        self.eval_batch_size = eval_params.eval_batch_size \

            if 'eval_batch_size' in eval_params else batch_size

        self.lr = lr

        self.loss_type = diffusion_params.diffusion_loss_type

        self.eval_epochs = eval_epochs

        self.visualize_sample_epoch = visualize_sample_epoch

        self.visualize_chain_epoch = visualize_chain_epoch

        self.eval_params = eval_params

        self.num_workers = num_workers

        self.augment_noise = augment_noise

        self.augment_rotation = augment_rotation

        self.dataset_info = dataset_params[dataset]

        self.T = diffusion_params.diffusion_steps

        self.clip_grad = clip_grad

        if clip_grad:

            self.gradnorm_queue = utils.Queue()

            # Add large value that will be flushed.

            self.gradnorm_queue.add(3000)

        self.lig_type_encoder = self.dataset_info['atom_encoder']

        self.lig_type_decoder = self.dataset_info['atom_decoder']

        self.pocket_type_encoder = self.dataset_info['aa_encoder'] \

            if self.pocket_representation == 'CA' \

            else self.dataset_info['atom_encoder']

        self.pocket_type_decoder = self.dataset_info['aa_decoder'] \

            if self.pocket_representation == 'CA' \

            else self.dataset_info['atom_decoder']

        smiles_list = None if eval_params.smiles_file is None \

            else np.load(eval_params.smiles_file)

        self.ligand_metrics = BasicMolecularMetrics(self.dataset_info,

                                                    smiles_list)

        self.molecule_properties = MoleculeProperties()

        self.ligand_type_distribution = CategoricalDistribution(

            self.dataset_info['atom_hist'], self.lig_type_encoder)

        if self.pocket_representation == 'CA':

            self.pocket_type_distribution = CategoricalDistribution(

                self.dataset_info['aa_hist'], self.pocket_type_encoder)

        else:

            self.pocket_type_distribution = None

        self.train_dataset = None

        self.val_dataset = None

        self.test_dataset = None

        self.virtual_nodes = virtual_nodes

        self.data_transform = None

        self.max_num_nodes = len(node_histogram) - 1

        if virtual_nodes:

            # symbol = 'virtual'

            symbol = 'Ne'  # visualize as Neon atoms

            self.lig_type_encoder[symbol] = len(self.lig_type_encoder)

            self.virtual_atom = self.lig_type_encoder[symbol]

            self.lig_type_decoder.append(symbol)

            self.data_transform = utils.AppendVirtualNodes(

                self.max_num_nodes, self.lig_type_encoder, symbol)

            # Update dataset_info dictionary. This is necessary for using the

            # visualization functions.

            self.dataset_info['atom_encoder'] = self.lig_type_encoder

            self.dataset_info['atom_decoder'] = self.lig_type_decoder

        self.atom_nf = len(self.lig_type_decoder)

        self.aa_nf = len(self.pocket_type_decoder)

        self.x_dims = 3

        net_dynamics = EGNNDynamics(

            atom_nf=self.atom_nf,

            residue_nf=self.aa_nf,

            n_dims=self.x_dims,

            joint_nf=egnn_params.joint_nf,

            device=egnn_params.device if torch.cuda.is_available() else 'cpu',

            hidden_nf=egnn_params.hidden_nf,

            act_fn=torch.nn.SiLU(),

            n_layers=egnn_params.n_layers,

            attention=egnn_params.attention,

            tanh=egnn_params.tanh,

            norm_constant=egnn_params.norm_constant,

            inv_sublayers=egnn_params.inv_sublayers,

            sin_embedding=egnn_params.sin_embedding,

            normalization_factor=egnn_params.normalization_factor,

            aggregation_method=egnn_params.aggregation_method,

            edge_cutoff_ligand=egnn_params.__dict__.get('edge_cutoff_ligand'),

            edge_cutoff_pocket=egnn_params.__dict__.get('edge_cutoff_pocket'),

            edge_cutoff_interaction=egnn_params.__dict__.get('edge_cutoff_interaction'),

            update_pocket_coords=(self.mode == 'joint'),

            reflection_equivariant=egnn_params.reflection_equivariant,

            edge_embedding_dim=egnn_params.__dict__.get('edge_embedding_dim'),

        )

        self.ddpm = ddpm_models[self.mode](

                dynamics=net_dynamics,

                atom_nf=self.atom_nf,

                residue_nf=self.aa_nf,

                n_dims=self.x_dims,

                timesteps=diffusion_params.diffusion_steps,

                noise_schedule=diffusion_params.diffusion_noise_schedule,

                noise_precision=diffusion_params.diffusion_noise_precision,

                loss_type=diffusion_params.diffusion_loss_type,

                norm_values=diffusion_params.normalize_factors,

                size_histogram=node_histogram,

                virtual_node_idx=self.lig_type_encoder[symbol] if virtual_nodes else None

        )

        self.auxiliary_loss = auxiliary_loss

        self.lj_rm = self.dataset_info['lennard_jones_rm']

        if self.auxiliary_loss:

            self.clamp_lj = loss_params.clamp_lj

            self.auxiliary_weight_schedule = WeightSchedule(

                T=diffusion_params.diffusion_steps,

                max_weight=loss_params.max_weight, mode=loss_params.schedule)

    def configure_optimizers(self):

        return torch.optim.AdamW(self.ddpm.parameters(), lr=self.lr,

                                 amsgrad=True, weight_decay=1e-12)

    def setup(self, stage: Optional[str] = None):

        if stage == 'fit':

            self.train_dataset = ProcessedLigandPocketDataset(

                Path(self.datadir, 'train.npz'), transform=self.data_transform)

            self.val_dataset = ProcessedLigandPocketDataset(

                Path(self.datadir, 'val.npz'), transform=self.data_transform)

        elif stage == 'test':

            self.test_dataset = ProcessedLigandPocketDataset(

                Path(self.datadir, 'test.npz'), transform=self.data_transform)

        else:

            raise NotImplementedError

    def train_dataloader(self):

        return DataLoader(self.train_dataset, self.batch_size, shuffle=True,

                          num_workers=self.num_workers,

                          collate_fn=self.train_dataset.collate_fn,

                          pin_memory=True)

    def val_dataloader(self):

        return DataLoader(self.val_dataset, self.batch_size, shuffle=False,

                          num_workers=self.num_workers,

                          collate_fn=self.val_dataset.collate_fn,

                          pin_memory=True)

    def test_dataloader(self):

        return DataLoader(self.test_dataset, self.batch_size, shuffle=False,

                          num_workers=self.num_workers,

                          collate_fn=self.test_dataset.collate_fn,

                          pin_memory=True)

    def get_ligand_and_pocket(self, data):

        ligand = {

            'x': data['lig_coords'].to(self.device, FLOAT_TYPE),

            'one_hot': data['lig_one_hot'].to(self.device, FLOAT_TYPE),

            'size': data['num_lig_atoms'].to(self.device, INT_TYPE),

            'mask': data['lig_mask'].to(self.device, INT_TYPE),

        }

        if self.virtual_nodes:

            ligand['num_virtual_atoms'] = data['num_virtual_atoms'].to(

                self.device, INT_TYPE)

        pocket = {

            'x': data['pocket_coords'].to(self.device, FLOAT_TYPE),

            'one_hot': data['pocket_one_hot'].to(self.device, FLOAT_TYPE),

            'size': data['num_pocket_nodes'].to(self.device, INT_TYPE),

            'mask': data['pocket_mask'].to(self.device, INT_TYPE)

        }

        return ligand, pocket

    def forward(self, data):

        ligand, pocket = self.get_ligand_and_pocket(data)

        # Note: \mathcal{L} terms in the paper represent log-likelihoods while

        # our loss terms are a negative(!) log-likelihoods

        delta_log_px, error_t_lig, error_t_pocket, SNR_weight, \

        loss_0_x_ligand, loss_0_x_pocket, loss_0_h, neg_log_const_0, \

        kl_prior, log_pN, t_int, xh_lig_hat, info = \

            self.ddpm(ligand, pocket, return_info=True)

        if self.loss_type == 'l2' and self.training:

            actual_ligand_size = ligand['size'] - ligand['num_virtual_atoms'] if self.virtual_nodes else ligand['size']

            # normalize loss_t

            denom_lig = self.x_dims * actual_ligand_size + \

                        self.ddpm.atom_nf * ligand['size']

            error_t_lig = error_t_lig / denom_lig

            denom_pocket = (self.x_dims + self.ddpm.residue_nf) * pocket['size']

            error_t_pocket = error_t_pocket / denom_pocket

            loss_t = 0.5 * (error_t_lig + error_t_pocket)

            # normalize loss_0

            loss_0_x_ligand = loss_0_x_ligand / (self.x_dims * actual_ligand_size)

            loss_0_x_pocket = loss_0_x_pocket / (self.x_dims * pocket['size'])

            loss_0 = loss_0_x_ligand + loss_0_x_pocket + loss_0_h

        # VLB objective or evaluation step

        else:

            # Note: SNR_weight should be negative

            loss_t = -self.T * 0.5 * SNR_weight * (error_t_lig + error_t_pocket)

            loss_0 = loss_0_x_ligand + loss_0_x_pocket + loss_0_h

            loss_0 = loss_0 + neg_log_const_0

        nll = loss_t + loss_0 + kl_prior

        # Correct for normalization on x.

        if not (self.loss_type == 'l2' and self.training):

            nll = nll - delta_log_px

            # always the same number of nodes if virtual nodes are added

            if not self.virtual_nodes:

                # Transform conditional nll into joint nll

                # Note:

                # loss = -log p(x,h|N) and log p(x,h,N) = log p(x,h|N) + log p(N)

                # Therefore, log p(x,h|N) = -loss + log p(N)

                # => loss_new = -log p(x,h,N) = loss - log p(N)

                nll = nll - log_pN

        # Add auxiliary loss term

        if self.auxiliary_loss and self.loss_type == 'l2' and self.training:

            x_lig_hat = xh_lig_hat[:, :self.x_dims]

            h_lig_hat = xh_lig_hat[:, self.x_dims:]

            weighted_lj_potential = \

                self.auxiliary_weight_schedule(t_int.long()) * \

                self.lj_potential(x_lig_hat, h_lig_hat, ligand['mask'])

            nll = nll + weighted_lj_potential

            info['weighted_lj'] = weighted_lj_potential.mean(0)

        info['error_t_lig'] = error_t_lig.mean(0)

        info['error_t_pocket'] = error_t_pocket.mean(0)

        info['SNR_weight'] = SNR_weight.mean(0)

        info['loss_0'] = loss_0.mean(0)

        info['kl_prior'] = kl_prior.mean(0)

        info['delta_log_px'] = delta_log_px.mean(0)

        info['neg_log_const_0'] = neg_log_const_0.mean(0)

        info['log_pN'] = log_pN.mean(0)

        return nll, info

    def lj_potential(self, atom_x, atom_one_hot, batch_mask):

        adj = batch_mask[:, None] == batch_mask[None, :]

        adj = adj ^ torch.diag(torch.diag(adj))  # remove self-edges

        edges = torch.where(adj)

        # Compute pair-wise potentials

        r = torch.sum((atom_x[edges[0]] - atom_x[edges[1]])**2, dim=1).sqrt()

        # Get optimal radii

        lennard_jones_radii = torch.tensor(self.lj_rm, device=r.device)

        # unit conversion pm -> A

        lennard_jones_radii = lennard_jones_radii / 100.0

        # normalization

        lennard_jones_radii = lennard_jones_radii / self.ddpm.norm_values[0]

        atom_type_idx = atom_one_hot.argmax(1)

        rm = lennard_jones_radii[atom_type_idx[edges[0]],

                                 atom_type_idx[edges[1]]]

        sigma = 2 ** (-1 / 6) * rm

        out = 4 * ((sigma / r) ** 12 - (sigma / r) ** 6)

        if self.clamp_lj is not None:

            out = torch.clamp(out, min=None, max=self.clamp_lj)

        # Compute potential per atom

        out = scatter_add(out, edges[0], dim=0, dim_size=len(atom_x))

        # Sum potentials of all atoms

        return scatter_add(out, batch_mask, dim=0)

    def log_metrics(self, metrics_dict, split, batch_size=None, **kwargs):

        for m, value in metrics_dict.items():

            self.log(f'{m}/{split}', value, batch_size=batch_size, **kwargs)

    def training_step(self, data, *args):

        if self.augment_noise > 0:

            raise NotImplementedError

            # Add noise eps ~ N(0, augment_noise) around points.

            eps = sample_center_gravity_zero_gaussian(x.size(), x.device)

            x = x + eps * args.augment_noise

        if self.augment_rotation:

            raise NotImplementedError

            x = utils.random_rotation(x).detach()

        try:

            nll, info = self.forward(data)

        except RuntimeError as e:

            # this is not supported for multi-GPU

            if self.trainer.num_devices < 2 and 'out of memory' in str(e):

                print('WARNING: ran out of memory, skipping to the next batch')

                return None

            else:

                raise e

        loss = nll.mean(0)

        info['loss'] = loss

        self.log_metrics(info, 'train', batch_size=len(data['num_lig_atoms']))

        return info

    def _shared_eval(self, data, prefix, *args):

        nll, info = self.forward(data)

        loss = nll.mean(0)

        info['loss'] = loss

        self.log_metrics(info, prefix, batch_size=len(data['num_lig_atoms']),

                         sync_dist=True)

        return info

    def validation_step(self, data, *args):

        self._shared_eval(data, 'val', *args)

    def test_step(self, data, *args):

        self._shared_eval(data, 'test', *args)

    def validation_epoch_end(self, validation_step_outputs):

        # Perform validation on single GPU

        if not self.trainer.is_global_zero:

            return

        suffix = '' if self.mode == 'joint' else '_given_pocket'

        if (self.current_epoch + 1) % self.eval_epochs == 0:

            tic = time()

            sampling_results = getattr(self, 'sample_and_analyze' + suffix)(

                self.eval_params.n_eval_samples, self.val_dataset,

                batch_size=self.eval_batch_size)

            self.log_metrics(sampling_results, 'val')

            print(f'Evaluation took {time() - tic:.2f} seconds')

        if (self.current_epoch + 1) % self.visualize_sample_epoch == 0:

            tic = time()

            getattr(self, 'sample_and_save' + suffix)(

                self.eval_params.n_visualize_samples)

            print(f'Sample visualization took {time() - tic:.2f} seconds')

        if (self.current_epoch + 1) % self.visualize_chain_epoch == 0:

            tic = time()

            getattr(self, 'sample_chain_and_save' + suffix)(

                self.eval_params.keep_frames)

            print(f'Chain visualization took {time() - tic:.2f} seconds')

    @torch.no_grad()

    def sample_and_analyze(self, n_samples, dataset=None, batch_size=None):

        print(f'Analyzing sampled molecules at epoch {self.current_epoch}...')

        batch_size = self.batch_size if batch_size is None else batch_size

        batch_size = min(batch_size, n_samples)

        # each item in molecules is a tuple (position, atom_type_encoded)

        molecules = []

        atom_types = []

        aa_types = []

        for i in range(math.ceil(n_samples / batch_size)):

            n_samples_batch = min(batch_size, n_samples - len(molecules))

            num_nodes_lig, num_nodes_pocket = \

                self.ddpm.size_distribution.sample(n_samples_batch)

            xh_lig, xh_pocket, lig_mask, _ = self.ddpm.sample(

                n_samples_batch, num_nodes_lig, num_nodes_pocket,

                device=self.device)

            x = xh_lig[:, :self.x_dims].detach().cpu()

            atom_type = xh_lig[:, self.x_dims:].argmax(1).detach().cpu()

            lig_mask = lig_mask.cpu()

            molecules.extend(list(

                zip(utils.batch_to_list(x, lig_mask),

                    utils.batch_to_list(atom_type, lig_mask))

            ))

            atom_types.extend(atom_type.tolist())

            aa_types.extend(

                xh_pocket[:, self.x_dims:].argmax(1).detach().cpu().tolist())

        return self.analyze_sample(molecules, atom_types, aa_types)

    def analyze_sample(self, molecules, atom_types, aa_types, receptors=None):

        # Distribution of node types

        kl_div_atom = self.ligand_type_distribution.kl_divergence(atom_types) \

            if self.ligand_type_distribution is not None else -1

        kl_div_aa = self.pocket_type_distribution.kl_divergence(aa_types) \

            if self.pocket_type_distribution is not None else -1

        # Convert into rdmols

        rdmols = [build_molecule(*graph, self.dataset_info) for graph in molecules]

        # Other basic metrics

        (validity, connectivity, uniqueness, novelty), (_, connected_mols) = \

            self.ligand_metrics.evaluate_rdmols(rdmols)

        qed, sa, logp, lipinski, diversity = \

            self.molecule_properties.evaluate_mean(connected_mols)

        out = {

            'kl_div_atom_types': kl_div_atom,

            'kl_div_residue_types': kl_div_aa,

            'Validity': validity,

            'Connectivity': connectivity,

            'Uniqueness': uniqueness,

            'Novelty': novelty,

            'QED': qed,

            'SA': sa,

            'LogP': logp,

            'Lipinski': lipinski,

            'Diversity': diversity

        }

        # Simple docking score

        if receptors is not None:

            # out['smina_score'] = np.mean(smina_score(rdmols, receptors))

            out['smina_score'] = np.mean(smina_score(connected_mols, receptors))

        return out

    def get_full_path(self, receptor_name):

        pdb, suffix = receptor_name.split('.')

        receptor_name = f'{pdb.upper()}-{suffix}.pdb'

        return Path(self.datadir, 'val', receptor_name)

    @torch.no_grad()

    def sample_and_analyze_given_pocket(self, n_samples, dataset=None,

                                        batch_size=None):

        print(f'Analyzing sampled molecules given pockets at epoch '

              f'{self.current_epoch}...')

        batch_size = self.batch_size if batch_size is None else batch_size

        batch_size = min(batch_size, n_samples)

        # each item in molecules is a tuple (position, atom_type_encoded)

        molecules = []

        atom_types = []

        aa_types = []

        receptors = []

        for i in range(math.ceil(n_samples / batch_size)):

            n_samples_batch = min(batch_size, n_samples - len(molecules))

            # Create a batch

            batch = dataset.collate_fn(

                [dataset[(i * batch_size + j) % len(dataset)]

                 for j in range(n_samples_batch)]

            )

            ligand, pocket = self.get_ligand_and_pocket(batch)

            receptors.extend([self.get_full_path(x) for x in batch['receptors']])

            if self.virtual_nodes:

                num_nodes_lig = self.max_num_nodes

            else:

                num_nodes_lig = self.ddpm.size_distribution.sample_conditional(

                    n1=None, n2=pocket['size'])

            xh_lig, xh_pocket, lig_mask, _ = self.ddpm.sample_given_pocket(

                pocket, num_nodes_lig)

            x = xh_lig[:, :self.x_dims].detach().cpu()

            atom_type = xh_lig[:, self.x_dims:].argmax(1).detach().cpu()

            lig_mask = lig_mask.cpu()

            if self.virtual_nodes:

                # Remove virtual nodes for analysis

                vnode_mask = (atom_type == self.virtual_atom)

                x = x[~vnode_mask, :]

                atom_type = atom_type[~vnode_mask]

                lig_mask = lig_mask[~vnode_mask]

            molecules.extend(list(

                zip(utils.batch_to_list(x, lig_mask),

                    utils.batch_to_list(atom_type, lig_mask))

            ))

            atom_types.extend(atom_type.tolist())

            aa_types.extend(

                xh_pocket[:, self.x_dims:].argmax(1).detach().cpu().tolist())

        return self.analyze_sample(molecules, atom_types, aa_types,

                                   receptors=receptors)

    def sample_and_save(self, n_samples):

        num_nodes_lig, num_nodes_pocket = \

            self.ddpm.size_distribution.sample(n_samples)

        xh_lig, xh_pocket, lig_mask, pocket_mask = \

            self.ddpm.sample(n_samples, num_nodes_lig, num_nodes_pocket,

                             device=self.device)

        if self.pocket_representation == 'CA':

            # convert residues into atom representation for visualization

            x_pocket, one_hot_pocket = utils.residues_to_atoms(

                xh_pocket[:, :self.x_dims], self.lig_type_encoder)

        else:

            x_pocket, one_hot_pocket = \

                xh_pocket[:, :self.x_dims], xh_pocket[:, self.x_dims:]

        x = torch.cat((xh_lig[:, :self.x_dims], x_pocket), dim=0)

        one_hot = torch.cat((xh_lig[:, self.x_dims:], one_hot_pocket), dim=0)

        outdir = Path(self.outdir, f'epoch_{self.current_epoch}')

        save_xyz_file(str(outdir) + '/', one_hot, x, self.lig_type_decoder,

                      name='molecule',

                      batch_mask=torch.cat((lig_mask, pocket_mask)))

        # visualize(str(outdir), dataset_info=self.dataset_info, wandb=wandb)

        visualize(str(outdir), dataset_info=self.dataset_info, wandb=None)

    def sample_and_save_given_pocket(self, n_samples):

        batch = self.val_dataset.collate_fn(

            [self.val_dataset[i] for i in torch.randint(len(self.val_dataset),

                                                        size=(n_samples,))]

        )

        ligand, pocket = self.get_ligand_and_pocket(batch)

        if self.virtual_nodes:

            num_nodes_lig = self.max_num_nodes

        else:

            num_nodes_lig = self.ddpm.size_distribution.sample_conditional(

                n1=None, n2=pocket['size'])

        xh_lig, xh_pocket, lig_mask, pocket_mask = \

            self.ddpm.sample_given_pocket(pocket, num_nodes_lig)

        if self.pocket_representation == 'CA':

            # convert residues into atom representation for visualization

            x_pocket, one_hot_pocket = utils.residues_to_atoms(

                xh_pocket[:, :self.x_dims], self.lig_type_encoder)

        else:

            x_pocket, one_hot_pocket = \

                xh_pocket[:, :self.x_dims], xh_pocket[:, self.x_dims:]

        x = torch.cat((xh_lig[:, :self.x_dims], x_pocket), dim=0)

        one_hot = torch.cat((xh_lig[:, self.x_dims:], one_hot_pocket), dim=0)

        outdir = Path(self.outdir, f'epoch_{self.current_epoch}')

        save_xyz_file(str(outdir) + '/', one_hot, x, self.lig_type_decoder,

                      name='molecule',

                      batch_mask=torch.cat((lig_mask, pocket_mask)))

        # visualize(str(outdir), dataset_info=self.dataset_info, wandb=wandb)

        visualize(str(outdir), dataset_info=self.dataset_info, wandb=None)

    def sample_chain_and_save(self, keep_frames):

        n_samples = 1

        num_nodes_lig, num_nodes_pocket = \

            self.ddpm.size_distribution.sample(n_samples)

        chain_lig, chain_pocket, _, _ = self.ddpm.sample(

            n_samples, num_nodes_lig, num_nodes_pocket,

            return_frames=keep_frames, device=self.device)

        chain_lig = utils.reverse_tensor(chain_lig)

        chain_pocket = utils.reverse_tensor(chain_pocket)

        # Repeat last frame to see final sample better.

        chain_lig = torch.cat([chain_lig, chain_lig[-1:].repeat(10, 1, 1)],

                              dim=0)

        chain_pocket = torch.cat(

            [chain_pocket, chain_pocket[-1:].repeat(10, 1, 1)], dim=0)

        # Prepare entire chain.

        x_lig = chain_lig[:, :, :self.x_dims]

        one_hot_lig = chain_lig[:, :, self.x_dims:]

        one_hot_lig = F.one_hot(

            torch.argmax(one_hot_lig, dim=2),

            num_classes=len(self.lig_type_decoder))

        x_pocket = chain_pocket[:, :, :self.x_dims]

        one_hot_pocket = chain_pocket[:, :, self.x_dims:]

        one_hot_pocket = F.one_hot(

            torch.argmax(one_hot_pocket, dim=2),

            num_classes=len(self.pocket_type_decoder))

        if self.pocket_representation == 'CA':

            # convert residues into atom representation for visualization

            x_pocket, one_hot_pocket = utils.residues_to_atoms(

                x_pocket, self.lig_type_encoder)

        x = torch.cat((x_lig, x_pocket), dim=1)

        one_hot = torch.cat((one_hot_lig, one_hot_pocket), dim=1)

        # flatten (treat frame (chain dimension) as batch for visualization)

        x_flat = x.view(-1, x.size(-1))

        one_hot_flat = one_hot.view(-1, one_hot.size(-1))

        mask_flat = torch.arange(x.size(0)).repeat_interleave(x.size(1))

        outdir = Path(self.outdir, f'epoch_{self.current_epoch}', 'chain')

        save_xyz_file(str(outdir), one_hot_flat, x_flat, self.lig_type_decoder,

                      name='/chain', batch_mask=mask_flat)

        visualize_chain(str(outdir), self.dataset_info, wandb=wandb)

    def sample_chain_and_save_given_pocket(self, keep_frames):

        n_samples = 1

        batch = self.val_dataset.collate_fn([

            self.val_dataset[torch.randint(len(self.val_dataset), size=(1,))]

        ])

        ligand, pocket = self.get_ligand_and_pocket(batch)

        if self.virtual_nodes:

            num_nodes_lig = self.max_num_nodes

        else:

            num_nodes_lig = self.ddpm.size_distribution.sample_conditional(

                n1=None, n2=pocket['size'])

        chain_lig, chain_pocket, _, _ = self.ddpm.sample_given_pocket(

            pocket, num_nodes_lig, return_frames=keep_frames)

        chain_lig = utils.reverse_tensor(chain_lig)

        chain_pocket = utils.reverse_tensor(chain_pocket)

        # Repeat last frame to see final sample better.

        chain_lig = torch.cat([chain_lig, chain_lig[-1:].repeat(10, 1, 1)],

                              dim=0)

        chain_pocket = torch.cat(

            [chain_pocket, chain_pocket[-1:].repeat(10, 1, 1)], dim=0)

        # Prepare entire chain.

        x_lig = chain_lig[:, :, :self.x_dims]

        one_hot_lig = chain_lig[:, :, self.x_dims:]

        one_hot_lig = F.one_hot(

            torch.argmax(one_hot_lig, dim=2),

            num_classes=len(self.lig_type_decoder))

        x_pocket = chain_pocket[:, :, :3]

        one_hot_pocket = chain_pocket[:, :, 3:]

        one_hot_pocket = F.one_hot(

            torch.argmax(one_hot_pocket, dim=2),

            num_classes=len(self.pocket_type_decoder))

        if self.pocket_representation == 'CA':

            # convert residues into atom representation for visualization

            x_pocket, one_hot_pocket = utils.residues_to_atoms(

                x_pocket, self.lig_type_encoder)

        x = torch.cat((x_lig, x_pocket), dim=1)

        one_hot = torch.cat((one_hot_lig, one_hot_pocket), dim=1)

        # flatten (treat frame (chain dimension) as batch for visualization)

        x_flat = x.view(-1, x.size(-1))

        one_hot_flat = one_hot.view(-1, one_hot.size(-1))

        mask_flat = torch.arange(x.size(0)).repeat_interleave(x.size(1))

        outdir = Path(self.outdir, f'epoch_{self.current_epoch}', 'chain')

        save_xyz_file(str(outdir), one_hot_flat, x_flat, self.lig_type_decoder,

                      name='/chain', batch_mask=mask_flat)

        visualize_chain(str(outdir), self.dataset_info, wandb=wandb)

    def prepare_pocket(self, biopython_residues, repeats=1):

        if self.pocket_representation == 'CA':

            pocket_coord = torch.tensor(np.array(

                [res['CA'].get_coord() for res in biopython_residues]),

                device=self.device, dtype=FLOAT_TYPE)

            pocket_types = torch.tensor(

                [self.pocket_type_encoder[three_to_one(res.get_resname())]

                 for res in biopython_residues], device=self.device)

        else:

            pocket_atoms = [a for res in biopython_residues

                            for a in res.get_atoms()

                            if (a.element.capitalize() in self.pocket_type_encoder or a.element != 'H')]

            pocket_coord = torch.tensor(np.array(

                [a.get_coord() for a in pocket_atoms]),

                device=self.device, dtype=FLOAT_TYPE)

            pocket_types = torch.tensor(

                [self.pocket_type_encoder[a.element.capitalize()]

                 for a in pocket_atoms], device=self.device)

        pocket_one_hot = F.one_hot(

            pocket_types, num_classes=len(self.pocket_type_encoder)

        )

        pocket_size = torch.tensor([len(pocket_coord)] * repeats,

                                   device=self.device, dtype=INT_TYPE)

        pocket_mask = torch.repeat_interleave(

            torch.arange(repeats, device=self.device, dtype=INT_TYPE),

            len(pocket_coord)

        )

        pocket = {

            'x': pocket_coord.repeat(repeats, 1),

            'one_hot': pocket_one_hot.repeat(repeats, 1),

            'size': pocket_size,

            'mask': pocket_mask

        }

        return pocket

    def generate_ligands(self, pdb_file, n_samples, pocket_ids=None,

                         ref_ligand=None, num_nodes_lig=None, sanitize=False,

                         largest_frag=False, relax_iter=0, timesteps=None,

                         n_nodes_bias=0, n_nodes_min=0, **kwargs):

        """

        Generate ligands given a pocket

        Args:

            pdb_file: PDB filename

            n_samples: number of samples

            pocket_ids: list of pocket residues in <chain>:<resi> format

            ref_ligand: alternative way of defining the pocket based on a

                reference ligand given in <chain>:<resi> format if the ligand is

                contained in the PDB file, or path to an SDF file that

                contains the ligand

            num_nodes_lig: number of ligand nodes for each sample (list of

                integers), sampled randomly if 'None'

            sanitize: whether to sanitize molecules or not

            largest_frag: only return the largest fragment

            relax_iter: number of force field optimization steps

            timesteps: number of denoising steps, use training value if None

            n_nodes_bias: added to the sampled (or provided) number of nodes

            n_nodes_min: lower bound on the number of sampled nodes

            kwargs: additional inpainting parameters

        Returns:

            list of molecules

        """

        assert (pocket_ids is None) ^ (ref_ligand is None)

        self.ddpm.eval()

        # Load PDB

        pdb_struct = PDBParser(QUIET=True).get_structure('', pdb_file)[0]

        if pocket_ids is not None:

            # define pocket with list of residues

            residues = [

                pdb_struct[x.split(':')[0]][(' ', int(x.split(':')[1]), ' ')]

                for x in pocket_ids]

        else:

            # define pocket with reference ligand

            residues = utils.get_pocket_from_ligand(pdb_struct, ref_ligand)

        pocket = self.prepare_pocket(residues, repeats=n_samples)

        # Pocket's center of mass

        pocket_com_before = scatter_mean(pocket['x'], pocket['mask'], dim=0)

        # Create dummy ligands

        if num_nodes_lig is None:

            num_nodes_lig = self.ddpm.size_distribution.sample_conditional(

                n1=None, n2=pocket['size'])

        # Add bias

        num_nodes_lig = num_nodes_lig + n_nodes_bias

        # Apply minimum ligand size

        num_nodes_lig = torch.clamp(num_nodes_lig, min=n_nodes_min)

        # Use inpainting

        if type(self.ddpm) == EnVariationalDiffusion:

            lig_mask = utils.num_nodes_to_batch_mask(

                len(num_nodes_lig), num_nodes_lig, self.device)

            ligand = {

                'x': torch.zeros((len(lig_mask), self.x_dims),

                                 device=self.device, dtype=FLOAT_TYPE),

                'one_hot': torch.zeros((len(lig_mask), self.atom_nf),

                                       device=self.device, dtype=FLOAT_TYPE),

                'size': num_nodes_lig,

                'mask': lig_mask

            }

            # Fix all pocket nodes but sample

            lig_mask_fixed = torch.zeros(len(lig_mask), device=self.device)

            pocket_mask_fixed = torch.ones(len(pocket['mask']),

                                           device=self.device)

            xh_lig, xh_pocket, lig_mask, pocket_mask = self.ddpm.inpaint(

                ligand, pocket, lig_mask_fixed, pocket_mask_fixed,

                timesteps=timesteps, **kwargs)

        # Use conditional generation

        elif type(self.ddpm) == ConditionalDDPM:

            xh_lig, xh_pocket, lig_mask, pocket_mask = \

                self.ddpm.sample_given_pocket(pocket, num_nodes_lig,

                                              timesteps=timesteps)

        else:

            raise NotImplementedError

        # Move generated molecule back to the original pocket position

        pocket_com_after = scatter_mean(

            xh_pocket[:, :self.x_dims], pocket_mask, dim=0)

        xh_pocket[:, :self.x_dims] += \

            (pocket_com_before - pocket_com_after)[pocket_mask]

        xh_lig[:, :self.x_dims] += \

            (pocket_com_before - pocket_com_after)[lig_mask]

        # Build mol objects

        x = xh_lig[:, :self.x_dims].detach().cpu()

        atom_type = xh_lig[:, self.x_dims:].argmax(1).detach().cpu()

        lig_mask = lig_mask.cpu()

        molecules = []

        for mol_pc in zip(utils.batch_to_list(x, lig_mask),

                          utils.batch_to_list(atom_type, lig_mask)):

            mol = build_molecule(*mol_pc, self.dataset_info, add_coords=True)

            mol = process_molecule(mol,

                                   add_hydrogens=False,

                                   sanitize=sanitize,

                                   relax_iter=relax_iter,

                                   largest_frag=largest_frag)

            if mol is not None:

                molecules.append(mol)

        return molecules

    def configure_gradient_clipping(self, optimizer, optimizer_idx,

                                    gradient_clip_val, gradient_clip_algorithm):

        if not self.clip_grad:

            return

        # Allow gradient norm to be 150% + 2 * stdev of the recent history.

        max_grad_norm = 1.5 * self.gradnorm_queue.mean() + \

                        2 * self.gradnorm_queue.std()

        # Get current grad_norm

        params = [p for g in optimizer.param_groups for p in g['params']]

        grad_norm = utils.get_grad_norm(params)

        # Lightning will handle the gradient clipping

        self.clip_gradients(optimizer, gradient_clip_val=max_grad_norm,

                            gradient_clip_algorithm='norm')

        if float(grad_norm) > max_grad_norm:

            self.gradnorm_queue.add(float(max_grad_norm))

        else:

            self.gradnorm_queue.add(float(grad_norm))

        if float(grad_norm) > max_grad_norm:

            print(f'Clipped gradient with value {grad_norm:.1f} '

                  f'while allowed {max_grad_norm:.1f}')

class WeightSchedule:

    def __init__(self, T, max_weight, mode='linear'):

        if mode == 'linear':

            self.weights = torch.linspace(max_weight, 0, T + 1)

        elif mode == 'constant':

            self.weights = max_weight * torch.ones(T + 1)

        else:

            raise NotImplementedError(f'{mode} weight schedule is not '

                                      f'available.')

    def __call__(self, t_array):

        """ all values in t_array are assumed to be integers in [0, T] """

        return self.weights[t_array].to(t_array.device)