#!/usr/bin/env python

# Copyright 2018 Earth Sciences Department, BSC-CNS
#
# This file is part of HERMESv3_GR.
#
# HERMESv3_GR is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# HERMESv3_GR is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with HERMESv3_GR. If not, see <http://www.gnu.org/licenses/>.


import os
import pandas as pd
import numpy as np

global refresh_log

global precision
precision = np.float64

global icomm
global comm
global rank
global size

global log_level
global df_times


def define_global_vars(in_log_level, mycomm):
    # TODO Documentation
    from mpi4py import MPI

    global icomm
    global comm
    global rank
    global size

    if mycomm is None:
        icomm = MPI.COMM_WORLD
    else:
        icomm = mycomm
    comm = icomm
    # comm = icomm.Split(color=0, key=0)
    rank = comm.Get_rank()
    size = comm.Get_size()

    global log_level
    log_level = in_log_level


def define_times_file():
    # TODO Documentation
    global df_times

    df_times = pd.DataFrame(columns=['Class', 'Function', rank])


def write_time(module, func, time, level=1):
    """
    Logs the execution time of a function to a DataFrame.

    This function appends a new row to the global `df_times` DataFrame, which contains the execution times of functions.
    The new row contains the name of the module and function, the rank of the process, and the execution time.
    The function only logs the execution time if the global log level is greater than or equal to the specified level.

    Parameters
    ----------
    module : str
        The name of the module where the function is defined.
    func : str
        The name of the function whose execution time is being logged.
    time : float
        The execution time of the function.
    level : int, optional
        The importance level of the message. The execution time will be logged if this level is less than or equal to
        the global log level. Default is 1.
    """
    global df_times
    if log_level >= level:
        new_row = pd.DataFrame({'Class': [module], 'Function': [func], rank: [time]})
        df_times = new_row if df_times.empty else pd.concat([df_times, new_row], ignore_index=True)


def finish_logs(output_dir, date):
    # TODO Documentation
    from functools import reduce

    global df_times
    df_times = df_times.groupby(['Class', 'Function']).sum().reset_index()
    data_frames = comm.gather(df_times, root=0)
    if rank == 0:
        times_path = os.path.join(output_dir, 'logs', 'HERMESv3_{0}_times_Procs{1}.csv'.format(
            date.strftime('%Y%m%d%H'), str(size).zfill(4)))
        if os.path.exists(times_path):
            os.remove(times_path)
        df_merged = reduce(lambda left, right: pd.merge(left, right, on=['Class', 'Function'], how='outer'),
                           data_frames)
        df_merged['min'] = df_merged.loc[:, list(range(size))].min(axis=1)
        df_merged['max'] = df_merged.loc[:, list(range(size))].max(axis=1)
        df_merged['mean'] = df_merged.loc[:, list(range(size))].mean(axis=1)

        df_merged.to_csv(times_path)
    comm.Barrier()
