src/model/lightning.py

import warnings
import tempfile
from typing import Optional, Union
from time import time
from pathlib import Path
from functools import partial
from itertools import accumulate
from argparse import Namespace

import numpy as np
import pandas as pd
from rdkit import Chem
import torch
from torch.utils.data import DataLoader, SubsetRandomSampler
from torch.distributions.categorical import Categorical
import pytorch_lightning as pl
from torch_scatter import scatter_mean

import src.utils as utils
from src.constants import atom_encoder, atom_decoder, aa_encoder, aa_decoder, \
    bond_encoder, bond_decoder, residue_encoder, residue_bond_encoder, \
    residue_decoder, residue_bond_decoder, aa_atom_index, aa_atom_mask
from src.data.dataset import ProcessedLigandPocketDataset, ClusteredDataset, get_wds
from src.data import data_utils
from src.data.data_utils import AppendVirtualNodesInCoM, center_data, Residues, TensorDict, randomize_tensors
from src.model.flows import CoordICFM, TorusICFM, CoordICFMPredictFinal, TorusICFMPredictFinal, SO3ICFM
from src.model.markov_bridge import UniformPriorMarkovBridge, MarginalPriorMarkovBridge
from src.model.dynamics import Dynamics
from src.model.dynamics_hetero import DynamicsHetero
from src.model.diffusion_utils import DistributionNodes
from src.model.loss_utils import TimestepWeights, clash_loss
from src.analysis.visualization_utils import pocket_to_rdkit, mols_to_pdbfile
from src.analysis.metrics import MoleculeValidity, CategoricalDistribution, MolecularProperties
from src.data.molecule_builder import build_molecule
from src.data.postprocessing import process_all
from src.sbdd_metrics.metrics import FullEvaluator
from src.sbdd_metrics.evaluation import VALIDITY_METRIC_NAME, aggregated_metrics, collection_metrics
from tqdm import tqdm

# derive additional constants
aa_atom_mask_tensor = torch.tensor([aa_atom_mask[aa] for aa in aa_decoder])
aa_atom_decoder = {aa: {v: k for k, v in aa_atom_index[aa].items()} for aa in aa_decoder}
aa_atom_type_tensor = torch.tensor([[atom_encoder.get(aa_atom_decoder[aa].get(i, '-')[0], -42)
                                     for i in range(14)] for aa in aa_decoder])


def set_default(namespace, key, default_val):
    val = vars(namespace).get(key, default_val)
    setattr(namespace, key, val)


