"""
Contains source code for loading in data and creating requisite PyTorch
data loader object
"""

import json
import pickle
import hashlib
import functools
import multiprocessing as mp
import os
import glob
import logging
from pathlib import Path
from typing import *

from matplotlib import pyplot as plt
import numpy as np

import torch
from torch import nn
from torch.utils.data import Dataset

LOCAL_DATA_DIR = Path(
    os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
)

CATH_DIR = LOCAL_DATA_DIR / "cath"
ALPHAFOLD_DIR = LOCAL_DATA_DIR / "alphafold"
SWISS_PROT_PDB = LOCAL_DATA_DIR / "swissprot_pdb_v6"
HOMO_DIR = LOCAL_DATA_DIR / "struct_token_bench/homo"
BINDINT_DIR = LOCAL_DATA_DIR / "struct_token_bench/interpro/binding"
BINDBIO_DIR = LOCAL_DATA_DIR / "struct_token_bench/biolip2/binding"
BINDSHAKE_DIR = LOCAL_DATA_DIR / "struct_token_bench/proteinshake"
REPEAT_DIR = LOCAL_DATA_DIR / "struct_token_bench/interpro/repeat"
EPT_DIR = LOCAL_DATA_DIR / "struct_token_bench/proteinglue"
ATLAS_DIR = LOCAL_DATA_DIR / "struct_token_bench/atlas"
FOLDSWITCHING_DIR = LOCAL_DATA_DIR / "struct_token_bench/foldswitching"
APOHOLO_DIR = LOCAL_DATA_DIR / "struct_token_bench/apoholo"
CATINT_DIR = LOCAL_DATA_DIR / "struct_token_bench/interpro/activesite"
CATBIO_DIR = LOCAL_DATA_DIR / "struct_token_bench/biolip2/catalytic"
CONSERVED_DIR = LOCAL_DATA_DIR / "struct_token_bench/interpro/conservedsite"
PRETRAIN_DIR = LOCAL_DATA_DIR / "vqvae_pretrain/train"
PRETRAIN_VALID_DIR = LOCAL_DATA_DIR / "vqvae_pretrain/validation"
CAMEO_DIR = LOCAL_DATA_DIR / "vqvae_pretrain/CAMEO"
CASP_DIR = LOCAL_DATA_DIR / "vqvae_pretrain/CASP14"

from geobpe.angles_and_coords import (
    canonical_distances_and_dihedrals,
    EXHAUSTIVE_ANGLES,
    EXHAUSTIVE_DISTS,
    extract_backbone_coords,
    extract_side_chain_coords,
    extract_backbone_residue_idxes,
    extract_aa_seq,
    extract_c_beta_coords
)
from geobpe.annotations import find_secondary_structures
from geobpe import utils

TRIM_STRATEGIES = Literal["leftalign", "randomcrop", "discard"]

FEATURE_SET_NAMES_TO_ANGULARITY = {
    "canonical": [False, False, False, True, True, True, True, True, True],
    "canonical-full-angles": [True, True, True, True, True, True],
    "canonical-minimal-angles": [True, True, True, True],
    "cart-coords": [False, False, False],
}
FEATURE_SET_NAMES_TO_FEATURE_NAMES = {
    "canonical": [
        "0C:1N",
        "N:CA",
        "CA:C",
        "phi",
        "psi",
        "omega",
        "tau",
        "CA:C:1N",
        "C:1N:1CA",
    ],
    "canonical-full-angles": [
        "phi",
        "psi",
        "omega",
        "tau",
        "CA:C:1N",
        "C:1N:1CA",
    ],
    "canonical-minimal-angles": ["phi", "psi", "omega", "tau"],
    "cart-coords": ["x", "y", "z"],
}

