import matplotlib.pyplot as plt
import numpy as np
from typing import List
import json


class Plot:

    def __init__(self, experiment_dirs: List[str]):
        self.experiment_dirs = experiment_dirs
        self.experiment_info = []
        self.range_counts = []

    def get_experiment_infos(self):

        self.times = []
        self.scores = []
        self.models = []
        self.accuracies = []
        self.stop_conditions = []
        self.stop_conditions_perc = []
        self.range_counts = []

        for experiment_dir in self.experiment_dirs:
            with open(f'{experiment_dir}/summary.json', 'r') as f:
                summary = json.load(f)
            self.models.append(summary['model'])
            
            range_counts_lists = sorted(summary['total_range_counts'].items()) # sorted by key, return a list of tuples
            
            self.range_counts.append(range_counts_lists)
            exp_times = []
            exp_scores = []
            exp_accuracies = []
            for _, results in summary['results'].items():
                exp_times.append(results['time'])
                exp_scores.append(results['score'])
                exp_accuracies.append(results['accuracy'])

            self.times.append(exp_times)
            self.scores.append(exp_scores)
            self.accuracies.append(exp_accuracies)
            if self.stop_conditions != []:
                assert len(list(summary['results'].keys())) == len(self.stop_conditions), \
                    "Experiments don't have a consistent number of stop conditions."
            else:
                self.stop_conditions = [int(k) for k in list(summary['results'].keys())]
            self.stop_conditions_perc = summary['stop_conditions_perc']

    def plot_experiments(self, save=False):

        colors = [
            (170/255, 0,       120/255),
            (150/255, 225/255, 0),
            (0/255,   0/255,   120/255),
            (0/255,   200/255, 195/255),
            (255/255, 205/255, 25/255),
            (255/255, 150/255, 50/255)
        ]
        lines = ['-', '--', '-.', (0, (1, 1)), (0, (1, 1)), '-.']
        markers = ['o', 's', 'D', '^', 'P', 'o']
        line_width = 2

        def save(fig, name):
            for exp_dir in self.experiment_dirs:
                fig.savefig(f'{exp_dir}/stopcond-scores-{name}.png')

        def create_single_plot(
            fig,
            ax,
            x,
            y,
            i,
            x_label,
            y_label,
            line_label,
            x_tick_labels=None,
            y_ticks=np.arange(0, 1.1, 0.1)
        ):

            ax.plot(
                x,
                y,
                markers[i % 5],
                linestyle=lines[i % 5],
                color=colors[i % 5],
                label=line_label,
                linewidth=line_width
            )
            ax.set_xlabel(x_label)
            ax.set_ylabel(y_label)

            if x_tick_labels:
                # TODO: Generalize better -- '-6' = the last 6 stop conditions that
                # can be displayed without overlaps on the graph.
                ax.set_xticks(x[-6:])
                ax.set_xticklabels(x_tick_labels[-6:])

            y_tick_labels = [str(round(v, 2))[1:] if round(v, 2) != 1.0 else round(v, 2) for v in y_ticks]
            ax.set_yticks(y_ticks)
            ax.set_yticklabels(y_tick_labels)
            ax.legend()
            ax.grid(b=True, which='major', color='k', linestyle='-', linewidth=0.2)
            return fig, ax

        # Plot 1 Recall
        fig, ax = plt.subplots(figsize=(7, 5), ncols=1, nrows=1)
        for i, (model, scores) in enumerate(zip(self.models, self.scores)):

            fig, ax = create_single_plot(
                fig,
                ax,
                x=self.stop_conditions,
                y=scores,
                i=i,
                x_label='% of the dataset searched',
                y_label='Recall',
                line_label=model,
                x_tick_labels=[f'{int(sc*100)}%' for sc in self.stop_conditions_perc]
            )
        if save:
            save(fig, "recall")
        else:
            fig.show()

        # Plot 2 Time
        fig, ax = plt.subplots(figsize=(7, 5), ncols=1, nrows=1)
        for i, (model, times, scores) in enumerate(zip(self.models, self.times, self.scores)):

            fig, ax = create_single_plot(
                fig,
                ax,
                x=times,
                y=scores,
                i=i,
                x_label='Time (s)',
                y_label='Recall',
                line_label=model
            )
        if save:
            save(fig, "time")
        else:
            fig.show()    
