"""
GradPose: PDB Superimposition using Gradient Descent
Authors: Daniel Rademaker, Kevin van Geemen
"""
import warnings
import glob
import time
import os
import math
from itertools import repeat
import numpy as np
import torch
import torch.multiprocessing as mp
import torch.nn.functional as F
from tqdm import tqdm, trange
from gradpose import util

# TODO: Consistent bar lengths (bar_format=)
# TODO: Argument to overwrite original files
# TODO: Check all of the docstrings, some arguments are be different!

class Rotator():
    """The Rotation module with all learnable parameters.
    """
    def __init__(self, xyz, presence_mask, center=None, quaternions=None,
        device='cpu', use_numpy=False):
        """Initializes the Rotator object.

        Args:
            xyz (torch.Tensor): Coordinates for each atom for each PDB.
            presence_mask (torch.Tensor): A mask to zero out missing atoms when optimizing.
            center (torch.Tensor, optional): Alternative center coordinates to use for each PDB.
                Shape=(No. of PDBs, 3) Defaults to None.
            quaternions (torch.Tensor, optional): Alternative quaternions to set for each PDB.
            device (str, optional): The device tensors will be stored on. Defaults to 'cpu'.
            use_numpy (bool, optional): Whether to use numpy instead of pytorch. Defaults to False.
        """

        self.device = device
        # Rotation is always over the center of the pdb structures
        self.center = self.get_center(xyz, presence_mask) \
            if center is None else center.to(self.device)
        self.xyz = xyz - self.center
        self.presence_mask = presence_mask
        self.nmb_samples  = self.xyz.shape[0]
        if use_numpy:
            self.quaternions = np.random.rand(self.nmb_samples, 4) \
                if quaternions is None else quaternions
        else:
            self.quaternions = torch.rand(self.nmb_samples, 4).to(self.device) \
                if quaternions is None else quaternions.to(self.device)
        self.quaternions.requires_grad = True

        # Normalize the scale of the proteins
        scale_factor = self.xyz[0].norm(dim=1).std()/100
        self.xyz = self.xyz.div(scale_factor)
        self.learning_rate = 100

        self.optimizer = torch.optim.SGD([self.quaternions], lr=self.learning_rate)
        self.use_numpy = use_numpy


    def get_center(self, xyz, presence_mask):
        """Get the centers of each PDB. Ignoring missing residues.

        Args:
            xyz (torch.Tensor): PDB coordinates.
            presence_mask (torch.Tensor): Loss mask for each PDB.

        Returns:
            torch.Tensor: Tensor with the center of each PDB.
        """
        return (xyz*presence_mask).sum(1).div(presence_mask.sum(1)).reshape(-1,1,3)

    def get_matrix(self):
        """Get the rotation matrix using the learned quaternions

        Returns:
            torch.Tensor: Rotation matrix.
        """
        # Normalize the q-values. The norm should equal to 1.
        q = self._normalize(self.quaternions)
        # init an empty matrix to fill
        if self.use_numpy:
            matrix = np.zeros((q.shape[0], 3, 3))
        else:
            matrix = torch.zeros(q.shape[0], 3, 3).to(self.device)
        # Fill the matrix
        matrix[:,0,0] = (2 * (q[:,0]**2     + q[:,1]**2)) - 1
        matrix[:,0,1] =  2 * (q[:,1]*q[:,2] - q[:,0]*q[:,3])
        matrix[:,0,2] =  2 * (q[:,1]*q[:,3] + q[:,0]*q[:,2])
        matrix[:,1,0] =  2 * (q[:,1]*q[:,2] + q[:,0]*q[:,3])
        matrix[:,1,1] = (2 * (q[:,0]**2     + q[:,2]**2)) - 1
        matrix[:,1,2] =  2 * (q[:,2]*q[:,3] - q[:,0]*q[:,1])
        matrix[:,2,0] =  2 * (q[:,1]*q[:,3] - q[:,0]*q[:,2])
        matrix[:,2,1] =  2 * (q[:,2]*q[:,3] + q[:,0]*q[:,1])
        matrix[:,2,2] = (2 * (q[:,0]**2     + q[:,3]**2)) - 1
        return matrix

    def get_rotated_xyz(self, pdb_index=None, xyz=None, rotation_matrix=None):
        """Returns the PDB coordinates after rotating by the Rotator's quaternions.

        Args:
            pdb_index (int, optional): If manually supplying xyz and rotation_matrix,
                the index to use. Defaults to None.
            xyz (torch.Tensor, optional): Override the coordinates matrix. Defaults to None.
            rotation_matrix (torch.Tensor, optional): Override the rotation matrix.
                Defaults to None.

        Returns:
            torch.Tensor: Rotated coordinates.
        """

        if rotation_matrix is None:
            # Retrieve rotation matrix
            rotation_matrix = self.get_matrix()

        # Apply the rotation matrix
        if pdb_index:
            if self.use_numpy:
                return np.matmul(xyz, rotation_matrix[pdb_index:pdb_index+1])
            return torch.matmul(xyz, rotation_matrix[pdb_index:pdb_index+1])
        if self.use_numpy:
            return np.matmul(self.xyz, rotation_matrix)
        return torch.matmul(self.xyz, rotation_matrix)

    def _normalize(self, quaternions):
        """Normalize quaternions.

        Args:
            quaternions (torch.Tensor): Matrix of quaternions.

        Returns:
            torch.Tensor: Normalized quaternions.
        """
        quaternions = quaternions / quaternions.norm(dim=1).reshape(-1,1)
        return quaternions

    def _loss(self):
        """Calculate the loss by comparing the rotated coordinates with the template's."""
        new_xyz = self.get_rotated_xyz()
        loss =  F.mse_loss(new_xyz, self.xyz[0], reduction='none') * self.presence_mask
        loss = loss.mean(2)
        loss = loss.mean(1)
        # self.history += [[i.item() for i in loss]]
        return loss.sum()

    def normalize_q(self):
        """Normalize the Rotator's quaternions."""
        self.quaternions.data = self._normalize(self.quaternions.data)

    # Take a gradient step
    def step(self):
        """Optimize the rotation matrix of the Rotator by one step.

        Returns:
            Float: The loss of the step.
        """
        loss = self._loss()
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()
        return loss.item()