def extract_pdb_code_and_chain(dataset_id):
    """
    Given an id like b'd1v4wb_', extract the PDB code and chain.
    Assumes the format is: b'd' + 4-character PDB code + chain + optional trailing characters.
    For example, 'd1v4wb_' -> PDB code '1V4W' (uppercase) and chain 'B'.
    """
    if isinstance(dataset_id, bytes):
        dataset_id = dataset_id.decode('utf-8')
    pdb_code = dataset_id[1:5].upper()  # characters 1-4 form the PDB code
    chain = dataset_id[5].upper()       # the 5th character is the chain identifier
    return pdb_code, chain


def featurize_one(
    fname,
    pfunc,
    coords_pfunc,
    full_coord_pfunc,
    full_atom_idx_map,
    side_chain_coords_pfunc,
    aa_seq_func,
    c_beta_func,
    secondary_pfunc=None,
):
    try:
        s = pfunc(fname)
        if s is None:
            logging.warning(f"Angles is None: {fname}")
            return None
        c = coords_pfunc(fname)
        if c is None:
            logging.warning(f"Coords is None: {fname}")
            return None
        c_full = full_coord_pfunc(fname)
        if c_full is None:
            logging.warning(f"Full coords is None: {fname}")
            return None
        idxes = full_atom_idx_map(fname)
        if idxes is None:
            logging.warning(f"Full idxes is None: {fname}")
            return None
        sc = side_chain_coords_pfunc(fname)
        aa = aa_seq_func(fname)
        if aa is None or len(aa) != len(s):
            logging.warning(f"AA sequence error (length mismatch or None): {fname}")
            return None
        c_beta = c_beta_func(fname)
        if c_beta is None or len(c_beta) != len(aa):
            logging.warning(f"C_beta error (length mismatch or None): {fname}")
            return None
        d = {
            "angles": s,
            "coords": c,
            "c_beta": c_beta,
            "full_coords": c_full,
            "full_idxes": idxes,
            "side_chain": sc,
            "aa": aa,
            "fname": fname,
        }
        if secondary_pfunc:
            try:
                sec = secondary_pfunc(fname)
                d["sec"] = sec
            except Exception as e:
                logging.warning(f"Secondary structure failed for {fname}: {e}")
        return d
    except Exception as e:
        logging.warning(f"Featurization failed for {fname}: {e}", exc_info=True)
        return None

    

