from turtle import color
import matplotlib.pyplot as plt


class PPOHistory(object):
    def __init__(self):
        self.meanRewards = []
        self.entropys = []
        self.actorLosses = []
        self.criticLosses = []

    def saveHis(self, rewards, entropys, aLosses, cLosses):
        self.meanRewards.extend([rewards])
        self.entropys.extend([entropys])
        self.actorLosses.extend(aLosses)
        self.criticLosses.extend(cLosses)

    def drawHis(self):
        plt.figure(figsize=(21, 13), facecolor="#011627")
        ax = plt.subplot(2, 2, 1)
        ax.set_facecolor("#011627")
        ax.spines["top"].set_color("#c9d2df")
        ax.spines["bottom"].set_color("#c9d2df")
        ax.spines["left"].set_color("#c9d2df")
        ax.spines["right"].set_color("#c9d2df")
        ax.plot(
            range(len(self.meanRewards)), self.meanRewards, color="#c9d2df", label="AverageRewards"
        )
        ax.set_title("meanRewards", color="#c9d2df")
        ax = plt.subplot(2, 2, 2)
        ax.set_facecolor("#011627")
        ax.spines["top"].set_color("#c9d2df")
        ax.spines["bottom"].set_color("#c9d2df")
        ax.spines["left"].set_color("#c9d2df")
        ax.spines["right"].set_color("#c9d2df")
        ax.plot(range(len(self.entropys)), self.entropys, color="#c9d2df", label="AverageEntropys")
        ax.set_title("entropys", color="#c9d2df")
        ax = plt.subplot(2, 2, 3)
        ax.set_facecolor("#011627")
        ax.spines["top"].set_color("#c9d2df")
        ax.spines["bottom"].set_color("#c9d2df")
        ax.spines["left"].set_color("#c9d2df")
        ax.spines["right"].set_color("#c9d2df")
        ax.plot(
            range(len(self.actorLosses)), self.actorLosses, color="#c9d2df", label="actorLosses"
        )
        ax.set_title("actorLosses", color="#c9d2df")
        ax = plt.subplot(2, 2, 4)
        ax.set_facecolor("#011627")
        ax.spines["top"].set_color("#c9d2df")
        ax.spines["bottom"].set_color("#c9d2df")
        ax.spines["left"].set_color("#c9d2df")
        ax.spines["right"].set_color("#c9d2df")
        ax.plot(
            range(len(self.criticLosses)), self.criticLosses, color="#c9d2df", label="criticLosses"
        )
        ax.set_title("criticLosses", color="#c9d2df")
        plt.show()