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

import numpy as np
try:
    import matplotlib.pyplot as plt
    import nglview as nv
    import IPython as ip
except ImportError:
    nv = None
    plt = None
    ip = None

try:
    import google.colab as gc
except ImportError:
    gc = None

import inductiva
from inductiva.utils import optional_deps

COMPRESSED_TRAJECTORY_FILE = "compressed_trajectory.xtc"
PROTEIN_TOPOLOGY_FILE = "protein.gro"
SYSTEM_TOPOLOGY_FILE = "solvated_protein.tpr"
RMSF_VALUES_FILE = "rmsf_values.npy"


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
        # universe will be loaded only when needed
        self.universe = None
        self.ipython_kernel = self.get_ipython_kernel()

    @optional_deps.needs_molecules_extra_deps
    def _load_universe(self):
        topology_path = os.path.join(self.sim_output_dir, PROTEIN_TOPOLOGY_FILE)
        trajectory_path = os.path.join(self.sim_output_dir,
                                       COMPRESSED_TRAJECTORY_FILE)

        if self.universe is None:
            self.universe = \
                inductiva.molecules.utils.unwrap_trajectory(
                    topology_path, trajectory_path)

    @optional_deps.needs_molecules_extra_deps
    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.
            """
        self.enable_vizualization()
        self._load_universe()

        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

    @optional_deps.needs_molecules_extra_deps
    def plot_rmsf_per_residue(self):
        """Plot 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."""
        rmsf_values_path = os.path.join(self.sim_output_dir, RMSF_VALUES_FILE)
        rmsf_values = np.load(rmsf_values_path)
        # 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()

        return rmsf_values

    @optional_deps.needs_molecules_extra_deps
    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.enable_vizualization()
        self._load_universe()

        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

    @optional_deps.needs_molecules_extra_deps
    def get_ipython_kernel(self):
        """Check if the current environment have an IPython kernel."""
        return ip.get_ipython()

    @optional_deps.needs_molecules_extra_deps
    def enable_vizualization(self):
        """Enable vizualization if IPython is available."""
        if self.ipython_kernel is None:
            raise ImportError("IPython is not available. Visualization is "
                              "only available in a python notebook.")

        if self.ipython_kernel.__module__.startswith("google.colab"):
            gc.output.enable_custom_widget_manager()