class PDBdataset():
    """Class that processes PDBs.
    """
    def __init__(self, pdbs, template_pdb, residues, chain,
            cores=mp.cpu_count(), output='result', device='cpu', verbosity=1):
        """_summary_

        Args:
            pdbs (str): Folder where the PDB files are located.
            template_pdb (str): Path to the template PDB.
            residues (list[int]): List of residue IDs to be used for alignment.
            chain (str): Chain ID to be used for alignment.
            cores (int, optional): Amount of CPU cores to use for multiprocessing.
                Defaults to the amount of cores on the system.
            output (str, optional): Path to output folder. Defaults to 'result'.
            device (str, optional): Device to store the tensors on. Defaults to 'cpu'.
            verbosity (int, optional): Verbosity level. Defaults to 1.
        """
        self.device = device
        self.verbosity = verbosity
        self.chain = chain

        if not residues or not chain:
            with open(template_pdb, 'r', encoding='utf-8') as template_file:
                template_file_data = template_file.read().split('\n')

                if not chain:
                    if self.verbosity > 0:
                        print("No chain selected, using longest chain: ", end='')

                    chain_dict = {}
                    for line in template_file_data:
                        if not line.startswith('ATOM  '):
                            continue
                        chain_id = line[21]
                        if not chain_id in chain_dict:
                            chain_dict[chain_id] = 0
                        chain_dict[chain_id] += 1
                    self.chain = sorted([
                        (count, chain_id)
                        for chain_id, count in chain_dict.items()
                    ], reverse=True)[0][1]

                    if self.verbosity > 0:
                        print(f"{self.chain}")

                if not residues:
                    if self.verbosity > 0:
                        print("No residues selected, determining from template chain: ", end='')
                    residues = list(set([
                        int(line[22:26])
                        for line in template_file_data
                        if line.startswith('ATOM ') and line[21] == self.chain
                    ]))

                    if self.verbosity > 0:
                        print(f"{residues[0]}:{residues[-1]}")


        self.residues = [int(i) for i in residues]
        self.residues_dict = {i:self.residues.index(i) for i in self.residues}
        self.pdbs = [template_pdb] + pdbs
        self.cores = cores
        self.output = output

        # Load CA coordinates for each PDB
        if self.verbosity > 1:
            print(f"Number of 3D Models: {len(self.pdbs)}")
            print('Start loading data')
        t0 = time.perf_counter()
        xyz, presence_mask = self.load_ca_atom_positions()
        if self.verbosity > 1:
            print(f'Retrieved data, time: {time.perf_counter()-t0:.2f} seconds')

        # Instantiate a Rotator
        self.rotator = Rotator(xyz, presence_mask, device=self.device)

    def _extract_pdb_atoms(self, file_name):
        """Extract all atom xyz coordinates

        Args:
            file_name (str): Path to PDB file.

        Returns:
            list: XYZ coordinates for each atom.
        """
        with open(file_name, encoding='utf-8') as pdb_file:
            return [[float(line[30:38]),float(line[38:46]),float(line[46:54])] \
                    for line in pdb_file.read().split('\n') \
                            if line.startswith('ATOM  ') or line.startswith('HETATM')]

    def _extract_pdb_ca(self, file_name):
        """Extract coordianates of all CA in specified chains and residues.

        Args:
            file_name (str): Path to PDB file.

        Returns:
            numpy.ndarray: Array that contains the XYZ coordinates of the CA atom per residue.
        """
        with open(file_name, encoding='utf-8') as pdb_file:
            pxyz = np.array(
                [
                    [
                        self.residues_dict[int(line[22:26])],
                        float(line[30:38]),
                        float(line[38:46]),
                        float(line[46:54])
                    ]
                    for line in pdb_file.read().split('\n') \
                    if line.startswith('ATOM ') and line[13:15] == 'CA' \
                        and line[21] == self.chain \
                        and int(line[22:26]) in self.residues
                ]
            )
        xyz = np.zeros((1, len(self.residues), 3))
        xyz[0][pxyz[:, 0].astype(int)] = pxyz[:, 1:]
        return xyz

    def load_ca_atom_positions(self):
        """Parse the PDBs to obtain the data needed for calculating the rotation matrices.

        Returns:
            torch.Tensor: CA atom coordinates for each PDB.
            torch.Tensor: The loss mask to exclude missing residues.
        """
        with mp.Pool(self.cores) as pool:
            ca_xyz = torch.tensor(
                np.concatenate(list(tqdm(
                    pool.imap(
                        self._extract_pdb_ca,
                        self.pdbs,
                        max(2, round(len(self.pdbs) / (self.cores * 16)))
                    ),
                    total=len(self.pdbs), desc="Extracting backbones",
                    disable=self.verbosity!=1
                )))).float().to(self.device)
        presence_mask = (ca_xyz.abs().sum(2) != 0).unsqueeze(2).float().to(self.device)
        return ca_xyz, presence_mask

    def _rotate_single_pdb(self, pdb_index, rotation_matrix):
        """Rotate a PDB by the quaternions of the Rotator and save the resulting PDB to a file.

        Args:
            pdb_index (int): The index of the PDB to be rotated.
            rotation_matrix (torch.Tensor): The rotation matrix to use for the rotator.
        """
        # TODO: Update docstring
        with torch.no_grad():
            pdb_file = self.pdbs[pdb_index]
            xyz = torch.Tensor(self._extract_pdb_atoms(pdb_file)).reshape(1, -1, 3)
            xyz = xyz - self.rotator.center[pdb_index]
            rotated_xyz = self.rotator.get_rotated_xyz(pdb_index, xyz, rotation_matrix).squeeze() \
                + self.rotator.center[0]

            with open(pdb_file, encoding='utf-8') as in_file:
                with open(os.path.join(self.output, os.path.basename(pdb_file)), 'w',
                        encoding='utf-8') as out_file:
                    if self.verbosity > 1:
                        print(f"Saving {out_file.name}")
                    coord_strings = (
                        '{:>8.3f}{:>8.3f}{:>8.3f}'.format(*atom.tolist()) for atom in rotated_xyz
                    )
                    pdb = [
                        line if not line.startswith('ATOM') or line.startswith('HETATM ')
                        else f'{line[:30]}{next(coord_strings)}{line[54:]}'
                        for line in in_file.readlines()
                    ]
                    out_file.write(''.join(pdb))


    def _rotate_single_pdb_pool(self, variables):
        """Intermediate function that unpacks the single tuple variable from imap
        as arguments to the true rotate function.

        Args:
            variables (Iterable): Iterable to be unpacked.
        """
        self._rotate_single_pdb(*variables)

    def rotate_all_pool(self):
        """Pool the rotating of PDBs."""
        with torch.no_grad():
            t0 = time.perf_counter()
            self.rotator.use_numpy = True  # Enable numpy for Rotator
            self.rotator.quaternions = self.rotator.quaternions.cpu()
            self.rotator.center = self.rotator.center.cpu()
            rotation_matrix = self.rotator.get_matrix()

            with mp.Pool(self.cores) as pool:
                list(tqdm(pool.imap(
                                self._rotate_single_pdb_pool,
                                zip(range(1, len(self.pdbs)), repeat(rotation_matrix)),
                                chunksize=max(2, round(len(self.pdbs) / (self.cores * 16)))
                                    ),
                        total=len(self.pdbs) - 1, desc="Rotating PDBs",
                        disable=self.verbosity == 0
                            ))
        if self.verbosity > 1:
            print(f'Rotating and saving data took: {time.perf_counter() - t0:.2f} seconds')

    def calc_rmsd_with_template(self, output_file):
        """Calculates the RMSD of each aligned PDB with the template.
        Only takes into account the selected residues and ignores any deletions.

        Args:
            output_file (str): Path to write to.
        """
        with torch.no_grad():
            rotated_xyz = self.rotator.get_rotated_xyz()
            pdb_range = torch.arange(start=1, end=len(self.pdbs))
            combined_masks = self.rotator.presence_mask[0] * self.rotator.presence_mask[pdb_range]
            rmsds = ((rotated_xyz[0] - rotated_xyz[pdb_range]) * combined_masks)\
                .pow(2).reshape(len(self.pdbs)-1, -1).sum(1)\
                .div(combined_masks.sum(1).squeeze()).sqrt()

            with open(output_file, 'a', encoding='utf-8') as rmsd_file:
                rmsd_file.write(f"Template\t{self.pdbs[0]}\n")
                rmsd_file.write("\n".join([
                    f"{os.path.basename(pdb)}\t{rmsds[i].item()}"
                    for i, pdb in enumerate(self.pdbs[1:])
                ]))

    def optimize(self):
        """Runs the optimization loop.
        Attempts to find the optimal rotation to superpose all PDBs to the template."""
        if self.verbosity > 1:
            print("Aligning backbones")
        t0 = time.perf_counter()
        label = self.rotator.xyz[0:1]
        shake_steps = max(75, -4*label.shape[1]+400)
        finetune_steps = 25
        step = 0
        for step in trange(0, shake_steps+finetune_steps, desc="Aligning backbones",
                disable=self.verbosity != 1):
            # Take a step
            loss = self.rotator.step()
            if step <= shake_steps:
                self.rotator.normalize_q()
            if step == shake_steps:
                self.rotator.optimizer = torch.optim.SGD([self.rotator.quaternions], lr=0.00001)
            if self.verbosity > 1:
                print(f'Step: {step+1} \t Loss: {loss}')
        if self.verbosity > 1:
            print(f'Steps: {step+1}, alignment-time: {time.perf_counter()-t0:.2f}')


