# code from https://github.com/gsurma/cartpole
from statistics import mean
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from collections import deque
import os
import csv
import numpy as np

AVERAGE_SCORE_TO_SOLVE = 195
CONSECUTIVE_RUNS_TO_SOLVE = 100

class ScoreLogger:
    def __init__(self, env_name, exp_name):
        self.scores = deque(maxlen=CONSECUTIVE_RUNS_TO_SOLVE)
        self.env_name = env_name

        self.SCORES_CSV_PATH = f"{os.getcwd()}/results/scores_{exp_name}.csv"
        self.SCORES_PNG_PATH = f"{os.getcwd()}/results/scores_{exp_name}.png"
        self.SOLVED_CSV_PATH = f"{os.getcwd()}/results/solved_{exp_name}.csv"
        self.SOLVED_PNG_PATH = f"{os.getcwd()}/results/solved_{exp_name}.png"

        if os.path.exists(self.SCORES_PNG_PATH):
            os.remove(self.SCORES_PNG_PATH)
        if os.path.exists(self.SCORES_CSV_PATH):
            os.remove(self.SCORES_CSV_PATH)

    def add_score(self, score, run):
        self._save_csv(self.SCORES_CSV_PATH, score)
        self._save_png(input_path=self.SCORES_CSV_PATH,
                       output_path=self.SCORES_PNG_PATH,
                       x_label="Episodes",
                       y_label="Total reward",
                       average_of_n_last=CONSECUTIVE_RUNS_TO_SOLVE,
                       show_goal=False,
                       show_trend=True,
                       show_legend=True)
        self.scores.append(score)
        mean_score = mean(self.scores)
        print(f"Scores: (score: {str(score)} min: {str(min(self.scores))}, avg: {str(mean_score)}, max:  {str(max(self.scores))})\n")
        # if mean_score >= AVERAGE_SCORE_TO_SOLVE and len(self.scores) >= CONSECUTIVE_RUNS_TO_SOLVE:
        #     solve_score = run-CONSECUTIVE_RUNS_TO_SOLVE
        #     print(f"Solved in {str(solve_score)} episodes, {str(run)} total episodes.")
        #     self._save_csv(self.SOLVED_CSV_PATH, solve_score)
        #     self._save_png(input_path=self.SOLVED_CSV_PATH,
        #                    output_path=self.SOLVED_PNG_PATH,
        #                    x_label="trials",
        #                    y_label="steps before solve",
        #                    average_of_n_last=None,
        #                    show_goal=False,
        #                    show_trend=False,
        #                    show_legend=False)
        #     return True

        return False

    def _save_png(self, input_path, output_path, x_label, y_label, average_of_n_last, show_goal, show_trend, show_legend):
        x = []
        y = []
        with open(input_path, "r") as scores:
            reader = csv.reader(scores)
            data = list(reader)
            for i in range(0, len(data)):
                x.append(int(i))
                y.append(float(data[i][0]))

        plt.subplots()
        plt.plot(x, y, label="score per episode")

        average_range = average_of_n_last if average_of_n_last is not None else len(x)
        plt.plot(x[-average_range:], [np.mean(y[-average_range:])] * len(y[-average_range:]), linestyle="--", label="last " + str(average_range) + " episodes average")

        if show_goal:
            plt.plot(x, [AVERAGE_SCORE_TO_SOLVE] * len(x), linestyle=":", label=str(AVERAGE_SCORE_TO_SOLVE) + " score average goal")

        if show_trend and len(x) > 1:
            trend_x = x[1:]
            z = np.polyfit(np.array(trend_x), np.array(y[1:]), 1)
            p = np.poly1d(z)
            plt.plot(trend_x, p(trend_x), linestyle="-.",  label="trend")

        plt.title(self.env_name)
        plt.xlabel(x_label)
        plt.ylabel(y_label)

        if show_legend:
            plt.legend(loc="upper left")

        plt.savefig(output_path, bbox_inches="tight")
        plt.close()

    def _save_csv(self, path, score):
        if not os.path.exists(path):
            with open(path, "w"):
                pass
        scores_file = open(path, "a")
        with scores_file:
            writer = csv.writer(scores_file)
            writer.writerow([score])
