import csv
import pickle


class Logger:
    def __init__(self):
        self.first_write = False
    
    def log(self, filename, steps, reward, total_eps):
        with open(filename, 'a', newline='') as csvfile:
            spamwriter = csv.writer(csvfile, delimiter=' ',
                                    quotechar='|', quoting=csv.QUOTE_MINIMAL)
            
            if not self.first_write:
                spamwriter.writerow(["steps", "reward", "total_eps"])
                self.first_write = True

            spamwriter.writerow([steps, reward, total_eps])

    class StateCounter:
        def __init__(self, size):
            self.width = size[0]
            self.height = size[1]

            self.count_dict = {}

        def inc_count(self, pos):
            pos_array = pos[0]
            pos_key = str(pos_array[0]) + str(pos_array[1])

            if pos_key not in self.count_dict:
                self.count_dict[pos_key] = 0
            self.count_dict[pos_key] += 1

        def save(self, filename):
            # open a file, where you ant to store the data
            file = open(filename, 'wb')

            # dump information to that file
            pickle.dump(self.count_dict, file)

            # close the file
            file.close()