class CathCanonicalAnglesDataset(Dataset):
    """
    Load in the dataset.

    All angles should be given between [-pi, pi]
    """

    feature_names = {
        "angles": [
            "0C:1N",
            "N:CA",
            "CA:C",
            "phi",
            "psi",
            "omega",
            "tau",
            "CA:C:1N",
            "C:1N:1CA",
        ],
        "coords": ["x", "y", "z"],
    }
    feature_is_angular = {
        "angles": [False, False, False, True, True, True, True, True, True],
        "coords": [False, False, False],
    }

    def __init__(
        self,
        pdbs: Union[
            Literal["cath", "alphafold", "homo", "ec"], str
        ] = "cath",  # Keyword or a directory
        split: Optional[Literal["train", "test", "validation"]] = None,
        pad: int = 512,
        min_length: int = 40,  # Set to 0 to disable
        trim_strategy: TRIM_STRATEGIES = "leftalign",
        toy: int = 0,
        zero_center: bool = False,  # Center the features to have 0 mean
        use_cache: bool = True,  # Use/build cached computations of dihedrals and angles
        cache_dir: Path = Path(os.path.dirname(os.path.abspath(__file__))),
        debug: bool = False
    ) -> None:
        super().__init__()
        assert pad > min_length
        self.trim_strategy = trim_strategy
        self.pad = pad
        self.min_length = min_length
        self.debug = debug
        
        # gather files
        self.pdbs_src = pdbs
        fnames = self.__get_pdb_fnames(pdbs)       
        self.fnames = fnames        

        # self.structures should be a list of dicts with keys (angles, coords, fname)
        # Define as None by default; allow for easy checking later
        self.structures = None
        codebase_hash = utils.md5_all_py_files(
            os.path.dirname(os.path.abspath(__file__))
        )
        # Default to false; assuming no cache, also doesn't match
        codebase_matches_hash = False
        self.use_cache = use_cache
        self.cache_dir = cache_dir
    # Always compute for toy; do not save
        if toy:
            if isinstance(toy, bool):
                toy = 150
            if isinstance(toy, str): # read from the file
                assert os.path.exists(toy)
                fnames = open(toy).readlines()
            else:
                fnames = fnames[:toy]
            logging.info(f"Loading toy dataset of {toy} structures")
            self.structures = self.__compute_featurization(fnames)
        elif use_cache and os.path.exists(self.cache_fname):
            logging.info(f"Loading cached full dataset from {self.cache_fname}")
            with open(self.cache_fname, "rb") as source:
                loaded_hash, loaded_structures = pickle.load(source)
                codebase_matches_hash = loaded_hash == codebase_hash
                if not codebase_matches_hash:
                    logging.warning(
                        "Mismatched hashes between codebase and cached values; updating cached values"
                    )
                else:
                    self.structures = loaded_structures
                    logging.info("Hash matches between codebase and cached values!")
        # We have not yet populated self.structures
        if self.structures is None:
            self.__clean_mismatched_caches()
            self.structures = self.__compute_featurization(fnames)
            if use_cache and not codebase_matches_hash:
                logging.info(f"Saving full dataset to cache at {self.cache_fname}")
                with open(self.cache_fname, "wb") as sink:
                    pickle.dump((codebase_hash, self.structures), sink)

        # If specified, remove sequences shorter than min_length
        if self.min_length:
            orig_len = len(self.structures)
            self.structures = [
                s for s in self.structures if s["angles"].shape[0] >= self.min_length
            ]
            len_delta = orig_len - len(self.structures)
            logging.info(
                f"Removing structures shorter than {self.min_length} residues excludes {len_delta}/{orig_len} --> {len(self.structures)} sequences"
            )
        if self.trim_strategy == "discard":
            orig_len = len(self.structures)
            self.structures = [
                s for s in self.structures if s["angles"].shape[0] <= self.pad
            ]
            len_delta = orig_len - len(self.structures)
            logging.info(
                f"Removing structures longer than {self.pad} produces {orig_len} - {len_delta} = {len(self.structures)} sequences"
            )

        # Split the dataset if requested. This is implemented here to maintain
        # functional parity with the original CATH dataset. Original CATH uses
        # a 80/10/10 split
        self.rng = np.random.default_rng(seed=6489)
        # Shuffle the sequences so contiguous splits acts like random splits
        self.rng.shuffle(self.structures)
        if split is not None:
            split_idx = int(len(self.structures) * 0.8)
            if split == "train":
                self.structures = self.structures[:split_idx]
            elif split == "validation":
                self.structures = self.structures[
                    split_idx : split_idx + int(len(self.structures) * 0.1)
                ]
            elif split == "test":
                self.structures = self.structures[
                    split_idx + int(len(self.structures) * 0.1) :
                ]
            else:
                raise ValueError(f"Unknown split: {split}")

            logging.info(f"Split {split} contains {len(self.structures)} structures")

        # if given, zero center the features
        self.means = None

        # Aggregate lengths
        self.all_lengths = [s["angles"].shape[0] for s in self.structures]
        self._length_rng = np.random.default_rng(seed=6489)
        logging.info(
            f"Length of angles: {np.min(self.all_lengths)}-{np.max(self.all_lengths)}, mean {np.mean(self.all_lengths)}"
        )

        # for ft in self.feature_names["angles"]:
        #     idx = self.feature_names["angles"].index(ft)
        #     is_angular = self.feature_is_angular["angles"][idx]
        #     logging.info(f"Feature {ft} is angular: {is_angular}")
        #     m, v = self.get_feature_mean_var(ft)
        #     logging.info(f"Feature {ft} mean, var: {m}, {v}")

    @staticmethod
    def __get_pdb_fnames(
        pdbs: Union[Literal["cath", "alphafold", "homo"], str, List[str], Tuple[str]]
    ) -> List[str]:
        """Return a list of filenames for PDB structures making up this dataset"""
        if isinstance(pdbs, (list, tuple)):
            # A list of PDBs
            for f in pdbs:
                assert os.path.isfile(f), f"Given file does not exist: {f}"
            fnames = pdbs
            logging.info(f"Given {len(fnames)} PDB files")
        elif Path(pdbs).is_dir():
            fnames = []
            for ext in [".pdb", ".pdb.gz"]:
                fnames.extend(glob.glob(os.path.join(pdbs, f"*{ext}")))
            assert fnames, f"No PDB files found in {pdbs}"
            logging.info(f"Found {len(fnames)} PDB files in {pdbs}")       
        else:  # Should be a keyword
            if pdbs == "all":
                ext = ".pdb"
                pat = os.path.join(LOCAL_DATA_DIR, "**/*.pdb")
                files = glob.glob(pat, recursive=True)
                c = Counter([Path(f).parent.relative_to(LOCAL_DATA_DIR) for f in files])
                logging.info(f"Counter {c}")
                num_uniq = len(set([Path(f).name for f in files]))                     
                logging.info(f"{num_uniq} unique")
                fnames = files
            elif pdbs == "test": # debug
                ext = ".pdb"
                pat = os.path.join(LOCAL_DATA_DIR, "**/*.pdb")
                files = glob.glob(pat, recursive=True)
                fnames = sorted(files)[:10]
            elif pdbs == "cath":
                fnames = glob.glob(os.path.join(CATH_DIR, "dompdb", "*"))
                assert fnames, f"No files found in {CATH_DIR}/dompdb"
            elif pdbs == "alphafold":
                fnames = glob.glob(os.path.join(ALPHAFOLD_DIR, "*.pdb.gz"))
                assert fnames, f"No files found in {ALPHAFOLD_DIR}"
            elif pdbs == "homo":
                fnames = glob.glob(os.path.join(HOMO_DIR, "*.pdb"))
            elif pdbs == "bindint":
                fnames = glob.glob(os.path.join(BINDINT_DIR, "*.pdb"))
            elif pdbs == "bindbio":
                fnames = glob.glob(os.path.join(BINDBIO_DIR, "*.pdb"))
            elif pdbs == "bindshake":
                fnames = glob.glob(os.path.join(BINDSHAKE_DIR, "*.pdb"))
            elif pdbs == "repeat":
                fnames = glob.glob(os.path.join(REPEAT_DIR, "*.pdb"))
            elif pdbs == "ept":
                fnames = glob.glob(os.path.join(EPT_DIR, "*.pdb"))
            elif pdbs == "catint":
                fnames = glob.glob(os.path.join(CATINT_DIR, "*.pdb"))
            elif pdbs == "catbio":
                fnames = glob.glob(os.path.join(CATBIO_DIR, "*.pdb"))
            elif pdbs == "conserved":
                fnames = glob.glob(os.path.join(CONSERVED_DIR, "*.pdb"))
            elif pdbs == "pretrain":
                fnames = glob.glob(os.path.join(PRETRAIN_DIR, "*.pdb"))
            elif pdbs == "prevalid":
                fnames = glob.glob(os.path.join(PRETRAIN_VALID_DIR, "*.pdb"))
            elif pdbs == "atlas":
                fnames = glob.glob(os.path.join(ATLAS_DIR, "*.pdb"))
            elif pdbs == "foldswitching":
                fnames = glob.glob(os.path.join(FOLDSWITCHING_DIR, "*.pdb"))
            elif pdbs == "apoholo":
                fnames = glob.glob(os.path.join(APOHOLO_DIR, "*.pdb"))                
            elif pdbs == "casp":
                fnames = glob.glob(os.path.join(CASP_DIR, "*.pdb"))
            elif pdbs == "cameo":
                fnames = glob.glob(os.path.join(CAMEO_DIR, "*.pdb"))
            elif pdbs == "swissprot":
                fnames = glob.glob(os.path.join(SWISS_PROT_PDB, "*.pdb.gz"))
            else:
                raise ValueError(f"Unknown pdb set: {pdbs}")
        # for debug
        # demo_pdb = './data/cath/dompdb/3w6sC00.pdb'
        demo_pdb = './data/remote_homology/test_superfamily_holdout_pdbs/2J7Q_A.pdb'
        # demo_pdb = './data/remote_homology/test_superfamily_holdout_pdbs/1QCR_D.pdb'
        if demo_pdb in fnames:
            fnames.remove(demo_pdb)
            fnames = [demo_pdb] + fnames
        return fnames

    @property
    def cache_fname(self) -> str:
        """Return the filename for the cache file"""
        if os.path.isdir(self.pdbs_src):
            k = os.path.basename(self.pdbs_src)
        else:
            k = self.pdbs_src

        # Create md5 of all the filenames (NOT their contents)
        hash_md5 = hashlib.md5()
        for fname in self.fnames:
            hash_md5.update(os.path.basename(fname).encode())
        filename_hash = hash_md5.hexdigest()

        return os.path.join(
            self.cache_dir, f"cache_canonical_structures_{k}_{filename_hash}.pkl"
        )

    def __clean_mismatched_caches(self) -> None:
        """Clean out mismatched cache files"""
        if not self.use_cache:
            logging.info("Not using cache -- skipping cache cleaning")
            return

        if os.path.isdir(self.pdbs_src):
            k = os.path.basename(self.pdbs_src)
        else:
            k = self.pdbs_src

        matches = glob.glob(
            os.path.join(self.cache_dir, f"cache_canonical_structures_{k}_*.pkl")
        )
        if not matches:
            logging.info(
                f"No cache files found matching {matches}, no cleaning necessary"
            )
        for fname in matches:
            if fname != self.cache_fname:
                logging.info(f"Removing old cache file {fname}")
                os.remove(fname)

    def __compute_featurization(
        self, fnames: Sequence[str]
    ) -> List[Dict[str, np.ndarray]]:
        """Get the featurization of the given fnames"""
        pfunc = functools.partial(
            canonical_distances_and_dihedrals,
            distances=EXHAUSTIVE_DISTS,
            angles=EXHAUSTIVE_ANGLES,
        )
        coords_pfunc = functools.partial(extract_backbone_coords, atoms=["CA"])
    
        logging.info(
            f"Computing full dataset of {len(fnames)} with {mp.cpu_count()} threads"
        )
        # Generate dihedral angles
        if not self.debug:
            pool = mp.Pool(processes=mp.cpu_count())
            struct_arrays = list(pool.map(pfunc, fnames, chunksize=250))
            coord_arrays = list(pool.map(coords_pfunc, fnames, chunksize=250))
            pool.close()
            pool.join()            
        else:
            struct_arrays = [pfunc(fname) for fname in fnames]        
            coord_arrays = [coords_pfunc(fname) for fname in fnames]
        
        # Contains only non-null structures
        structures = []
        for fname, s, c in zip(fnames, struct_arrays, coord_arrays):
            if s is None:
                continue 
            structures.append(
                {
                    "angles": s,
                    "coords": c,
                    "fname": fname,
                }
            )
        return structures

    def sample_length(self, n: int = 1) -> Union[int, List[int]]:
        """
        Sample a observed length of a sequence
        """
        assert n > 0
        if n == 1:
            l = self._length_rng.choice(self.all_lengths)
        else:
            l = self._length_rng.choice(self.all_lengths, size=n, replace=True).tolist()
        return l

    def get_masked_means(self) -> np.ndarray:
        """Return the means subset to the actual features used"""
        if self.means is None:
            return None
        return np.copy(self.means)

    @functools.cached_property
    def filenames(self) -> List[str]:
        """Return the filenames that constitute this dataset"""
        return [s["fname"] for s in self.structures]

    def __len__(self) -> int:
        return len(self.structures)

    def __getitem__(
        self, index, ignore_zero_center: bool = False
    ) -> Dict[str, torch.Tensor]:
        if not 0 <= index < len(self):
            raise IndexError("Index out of range")

        angles = self.structures[index]["angles"]
        # NOTE coords are NOT shifted or wrapped, has same length as angles
        coords = self.structures[index]["coords"]
        assert angles.shape[0] == coords.shape[0]

        # If given, offset the angles with mean
        if self.means is not None and not ignore_zero_center:
            assert (
                self.means.shape[0] == angles.shape[1]
            ), f"Mismatched shapes for mean offset: {self.means.shape} != {angles.shape}"
            angles = angles - self.means

            # The distance features all contain a single ":"
            colon_count = np.array([c.count(":") for c in angles.columns])
            # WARNING this uses a very hacky way to find the angles
            angular_idx = np.where(colon_count != 1)[0]
            angles.iloc[:, angular_idx] = utils.modulo_with_wrapped_range(
                angles.iloc[:, angular_idx], -np.pi, np.pi
            )

        # Subset angles to ones we are actaully using as features
        angles = angles.loc[
            :, CathCanonicalAnglesDataset.feature_names["angles"]
        ].values
        assert angles is not None
        assert angles.shape[1] == len(
            CathCanonicalAnglesDataset.feature_is_angular["angles"]
        ), f"Mismatched shapes for angles: {angles.shape[1]} != {len(CathCanonicalAnglesDataset.feature_is_angular['angles'])}"

        # Replace nan values with zero
        np.nan_to_num(angles, copy=False, nan=0)

        # Create attention mask. 0 indicates masked
        l = min(self.pad, angles.shape[0])
        attn_mask = torch.zeros(size=(self.pad,))
        attn_mask[:l] = 1.0

        # Additionally, mask out positions that are nan
        # is_nan = np.where(np.any(np.isnan(angles), axis=1))[0]
        # attn_mask[is_nan] = 0.0  # Mask out the nan positions

        # Perform padding/trimming
        if angles.shape[0] < self.pad:
            angles = np.pad(
                angles,
                ((0, self.pad - angles.shape[0]), (0, 0)),
                mode="constant",
                constant_values=0,
            )
            coords = np.pad(
                coords,
                ((0, self.pad - coords.shape[0]), (0, 0)),
                mode="constant",
                constant_values=0,
            )
        elif angles.shape[0] > self.pad:
            if self.trim_strategy == "leftalign":
                angles = angles[: self.pad]
                coords = coords[: self.pad]
            elif self.trim_strategy == "randomcrop":
                # Randomly crop the sequence to
                start_idx = self.rng.integers(0, angles.shape[0] - self.pad)
                end_idx = start_idx + self.pad
                assert end_idx < angles.shape[0]
                angles = angles[start_idx:end_idx]
                coords = coords[start_idx:end_idx]
                assert angles.shape[0] == coords.shape[0] == self.pad
            else:
                raise ValueError(f"Unknown trim strategy: {self.trim_strategy}")

        # Create position IDs
        position_ids = torch.arange(start=0, end=self.pad, step=1, dtype=torch.long)

        angular_idx = np.where(CathCanonicalAnglesDataset.feature_is_angular["angles"])[
            0
        ]
        assert utils.tolerant_comparison_check(
            angles[:, angular_idx], ">=", -np.pi
        ), f"Illegal value: {np.min(angles[:, angular_idx])}"
        assert utils.tolerant_comparison_check(
            angles[:, angular_idx], "<=", np.pi
        ), f"Illegal value: {np.max(angles[:, angular_idx])}"
        angles = torch.from_numpy(angles).float()
        coords = torch.from_numpy(coords).float()

        retval = {
            "angles": angles,
            "coords": coords,
            "attn_mask": attn_mask,
            "position_ids": position_ids,
            "lengths": torch.tensor(l, dtype=torch.int64),
        }
        return retval

    def get_feature_mean_var(self, ft_name: str) -> Tuple[float, float]:
        """
        Return the mean and variance associated with a given feature
        """
        assert ft_name in self.feature_names["angles"], f"Unknown feature {ft_name}"
        idx = self.feature_names["angles"].index(ft_name)
        logging.info(f"Computing metrics for {ft_name} - idx {idx}")

        all_vals = []
        for i in range(len(self)):
            item = self[i]
            attn_idx = torch.where(item["attn_mask"] == 1.0)[0]
            vals = item["angles"][attn_idx, idx]
            all_vals.append(vals)
        all_vals = torch.cat(all_vals)
        assert all_vals.ndim == 1
        return torch.var_mean(all_vals)[::-1]  # Default is (var, mean)


