#!/usr/bin/env python3
# -*- coding: utf-8 -*-

__author__ = 'MICROYU'

import csv
import os
import json, time
import numpy as np
from rl_plotter import log_utils as lu

class Logger():
    def __init__(self, log_dir="./logs", exp_name=None, env_name=None, seed=0, config=None, filename='evaluator.csv'):
        self.exp_name = exp_name
        self.env_name = env_name
        self.seed = seed
        self.previous_log_time = time.time()
        self.start_log_time = self.previous_log_time

        num_exps = 0
        self.log_dir = f"./{log_dir}/{exp_name.replace('-', '_')}_{env_name.replace('-', '_')}_seed{seed}"
        while True:
            if os.path.exists(f"{self.log_dir}-{str(num_exps)}/"):
                num_exps += 1
            else:
                self.log_dir += f"-{str(num_exps)}/"
                os.makedirs(self.log_dir)
                break
        self.csv_file = open(self.log_dir + '/' + filename, 'w', encoding='utf8')
        # header={"t_start": time.time(), 'env_id' : env_name, 'exp_name': exp_name, 'seed': seed}
        # header = '# {} \n'.format(json.dumps(header))
        # self.csv_file.write(header)
        self.logger = csv.DictWriter(self.csv_file, fieldnames=('mean_score', 'total_steps', 'std_score', 'max_score', 'min_score'))
        self.logger.writeheader()
        self.csv_file.flush()

        if config != None:
            lu.save_config(exp_name, config, self.log_dir)

    def monitor_env(self, env):
        from baselines import bench
        env = bench.Monitor(env, self.log_dir, allow_early_resets=True)
        return env

    def update(self, score, total_steps):
        '''
            Score is a list
        '''
        current_log_time = time.time()
        avg_score = np.mean(score)
        std_score = np.std(score)
        max_score = np.max(score)
        min_score = np.min(score)

        print(lu.colorize(f"\nTime: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}, Time spent from previous logger: {(current_log_time - self.previous_log_time):.3f} s", 'yellow', bold=True))
        print(lu.colorize(f"Evaluation over {len(score)} episodes after {total_steps}:", 'yellow', bold=True))
        print(lu.colorize(f"Avg: {avg_score:.3f} Std: {std_score:.3f} Max: {max_score:.3f} Min: {min_score:.3f}\n", 'yellow', bold=True))
        self.previous_log_time = current_log_time
        
        epinfo = {"mean_score": avg_score, "total_steps": total_steps, "std_score": std_score, "max_score": max_score, "min_score": min_score}
        self.logger.writerow(epinfo)
        self.csv_file.flush()
    
    def new_custom_logger(self, filename=None, fieldnames=[]):
        custom_logger = CustomLogger(self.log_dir, self.exp_name, self.env_name, self.seed, filename, fieldnames)
        return custom_logger
    
    def new_eval_logger(self, filename=None):
        eval_logger = EvalLogger(self.log_dir, self.exp_name, self.env_name, self.seed, filename)
        return eval_logger

class CustomLogger():
    def __init__(self, log_dir="./logs", exp_name=None, env_name=None, seed=0, filename="logger.csv", fieldnames=[]):
        self.previous_log_time = time.time()
        self.start_log_time = self.previous_log_time

        self.fieldnames = ["total_steps"] + fieldnames
        self.csv_file = open(log_dir + '/' + filename, 'w', encoding='utf8')
        # header={"t_start": time.time(), 'env_id' : env_name, 'exp_name': exp_name, 'seed': seed}
        # header = '# {} \n'.format(json.dumps(header))
        # self.csv_file.write(header)
        self.logger = csv.DictWriter(self.csv_file, fieldnames=(self.fieldnames))
        self.logger.writeheader()
        self.csv_file.flush()

    def update(self, fieldvalues, total_steps):
        epinfo = {}
        fieldvalues = [total_steps] + fieldvalues

        current_log_time = time.time()
        print(lu.colorize(f"\nTime: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}, Time spent from previous logger: {(current_log_time - self.previous_log_time):.3f} s", 'blue', bold=True))
        print(lu.colorize(f"CustomLogger with fileds: {self.fieldnames}", 'blue', bold=True))
        print(lu.colorize(f"fieldvalues: {fieldvalues}\n", 'blue', bold=True))
        self.previous_log_time = current_log_time
       
        for filedname, filedvalue in zip(self.fieldnames, fieldvalues):
            epinfo.update({filedname: filedvalue})
        self.logger.writerow(epinfo)
        self.csv_file.flush()

class EvalLogger():
    def __init__(self, log_dir="./logs", exp_name=None, env_name=None, seed=0, filename="logger.csv"):
        self.previous_log_time = time.time()
        self.start_log_time = self.previous_log_time

        self.csv_file = open(log_dir + '/' + filename, 'w', encoding='utf8')
        # header={"t_start": time.time(), 'env_id' : env_name, 'exp_name': exp_name, 'seed': seed}
        # header = '# {} \n'.format(json.dumps(header))
        # self.csv_file.write(header)
        self.logger = csv.DictWriter(self.csv_file, fieldnames=('mean_score', 'total_steps', 'std_score', 'max_score', 'min_score'))
        self.logger.writeheader()
        self.csv_file.flush()

    def update(self, score, total_steps):
        '''
            Score is a list
        '''
        current_log_time = time.time()
        avg_score = np.mean(score)
        std_score = np.std(score)
        max_score = np.max(score)
        min_score = np.min(score)

        print(lu.colorize(f"\nTime: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}, Time spent from previous logger: {(current_log_time - self.previous_log_time):.3f} s", 'yellow', bold=True))
        print(lu.colorize(f"Evaluation over {len(score)} episodes after {total_steps}:", 'yellow', bold=True))
        print(lu.colorize(f"Avg: {avg_score:.3f} Std: {std_score:.3f} Max: {max_score:.3f} Min: {min_score:.3f}\n", 'yellow', bold=True))
        self.previous_log_time = current_log_time
        
        epinfo = {"mean_score": avg_score, "total_steps": total_steps, "std_score": std_score, "max_score": max_score, "min_score": min_score}
        self.logger.writerow(epinfo)
        self.csv_file.flush()
    
    def update(self, avg_score, std_score, max_score, min_score, total_steps):
        
        epinfo = {"mean_score": avg_score, "total_steps": total_steps, "std_score": std_score, "max_score": max_score, "min_score": min_score}
        self.logger.writerow(epinfo)
        self.csv_file.flush()
    