#!/usr/bin/env python3

# Copyright 2017-2020 Earth Sciences Department, BSC-CNS

# This file is part of Autosubmit.

# Autosubmit 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.

# Autosubmit 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 Autosubmit.  If not, see <http://www.gnu.org/licenses/>.

import itertools
import math
import traceback

import matplotlib as mtp

mtp.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.patches as mpatches
from autosubmit.statistics.statistics import Statistics
from autosubmit.job.job import Job
from log.log import Log
from datetime import datetime
from typing import List, Dict

Log.get_logger("Autosubmit")

# Autosubmit stats constants
RATIO = 4
MAX_JOBS_PER_PLOT = 12.0
MAX_NUM_PLOTS = 40



def _seq(start, end, step):
    """From: https://pynative.com/python-range-for-float-numbers/"""
    sample_count = int(abs(end - start) / step)
    return itertools.islice(itertools.count(start, step), sample_count)

def create_bar_diagram(experiment_id, jobs_list, general_stats, output_file, period_ini=None, period_fi=None,
                       queue_time_fixes=None):
    # type: (str, List[Job], List[str], str, datetime, datetime, Dict[str, int]) -> None
    """
    Creates a bar diagram of the statistics.

    :param queue_time_fixes:
    :param experiment_id: experiment's identifier
    :type experiment_id: str
    :param jobs_list: list of jobs
    :type jobs_list: List[Job]

    :param general_stats: list of sections and options in the %DEFAULT.EXPID%_GENERAL_STATS file
    :type general_stats: list of tuples  
    :param output_file: path to the output file  
    :type output_file: str  
    :param period_ini: starting date and time
    :type period_ini: datetime  
    :param period_fi: finish date and time
    :type period_fi: datetime  
    """
    # Error prevention
    plt.close('all')
    normal_plots_count = 0
    failed_jobs_plots_count = 0
    exp_stats = None
    try:
        exp_stats = Statistics(jobs_list, period_ini, period_fi, queue_time_fixes)
        exp_stats.calculate_statistics()
        exp_stats.calculate_summary()
        exp_stats.make_old_format()

        failed_jobs_dict = exp_stats.build_failed_jobs_only_list()
        # Stats variables definition
        normal_plots_count = int(math.ceil(len(exp_stats.jobs_stat) / MAX_JOBS_PER_PLOT))
        failed_jobs_plots_count = int(math.ceil(len(failed_jobs_dict) / MAX_JOBS_PER_PLOT))
    except Exception as exp:
        print(exp)
        print((traceback.format_exc()))

    # Plotting
    total_plots_count = normal_plots_count + failed_jobs_plots_count
    # num_plots = norma
    # ind = np.arrange(int(MAX_JOBS_PER_PLOT))
    width = 0.16
    # Creating stats figure + sanity check
    plot = True
    err_message = "The results are too large to be shown, try narrowing your query.\nUse a filter like -ft where you supply a list of job types, e.g. INI, SIM or use the flag -fp where you supply an integer that represents the number of hours into the past that should be queried:\nSuppose it is noon, if you supply -fp 5 the query will consider changes starting from 7:00 am. If you really wish to query the whole experiment, refer to Autosubmit GUI."
    if total_plots_count > MAX_NUM_PLOTS:
        Log.info(err_message)
        plot = False
    else:
        fig = plt.figure(figsize=(RATIO * 4, 3 * RATIO * total_plots_count))
        fig.suptitle('STATS - ' + experiment_id, fontsize=24, fontweight='bold')
        # Variables initialization
        ax, ax2 = [], []
        rects = [None] * 5
        # print("Normal plots: {}".format(normal_plots_count))
        # print("Failed jobs plots: {}".format(failed_jobs_plots_count))
        # print("Total plots: {}".format(total_plots_count))
        grid_spec = gridspec.GridSpec(RATIO * total_plots_count + 2, 1)
        i_plot = 0
        for plot in range(1, normal_plots_count + 1):
            try:
                # Calculating jobs inside the given plot
                l1 = int((plot - 1) * MAX_JOBS_PER_PLOT)
                l2 = min(int(plot * MAX_JOBS_PER_PLOT), len(exp_stats.jobs_stat))
                if l2 - l1 <= 0:
                    continue
                ind = range(l2 - l1)
                ind_width = [x + width for x in ind]
                ind_width_3 = [x + width * 3 for x in ind]
                ind_width_4 = [x + width * 4 for x in ind]
                # Building plot axis
                ax.append(fig.add_subplot(grid_spec[RATIO * plot - RATIO + 2:RATIO * plot + 1]))
                ax[plot - 1].set_ylabel('hours')
                ax[plot - 1].set_xticks(ind_width)
                ax[plot - 1].set_xticklabels(
                    [job.name for job in jobs_list[l1:l2]], rotation='vertical')
                ax[plot - 1].set_title(experiment_id, fontsize=20)
                upper_limit = round(1.10 * exp_stats.max_time, 4)
                step = round(upper_limit / 10, 4)
                # Here we use ``upper_limit + step`` as np.arange is inclusive at the end,
                # ``islice`` is not.
                y_ticks = [round(x, 4) for x in _seq(0, upper_limit + step, step)]
                # ax[plot - 1].set_yticks(np.arange(0, upper_limit, round(upper_limit / 10, 4)))
                ax[plot - 1].set_yticks(y_ticks)
                ax[plot - 1].set_ylim(0, float(1.10 * exp_stats.max_time))
                # Building reacts
                rects[0] = ax[plot - 1].bar(ind, exp_stats.queued[l1:l2], width, color='lightpink')
                rects[1] = ax[plot - 1].bar(ind_width, exp_stats.run[l1:l2], width, color='green')
                rects[2] = ax[plot - 1].bar(ind_width_3, exp_stats.fail_queued[l1:l2], width, color='lightsalmon')
                rects[3] = ax[plot - 1].bar(ind_width_4, exp_stats.fail_run[l1:l2], width, color='salmon')
                rects[4] = ax[plot - 1].plot([0., width * 6 * MAX_JOBS_PER_PLOT],
                                             [exp_stats.threshold, exp_stats.threshold], "k--", label='wallclock sim')
                # Building legend
                i_plot = plot
            except Exception as exp:
                print((traceback.format_exc()))
                print(exp)

        job_names_in_failed = [name for name in exp_stats.failed_jobs_dict]
        failed_jobs_rects = [None]
        for j_plot in range(1, failed_jobs_plots_count + 1):
            try:
                l1 = int((j_plot - 1) * MAX_JOBS_PER_PLOT)
                l2 = min(int(j_plot * MAX_JOBS_PER_PLOT), len(job_names_in_failed))
                if l2 - l1 <= 0:
                    continue
                ind = range(l2 - l1)
                ind_width = [x + width for x in ind]
                ind_width_2 = [x + width * 2 for x in ind]
                plot = i_plot + j_plot
                ax.append(fig.add_subplot(grid_spec[RATIO * plot - RATIO + 2:RATIO * plot + 1]))
                ax[plot - 1].set_ylabel('# failed attempts')
                ax[plot - 1].set_xticks(ind_width)
                ax[plot - 1].set_xticklabels([name for name in job_names_in_failed[l1:l2]], rotation='vertical')
                ax[plot - 1].set_title(experiment_id, fontsize=20)
                ax[plot - 1].set_ylim(0, float(1.10 * exp_stats.max_fail))
                ax[plot - 1].set_yticks(range(0, exp_stats.max_fail + 2))
                failed_jobs_rects[0] = ax[plot - 1].bar(ind_width_2, [exp_stats.failed_jobs_dict[name] for name in
                                                                          job_names_in_failed[l1:l2]], width, color='red')
            except Exception as exp:
                print((traceback.format_exc()))
                print(exp)

        # Building legends subplot
        legends_plot = fig.add_subplot(grid_spec[0, 0])
        legends_plot.set_frame_on(False)
        legends_plot.axes.get_xaxis().set_visible(False)
        legends_plot.axes.get_yaxis().set_visible(False)

        try:
            # Building legends
            # print("Legends")
            build_legends(legends_plot, rects, exp_stats, general_stats)
            # Saving output figure
            grid_spec.tight_layout(fig, rect=[0, 0.03, 1, 0.97])
            plt.savefig(output_file)
        except Exception as exp:
            print(exp)
            print((traceback.format_exc()))
    try:
        create_csv_stats(exp_stats, jobs_list, output_file)
    except Exception as exp:
        Log.info(f'Error while creating csv stats:\n{err_message}')
    return plot