class FullCathCanonicalCoordsDataset(CathCanonicalAnglesDataset):
    """
    Building on the CATH dataset, return the XYZ coordaintes of each alpha carbon
    """

    feature_names = {"coords": list("xyz")}
    feature_is_angular = {"coords": [False, False, False]}
    custom_kwargs = ["secondary"]

    def __init__(self, *args, **kwargs) -> None:        
        for kwarg in self.custom_kwargs:
            if kwarg in kwargs: 
                print(kwarg, kwargs[kwarg])               
                setattr(self, kwarg, kwargs[kwarg])
                kwargs.pop(kwarg)
        super().__init__(*args, **kwargs)


    @staticmethod
    def single_compute_featurization(fname):
        featurizer = functools.partial(
            featurize_one,
            pfunc=functools.partial(
                canonical_distances_and_dihedrals, distances=EXHAUSTIVE_DISTS, angles=EXHAUSTIVE_ANGLES
            ),
            coords_pfunc=functools.partial(extract_backbone_coords, atoms=["CA"]),
            full_coord_pfunc=functools.partial(extract_backbone_coords, atoms=["N", "CA", "C"]),
            full_atom_idx_map=functools.partial(extract_backbone_residue_idxes, atoms=["N", "CA", "C"]),
            side_chain_coords_pfunc=extract_side_chain_coords,
            aa_seq_func=extract_aa_seq,
            c_beta_func=extract_c_beta_coords,
            secondary_pfunc=None
        )
        return featurizer(fname)

    # due to name mangling
    def _CathCanonicalAnglesDataset__compute_featurization(self, fnames):
        import functools
        import multiprocessing

        # Build partial with parameters
        featurizer = functools.partial(
            featurize_one,
            pfunc=functools.partial(
                canonical_distances_and_dihedrals, distances=EXHAUSTIVE_DISTS, angles=EXHAUSTIVE_ANGLES
            ),
            coords_pfunc=functools.partial(extract_backbone_coords, atoms=["CA"]),
            full_coord_pfunc=functools.partial(extract_backbone_coords, atoms=["N", "CA", "C"]),
            full_atom_idx_map=functools.partial(extract_backbone_residue_idxes, atoms=["N", "CA", "C"]),
            side_chain_coords_pfunc=extract_side_chain_coords,
            aa_seq_func=extract_aa_seq,
            c_beta_func=extract_c_beta_coords,
            secondary_pfunc=find_secondary_structures if (hasattr(self, "secondary") and self.secondary) else None
        )

        logging.info(f"Computing full dataset of {len(fnames)} with {multiprocessing.cpu_count()} threads")
        if not self.debug:
            with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
                structures = list(filter(None, pool.map(featurizer, fnames, chunksize=50)))
        else:
            structures = []
            for fname in fnames:
                out = featurizer(fname)
                if out is not None:
                    structures.append(out)
        return structures

    def __getitem__(
        self, index, ignore_zero_center: bool = True
    ) -> Dict[str, torch.Tensor]:
        return_dict = super().__getitem__(index, ignore_zero_center=ignore_zero_center)
        return_dict["full_coords"] = self.structures[index]['full_coords'] # ignore pad
        return return_dict
