"""PDB dataset loader."""
import math
from typing import Optional

import torch
import torch.distributed as dist
import json
import tree
import numpy as np
import torch
import pandas as pd
import logging
import random
import functools as fn

from torch.utils import data
from data import utils as du
from analysis import utils as au
from openfold.data import data_transforms
from openfold.np import residue_constants
from openfold.utils import rigid_utils

from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
import itertools

def _rog_quantile_curve(df, quantile, eval_x):
    y_quant = pd.pivot_table(
        df,
        values='radius_gyration',
        index='modeled_seq_len',
        aggfunc=lambda x: np.quantile(x, quantile)
    )
    x_quant = y_quant.index.to_numpy()
    y_quant = y_quant.radius_gyration.to_numpy()

    # Fit polynomial regressor
    poly = PolynomialFeatures(degree=4, include_bias=True)
    poly_features = poly.fit_transform(x_quant[:, None])
    poly_reg_model = LinearRegression()
    poly_reg_model.fit(poly_features, y_quant)

    # Calculate cutoff for all sequence lengths
    pred_poly_features = poly.fit_transform(eval_x[:, None])
    # Add a little more.
    pred_y = poly_reg_model.predict(pred_poly_features) + 0.1
    return pred_y


class MotifPdbDataset(data.Dataset):
    def __init__(
            self,
            *,
            data_conf,
            motif_type,
            num_motif=1,
        ):
        self._log = logging.getLogger(__name__)
        self._data_conf = data_conf
        self._motif_type = motif_type
        self.num_motif = num_motif
        self._init_metadata()
        self.comb_motif = list(itertools.combinations(range(len(self.csv)), num_motif))
        self.motif_keys = {
            'motif_seq_masks', # [N_backbone_res, N_motifs], np.bool, True indicates motif
            'motif_aatype', # [N_motifs, N_max_motif_len(80)], torch.int64
            'motif_rigids_0', # [N_motifs, N_max_motif_len(80), 7], torch.float32, quat rot + trans
            'motif_masks', # [N_motifs, N_max_motif_len(80)], torch.bool, False for padding
            'motif_torsion_angles_sin_cos', # [N_motifs, N_max_motif_len(80), 7, 2], torch.float32
        }
        

    @property
    def data_conf(self):
        return self._data_conf

    def _init_metadata(self):
        """Initialize metadata."""
        pdb_csv = pd.read_csv(self.data_conf.motif_csv_path)
        self.raw_csv = pdb_csv
        pdb_csv = pdb_csv.sort_values('modeled_seq_len', ascending=False)
        self.csv = pdb_csv
        self.motif_mask_file = json.load(open(self.data_conf.motif_mask_path, 'r'))
        if self._motif_type == 'rfdiffusion_motif':
            self.motif_mask_dict = {k.split('_')[0]: np.array(v) for k,v in self.motif_mask_file.items()}
        elif 'sz_motif' in self._motif_type:
            self.motif_mask_dict = {f"{v['PDB']}_{v['Chain']}_{v['Start']}_{v['Stop']}": np.array(v['mask']) for idx, v in enumerate(self.motif_mask_file)}
        else:
            raise NotImplementedError('motif type not implemented')
    
    # cache make the same sample in same batch 
    @fn.lru_cache(maxsize=100)
    def _process_csv_row(self, processed_file_path):
        processed_feats = du.read_pkl(processed_file_path)
        processed_feats = du.parse_chain_feats(processed_feats)

        # Only take modeled residues.
        modeled_idx = processed_feats['modeled_idx']
        min_idx = np.min(modeled_idx)
        max_idx = np.max(modeled_idx)
        del processed_feats['modeled_idx']
        processed_feats = tree.map_structure(
            lambda x: x[min_idx:(max_idx+1)], processed_feats)

        # Run through OpenFold data transforms.
        chain_feats = {
            'aatype': torch.tensor(processed_feats['aatype']).long(),
            'all_atom_positions': torch.tensor(processed_feats['atom_positions']).double(),
            'all_atom_mask': torch.tensor(processed_feats['atom_mask']).double()
        }
        chain_feats = data_transforms.atom37_to_frames(chain_feats)
        chain_feats = data_transforms.make_atom14_masks(chain_feats)
        chain_feats = data_transforms.make_atom14_positions(chain_feats)
        chain_feats = data_transforms.atom37_to_torsion_angles()(chain_feats)

        # Re-number residue indices for each chain such that it starts from 1.
        # Randomize chain indices.
        chain_idx = processed_feats["chain_index"]
        res_idx = processed_feats['residue_index']
        new_res_idx = np.zeros_like(res_idx)
        new_chain_idx = np.zeros_like(res_idx)
        all_chain_idx = np.unique(chain_idx).tolist()
        shuffled_chain_idx = np.array(
            random.sample(all_chain_idx, len(all_chain_idx))) - np.min(all_chain_idx) + 1
        for i,chain_id in enumerate(all_chain_idx):
            chain_mask = (chain_idx == chain_id).astype(int)
            chain_min_idx = np.min(res_idx + (1 - chain_mask) * 1e3).astype(int)
            new_res_idx = new_res_idx + (res_idx - chain_min_idx + 1) * chain_mask

            # Shuffle chain_index
            replacement_chain_id = shuffled_chain_idx[i]
            new_chain_idx = new_chain_idx + replacement_chain_id * chain_mask

        # To speed up processing, only take necessary features
        final_feats = {
            'aatype': chain_feats['aatype'],
            'seq_idx': new_res_idx,
            'chain_idx': new_chain_idx,
            'residx_atom14_to_atom37': chain_feats['residx_atom14_to_atom37'],
            'residue_index': processed_feats['residue_index'],
            'res_mask': processed_feats['bb_mask'],
            'atom37_pos': chain_feats['all_atom_positions'],
            'atom37_mask': chain_feats['all_atom_mask'],
            'atom14_pos': chain_feats['atom14_gt_positions'],
            'rigidgroups_0': chain_feats['rigidgroups_gt_frames'],
            'torsion_angles_sin_cos': chain_feats['torsion_angles_sin_cos'],
        }
        return final_feats

    def __len__(self):
        return len(self.comb_motif)

    def _get_single_motif(self, example_idx):
        # Sample data example.
        csv_row = self.csv.iloc[example_idx]
        if 'pdb_name' in csv_row:
            pdb_name = csv_row['pdb_name']
        elif 'chain_name' in csv_row:
            pdb_name = csv_row['chain_name']
        else:
            raise ValueError('Need chain identifier.')
        processed_file_path = csv_row['processed_path']
        chain_feats = self._process_csv_row(processed_file_path)

        gt_bb_rigid = rigid_utils.Rigid.from_tensor_4x4(
            chain_feats['rigidgroups_0'])[:, 0]
        diffused_mask = np.ones_like(chain_feats['res_mask'])
        if np.sum(diffused_mask) < 1:
            raise ValueError('Must be diffused')
        fixed_mask = 1 - diffused_mask
        chain_feats['fixed_mask'] = fixed_mask
        chain_feats['rigids_0'] = gt_bb_rigid.to_tensor_7()
        chain_feats['sc_ca_t'] = torch.zeros_like(gt_bb_rigid.get_trans())
        
        motif_seq_mask = self.motif_mask_dict[pdb_name]
        chain_feats['motif_seq_masks'] = motif_seq_mask
        chain_feats['motif_aatype'] = chain_feats['aatype']
        chain_feats['motif_masks'] = torch.ones_like(chain_feats['motif_aatype'], dtype=torch.bool)
        chain_feats['motif_rigids_0'] = chain_feats['rigids_0']
        chain_feats['motif_torsion_angles_sin_cos'] = chain_feats['torsion_angles_sin_cos']
            
        # Convert all features to tensors.
        final_feats = tree.map_structure(
            lambda x: x if torch.is_tensor(x) else torch.tensor(x), chain_feats)
        # final_feats = du.pad_feats(final_feats, csv_row['modeled_seq_len'])
        
        final_feats = dict(filter(lambda item: item[0] in self.motif_keys, final_feats.items()))
            
        return final_feats, pdb_name

    def __getitem__(self, idx):
        motif_feats_list, motif_name_list = zip(*[self._get_single_motif(example_idx) for example_idx in self.comb_motif[idx]])
        max_motif_len = max([feats['motif_masks'].shape[-1] for feats in motif_feats_list])
        max_bb_len = max([feats['motif_seq_masks'].shape[0] for feats in motif_feats_list])
        pad_motif_feats_list = [du.pad_motif_feats(feats, max_motif_len, max_bb_len, use_torch=True) for feats in motif_feats_list] 
        final_feats = {key: torch.stack(list(map(lambda d: d[key], pad_motif_feats_list))) for key in self.motif_keys}
        final_feats['motif_seq_masks'] = final_feats['motif_seq_masks']
        return final_feats, '+'.join(motif_name_list)