def create_csv_stats(exp_stats, jobs_list, output_file):
    job_names = [job.name for job in exp_stats.jobs_stat]
    start_times = exp_stats.start_times
    end_times = exp_stats.end_times
    queuing_times = exp_stats.queued
    running_times = exp_stats.run

    output_file = output_file.replace('pdf', 'csv')
    with open(output_file, 'w') as file:
        file.write(
            "Job,Started,Ended,Queuing time (hours),Running time (hours)\n")
        # In the other function, job_names,start_times... etc is only filled if the job has completed retrials
        # So I'll change this one to do the same
        for i in range(len([ job for job in jobs_list if job.get_last_retrials() ])):
            file.write("{0},{1},{2},{3},{4}\n".format(
                job_names[i], start_times[i], end_times[i], queuing_times[i], running_times[i]))


def build_legends(plot, rects, experiment_stats, general_stats):
    # type: (plt.figure, List[plt.bar], Statistics, List[str]) -> None
    # Main legend with colourful rectangles

    legend_rects = [[rect[0] for rect in rects]]

    legend_titles = [
        ['Queued (h)', 'Run (h)', 'Fail Queued (h)', 'Fail Run (h)', 'Max wallclock (h)']
    ]
    legend_locs = ["upper right"]
    legend_handlelengths = [None]

    # General stats legends, if exists
    if len(general_stats) > 0:
        legend_rects.append(get_whites_array(len(general_stats)))
        legend_titles.append([str(key) + ': ' + str(value) for key, value in general_stats])
        legend_locs.append("upper center")
        legend_handlelengths.append(0)

    # Total stats legend
    stats_summary_as_list = experiment_stats.get_summary_as_list()
    legend_rects.append(get_whites_array(len(stats_summary_as_list)))
    legend_titles.append(stats_summary_as_list)
    legend_locs.append("upper left")
    legend_handlelengths.append(0)

    # Creating the legends
    legends = create_legends(plot, legend_rects, legend_titles, legend_locs, legend_handlelengths)
    for legend in legends:
        plt.gca().add_artist(legend)


def create_legends(plot, rects, titles, locs, handlelengths):
    legends = []
    for i in range(len(rects)):
        legends.append(create_legend(
            plot, rects[i], titles[i], locs[i], handlelengths[i]))
    return legends


def create_legend(plot, rects, titles, loc, handlelength=None):
    return plot.legend(rects, titles, loc=loc, handlelength=handlelength)


def get_whites_array(length):
    white = mpatches.Rectangle((0, 0), 0, 0, alpha=0.0)
    return [white for _ in range(length)]