class DrugFlow(pl.LightningModule):
    def __init__(
            self,
            pocket_representation: str,
            train_params: Namespace,
            loss_params: Namespace,
            eval_params: Namespace,
            predictor_params: Namespace,
            simulation_params: Namespace,
            virtual_nodes: Union[list, None],
            flexible: bool,
            flexible_bb: bool = False,
            debug: bool = False,
            overfit: bool = False,
    ):
        super(DrugFlow, self).__init__()
        self.save_hyperparameters()

        # Set default parameters
        set_default(train_params, "sharded_dataset", False)
        set_default(train_params, "sample_from_clusters", False)
        set_default(train_params, "lr_step_size", None)
        set_default(train_params, "lr_gamma", None)
        set_default(train_params, "gnina", None)
        set_default(loss_params, "lambda_x", 1.0)
        set_default(loss_params, "lambda_clash", None)
        set_default(loss_params, "reduce", "mean")
        set_default(loss_params, "regularize_uncertainty", None)
        set_default(eval_params, "n_loss_per_sample", 1)
        set_default(eval_params, "n_sampling_steps", simulation_params.n_steps)
        set_default(predictor_params, "transform_sc_pred", False)
        set_default(predictor_params, "add_chi_as_feature", False)
        set_default(predictor_params, "augment_residue_sc", False)
        set_default(predictor_params, "augment_ligand_sc", False)
        set_default(predictor_params, "add_all_atom_diff", False)
        set_default(predictor_params, "angle_act_fn", None)
        set_default(simulation_params, "predict_confidence", False)
        set_default(simulation_params, "predict_final", False)
        set_default(simulation_params, "scheduler_chi", None)

        # Check for invalid configurations
        assert pocket_representation in {'side_chain_bead', 'CA+'}
        self.pocket_representation = pocket_representation

        assert flexible or not predictor_params.augment_residue_sc
        self.augment_residue_sc = predictor_params.augment_residue_sc \
            if 'augment_residue_sc' in predictor_params else False
        self.augment_ligand_sc = predictor_params.augment_ligand_sc \
            if 'augment_ligand_sc' in predictor_params else False

        assert not (flexible_bb and predictor_params.normal_modes), \
            "Normal mode eigenvectors are only meaningful for fixed backbones"
        assert (not flexible_bb) or flexible, \
            "Currently atom vectors aren't updated if flexible=False"

        assert not (simulation_params.predict_confidence and
                    (not predictor_params.heterogeneous_graph or simulation_params.predict_final))

        # Set parameters
        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None
        self.virtual_nodes = virtual_nodes
        self.flexible = flexible
        self.flexible_bb = flexible_bb
        self.debug = debug
        self.overfit = overfit
        self.predict_confidence = simulation_params.predict_confidence

        if self.virtual_nodes:
            self.add_virtual_min = virtual_nodes[0]
            self.add_virtual_max = virtual_nodes[1]

        # Training parameters
        self.datadir = train_params.datadir
        self.receptor_dir = train_params.datadir
        self.batch_size = train_params.batch_size
        self.lr = train_params.lr
        self.lr_step_size = train_params.lr_step_size
        self.lr_gamma = train_params.lr_gamma
        self.num_workers = train_params.num_workers
        self.sample_from_clusters = train_params.sample_from_clusters
        self.sharded_dataset = train_params.sharded_dataset
        self.clip_grad = train_params.clip_grad
        if self.clip_grad:
            self.gradnorm_queue = utils.Queue()
            # Add large value that will be flushed.
            self.gradnorm_queue.add(3000)

        # Evaluation parameters
        self.outdir = eval_params.outdir
        self.eval_batch_size = eval_params.eval_batch_size
        self.eval_epochs = eval_params.eval_epochs
        # assert eval_params.visualize_sample_epoch % self.eval_epochs == 0
        self.visualize_sample_epoch = eval_params.visualize_sample_epoch
        self.visualize_chain_epoch = eval_params.visualize_chain_epoch
        self.sample_with_ground_truth_size = eval_params.sample_with_ground_truth_size
        self.n_loss_per_sample = eval_params.n_loss_per_sample
        self.n_eval_samples = eval_params.n_eval_samples
        self.n_visualize_samples = eval_params.n_visualize_samples
        self.keep_frames = eval_params.keep_frames
        self.gnina = train_params.gnina

        # Feature encoders/decoders
        self.atom_encoder = atom_encoder
        self.atom_decoder = atom_decoder
        self.bond_encoder = bond_encoder
        self.bond_decoder = bond_decoder
        self.aa_encoder = aa_encoder
        self.aa_decoder = aa_decoder
        self.residue_encoder = residue_encoder
        self.residue_decoder = residue_decoder
        self.residue_bond_encoder = residue_bond_encoder
        self.residue_bond_decoder = residue_bond_decoder

        self.atom_nf = len(self.atom_decoder)
        self.residue_nf = len(self.aa_decoder)
        if self.pocket_representation == 'side_chain_bead':
            self.residue_nf += len(self.residue_encoder)
        if self.pocket_representation == 'CA+':
            self.aa_atom_index = aa_atom_index
            self.n_atom_aa = max([x for aa in aa_atom_index.values() for x in aa.values()]) + 1
            self.residue_nf = (self.residue_nf, self.n_atom_aa)  # (s, V)
        self.bond_nf = len(self.bond_decoder)
        self.pocket_bond_nf = len(self.residue_bond_decoder)
        self.x_dim = 3

        # Set up the neural network
        self.dynamics = self.init_model(predictor_params)

        # Initialize objects for each variable type
        if simulation_params.predict_final:
            self.module_x = CoordICFMPredictFinal(None)
            self.module_chi = TorusICFMPredictFinal(None, 5) if self.flexible else None
            if self.flexible_bb:
                raise NotImplementedError()
        else:
            self.module_x = CoordICFM(None)
            # self.module_chi = AngleICFM(None, 5) if self.flexible else None
            scheduler_args = None if simulation_params.scheduler_chi is None else vars(simulation_params.scheduler_chi)
            self.module_chi = TorusICFM(None, 5, scheduler_args) if self.flexible else None
            self.module_trans = CoordICFM(None) if self.flexible_bb else None
            self.module_rot = SO3ICFM(None) if self.flexible_bb else None

        if simulation_params.prior_h == 'uniform':
            self.module_h = UniformPriorMarkovBridge(self.atom_nf, loss_type=loss_params.discrete_loss)
        elif simulation_params.prior_h == 'marginal':
            self.register_buffer('prior_h', self.get_categorical_prop('atom'))  # add to module
            self.module_h = MarginalPriorMarkovBridge(self.atom_nf, self.prior_h, loss_type=loss_params.discrete_loss)

        if simulation_params.prior_e == 'uniform':
            self.module_e = UniformPriorMarkovBridge(self.bond_nf, loss_type=loss_params.discrete_loss)
        elif simulation_params.prior_e == 'marginal':
            self.register_buffer('prior_e', self.get_categorical_prop('bond'))  # add to module
            self.module_e = MarginalPriorMarkovBridge(self.bond_nf, self.prior_e, loss_type=loss_params.discrete_loss)


        # Loss parameters
        self.loss_reduce = loss_params.reduce
        self.lambda_x = loss_params.lambda_x
        self.lambda_h = loss_params.lambda_h
        self.lambda_e = loss_params.lambda_e
        self.lambda_chi = loss_params.lambda_chi if self.flexible else None
        self.lambda_trans = loss_params.lambda_trans if self.flexible_bb else None
        self.lambda_rot = loss_params.lambda_rot if self.flexible_bb else None
        self.lambda_clash = loss_params.lambda_clash
        self.regularize_uncertainty = loss_params.regularize_uncertainty

        if loss_params.timestep_weights is not None:
            weight_type = loss_params.timestep_weights.split('_')[0]
            kwargs = loss_params.timestep_weights.split('_')[1:]
            kwargs = {x.split('=')[0]: float(x.split('=')[1]) for x in kwargs}
            self.timestep_weights = TimestepWeights(weight_type, **kwargs)
        else:
            self.timestep_weights = None


        # Sampling
        self.T_sampling = eval_params.n_sampling_steps
        self.train_step_size = 1 / simulation_params.n_steps
        self.size_distribution = None  # initialized only if needed


        # Metrics, initialized only if needed
        self.train_smiles = None
        self.ligand_metrics = None
        self.molecule_properties = None
        self.evaluator = None
        self.ligand_atom_type_distribution = None
        self.ligand_bond_type_distribution = None

        # containers for metric aggregation
        self.training_step_outputs = []
        self.validation_step_outputs = []

    def on_load_checkpoint(self, checkpoint):
        """
        This hook is only used for backward compatibility with checkpoints that
        did not save prior_h and prior_e in state_dict in the past
        """
        if hasattr(self, "prior_h") and "prior_h" not in checkpoint["state_dict"]:
            checkpoint["state_dict"]["prior_h"] = self.get_categorical_prop('atom')
        if hasattr(self, "prior_e") and "prior_e" not in checkpoint["state_dict"]:
            checkpoint["state_dict"]["prior_e"] = self.get_categorical_prop('bond')
        if "prior_e" in checkpoint["state_dict"] and not hasattr(self, "prior_e"):
            # NOTE: a very exotic case that happened to one model. Potentially can be removed in the future
            self.register_buffer("prior_e", self.get_categorical_prop('bond'))

    def init_model(self, predictor_params):

        model_type = predictor_params.backbone

        if 'heterogeneous_graph' in predictor_params and predictor_params.heterogeneous_graph:
            return DynamicsHetero(
                atom_nf=self.atom_nf,
                residue_nf=self.residue_nf,
                bond_dict=self.bond_encoder,
                pocket_bond_dict=self.residue_bond_encoder,
                model=model_type,
                num_rbf_time=predictor_params.__dict__.get('num_rbf_time'),
                model_params=getattr(predictor_params, model_type + '_params'),
                edge_cutoff_ligand=predictor_params.edge_cutoff_ligand,
                edge_cutoff_pocket=predictor_params.edge_cutoff_pocket,
                edge_cutoff_interaction=predictor_params.edge_cutoff_interaction,
                predict_angles=self.flexible,
                predict_frames=self.flexible_bb,
                add_cycle_counts=predictor_params.cycle_counts,
                add_spectral_feat=predictor_params.spectral_feat,
                add_nma_feat=predictor_params.normal_modes,
                reflection_equiv=predictor_params.reflection_equivariant,
                d_max=predictor_params.d_max,
                num_rbf_dist=predictor_params.num_rbf,
                self_conditioning=predictor_params.self_conditioning,
                augment_residue_sc=self.augment_residue_sc,
                augment_ligand_sc=self.augment_ligand_sc,
                add_chi_as_feature=predictor_params.add_chi_as_feature,
                angle_act_fn=predictor_params.angle_act_fn,
                add_all_atom_diff=predictor_params.add_all_atom_diff,
                predict_confidence=self.predict_confidence,
            )

        else:
            if predictor_params.__dict__.get('num_rbf_time') is not None:
                raise NotImplementedError("RBF time embedding not yet implemented")

            return Dynamics(
                atom_nf=self.atom_nf,
                residue_nf=self.residue_nf,
                joint_nf=predictor_params.joint_nf,
                bond_dict=self.bond_encoder,
                pocket_bond_dict=self.residue_bond_encoder,
                edge_nf=predictor_params.edge_nf,
                hidden_nf=predictor_params.hidden_nf,
                model=model_type,
                model_params=getattr(predictor_params, model_type + '_params'),
                edge_cutoff_ligand=predictor_params.edge_cutoff_ligand,
                edge_cutoff_pocket=predictor_params.edge_cutoff_pocket,
                edge_cutoff_interaction=predictor_params.edge_cutoff_interaction,
                predict_angles=self.flexible,
                predict_frames=self.flexible_bb,
                add_cycle_counts=predictor_params.cycle_counts,
                add_spectral_feat=predictor_params.spectral_feat,
                add_nma_feat=predictor_params.normal_modes,
                self_conditioning=predictor_params.self_conditioning,
                augment_residue_sc=self.augment_residue_sc,
                augment_ligand_sc=self.augment_ligand_sc,
                add_chi_as_feature=predictor_params.add_chi_as_feature,
                angle_act_fn=predictor_params.angle_act_fn,
            )

    def _load_histogram(self, type):
        """
        Load empirical categorical distributions of atom or bond types from disk.
        Returns None if the required file is not found.
        """
        assert type in {"atom", "bond"}
        filename = 'ligand_type_histogram.npy' if type == 'atom' else 'ligand_bond_type_histogram.npy'
        encoder = self.atom_encoder if type == 'atom' else self.bond_encoder
        hist_file = Path(self.datadir, filename)
        if not hist_file.exists():
            return None
        hist = np.load(hist_file, allow_pickle=True).item()
        return CategoricalDistribution(hist, encoder)

    def get_categorical_prop(self, type):
        hist = self._load_histogram(type)
        encoder = self.atom_encoder if type == 'atom' else self.bond_encoder
        # Note: default value ensures that code will crash if prior is not
        #  read from disk or loaded from checkpoint later on
        return torch.zeros(len(encoder)) * float("nan") if hist is None else torch.tensor(hist.p)

    def configure_optimizers(self):
        optimizers = [
            torch.optim.AdamW(self.parameters(), lr=self.lr, amsgrad=True, weight_decay=1e-12),
        ]

        if self.lr_step_size is None or self.lr_gamma is None:
            lr_schedulers = []
        else:
            lr_schedulers = [
                torch.optim.lr_scheduler.StepLR(optimizers[0], step_size=self.lr_step_size, gamma=self.lr_gamma),
            ]
        return optimizers, lr_schedulers

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

        self.setup_sampling()

        if stage == 'fit':
            self.train_dataset = self.get_dataset(stage='train')
            self.val_dataset = self.get_dataset(stage='val')
            self.setup_metrics()
        elif stage == 'val':
            self.val_dataset = self.get_dataset(stage='val')
            self.setup_metrics()
        elif stage == 'test':
            self.test_dataset = self.get_dataset(stage='test')
            self.setup_metrics()
        elif stage == 'generation':
            pass
        else:
            raise NotImplementedError

    def get_dataset(self, stage, pocket_transform=None):

        # when sampling we don't append virtual nodes as we might need access to the ground truth size
        if self.virtual_nodes and stage == "train":
            ligand_transform = AppendVirtualNodesInCoM(
                atom_encoder, bond_encoder, add_min=self.add_virtual_min, add_max=self.add_virtual_max)
        else:
            ligand_transform = None

        # we want to know if something goes wrong on the validation or test set
        catch_errors = stage == "train"

        if self.sharded_dataset:
            return get_wds(
                data_path=self.datadir,
                stage='val' if self.debug else stage,
                ligand_transform=ligand_transform,
                pocket_transform=pocket_transform,
            )

        if self.sample_from_clusters and stage == "train":  # val/test should be deterministic
            return ClusteredDataset(
                pt_path=Path(self.datadir, 'val.pt' if self.debug else f'{stage}.pt'),
                ligand_transform=ligand_transform,
                pocket_transform=pocket_transform,
                catch_errors=catch_errors
            )

        return ProcessedLigandPocketDataset(
            pt_path=Path(self.datadir, 'val.pt' if self.debug else f'{stage}.pt'),
            ligand_transform=ligand_transform,
            pocket_transform=pocket_transform,
            catch_errors=catch_errors
        )

    def setup_sampling(self):
        # distribution of nodes
        histogram_file = Path(self.datadir, 'size_distribution.npy')  # TODO: store this in model checkpoint so that we can sample without this file
        size_histogram = np.load(histogram_file).tolist()
        self.size_distribution = DistributionNodes(size_histogram)

    def setup_metrics(self):
        # For metrics
        smiles_file = Path(self.datadir, 'train_smiles.npy')
        self.train_smiles = None if not smiles_file.exists() else np.load(smiles_file)

        self.ligand_metrics = MoleculeValidity()
        self.molecule_properties = MolecularProperties()
        self.evaluator = FullEvaluator(gnina=self.gnina, exclude_evaluators=['geometry', 'ring_count'])
        self.ligand_atom_type_distribution = self._load_histogram('atom')
        self.ligand_bond_type_distribution = self._load_histogram('bond')

    def train_dataloader(self):
        shuffle = None if self.overfit else False if self.sharded_dataset else True
        return DataLoader(self.train_dataset, self.batch_size, shuffle=shuffle,
                          sampler=SubsetRandomSampler([0]) if self.overfit else None,
                          num_workers=self.num_workers,
                          collate_fn=self.train_dataset.collate_fn,
                          # collate_fn=partial(self.train_dataset.collate_fn, ligand_transform=batch_transform),
                          pin_memory=True)

    def val_dataloader(self):
        if self.overfit:
            return self.train_dataloader()

        return DataLoader(self.val_dataset, self.eval_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.eval_batch_size, shuffle=False,
                          num_workers=self.num_workers,
                          collate_fn=self.test_dataset.collate_fn,
                          pin_memory=True)

    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 aggregate_metrics(self, step_outputs, prefix):
        if 'timestep' in step_outputs[0]:
            timesteps = torch.cat([x['timestep'] for x in step_outputs]).squeeze()

        if 'loss_per_sample' in step_outputs[0]:
            losses = torch.cat([x['loss_per_sample'] for x in step_outputs])
            pearson_corr = torch.corrcoef(torch.stack([timesteps, losses], dim=0))[0, 1]
            self.log(f'corr_loss_timestep/{prefix}', pearson_corr, prog_bar=False)

        if 'eps_hat_norm' in step_outputs[0]:
            eps_norm = torch.cat([x['eps_hat_norm'] for x in step_outputs])
            pearson_corr = torch.corrcoef(torch.stack([timesteps, eps_norm], dim=0))[0, 1]
            self.log(f'corr_eps_timestep/{prefix}', pearson_corr, prog_bar=False)

    def on_train_epoch_end(self):
        self.aggregate_metrics(self.training_step_outputs, 'train')
        self.training_step_outputs.clear()

    # TODO: doesn't work in multi-GPU mode
    # def on_before_batch_transfer(self, batch, dataloader_idx):
    #     """
    #     Performs operations on data before it is transferred to the GPU.
    #     Hence, supports multiple dataloaders for speedup.
    #     """
    #     batch['pocket'] = Residues(**batch['pocket'])
    #     return batch

    # # TODO: try if this is compatible with DDP
    # def on_after_batch_transfer(self, batch, dataloader_idx):
    #     """
    #     Performs operations on data after it is transferred to the GPU.
    #     """
    #     batch['pocket'] = Residues(**batch['pocket'])
    #     batch['ligand'] = TensorDict(**batch['ligand'])
    #     return batch

    def get_sc_transform_fn(self, zt_chi, zt_x, t, z0_chi, ligand_mask, pocket):
        sc_transform = {}

        if self.augment_residue_sc:
            def pred_all_atom(pred_chi, pred_trans=None, pred_rot=None):
                temp_pocket = pocket.deepcopy()

                if pred_trans is not None and pred_rot is not None:
                    zt_trans = pocket['x']
                    zt_rot = pocket['axis_angle']
                    z1_trans_pred = self.module_trans.get_z1_given_zt_and_pred(
                        zt_trans, pred_trans, None, t, pocket['mask'])
                    z1_rot_pred = self.module_rot.get_z1_given_zt_and_pred(
                        zt_rot, pred_rot, None, t, pocket['mask'])
                    temp_pocket.set_frame(z1_trans_pred, z1_rot_pred)

                z1_chi_pred = self.module_chi.get_z1_given_zt_and_pred(
                    zt_chi[..., :5], pred_chi, z0_chi, t, pocket['mask'])
                temp_pocket.set_chi(z1_chi_pred)

                all_coord = temp_pocket['v'] + temp_pocket['x'].unsqueeze(1)
                return all_coord - pocket['x'].unsqueeze(1)

            sc_transform['residues'] = pred_all_atom

        if self.augment_ligand_sc:
            # sc_transform['atoms'] = partial(self.module_x.get_z1_given_zt_and_pred, zt=zs_x, z0=None, t=t, batch_mask=lig_mask)
            sc_transform['atoms'] = lambda pred: (self.module_x.get_z1_given_zt_and_pred(
                zt_x, pred.squeeze(1), None, t, ligand_mask) - zt_x).unsqueeze(1)

        return sc_transform

    def compute_loss(self, ligand, pocket, return_info=False):
        """
        Samples time steps and computes network predictions
        """
        # TODO: move somewhere else (like collate_fn)
        pocket = Residues(**pocket)

        # Center sample
        ligand, pocket = center_data(ligand, pocket)
        if pocket['x'].numel() > 0:
            pocket_com = scatter_mean(pocket['x'], pocket['mask'], dim=0)
        else:
            pocket_com = scatter_mean(ligand['x'], ligand['mask'], dim=0)

        # # Normalize pocket coordinates
        # pocket['x'] = self.module_x.normalize(pocket['x'])

        # Sample a timestep t for each example in batch
        t = torch.rand(ligand['size'].size(0), device=ligand['x'].device).unsqueeze(-1)

        # Noise
        z0_x = self.module_x.sample_z0(pocket_com, ligand['mask'])
        z0_h = self.module_h.sample_z0(ligand['mask'])
        z0_e = self.module_e.sample_z0(ligand['bond_mask'])
        zt_x = self.module_x.sample_zt(z0_x, ligand['x'], t, ligand['mask'])
        zt_h = self.module_h.sample_zt(z0_h, ligand['one_hot'], t, ligand['mask'])
        zt_e = self.module_e.sample_zt(z0_e, ligand['bond_one_hot'], t, ligand['bond_mask'])

        if self.flexible_bb:
            z0_trans = self.module_trans.sample_z0(pocket_com, pocket['mask'])
            z1_trans = pocket['x'].detach().clone()
            zt_trans = self.module_trans.sample_zt(z0_trans, z1_trans, t, pocket['mask'])

            z0_rot = self.module_rot.sample_z0(pocket['mask'])
            z1_rot = pocket['axis_angle'].detach().clone()
            zt_rot = self.module_rot.sample_zt(z0_rot, z1_rot, t, pocket['mask'])

            # update pocket
            pocket.set_frame(zt_trans, zt_rot)

        z0_chi, zt_chi = None, None
        if self.flexible:
            # residues = [data_utils.residue_from_internal_coord(ic) for ic in pocket['residues']]
            # residues = pocket['residues']
            # z1_chi = torch.stack([data_utils.get_torsion_angles(r, device=self.device) for r in residues], dim=0)
            z1_chi = pocket['chi'][:, :5].detach().clone()

            z0_chi = self.module_chi.sample_z0(pocket['mask'])
            zt_chi = self.module_chi.sample_zt(z0_chi, z1_chi, t, pocket['mask'])

            # internal to external coordinates
            pocket.set_chi(zt_chi)
        if pocket['x'].numel() == 0:
            pocket.set_empty_v()

        # Predict denoising
        sc_transform = self.get_sc_transform_fn(zt_chi, zt_x, t, z0_chi, ligand['mask'], pocket)
        # sc_transform = None
        pred_ligand, pred_residues = self.dynamics(
            zt_x, zt_h, ligand['mask'], pocket, t,
            bonds_ligand=(ligand['bonds'], zt_e), sc_transform=sc_transform
        )

        # Compute L2 loss
        if self.predict_confidence:
            loss_x = self.module_x.compute_loss(pred_ligand['vel'], z0_x, ligand['x'], t, ligand['mask'], reduce='none')

            # compute confidence regularization
            k = self.module_x.dim  # pred.size(-1)
            sigma = pred_ligand['uncertainty_vel']
            loss_x = loss_x / (2 * sigma ** 2) + k * torch.log(sigma)

            if self.regularize_uncertainty is not None:
                loss_x = loss_x + self.regularize_uncertainty * (pred_ligand['uncertainty_vel'] - 1) ** 2

            loss_x = self.module_x.reduce_loss(loss_x, ligand['mask'], reduce=self.loss_reduce)
        else:
            loss_x = self.module_x.compute_loss(pred_ligand['vel'], z0_x, ligand['x'], t, ligand['mask'], reduce=self.loss_reduce)

        # Loss for categorical variables
        t_next = torch.clamp(t + self.train_step_size, max=1.0)
        loss_h = self.module_h.compute_loss(pred_ligand['logits_h'], zt_h, ligand['one_hot'], ligand['mask'], t, t_next, reduce=self.loss_reduce)
        loss_e = self.module_e.compute_loss(pred_ligand['logits_e'], zt_e, ligand['bond_one_hot'], ligand['bond_mask'], t, t_next, reduce=self.loss_reduce)

        loss = self.lambda_x * loss_x + self.lambda_h * loss_h + self.lambda_e * loss_e
        if self.flexible:
            loss_chi = self.module_chi.compute_loss(pred_residues['chi'], z0_chi, z1_chi, zt_chi, t, pocket['mask'], reduce=self.loss_reduce)
            loss = loss + self.lambda_chi * loss_chi

        if self.flexible_bb:
            loss_trans = self.module_trans.compute_loss(pred_residues['trans'], z0_trans, z1_trans, t, pocket['mask'], reduce=self.loss_reduce)
            loss_rot = self.module_rot.compute_loss(pred_residues['rot'], z0_rot, z1_rot, zt_rot, t, pocket['mask'], reduce=self.loss_reduce)
            loss = loss + self.lambda_trans * loss_trans + self.lambda_rot * loss_rot

        if self.lambda_clash is not None and self.lambda_clash > 0:

            if self.flexible_bb:
                pred_z1_trans = self.module_trans.get_z1_given_zt_and_pred(zt_trans, pred_residues['trans'], z0_trans, t, pocket['mask'])
                pred_z1_rot = self.module_rot.get_z1_given_zt_and_pred(zt_rot, pred_residues['rot'], z0_rot, t, pocket['mask'])
                pocket.set_frame(pred_z1_trans, pred_z1_rot)

            if self.flexible:
                # internal to external coordinates
                pred_z1_chi = self.module_chi.get_z1_given_zt_and_pred(zt_chi, pred_residues['chi'], z0_chi, t, pocket['mask'])
                pocket.set_chi(pred_z1_chi)

            pocket_coord = pocket['x'].unsqueeze(1) + pocket['v']
            pocket_types = aa_atom_type_tensor[pocket['one_hot'].argmax(dim=-1)]
            pocket_mask = pocket['mask'].unsqueeze(-1).repeat((1, pocket['v'].size(1)))

            # Extract only existing atoms
            atom_mask = aa_atom_mask_tensor[pocket['one_hot'].argmax(dim=-1)]
            pocket_coord = pocket_coord[atom_mask]
            pocket_types = pocket_types[atom_mask]
            pocket_mask = pocket_mask[atom_mask]

            # pred_z1_x = pred_x + z0_x
            pred_z1_x = self.module_x.get_z1_given_zt_and_pred(zt_x, pred_ligand['vel'], z0_x, t, ligand['mask'])
            pred_z1_h = pred_ligand['logits_h'].argmax(dim=-1)
            loss_clash = clash_loss(pred_z1_x, pred_z1_h, ligand['mask'],
                                    pocket_coord, pocket_types, pocket_mask)
            loss = loss + self.lambda_clash * loss_clash

        if self.timestep_weights is not None:
            w_t = self.timestep_weights(t).squeeze()
            loss = w_t * loss

        loss = loss.mean(0)

        info = {
            'loss_x': loss_x.mean().item(),
            'loss_h': loss_h.mean().item(),
            'loss_e': loss_e.mean().item(),
        }
        if self.flexible:
            info['loss_chi'] = loss_chi.mean().item()
        if self.flexible_bb:
            info['loss_trans'] = loss_trans.mean().item()
            info['loss_rot'] = loss_rot.mean().item()
        if self.lambda_clash is not None:
            info['loss_clash'] = loss_clash.mean().item()
        if self.predict_confidence:
            sigma_x_mol = scatter_mean(pred_ligand['uncertainty_vel'], ligand['mask'], dim=0)
            info['pearson_sigma_x'] = torch.corrcoef(torch.stack([sigma_x_mol.detach(), t.squeeze()]))[0, 1].item()
            info['mean_sigma_x'] = sigma_x_mol.mean().item()
            entropy_h = Categorical(logits=pred_ligand['logits_h']).entropy()
            entropy_h_mol = scatter_mean(entropy_h, ligand['mask'], dim=0)
            info['pearson_entropy_h'] = torch.corrcoef(torch.stack([entropy_h_mol.detach(), t.squeeze()]))[0, 1].item()
            info['mean_entropy_h'] = entropy_h_mol.mean().item()
            entropy_e = Categorical(logits=pred_ligand['logits_e']).entropy()
            entropy_e_mol = scatter_mean(entropy_e, ligand['bond_mask'], dim=0)
            info['pearson_entropy_e'] = torch.corrcoef(torch.stack([entropy_e_mol.detach(), t.squeeze()]))[0, 1].item()
            info['mean_entropy_e'] = entropy_e_mol.mean().item()

        return (loss, info) if return_info else loss

    def training_step(self, data, *args):
        ligand, pocket = data['ligand'], data['pocket']
        try:
            loss, info = self.compute_loss(ligand, pocket, return_info=True)
        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

        log_dict = {k: v for k, v in info.items() if isinstance(v, float)
                    or torch.numel(v) <= 1}
        # if self.learn_nu:
        #     log_dict['nu_x'] = self.noise_schedules['x'].nu.item()
        #     log_dict['nu_h'] = self.noise_schedules['h'].nu.item()
        #     log_dict['nu_e'] = self.noise_schedules['e'].nu.item()

        self.log_metrics({'loss': loss, **log_dict}, 'train',
                         batch_size=len(ligand['size']))

        out = {'loss': loss, **info}
        self.training_step_outputs.append(out)
        return out

    def validation_step(self, data, *args):

        # Compute the loss N times and average to get a better estimate
        loss_list, info_list = [], []
        self.dynamics.train()  # TODO: this is currently necessary to make self-conditioning work
        for _ in range(self.n_loss_per_sample):
            loss, info = self.compute_loss(data['ligand'].copy(),
                                           data['pocket'].copy(),
                                           return_info=True)
            loss_list.append(loss.item())
            info_list.append(info)
        self.dynamics.eval()
        if len(loss_list) >= 1:
            loss = np.mean(loss_list)
            info = {k: np.mean([x[k] for x in info_list]) for k in info_list[0]}
            self.log_metrics({'loss': loss, **info}, 'val', batch_size=len(data['ligand']['size']))

        # Sample
        rdmols, rdpockets, _ = self.sample(
            data=data,
            n_samples=self.n_eval_samples,
            num_nodes="ground_truth" if self.sample_with_ground_truth_size else None,
        )

        out = {
            'ligands': rdmols,
            'pockets': rdpockets,
            'receptor_files': [Path(self.receptor_dir, 'val', x) for x in data['pocket']['name']]
        }
        self.validation_step_outputs.append(out)
        return out

    # def test_step(self, data, *args):
    #     self._shared_eval(data, 'test', *args)

    def on_validation_epoch_end(self):

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

        rdmols = [m for x in self.validation_step_outputs for m in x['ligands']]
        rdpockets = [p for x in self.validation_step_outputs for p in x['pockets']]
        receptors = [r for x in self.validation_step_outputs for r in x['receptor_files']]
        self.validation_step_outputs.clear()

        ligand_atom_types = [atom_encoder[a.GetSymbol()] for m in rdmols for a in m.GetAtoms()]
        ligand_bond_types = []
        for m in rdmols:
            bonds = m.GetBonds()
            no_bonds = m.GetNumAtoms() * (m.GetNumAtoms() - 1) // 2 - m.GetNumBonds()
            ligand_bond_types += [bond_encoder['NOBOND']] * no_bonds
            for b in bonds:
                ligand_bond_types.append(bond_encoder[b.GetBondType().name])

        tic = time()
        results = self.analyze_sample(
            rdmols, ligand_atom_types, ligand_bond_types, receptors=(rdpockets if len(rdpockets) != 0 else None)
        )
        self.log_metrics(results, 'val')
        print(f'Evaluation took {time() - tic:.2f} seconds')

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

            outdir.mkdir(exist_ok=True, parents=True)

            # center for better visualization
            rdmols = rdmols[:self.n_visualize_samples]
            rdpockets = rdpockets[:self.n_visualize_samples]
            for m, p in zip(rdmols, rdpockets):
                center = m.GetConformer().GetPositions().mean(axis=0)
                for i in range(m.GetNumAtoms()):
                    x, y, z = m.GetConformer().GetPositions()[i] - center
                    m.GetConformer().SetAtomPosition(i, (x, y, z))
                for i in range(p.GetNumAtoms()):
                    x, y, z = p.GetConformer().GetPositions()[i] - center
                    p.GetConformer().SetAtomPosition(i, (x, y, z))

            # save molecule
            utils.write_sdf_file(Path(outdir, 'molecules.sdf'), rdmols)

            # save pocket
            utils.write_sdf_file(Path(outdir, 'pockets.sdf'), rdpockets)

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

        if (self.current_epoch + 1) % self.visualize_chain_epoch == 0:
            tic = time()
            outdir.mkdir(exist_ok=True, parents=True)

            if self.sharded_dataset:
                index = torch.randint(len(self.val_dataset), size=(1,)).item()
                for i, x in enumerate(self.val_dataset):
                    if i == index:
                        break
                batch = self.val_dataset.collate_fn([x])
            else:
                batch = self.val_dataset.collate_fn([self.val_dataset[torch.randint(len(self.val_dataset), size=(1,))]])
            batch['pocket'] = Residues(**batch['pocket']).to(self.device)
            pocket_copy = batch['pocket'].copy()

            if len(batch['pocket']['x']) > 0:
                ligand_chain, pocket_chain, info = self.sample_chain(batch['pocket'], self.keep_frames)
            else:
                num_nodes, _ = self.size_distribution.sample()
                ligand_chain, pocket_chain, info = self.sample_chain(batch['pocket'], self.keep_frames, num_nodes=num_nodes)

            # utils.write_sdf_file(Path(outdir, 'chain_pocket.sdf'), pocket_chain)
            # utils.write_chain(Path(outdir, 'chain_pocket.xyz'), pocket_chain)
            if self.flexible or self.flexible_bb:
                # insert ground truth at the beginning so that it's used by PyMOL to determine the connectivity
                ground_truth_pocket = pocket_to_rdkit(
                    pocket_copy, self.pocket_representation,
                    self.atom_encoder, self.atom_decoder,
                    self.aa_decoder, self.residue_decoder,
                    self.aa_atom_index
                )[0]
                ground_truth_ligand = build_molecule(
                    batch['ligand']['x'], batch['ligand']['one_hot'].argmax(1),
                    bonds=batch['ligand']['bonds'],
                    bond_types=batch['ligand']['bond_one_hot'].argmax(1),
                    atom_decoder=self.atom_decoder,
                    bond_decoder=self.bond_decoder
                )
                pocket_chain.insert(0, ground_truth_pocket)
                ligand_chain.insert(0, ground_truth_ligand)
                # pocket_chain.insert(0, pocket_chain[-1])
                # ligand_chain.insert(0, ligand_chain[-1])

            # save molecules
            utils.write_sdf_file(Path(outdir, 'chain_ligand.sdf'), ligand_chain)

            # save pocket
            mols_to_pdbfile(pocket_chain, Path(outdir, 'chain_pocket.pdb'))

            self.log_metrics(info, 'val')
            print(f'Chain visualization took {time() - tic:.2f} seconds')


    # NOTE: temporary fix of this Lightning bug:
    # https://github.com/Lightning-AI/pytorch-lightning/discussions/18110
    # Without it resume training has a strange behavior and fails
    @property
    def total_batch_idx(self) -> int:
        """Returns the current batch index (across epochs)"""
        # use `ready` instead of `completed` in case this is accessed after `completed` has been increased
        # but before the next `ready` increase
        return max(0, self.batch_progress.total.ready - 1)

    @property
    def batch_idx(self) -> int:
        """Returns the current batch index (within this epoch)"""
        # use `ready` instead of `completed` in case this is accessed after `completed` has been increased
        # but before the next `ready` increase
        return max(0, self.batch_progress.current.ready - 1)

    # def analyze_sample(self, rdmols, atom_types, bond_types, aa_types=None, receptors=None):
    #     out = {}

    #     # Distribution of node types
    #     kl_div_atom = self.ligand_atom_type_distribution.kl_divergence(atom_types) \
    #         if self.ligand_atom_type_distribution is not None else -1
    #     out['kl_div_atom_types'] = kl_div_atom

    #     # Distribution of edge types
    #     kl_div_bond = self.ligand_bond_type_distribution.kl_divergence(bond_types) \
    #         if self.ligand_bond_type_distribution is not None else -1
    #     out['kl_div_bond_types'] = kl_div_bond

    #     if aa_types is not None:
    #         kl_div_aa = self.pocket_type_distribution.kl_divergence(aa_types) \
    #             if self.pocket_type_distribution is not None else -1
    #         out['kl_div_residue_types'] = kl_div_aa

    #     # Post-process sample
    #     processed_mols = [process_all(m) for m in rdmols]

    #     # Other basic metrics
    #     results = self.ligand_metrics(rdmols)
    #     out['n_samples'] = results['n_total']
    #     out['Validity'] = results['validity']
    #     out['Connectivity'] = results['connectivity']
    #     out['valid_and_connected'] = results['valid_and_connected']

    #     # connected_mols = [get_largest_fragment(m) for m in rdmols]
    #     connected_mols = [process_all(m, largest_frag=True, adjust_aromatic_Ns=False, relax_iter=0) for m in rdmols]
    #     connected_mols = [m for m in connected_mols if m is not None]
    #     out.update(self.molecule_properties(connected_mols))

    #     # Repeat after post-processing
    #     results = self.ligand_metrics(processed_mols)
    #     out['validity_processed'] = results['validity']
    #     out['connectivity_processed'] = results['connectivity']
    #     out['valid_and_connected_processed'] = results['valid_and_connected']

    #     processed_mols = [m for m in processed_mols if m is not None]
    #     for k, v in self.molecule_properties(processed_mols).items():
    #         out[f"{k}_processed"] = v

    #     # Simple docking score
    #     if receptors is not None and self.gnina is not None:
    #         assert len(receptors) == len(rdmols)
    #         docking_results = compute_gnina_scores(rdmols, receptors, gnina=self.gnina)
    #         out.update(docking_results)

    #     # Clash score
    #     if receptors is not None:
    #         assert len(receptors) == len(rdmols)
    #         clashes = {
    #             'ligands': [legacy_clash_score(m) for m in rdmols],
    #             'pockets': [legacy_clash_score(p) for p in receptors],
    #             'between': [legacy_clash_score(m, p) for m, p in zip(rdmols, receptors)],
    #             'v2_ligands': [clash_score(m) for m in rdmols],
    #             'v2_pockets': [clash_score(p) for p in receptors],
    #             'v2_between': [clash_score(m, p) for m, p in zip(rdmols, receptors)]
    #         }
    #         for k, v in clashes.items():
    #             out[f'mean_clash_score_{k}'] = np.mean(v)
    #             out[f'frac_no_clashes_{k}'] = np.mean(np.array(v) <= 0.0)

    #     return out

    def analyze_sample(self, rdmols, atom_types, bond_types, aa_types=None, receptors=None):
        out = {}

        # Distribution of node types
        kl_div_atom = self.ligand_atom_type_distribution.kl_divergence(atom_types) \
            if self.ligand_atom_type_distribution is not None else -1
        out['kl_div_atom_types'] = kl_div_atom

        # Distribution of edge types
        kl_div_bond = self.ligand_bond_type_distribution.kl_divergence(bond_types) \
            if self.ligand_bond_type_distribution is not None else -1
        out['kl_div_bond_types'] = kl_div_bond

        if aa_types is not None:
            kl_div_aa = self.pocket_type_distribution.kl_divergence(aa_types) \
                if self.pocket_type_distribution is not None else -1
            out['kl_div_residue_types'] = kl_div_aa

        # Evaluation
        results = []
        if receptors is not None:
            with tempfile.TemporaryDirectory() as tmpdir:
                for mol, receptor in zip(tqdm(rdmols, desc='FullEvaluator'), receptors):
                    receptor_path = Path(tmpdir, 'receptor.pdb')
                    Chem.MolToPDBFile(receptor, str(receptor_path))
                    results.append(self.evaluator(mol, receptor_path))
        else:
            for mol in tqdm(rdmols, desc='FullEvaluator'):
                self.evaluator = FullEvaluator(pb_conf='mol')
                results.append(self.evaluator(mol))

        results = pd.DataFrame(results)
        agg_results = aggregated_metrics(results, self.evaluator.dtypes, VALIDITY_METRIC_NAME).fillna(0)
        agg_results['metric'] = agg_results['metric'].str.replace('.', '/')

        col_results = collection_metrics(results, self.train_smiles, VALIDITY_METRIC_NAME, exclude_evaluators='fcd')
        col_results['metric'] = 'collection/' + col_results['metric']

        all_results = pd.concat([agg_results, col_results])
        out.update(**dict(all_results[['metric', 'value']].values))

        return out

    def sample_zt_given_zs(self, zs_ligand, zs_pocket, s, t, delta_eps_x=None, uncertainty=None):

        sc_transform = self.get_sc_transform_fn(zs_pocket.get('chi'), zs_ligand['x'], s, None, zs_ligand['mask'], zs_pocket)
        pred_ligand, pred_residues = self.dynamics(
            zs_ligand['x'], zs_ligand['h'], zs_ligand['mask'], zs_pocket, s, bonds_ligand=(zs_ligand['bonds'], zs_ligand['e']),
            sc_transform=sc_transform
        )

        if delta_eps_x is not None:
            pred_ligand['vel'] = pred_ligand['vel'] + delta_eps_x

        zt_ligand = zs_ligand.copy()
        zt_ligand['x'] = self.module_x.sample_zt_given_zs(zs_ligand['x'], pred_ligand['vel'], s, t, zs_ligand['mask'])

        zt_ligand['h'] = self.module_h.sample_zt_given_zs(zs_ligand['h'], pred_ligand['logits_h'], s, t, zs_ligand['mask'])
        zt_ligand['e'] = self.module_e.sample_zt_given_zs(zs_ligand['e'], pred_ligand['logits_e'], s, t, zs_ligand['edge_mask'])

        zt_pocket = zs_pocket.copy()
        if self.flexible_bb:
            zt_trans_pocket = self.module_trans.sample_zt_given_zs(zs_pocket['x'], pred_residues['trans'], s, t, zs_pocket['mask'])
            zt_rot_pocket = self.module_rot.sample_zt_given_zs(zs_pocket['axis_angle'], pred_residues['rot'], s, t, zs_pocket['mask'])

            # update pocket in-place
            zt_pocket.set_frame(zt_trans_pocket, zt_rot_pocket)

        if self.flexible:
            zt_chi_pocket = self.module_chi.sample_zt_given_zs(zs_pocket['chi'][..., :5], pred_residues['chi'], s, t, zs_pocket['mask'])

            # update pocket in-place
            zt_pocket.set_chi(zt_chi_pocket)

        if self.predict_confidence:
            assert uncertainty is not None
            dt = (t - s).view(-1)[zt_ligand['mask']]
            uncertainty['sigma_x_squared'] += (dt * pred_ligand['uncertainty_vel']**2)
            uncertainty['entropy_h'] += (dt * Categorical(logits=pred_ligand['logits_h']).entropy())

        return zt_ligand, zt_pocket

    def simulate(self, ligand, pocket, timesteps, t_start, t_end=1.0,
                 return_frames=1, guide_log_prob=None):
        """
        Take a version of the ligand and pocket (at any time step t_start) and
        simulate the generative process from t_start to t_end.
        """

        assert 0 < return_frames <= timesteps
        assert timesteps % return_frames == 0
        assert 0.0 <= t_start < 1.0
        assert 0 < t_end <= 1.0
        assert t_start < t_end

        device = ligand['x'].device
        n_samples = len(pocket['size'])
        delta_t = (t_end - t_start) / timesteps

        # Initialize output tensors
        out_ligand = {
            'x': torch.zeros((return_frames, len(ligand['mask']), self.x_dim), device=device),
            'h': torch.zeros((return_frames, len(ligand['mask']), self.atom_nf), device=device),
            'e': torch.zeros((return_frames, len(ligand['edge_mask']), self.bond_nf), device=device)
        }
        if self.predict_confidence:
            out_ligand['sigma_x'] = torch.zeros((return_frames, len(ligand['mask'])), device=device)
            out_ligand['entropy_h'] = torch.zeros((return_frames, len(ligand['mask'])), device=device)
        out_pocket = {
            'x': torch.zeros((return_frames, len(pocket['mask']), 3), device=device),  # CA-coord
            'v': torch.zeros((return_frames, len(pocket['mask']), self.n_atom_aa, 3), device=device)  # difference vectors to all other atoms
        }

        cumulative_uncertainty = {
            'sigma_x_squared': torch.zeros(len(ligand['mask']), device=device),
            'entropy_h': torch.zeros(len(ligand['mask']), device=device)
        } if self.predict_confidence else None

        for i, t in enumerate(torch.linspace(t_start, t_end - delta_t, timesteps)):
            t_array = torch.full((n_samples, 1), fill_value=t, device=device)

            if guide_log_prob is not None:
                raise NotImplementedError('Not yet implemented for flow matching model')
                alpha_t = self.diffusion_x.schedule.alpha(self.gamma_x(t_array))

                with torch.enable_grad():
                    zt_x_ligand.requires_grad = True
                    g = guide_log_prob(t_array, x=ligand['x'], h=ligand['h'], batch_mask=ligand['mask'],
                                       bonds=ligand['bonds'], bond_types=ligand['e'])

                    # Compute gradient w.r.t. coordinates
                    grad_x_lig = torch.autograd.grad(g.sum(), inputs=ligand['x'])[0]

                    # clip gradients
                    g_max = 1.0
                    clip_mask = (grad_x_lig.norm(dim=-1) > g_max)
                    grad_x_lig[clip_mask] = \
                        grad_x_lig[clip_mask] / grad_x_lig[clip_mask].norm(
                            dim=-1, keepdim=True) * g_max

                delta_eps_lig = -1 * (1 - alpha_t[lig_mask]).sqrt() * grad_x_lig
            else:
                delta_eps_lig = None

            ligand, pocket = self.sample_zt_given_zs(
                ligand, pocket, t_array, t_array + delta_t, delta_eps_lig, cumulative_uncertainty)

            # save frame
            if (i + 1) % (timesteps // return_frames) == 0:
                idx = (i + 1) // (timesteps // return_frames)
                idx = idx - 1

                out_ligand['x'][idx] = ligand['x'].detach()
                out_ligand['h'][idx] = ligand['h'].detach()
                out_ligand['e'][idx] = ligand['e'].detach()
                if pocket['x'].numel() > 0:
                    out_pocket['x'][idx] = pocket['x'].detach()
                    out_pocket['v'][idx] = pocket['v'][:, :self.n_atom_aa, :].detach()
                if self.predict_confidence:
                    out_ligand['sigma_x'][idx] = cumulative_uncertainty['sigma_x_squared'].sqrt().detach()
                    out_ligand['entropy_h'][idx] = cumulative_uncertainty['entropy_h'].detach()

        # remove frame dimension if only the final molecule is returned
        out_ligand = {k: v.squeeze(0) for k, v in out_ligand.items()}
        out_pocket = {k: v.squeeze(0) for k, v in out_pocket.items()}

        return out_ligand, out_pocket

    def init_ligand(self, num_nodes_lig, pocket):
        device = pocket['x'].device

        n_samples = len(pocket['size'])
        lig_mask = utils.num_nodes_to_batch_mask(n_samples, num_nodes_lig, device)

        # only consider upper triangular matrix for symmetry
        lig_bonds = torch.stack(torch.where(torch.triu(
            lig_mask[:, None] == lig_mask[None, :], diagonal=1)), dim=0)
        lig_edge_mask = lig_mask[lig_bonds[0]]

        # Sample from Normal distribution in the pocket center
        pocket_com = scatter_mean(pocket['x'], pocket['mask'], dim=0)
        z0_x = self.module_x.sample_z0(pocket_com, lig_mask)
        z0_h = self.module_h.sample_z0(lig_mask)
        z0_e = self.module_e.sample_z0(lig_edge_mask)

        return TensorDict(**{
            'x': z0_x, 'h': z0_h, 'e': z0_e, 'mask': lig_mask,
            'bonds': lig_bonds, 'edge_mask': lig_edge_mask
        })

    def init_pocket(self, pocket):

        if self.flexible_bb:
            pocket_com = scatter_mean(pocket['x'], pocket['mask'], dim=0)
            z0_trans = self.module_trans.sample_z0(pocket_com, pocket['mask'])
            z0_rot = self.module_rot.sample_z0(pocket['mask'])

            # update pocket in-place
            pocket.set_frame(z0_trans, z0_rot)

        if self.flexible:
            z0_chi = self.module_chi.sample_z0(pocket['mask'])

            # # DEBUG ##
            # z0_chi = torch.stack([data_utils.get_torsion_angles(r, device=self.device) for r in pocket['residues']], dim=0)
            # ####

            # internal to external coordinates
            pocket.set_chi(z0_chi)

        if pocket['x'].numel() == 0:
            pocket.set_empty_v()

        return pocket

    def parse_num_nodes_spec(self, batch, spec=None, size_model=None):

        if spec == "2d_histogram" or spec is None:  # default option
            assert "pocket" in batch
            num_nodes = self.size_distribution.sample_conditional(
                n1=None, n2=batch['pocket']['size'])

            # make sure there is at least one potential bond
            num_nodes[num_nodes < 2] = 2

        elif isinstance(spec, (int, torch.Tensor)):
            num_nodes = spec

        elif spec == "ground_truth":
            assert "ligand" in batch
            num_nodes = batch['ligand']['size']

        elif spec == "nn_prediction":
            assert size_model is not None
            assert "pocket" in batch
            predictions = size_model.forward(batch['pocket'])
            predictions = torch.softmax(predictions, dim=-1)
            predictions[:, :5] = 0.0
            probabilities = predictions / predictions.sum(dim=1, keepdims=True)
            num_nodes = torch.distributions.Categorical(probabilities).sample()

        elif isinstance(spec, str) and spec.startswith("uniform"):
            # expected format: uniform_low_high
            assert "pocket" in batch
            left, right = map(int, spec.split("_")[1:])
            shape = batch['pocket']['size'].shape
            num_nodes = torch.randint(left, right + 1, shape, dtype=torch.long)

        else:
            raise NotImplementedError(f"Invalid size specification {spec}")

        if self.virtual_nodes:
            num_nodes += self.add_virtual_max

        return num_nodes

    @torch.no_grad()
    def sample(self, data, n_samples, num_nodes=None, timesteps=None,
               guide_log_prob=None, size_model=None, **kwargs):

        # TODO: move somewhere else (like collate_fn)
        data['pocket'] = Residues(**data['pocket'])

        timesteps = self.T_sampling if timesteps is None else timesteps

        if len(data['pocket']['x']) > 0:
            pocket = data_utils.repeat_items(data['pocket'], n_samples)
        else:
            pocket = Residues(**{key: value for key, value in data['pocket'].items()})
            pocket['name'] = pocket['name'] * n_samples
            pocket['size'] = pocket['size'].repeat(n_samples)
            pocket['n_bonds'] = pocket['n_bonds'].repeat(n_samples)

        _ligand = data_utils.repeat_items(data['ligand'], n_samples)
        # _ligand = randomize_tensors(_ligand, exclude_keys=['size', 'name'])  # avoid data leakage

        batch = {"ligand": _ligand, "pocket": pocket}
        num_nodes = self.parse_num_nodes_spec(batch, spec=num_nodes, size_model=size_model)

        # Sample from prior
        if pocket['x'].numel() > 0:
            ligand = self.init_ligand(num_nodes, pocket)
        else:
            ligand = self.init_ligand(num_nodes, _ligand)
        pocket = self.init_pocket(pocket)

        # return prior samples
        if timesteps == 0:
            # Convert into rdmols
            rdmols = [build_molecule(coords=m['x'], 
                atom_types=m['h'].argmax(1), 
                bonds=m['bonds'], 
                bond_types=m['e'].argmax(1), 
                atom_decoder=self.atom_decoder, bond_decoder=self.bond_decoder) 
                for m in data_utils.split_entity(ligand.detach().cpu(), edge_types={"e", "edge_mask"}, edge_mask=ligand["edge_mask"])]

            rdpockets = pocket_to_rdkit(pocket, self.pocket_representation,
                                        self.atom_encoder, self.atom_decoder,
                                        self.aa_decoder, self.residue_decoder,
                                        self.aa_atom_index)

            return rdmols, rdpockets, _ligand['name']

        out_tensors_ligand, out_tensors_pocket = self.simulate(
            ligand, pocket, timesteps, 0.0, 1.0,
            guide_log_prob=guide_log_prob
        )

        # Build mol objects
        x = out_tensors_ligand['x'].detach().cpu()
        ligand_type = out_tensors_ligand['h'].argmax(1).detach().cpu()
        edge_type = out_tensors_ligand['e'].argmax(1).detach().cpu()
        lig_mask = ligand['mask'].detach().cpu()
        lig_bonds = ligand['bonds'].detach().cpu()
        lig_edge_mask = ligand['edge_mask'].detach().cpu()
        sizes = torch.unique(ligand['mask'], return_counts=True)[1].tolist()
        offsets = list(accumulate(sizes[:-1], initial=0))
        mol_kwargs = {
            'coords': utils.batch_to_list(x, lig_mask),
            'atom_types': utils.batch_to_list(ligand_type, lig_mask),
            'bonds': utils.batch_to_list_for_indices(lig_bonds, lig_edge_mask, offsets),
            'bond_types': utils.batch_to_list(edge_type, lig_edge_mask)
        }
        if self.predict_confidence:
            sigma_x = out_tensors_ligand['sigma_x'].detach().cpu()
            entropy_h = out_tensors_ligand['entropy_h'].detach().cpu()
            mol_kwargs['atom_props'] = [
                {'sigma_x': x[0], 'entropy_h': x[1]}
                for x in zip(utils.batch_to_list(sigma_x, lig_mask),
                             utils.batch_to_list(entropy_h, lig_mask))
            ]
        mol_kwargs = [{k: v[i] for k, v in mol_kwargs.items()}
                      for i in range(len(mol_kwargs['coords']))]

        # Convert into rdmols
        rdmols = [build_molecule(
            **m, atom_decoder=self.atom_decoder, bond_decoder=self.bond_decoder)
            for m in mol_kwargs
        ]

        out_pocket = pocket.copy()
        out_pocket['x'] = out_tensors_pocket['x']
        out_pocket['v'] = out_tensors_pocket['v']
        rdpockets = pocket_to_rdkit(out_pocket, self.pocket_representation,
                                    self.atom_encoder, self.atom_decoder,
                                    self.aa_decoder, self.residue_decoder,
                                    self.aa_atom_index)

        return rdmols, rdpockets, _ligand['name']

    @torch.no_grad()
    def sample_chain(self, pocket, keep_frames, num_nodes=None, timesteps=None,
                     guide_log_prob=None, **kwargs):

        # TODO: move somewhere else (like collate_fn)
        pocket = Residues(**pocket)

        info = {}

        timesteps = self.T_sampling if timesteps is None else timesteps

        # n_samples = 1
        # TODO: get batch_size differently
        assert len(pocket['mask'].unique()) <= 1, "sample_chain only supports a single sample"

        # # Pocket's initial center of mass
        # pocket_com_before = scatter_mean(pocket['x'], pocket['mask'], dim=0)

        num_nodes = self.parse_num_nodes_spec(batch={"pocket": pocket}, spec=num_nodes)

        # Sample from prior
        if pocket['x'].numel() > 0:
            ligand = self.init_ligand(num_nodes, pocket)
        else:
            dummy_pocket = Residues.empty(pocket['x'].device)
            ligand = self.init_ligand(num_nodes, dummy_pocket)

        pocket = self.init_pocket(pocket)

        out_tensors_ligand, out_tensors_pocket = self.simulate(
            ligand, pocket, timesteps, 0.0, 1.0, guide_log_prob=guide_log_prob, return_frames=keep_frames)

        # chain_lig = utils.reverse_tensor(chain_lig)
        # chain_pocket = utils.reverse_tensor(chain_pocket)
        # chain_bond = utils.reverse_tensor(chain_bond)

        info['traj_displacement_lig'] = torch.norm(out_tensors_ligand['x'][-1] - out_tensors_ligand['x'][0], dim=-1).mean()
        info['traj_rms_lig'] = out_tensors_ligand['x'].std(dim=0).mean()

        # # 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)
        # chain_bond = torch.cat([chain_bond, chain_bond[-1:].repeat(10, 1, 1)], dim=0)

        # Flatten
        assert keep_frames == out_tensors_ligand['x'].size(0) == out_tensors_pocket['x'].size(0)
        n_atoms = out_tensors_ligand['x'].size(1)
        n_bonds = out_tensors_ligand['e'].size(1)
        n_residues = out_tensors_pocket['x'].size(1)
        device = out_tensors_ligand['x'].device

        def flatten_tensor(chain):
            if len(chain.size()) == 3:  # l=0 values
                return chain.view(-1, chain.size(-1))
            elif len(chain.size()) == 4:  # vectors
                return chain.view(-1, chain.size(-2), chain.size(-1))
            else:
                warnings.warn(f"Could not flatten frame dimension of tensor with shape {list(chain.size())}")
                return chain

        out_tensors_ligand_flat = {k: flatten_tensor(chain) for k, chain in out_tensors_ligand.items()}
        out_tensors_pocket_flat = {k: flatten_tensor(chain) for k, chain in out_tensors_pocket.items()}
        # ligand_flat = chain_lig.view(-1, chain_lig.size(-1))
        # ligand_mask_flat = torch.arange(chain_lig.size(0)).repeat_interleave(chain_lig.size(1)).to(chain_lig.device)
        ligand_mask_flat = torch.arange(keep_frames).repeat_interleave(n_atoms).to(device)

        # # pocket_flat = chain_pocket.view(-1, chain_pocket.size(-1))
        # # pocket_v_flat = pocket['v'].repeat(100, 1, 1)
        # pocket_flat = chain_pocket.view(-1, chain_pocket.size(-2), chain_pocket.size(-1))
        # pocket_mask_flat = torch.arange(chain_pocket.size(0)).repeat_interleave(chain_pocket.size(1)).to(chain_pocket.device)
        pocket_mask_flat = torch.arange(keep_frames).repeat_interleave(n_residues).to(device)

        # bond_flat = chain_bond.view(-1, chain_bond.size(-1))
        # bond_mask_flat = torch.arange(chain_bond.size(0)).repeat_interleave(chain_bond.size(1)).to(chain_bond.device)
        bond_mask_flat = torch.arange(keep_frames).repeat_interleave(n_bonds).to(device)
        edges_flat = ligand['bonds'].repeat(1, keep_frames)

        # # Move generated molecule back to the original pocket position
        # pocket_com_after = scatter_mean(pocket_flat[:, 0, :], pocket_mask_flat, dim=0)
        # ligand_flat[:, :self.x_dim] += (pocket_com_before - pocket_com_after)[ligand_mask_flat]
        #
        # # Move pocket back as well (for visualization purposes)
        # pocket_flat[:, 0, :] += (pocket_com_before - pocket_com_after)[pocket_mask_flat]

        # Build ligands
        x = out_tensors_ligand_flat['x'].detach().cpu()
        ligand_type = out_tensors_ligand_flat['h'].argmax(1).detach().cpu()
        ligand_mask_flat = ligand_mask_flat.detach().cpu()
        bond_mask_flat = bond_mask_flat.detach().cpu()
        edges_flat = edges_flat.detach().cpu()
        edge_type = out_tensors_ligand_flat['e'].argmax(1).detach().cpu()
        offsets = torch.zeros(keep_frames, dtype=int)  # edges_flat is already zero-based
        molecules = list(
            zip(utils.batch_to_list(x, ligand_mask_flat),
                utils.batch_to_list(ligand_type, ligand_mask_flat),
                utils.batch_to_list_for_indices(edges_flat, bond_mask_flat, offsets),
                utils.batch_to_list(edge_type, bond_mask_flat)
                )
        )

        # Convert into rdmols
        ligand_chain = [build_molecule(
            *graph, atom_decoder=self.atom_decoder,
            bond_decoder=self.bond_decoder) for graph in molecules
        ]

        # Build pockets
        # as long as the pocket does not change during sampling, we can ust
        # write it once
        out_pocket = {
            'x': out_tensors_pocket_flat['x'],
            'one_hot': pocket['one_hot'].repeat(keep_frames, 1),
            'mask': pocket_mask_flat,
            'v': out_tensors_pocket_flat['v'],
            'atom_mask': pocket['atom_mask'].repeat(keep_frames, 1),
        } if self.flexible else pocket
        pocket_chain = pocket_to_rdkit(out_pocket, self.pocket_representation,
                                       self.atom_encoder, self.atom_decoder,
                                       self.aa_decoder, self.residue_decoder,
                                       self.aa_atom_index)

        return ligand_chain, pocket_chain, info

    # def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_val, gradient_clip_algorithm):
    # def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm):
    def configure_gradient_clipping(self, optimizer, *args, **kwargs):

        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()

        # hard upper limit
        max_grad_norm = min(max_grad_norm, 10.0)

        # 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:
            print(f'Clipped gradient with value {grad_norm:.1f} '
                  f'while allowed {max_grad_norm:.1f}')
            grad_norm = max_grad_norm

        self.gradnorm_queue.add(float(grad_norm))