class PdbDataset(data.Dataset):
    def __init__(
            self,
            *,
            data_conf,
            diffuser,
            is_training,
            is_save_virtual_motif=False,   
        ):
        self._log = logging.getLogger(__name__)
        self._is_training = is_training
        self.is_save_virtual_motif = is_save_virtual_motif
        self._data_conf = data_conf
        self._init_metadata()
        self._diffuser = diffuser

    @property
    def is_training(self):
        return self._is_training

    @property
    def diffuser(self):
        return self._diffuser

    @property
    def data_conf(self):
        return self._data_conf

    def _init_metadata(self):
        """Initialize metadata."""

        # Process CSV with different filtering criterions.
        filter_conf = self.data_conf.filtering
        pdb_csv = pd.read_csv(self.data_conf.csv_path)
        self.raw_csv = pdb_csv
        if filter_conf.allowed_oligomer is not None and len(filter_conf.allowed_oligomer) > 0:
            pdb_csv = pdb_csv[pdb_csv.oligomeric_detail.isin(
                filter_conf.allowed_oligomer)]
        if filter_conf.max_len is not None:
            pdb_csv = pdb_csv[pdb_csv.modeled_seq_len <= filter_conf.max_len]
        if filter_conf.min_len is not None:
            pdb_csv = pdb_csv[pdb_csv.modeled_seq_len >= filter_conf.min_len]
        if filter_conf.max_helix_percent is not None:
            pdb_csv = pdb_csv[
                pdb_csv.helix_percent < filter_conf.max_helix_percent]
        if filter_conf.max_loop_percent is not None:
            pdb_csv = pdb_csv[
                pdb_csv.coil_percent < filter_conf.max_loop_percent]
        if filter_conf.min_beta_percent is not None:
            pdb_csv = pdb_csv[
                pdb_csv.strand_percent > filter_conf.min_beta_percent]
        if filter_conf.rog_quantile is not None \
            and filter_conf.rog_quantile > 0.0:
            prot_rog_low_pass = _rog_quantile_curve(
                pdb_csv, 
                filter_conf.rog_quantile,
                np.arange(filter_conf.max_len))
            row_rog_cutoffs = pdb_csv.modeled_seq_len.map(
                lambda x: prot_rog_low_pass[x-1])
            pdb_csv = pdb_csv[pdb_csv.radius_gyration < row_rog_cutoffs]
        if filter_conf.subset is not None:
            pdb_csv = pdb_csv[:filter_conf.subset]
        pdb_csv = pdb_csv.sort_values('modeled_seq_len', ascending=False)
        self._create_split(pdb_csv)

    def _create_split(self, pdb_csv):
        # Training or validation specific logic.
        if self.is_training:
            self.csv = pdb_csv
            self._log.info(
                f'Training: {len(self.csv)} examples')
        else:
            all_lengths = np.sort(pdb_csv.modeled_seq_len.unique())
            length_indices = (len(all_lengths) - 1) * np.linspace(
                0.0, 1.0, self._data_conf.num_eval_lengths)
            length_indices = length_indices.astype(int)
            eval_lengths = all_lengths[length_indices]
            eval_csv = pdb_csv[pdb_csv.modeled_seq_len.isin(eval_lengths)]
            # Fix a random seed to get the same split each time.
            eval_csv = eval_csv.groupby('modeled_seq_len').sample(
                self._data_conf.samples_per_eval_length, replace=True, random_state=123)
            eval_csv = eval_csv.sort_values('modeled_seq_len', ascending=False)
            self.csv = eval_csv
            self._log.info(
                f'Validation: {len(self.csv)} examples with lengths {eval_lengths}')
    # cache make the same sample in same batch 
    @fn.lru_cache(maxsize=100)
    def _process_csv_row(self, processed_file_path):
        processed_feats = du.read_pkl(processed_file_path)
        processed_feats = du.parse_chain_feats(processed_feats)

        # Only take modeled residues.
        modeled_idx = processed_feats['modeled_idx']
        min_idx = np.min(modeled_idx)
        max_idx = np.max(modeled_idx)
        del processed_feats['modeled_idx']
        processed_feats = tree.map_structure(
            lambda x: x[min_idx:(max_idx+1)], processed_feats)

        # Run through OpenFold data transforms.
        chain_feats = {
            'aatype': torch.tensor(processed_feats['aatype']).long(),
            'all_atom_positions': torch.tensor(processed_feats['atom_positions']).double(),
            'all_atom_mask': torch.tensor(processed_feats['atom_mask']).double()
        }
        chain_feats = data_transforms.atom37_to_frames(chain_feats)
        chain_feats = data_transforms.make_atom14_masks(chain_feats)
        chain_feats = data_transforms.make_atom14_positions(chain_feats)
        chain_feats = data_transforms.atom37_to_torsion_angles()(chain_feats)

        # Re-number residue indices for each chain such that it starts from 1.
        # Randomize chain indices.
        chain_idx = processed_feats["chain_index"]
        res_idx = processed_feats['residue_index']
        new_res_idx = np.zeros_like(res_idx)
        new_chain_idx = np.zeros_like(res_idx)
        all_chain_idx = np.unique(chain_idx).tolist()
        shuffled_chain_idx = np.array(
            random.sample(all_chain_idx, len(all_chain_idx))) - np.min(all_chain_idx) + 1
        for i,chain_id in enumerate(all_chain_idx):
            chain_mask = (chain_idx == chain_id).astype(int)
            chain_min_idx = np.min(res_idx + (1 - chain_mask) * 1e3).astype(int)
            new_res_idx = new_res_idx + (res_idx - chain_min_idx + 1) * chain_mask

            # Shuffle chain_index
            replacement_chain_id = shuffled_chain_idx[i]
            new_chain_idx = new_chain_idx + replacement_chain_id * chain_mask

        # To speed up processing, only take necessary features
        final_feats = {
            'aatype': chain_feats['aatype'],
            'seq_idx': new_res_idx,
            'chain_idx': new_chain_idx,
            'residx_atom14_to_atom37': chain_feats['residx_atom14_to_atom37'],
            'residue_index': processed_feats['residue_index'],
            'res_mask': processed_feats['bb_mask'],
            'atom37_pos': chain_feats['all_atom_positions'],
            'atom37_mask': chain_feats['all_atom_mask'],
            'atom14_pos': chain_feats['atom14_gt_positions'],
            'rigidgroups_0': chain_feats['rigidgroups_gt_frames'],
            'torsion_angles_sin_cos': chain_feats['torsion_angles_sin_cos'],
        }
        return final_feats

    def _create_diffused_masks(self, atom37_pos, rng, row):
        bb_pos = atom37_pos[:, residue_constants.atom_order['CA']]
        dist2d = np.linalg.norm(bb_pos[:, None, :] - bb_pos[None, :, :], axis=-1)

        # Randomly select residue then sample a distance cutoff
        # TODO: Use a more robust diffuse mask sampling method.
        diff_mask = np.zeros_like(bb_pos)
        attempts = 0
        while np.sum(diff_mask) < 1:
            crop_seed = rng.integers(dist2d.shape[0])
            seed_dists = dist2d[crop_seed]
            max_scaffold_size = min(
                self._data_conf.scaffold_size_max,
                seed_dists.shape[0] - self._data_conf.motif_size_min
            )
            scaffold_size = rng.integers(
                low=self._data_conf.scaffold_size_min,
                high=max_scaffold_size
            )
            dist_cutoff = np.sort(seed_dists)[scaffold_size]
            diff_mask = (seed_dists < dist_cutoff).astype(float)
            attempts += 1
            if attempts > 100:
                raise ValueError(
                    f'Unable to generate diffusion mask for {row}')
        return diff_mask

    def _generate_motif_masks(self, seq_len, min_len=50, max_len=80, num_motifs=2, type='sequential'):
        # sequence maybe incomplete, seq_idx maybe discontinuous
        # seq_idx = torch.arange(0, seq_len)
        if type == 'sequential':
            if min_len > seq_len:
                raise ValueError("Not enough space to extract the specified number of motifs with the minimum length.")
            all_start_positions = torch.arange(seq_len - min_len + 1)

            # # random select in possible start positions
            # selected_start_positions = all_start_positions[torch.randperm(len(all_start_positions))[:num_motifs]]
            # selected_start_positions, _ = torch.sort(selected_start_positions)

            # seg num_motifs fragments
            segment_length = len(all_start_positions) // num_motifs
            selected_start_positions = torch.randint(0, segment_length-1, (num_motifs,)) + torch.cumsum(torch.full((num_motifs,),segment_length), dim=0) - segment_length

            masks = []
            for start_pos in selected_start_positions:
                # make sure the length of motif will not exceed the seq_len
                max_possible_len = min(max_len, seq_len - start_pos)
                # random motif length
                length = torch.randint(min_len, max_possible_len + 1, (1,)).item()
                # sub_seq = seq_idx[start_pos:start_pos + length]
                # Set the corresponding positions in the mask to one
                mask = torch.zeros((seq_len,))
                mask[start_pos:start_pos + length] = 1
                masks.append(mask)
            masks = torch.stack(masks, dim=1)
            return masks
        elif type == 'spatial':
            raise NotImplementedError('spatial motif generate type')
        else:
            raise NotImplementedError('ilegal motif generate type')


    def _generate_motif(self, rigid_0, seq_len, aatype, torsion_angles_sin_cos, min_len=20, max_len=80, num_motifs=2, type='sequential'):
        # sequence maybe incomplete, seq_idx maybe discontinuous
        # seq_idx = torch.arange(0, seq_len)
        if type == 'sequential':
            if min_len > seq_len:
                raise ValueError("Not enough space to extract the specified number of motifs with the minimum length.")
            #all_start_positions = np.arange(seq_len - min_len + 1)

            # # random select in possible start positions
            # selected_start_positions = all_start_positions[torch.randperm(len(all_start_positions))[:num_motifs]]
            # selected_start_positions, _ = torch.sort(selected_start_positions)

            # seg num_motifs fragments
            #segment_length = len(all_start_positions) // num_motifs
            #selected_start_positions = np.random.randint(0, segment_length-1, (num_motifs,)) + np.cumsum(np.full((num_motifs,),segment_length), dim=0) - segment_length
            #selected_start_positions = np.random.randint(0, segment_length-1, (num_motifs,)) + np.cumsum(np.full((num_motifs,),segment_length)) - segment_length
            seq_for_motif_len = seq_len // num_motifs
            position_len = seq_for_motif_len - min_len 
            if position_len <= 0:
                position_len = 1
            selected_start_positions = []
            for i in range(num_motifs):
                selected_start_positions.append(np.random.randint(0, position_len)+i*seq_for_motif_len)
            seq_masks = []
            motif_masks = []
            motif_aatype = []
            motif_rigids_0 = []
            for i, start_pos in enumerate(selected_start_positions):
                # make sure the length of motif will not exceed the seq_len
                if i < (len(selected_start_positions)-1):
                    possible_len = selected_start_positions[i+1] - start_pos
                else:
                    possible_len = seq_len - start_pos
                max_possible_len = min(max_len, possible_len)
                # random motif length
                #print(min_len, max_possible_len)
                length = np.random.randint(min_len, max_possible_len + 1, (1,))
                # sub_seq = seq_idx[start_pos:start_pos + length]
                # Set the corresponding positions in the mask to one
                mask = np.zeros((seq_len,), dtype=np.bool_)

                mask[start_pos:start_pos + length[0]] = True
                seq_masks.append(mask)
                motif_rigids_0.append(rigid_0[mask])
                motif_aatype.append(aatype[mask])

            # padding & merge, pad motif_aatype and motif_rigids_0 to the same length
            max_motif_len = max_len
            for idx in range(len(motif_aatype)):
                if max_motif_len > len(motif_aatype[idx]):
                    motif_mask = np.concatenate((np.ones((len(motif_aatype[idx]),), dtype=np.bool_), np.zeros((max_motif_len-len(motif_aatype[idx]),), dtype=np.bool_)))
                else:
                    motif_mask = np.zeros((max_motif_len,), dtype=np.bool_)
                motif_masks.append(motif_mask)

                motif_aatype[idx] = du.pad(motif_aatype[idx], max_motif_len, use_torch=False)
                motif_rigids_0[idx] = du.pad_rigid(motif_rigids_0[idx], max_motif_len)
            return np.stack(motif_aatype), np.stack(motif_masks), np.stack(seq_masks), np.stack(motif_rigids_0)
        elif type == 'spatial':
            raise NotImplementedError('spatial motif generate type')
        else:
            raise NotImplementedError('ilegal motif generate type')

    def __len__(self):
        return len(self.csv)

    def ind_reorder(self, seq_idx, motif_seq_mask):
        seq_mask = np.ones(motif_seq_mask.shape[-1], dtype=np.bool_)
        for i in range(motif_seq_mask.shape[0]):
            seq_mask *= (~motif_seq_mask[0])
        middle = int(seq_mask.sum()//2)
        idx = 1
        for i in range(seq_idx.shape[0]):
            if seq_mask[i]:
                if idx == middle:
                    idx += 1
                seq_idx[i] = idx
                idx += 1
            else:
                seq_idx[i] = middle
        return seq_idx

    def __getitem__(self, idx):

        # Sample data example.
        example_idx = idx
        csv_row = self.csv.iloc[example_idx]
        if 'pdb_name' in csv_row:
            pdb_name = csv_row['pdb_name']
        elif 'chain_name' in csv_row:
            pdb_name = csv_row['chain_name']
        else:
            raise ValueError('Need chain identifier.')
        processed_file_path = csv_row['processed_path']
        chain_feats = self._process_csv_row(processed_file_path)

        # Use a fixed seed for evaluation.
        if self.is_training:
            rng = np.random.default_rng(None)
        else:
            rng = np.random.default_rng(idx)
        gt_bb_rigid = rigid_utils.Rigid.from_tensor_4x4(
            chain_feats['rigidgroups_0'])[:, 0]
        diffused_mask = np.ones_like(chain_feats['res_mask'])
        if np.sum(diffused_mask) < 1:
            raise ValueError('Must be diffused')
        fixed_mask = 1 - diffused_mask
        chain_feats['fixed_mask'] = fixed_mask
        chain_feats['rigids_0'] = gt_bb_rigid.to_tensor_7()
        chain_feats['sc_ca_t'] = torch.zeros_like(gt_bb_rigid.get_trans())
        chain_feats['one_hot'] = torch.nn.functional.one_hot(chain_feats['aatype'], 22)
        if 'motif' in self._data_conf:
            motif_aatype, motif_masks, seq_masks, motif_rigids_0 = self._generate_motif(chain_feats['rigids_0'], csv_row['modeled_seq_len'], chain_feats['aatype'], chain_feats['torsion_angles_sin_cos'], self._data_conf.motif.min_len, self._data_conf.motif.max_len, self._data_conf.motif.num_motifs, self._data_conf.motif.type)
            chain_feats['motif_masks'] = motif_masks
            chain_feats['motif_seq_masks'] = seq_masks
            chain_feats['motif_seq_masks_for_loss'] = seq_masks
            #print("#############1212121", chain_feats['motif_seq_masks'].shape, chain_feats['rigids_0'].shape)
            chain_feats['motif_aatype'] = motif_aatype
            chain_feats['motif_rigids_0'] = motif_rigids_0
        # Sample t and diffuse.
        if self.is_training:
            t = rng.uniform(self._data_conf.min_t, 1.0)
            diff_feats_t = self._diffuser.forward_marginal(
                rigids_0=gt_bb_rigid,
                representation = chain_feats['one_hot'],
                t=t,
                diffuse_mask=None,
                motif_seq_mask=seq_masks
            )
        else:
            t = 1.0
            diff_feats_t = self.diffuser.sample_ref(
                n_samples=gt_bb_rigid.shape[0],
                impute=gt_bb_rigid,
                diffuse_mask=None,
                as_tensor_7=True,
                motif_dict=chain_feats,
                motif_mask = chain_feats['motif_seq_masks']
            )
        
        chain_feats.update(diff_feats_t)
        chain_feats['t'] = t
        #chain_feats['seq_idx'] = self.ind_reorder(chain_feats['seq_idx'], chain_feats['motif_seq_masks'])
        chain_feats['motif_position'] = np.zeros_like(chain_feats['motif_seq_masks'].shape[-1], dtype=np.bool_)
        for i in range(chain_feats['motif_seq_masks'].shape[0]):
            chain_feats['motif_position'] = (chain_feats['motif_position'] | chain_feats['motif_seq_masks'][i,:])
            # motif_masks = self._generate_motif_masks(csv_row['modeled_seq_len'], self._data_conf.motif.min_len, self._data_conf.motif.max_len, self._data_conf.motif.num_motifs, self._data_conf.motif.type)
            # chain_feats['motif_masks'] = motif_masks 
            # chain_feats['motif_aatype'] = chain_feats['aatype'][..., None] * motif_masks
            # chain_feats['motif_rigids_0'] = chain_feats['rigids_0'][..., None, :] * motif_masks[..., None] 
            # chain_feats['motif_aatype'] = chain_feats['aatype'][..., None].repeat(1,self._data_conf.motif.num_motifs)
            # chain_feats['motif_rigids_0'] = chain_feats['rigids_0'][..., None, :].repeat(1,self._data_conf.motif.num_motifs,1)
        
        # Convert all features to tensors.
        final_feats = tree.map_structure(
            lambda x: x if torch.is_tensor(x) else torch.tensor(x), chain_feats)
        for key in diff_feats_t.keys():
            final_feats[key] =final_feats[key].to(dtype=final_feats['rigids_t'].dtype)
        final_feats = du.pad_feats(final_feats, csv_row['modeled_seq_len'])
        if self.is_training and not self.is_save_virtual_motif:
            return final_feats
        else:
            return final_feats, pdb_name


class TrainSampler(data.Sampler):

    def __init__(
            self,
            *,
            data_conf,
            dataset,
            batch_size,
            sample_mode,
        ):
        self._log = logging.getLogger(__name__)
        self._data_conf = data_conf
        self._dataset = dataset
        self._data_csv = self._dataset.csv
        self._dataset_indices = list(range(len(self._data_csv)))
        self._data_csv['index'] = self._dataset_indices
        self._batch_size = batch_size
        self.epoch = 0
        self._sample_mode = sample_mode
        self.sampler_len = len(self._dataset_indices) * self._batch_size

        if self._sample_mode in ['cluster_length_batch', 'cluster_time_batch']:
            self._pdb_to_cluster = self._read_clusters()
            self._max_cluster = max(self._pdb_to_cluster.values())
            self._log.info(f'Read {self._max_cluster} clusters.')
            self._missing_pdbs = 0
            def cluster_lookup(pdb):
                pdb = pdb.upper()
                if pdb not in self._pdb_to_cluster:
                    self._pdb_to_cluster[pdb] = self._max_cluster + 1
                    self._max_cluster += 1
                    self._missing_pdbs += 1
                return self._pdb_to_cluster[pdb]
            self._data_csv['cluster'] = self._data_csv['pdb_name'].map(cluster_lookup)
            num_clusters = len(set(self._data_csv['cluster']))
            self.sampler_len = num_clusters * self._batch_size
            self._log.info(
                f'Training on {num_clusters} clusters. PDBs without clusters: {self._missing_pdbs}'
            )

    def _read_clusters(self):
        pdb_to_cluster = {}
        with open(self._data_conf.cluster_path, "r") as f:
            for i,line in enumerate(f):
                for chain in line.split(' '):
                    pdb = chain.split('_')[0]
                    pdb_to_cluster[pdb.upper()] = i
        return pdb_to_cluster

    def __iter__(self):
        if self._sample_mode == 'length_batch':
            # Each batch contains multiple proteins of the same length.
            sampled_order = self._data_csv.groupby('modeled_seq_len').sample(
                self._batch_size, replace=True, random_state=self.epoch)
            return iter(sampled_order['index'].tolist())
        elif self._sample_mode == 'time_batch':
            # Each batch contains multiple time steps of the same protein.
            random.shuffle(self._dataset_indices)
            repeated_indices = np.repeat(self._dataset_indices, self._batch_size)
            return iter(repeated_indices)
        elif self._sample_mode == 'cluster_length_batch':
            # Each batch contains multiple clusters of the same length.
            sampled_clusters = self._data_csv.groupby('cluster').sample(
                1, random_state=self.epoch)
            sampled_order = sampled_clusters.groupby('modeled_seq_len').sample(
                self._batch_size, replace=True, random_state=self.epoch)
            return iter(sampled_order['index'].tolist())
        elif self._sample_mode == 'cluster_time_batch':
            # Each batch contains multiple time steps of a protein from a cluster.
            sampled_clusters = self._data_csv.groupby('cluster').sample(
                1, random_state=self.epoch)
            dataset_indices = sampled_clusters['index'].tolist()
            repeated_indices = np.repeat(dataset_indices, self._batch_size)
            return iter(repeated_indices.tolist())
        else:
            raise ValueError(f'Invalid sample mode: {self._sample_mode}')

    def set_epoch(self, epoch):
        self.epoch = epoch

    def __len__(self):
        return self.sampler_len


# modified from torch.utils.data.distributed.DistributedSampler
# key points: shuffle of each __iter__ is determined by epoch num to ensure the same shuffle result for each proccessor
class DistributedTrainSampler(data.Sampler):
    r"""Sampler that restricts data loading to a subset of the dataset.

    modified from torch.utils.data.distributed import DistributedSampler

    .. note::
        Dataset is assumed to be of constant size and that any instance of it always
        returns the same elements in the same order.

    Args:
        dataset: Dataset used for sampling.
        num_replicas (int, optional): Number of processes participating in
            distributed training. By default, :attr:`world_size` is retrieved from the
            current distributed group.
        rank (int, optional): Rank of the current process within :attr:`num_replicas`.
            By default, :attr:`rank` is retrieved from the current distributed
            group.
        shuffle (bool, optional): If ``True`` (default), sampler will shuffle the
            indices.
        seed (int, optional): random seed used to shuffle the sampler if
            :attr:`shuffle=True`. This number should be identical across all
            processes in the distributed group. Default: ``0``.
        drop_last (bool, optional): if ``True``, then the sampler will drop the
            tail of the data to make it evenly divisible across the number of
            replicas. If ``False``, the sampler will add extra indices to make
            the data evenly divisible across the replicas. Default: ``False``.

    .. warning::
        In distributed mode, calling the :meth:`set_epoch` method at
        the beginning of each epoch **before** creating the :class:`DataLoader` iterator
        is necessary to make shuffling work properly across multiple epochs. Otherwise,
        the same ordering will be always used.

    Example::

        >>> # xdoctest: +SKIP
        >>> sampler = DistributedSampler(dataset) if is_distributed else None
        >>> loader = DataLoader(dataset, shuffle=(sampler is None),
        ...                     sampler=sampler)
        >>> for epoch in range(start_epoch, n_epochs):
        ...     if is_distributed:
        ...         sampler.set_epoch(epoch)
        ...     train(loader)
    """

    def __init__(self, 
                *,
                data_conf,
                dataset,
                batch_size,
                num_replicas: Optional[int] = None,
                rank: Optional[int] = None, shuffle: bool = True,
                seed: int = 0, drop_last: bool = False) -> None:
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()
        if rank >= num_replicas or rank < 0:
            raise ValueError(
                "Invalid rank {}, rank should be in the interval"
                " [0, {}]".format(rank, num_replicas - 1))
        self._data_conf = data_conf
        self._dataset = dataset
        self._data_csv = self._dataset.csv
        self._dataset_indices = list(range(len(self._data_csv)))
        self._data_csv['index'] = self._dataset_indices
        # _repeated_size is the size of the dataset multiply by batch size
        self._repeated_size = batch_size * len(self._data_csv)
        self._batch_size = batch_size
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.drop_last = drop_last
        # If the dataset length is evenly divisible by # of replicas, then there
        # is no need to drop any data, since the dataset will be split equally.
        if self.drop_last and self._repeated_size % self.num_replicas != 0:  # type: ignore[arg-type]
            # Split to nearest available length that is evenly divisible.
            # This is to ensure each rank receives the same amount of data when
            # using this Sampler.
            self.num_samples = math.ceil(
                (self._repeated_size - self.num_replicas) / self.num_replicas  # type: ignore[arg-type]
            )
        else:
            self.num_samples = math.ceil(self._repeated_size / self.num_replicas)  # type: ignore[arg-type]
        self.total_size = self.num_samples * self.num_replicas
        self.shuffle = shuffle
        self.seed = seed

    def __iter__(self) :
        if self.shuffle:
            # deterministically shuffle based on epoch and seed
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            indices = torch.randperm(len(self._data_csv), generator=g).tolist()  # type: ignore[arg-type]
        else:
            indices = list(range(len(self._data_csv)))  # type: ignore[arg-type]

        # indices is expanded by self._batch_size times
        indices = np.repeat(indices, self._batch_size)
        
        if not self.drop_last:
            # add extra samples to make it evenly divisible
            padding_size = self.total_size - len(indices)
            if padding_size <= len(indices):
                indices = np.concatenate((indices, indices[:padding_size]), axis=0)
            else:
                indices = np.concatenate((indices, np.repeat(indices, math.ceil(padding_size / len(indices)))[:padding_size]), axis=0)

        else:
            # remove tail of data to make it evenly divisible.
            indices = indices[:self.total_size]
        assert len(indices) == self.total_size

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        # 
        assert len(indices) == self.num_samples

        return iter(indices)

    def __len__(self) -> int:
        return self.num_samples

    def set_epoch(self, epoch: int) -> None:
        r"""
        Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
        use a different random ordering for each epoch. Otherwise, the next iteration of this
        sampler will yield the same ordering.

        Args:
            epoch (int): Epoch number.
        """
        self.epoch = epoch