def superpose(pdbs_list, template, output=None, residues=None, chain=None, cores=mp.cpu_count(),
    batch_size=50000, gpu=False, rmsd_path=None, verbosity=1):
    """Runs all the steps to superpose a list of PDBs.

    Args:
        pdbs_list (list): List of paths to PDB files.
        template (str): Template PDB path.
        output (str, optional): Folder to write the rotated PDBs. Defaults to None.
        residues (list, optional): List of residue indices to align to. Defaults to None.
        chain (str, optional): Chain to align to.. Defaults to None.
        cores (int, optional): The number of CPU cores to use for multiprocessing.
            Defaults to the amount of cores on the system.
        batch_size (int, optional): The batch size of the alignment. Defaults to 50000.
        gpu (bool, optional): Flag to indicate whether GPU should be utilized. Defaults to False.
        rmsd_path (str, optional): If set, will write RMSDs of all models
            compared to the template. Defaults to None.
        verbosity (int, optional): Verbosity level. 0 = Silent, 1 = Normal, 2 = Verbose.
            Defaults to 1.
    """
    if verbosity > 0:
        start_time = time.perf_counter()
        print("\n=== Superimpose ===")

    torch.set_num_threads(cores)

    # TODO: Do we keep this?
    # for REPRODUCIBILITY
    warnings.filterwarnings("ignore")

    # Make output folder
    os.makedirs(output, exist_ok=True)

    if gpu:
        device = torch.device("cuda")
        mp.set_start_method('spawn')
    else:
        device = torch.device("cpu")

    num_batches = math.ceil(len(pdbs_list)/batch_size)
    for i in range(num_batches):
        batch_start = i * batch_size
        if verbosity > 0:
            print(f"Processing batch {i+1}/{num_batches}")
        data_processor = PDBdataset(
            pdbs_list[batch_start:batch_start+batch_size],
            template,
            residues,
            chain,
            cores,
            output,
            device,
            verbosity
        )
        data_processor.optimize()
        data_processor.rotate_all_pool()
        if rmsd_path:
            data_processor.calc_rmsd_with_template(rmsd_path)

    if verbosity > 0:
        end_time = time.perf_counter()
        print(f"=== Superimposing completed in {end_time - start_time:.1f}s ===\n")


def main():
    """Runs the PDB superposition using command line arguments."""
    arguments = util.parse_args()

    pdbs_list = []
    if arguments.subfolders:
        for root, _, files in os.walk(arguments.input):
            for file in [file for file in files if file.endswith('.pdb')]:
                pdbs_list.append(str(os.path.join(root, file)))
    else:
        pdbs_list += glob.glob(str(os.path.join(arguments.input, '*.pdb')))

    if not pdbs_list:
        raise Exception("Specified folder does not contain .pdb files.")

    # Sort the files alphabetically
    pdbs_list.sort(key=str.lower)

    if not arguments.template:
        arguments.template = pdbs_list[0]
        if arguments.verbosity > 0:
            print(f"No template assigned. Using: {arguments.template}")

    if arguments.rmsd:
        arguments.rmsd = os.path.join(arguments.output, "rmsd.tsv")

    superpose(
        pdbs_list,
        arguments.template,
        arguments.output,
        arguments.residues,
        arguments.chain,
        arguments.n_cores,
        arguments.batch_size,
        arguments.gpu,
        arguments.rmsd,
        arguments.verbosity
    )



if __name__ == '__main__':
    main()
