"Postprocessing steps for the MDWaterBox scenario."
import os
import pathlib
import time
from typing import Literal

import matplotlib.pyplot as plt
import MDAnalysis as mda
import nglview as nv
import numpy as np

import inductiva

FULL_TRAJECTORY_FILE = "full_trajectory.trr"
COMPRESSED_TRAJECTORY_FILE = "compressed_trajectory.xtc"
PROTEIN_TOPOLOGY_FILE = "protein.gro"
SYSTEM_TOPOLOGY_FILE = "solvated_protein.tpr"


class ProteinSolvationOutput:
    """Post process the simulation output of a ProteinSolvation scenario."""

    def __init__(self, sim_output_path: pathlib.Path = None):
        """Initializes a `ProteinSolvationOutput` object.

        Given a simulation output directory that contains the standard files
        generated by a ProteinSolvation simulation run, this class provides 
        methods to visualize the simulation outputs in a notebook interactively.

        Args:
            sim_output_path: Path to the simulation output directory."""

        self.sim_output_dir = sim_output_path
        topology_path = os.path.join(self.sim_output_dir, PROTEIN_TOPOLOGY_FILE)
        trajectory_path = os.path.join(self.sim_output_dir,
                                       COMPRESSED_TRAJECTORY_FILE)
        self.universe = inductiva.molecules.scenarios.utils.unwrap_trajectory(
            topology_path, trajectory_path)

    def render_interactive(self,
                           representation: Literal["cartoon", "ball+stick",
                                                   "line", "point", "surface",
                                                   "ribbon"] = "ball+stick",
                           selection: str = "protein",
                           add_backbone: bool = True):
        """
        Render the simulation outputs in an interactive visualization.
        Args: 
            representation: The protein representation to use for the 
            visualization.
            selection: The selection to use for the representation. Check 
            https://nglviewer.org/ngl/api/manual/usage/selection-language.html 
            for details.
            add_backbone: Whether to add the protein backbone to the 
            visualization.
            """
        view = nv.show_mdanalysis(self.universe)
        view.add_representation(representation, selection=selection)
        if add_backbone:  #hardcoding the backbone as a cartoon representation
            view.add_representation("cartoon", selection="protein")
        view.center()

        print("System Information: ")
        print(f"Number of protein atoms: {len(self.universe.atoms)}")
        print(f"Number of amino acids: "
              f"{self.universe.select_atoms('protein').n_residues}")
        print(f"Number of trajectory frames: {len(self.universe.trajectory)}")
        return view

    def calculate_rmsf_trajectory(self):
        """Calculate the root mean square fluctuation (RMSF) over a trajectory.  

        It is typically calculated for the alpha carbon atom of each residue. 
        These atoms make the backbone of the protein.The RMSF is the square root 
        of the variance of the fluctuation around the average position:
        &rhoi = √⟨(xi - ⟨xi⟩)²⟩
        It quantifies how much a structure diverges from the average structure 
        over time, the RSMF can reveal which areas of the system are the most 
        mobile. Check 
        https://userguide.mdanalysis.org/stable/examples/analysis/alignment_and_rms/rmsf.html 
        for more details.

        Args:
            nglview_visualization: Whether to return visualization of the 
            RMSF using nglview or not."""
        start_time = time.time()
        topology_path = os.path.join(self.sim_output_dir, SYSTEM_TOPOLOGY_FILE)
        full_trajectory_path = os.path.join(self.sim_output_dir,
                                            FULL_TRAJECTORY_FILE)
        full_precision_universe = mda.Universe(topology_path,
                                               full_trajectory_path)

        aligned_trajectory_path = os.path.join(self.sim_output_dir,
                                               "aligned_traj.dcd")
        inductiva.molecules.scenarios.utils.align_trajectory_to_average(
            full_precision_universe, aligned_trajectory_path)
        align_universe = mda.Universe(topology_path, aligned_trajectory_path)

        # Calculate RMSF for carbon alpha atoms
        c_alphas = align_universe.select_atoms("protein and name CA")
        rmsf = mda.analysis.rms.RMSF(c_alphas).run()
        rmsf_values = rmsf.results.rmsf
        duration = time.time() - start_time
        print(f"RMSF calculation took {duration:.2f} seconds.")
        return rmsf_values

    def plot_rmsf_per_residue(self, rmsf_values: np.array):
        "Plot RMSF values per residue."
        # Plot the data
        residue_number = np.arange(len(rmsf_values))
        plt.plot(residue_number, rmsf_values)
        plt.xlabel("Residue Number")
        plt.ylabel("RMSF")
        plt.title("RMSF per residue")
        plt.grid(True)
        plt.show()

    def render_attribute_per_residue(
        self,
        residue_attributes: np.ndarray,
        representation: Literal["cartoon", "ball+stick", "line", "point",
                                "surface", "ribbon"] = "cartoon"):
        """Render a specific protein attribute in an interactive visualization.
        Args: 
            residue_attributes: The per residue values of the attribute you want 
            to visualize.
            representation: The protein representation to use for the 
            visualization.
            """
        self.universe.add_TopologyAttr("tempfactors")
        protein = self.universe.select_atoms("protein")
        for residue, value in zip(protein.residues, residue_attributes):
            residue.atoms.tempfactors = value
        view = nv.show_mdanalysis(self.universe)
        view.add_representation(representation, selection="protein")
        view.update_representation(color_scheme="bfactor")
        view.center()
        return view
