####################################################################
#                                                                  #
#    ensemble_md,                                                  #
#    a python package for running GROMACS simulation ensembles     #
#                                                                  #
#    Written by Wei-Tse Hsu <wehs7661@colorado.edu>                #
#    Copyright (c) 2022 University of Colorado Boulder             #
#                                                                  #
####################################################################
"""
The :obj:`.utils` module provides useful utility functions.
"""
import os
import sys
import natsort
import datetime
import numpy as np


class Logger:
    """
    Redirects the STDOUT and STDERR to a specified output file while preserving them on screen.

    Parameters
    ----------
    logfile : str
        Name of the output file to write the logged messages.

    Attributes
    ----------
    terminal : file object
        The file object that represents the original STDOUT (i.e., the screen).
    log : file object
        The file object that represents the logfile where messages will be written.
    """

    def __init__(self, logfile):
        """
        Initializes a Logger instance.

        Parameters
        ----------
        logfile : str
            Name of the output file to write the logged messages.
        """
        self.terminal = sys.stdout
        self.log = open(logfile, "a")

    def write(self, message):
        """
        Writes the given message to both the STDOUT and the logfile.

        Parameters
        ----------
        message : str
            The message to be written to STDOUT and logfile.
        """
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        """
        This method is needed for Python 3 compatibility. This handles the flush command by doing nothing.
        You might want to specify some extra behavior here.
        """
        # self.terminal.log()
        pass


def gmx_output(gmx_obj):
    """
    Tells if a command launched by ``gmxapi.commandline_operation`` failed or not.
    If the command failed, the function will print out STDOUT and STDERR of the process.

    Parameters
    ----------
    gmx_obj : obj
        An object generated by gmx.commandline_operation.
    """
    try:
        rtn_code = gmx_obj.output.returncode.result()
        if type(rtn_code) == list:  # simulation ensemble
            if sum(rtn_code) == 0:
                print(" DONE")
            else:
                print(f" Return codes: {rtn_code}")
                for i in range(len(rtn_code)):
                    if rtn_code[i] != 0:
                        print(f"\nSTDERR of the process:\n\n {gmx_obj.output.stderr.result()[i]}\n")
        else:  # single simulation
            if rtn_code == 0:
                print(" DONE")
            else:
                print(f"\nSTDERR of the process:\n\n {gmx_obj.output.stderr.result()}\n")
    except AttributeError:
        raise RuntimeError("The input object is not an object generated by commandline_operation.")


def clean_up(dir_before, dir_after, verbose=False):
    """
    Removes newly generated folders that are empty after a command launched by gmxapi.

    Parameters
    ----------
    dir_before : list
        The list of directories existed before the command was executed.
    dir_after : list
        The list of directories existed after the command was executed. This helps
        figure out which directories are newly generated by command.
    verbose : bool
        If details should be printed.
    """
    new_dirs = natsort.natsorted([i for i in dir_after if i not in dir_before])
    if len(new_dirs) != 0:
        if verbose is True:
            print("\nCleaning up/restructuring the directories ...")
    for i in new_dirs:
        if len(os.listdir(i)) == 0:
            if verbose is True:
                print(f"  Removing the empty folder {i} ...")
            os.rmdir(i)


def format_time(t):
    """
    Converts time in seconds to the "most readable" format.

    Parameters
    ----------
    t : float
        The time in seconds.

    Returns
    -------
    t_str : sttr
        A string in the format of "XX day XX hour(s) XX minute(s) XX second(s)".
    """
    hh_mm_ss = str(datetime.timedelta(seconds=t)).split(":")

    if "day" in hh_mm_ss[0]:
        # hh_mm_ss[0] will contain "day" and cannot be converted to float
        hh, mm, ss = hh_mm_ss[0], float(hh_mm_ss[1]), float(hh_mm_ss[2])
        t_str = f"{hh} hour(s) {mm:.0f} minute(s) {ss:.0f} second(s)"
    else:
        hh, mm, ss = float(hh_mm_ss[0]), float(hh_mm_ss[1]), float(hh_mm_ss[2])
        if hh == 0:
            if mm == 0:
                t_str = f"{ss:.1f} second(s)"
            else:
                t_str = f"{mm:.0f} minute(s) {ss:.0f} second(s)"
        else:
            t_str = f"{hh:.0f} hour(s) {mm:.0f} minute(s) {ss:.0f} second(s)"

    return t_str


def autoconvert(s):
    """
    Converts input to a numerical type if possible. Used for the MDP parser.
    Modified from `utilities.py in GromacsWrapper <https://github.com/Becksteinlab/GromacsWrapper>`_.
    Copyright (c) 2009 Oliver Beckstein <orbeckst@gmail.com>

    Parameters
    ----------
    s : str or any
        The input value to be converted to a numerical type if possible. If :code:`s` is not a string,
        it is returned as is.

    Returns
    -------
    numerical : int, float, numpy.ndarray, or any
        The converted numerical value. If :code:`s` can be converted to a single numerical value,
        that value is returned as an :code:`int` or :code:`float`. If :code:`s` can be converted to
        multiple numerical values, a :code:`numpy.ndarray` containing those values is returned.
        If :code:`s` cannot be converted to a numerical value, :code:`s` is returned as is.

    Raises
    ------
    ValueError
        If :code:`s` cannot be converted to a numerical value.
    """
    if type(s) is not str:
        return s
    for converter in int, float, str:  # try them in increasing order of lenience
        try:
            s = [converter(i) for i in s.split()]
            if len(s) == 1:
                return s[0]
            else:
                return np.array(s)
        except (ValueError, AttributeError):
            pass
    raise ValueError("Failed to autoconvert {0!r}".format(s))


def get_subplot_dimension(n_panels):
    """
    Gets the numbers of rows and columns in a subplot such that
    the arrangement of the .

    Parameters
    ----------
    n_panels : int
        The number of panels in the subplot.
    """
    if int(np.sqrt(n_panels) + 0.5) ** 2 == n_panels:
        # perfect square number
        n_cols = int(np.sqrt(n_panels))
    else:
        n_cols = int(np.floor(np.sqrt(n_panels))) + 1

    if n_panels % n_cols == 0:
        n_rows = int(np.floor(n_panels / n_cols))
    else:
        n_rows = int(np.floor(n_panels / n_cols)) + 1

    return n_rows, n_cols


def weighted_mean(vals, errs):
    """
    Calculates the inverse-variance-weighted mean.

    Parameters
    ----------
    vals : list
        A list of values to be averaged.
    errs : list
        A list of errors corresponding to the given values

    Returns
    -------
    mean : float
        The inverse-variance-weighted mean.
    err : float
        The propgated error of the mean.
    """
    w = [1 / (i ** 2) for i in errs]
    wx = [w[i] * vals[i] for i in range(len(vals))]
    mean = np.sum(wx) / np.sum(w)
    err = np.sqrt(1 / np.sum(w))

    return mean, err
