import numpy as np
import time

"""
Please do not make changes to this file. 
This is only a reference script provided to allow you 
to do local evaluation. The evaluator **DOES NOT** 
use this script for orchestrating the evaluations. 
"""

# to avoid crashes but might cause results to be different 
# https://github.com/dmlc/xgboost/issues/1715

# import os
# os.environ['KMP_DUPLICATE_LIB_OK']='True'


from agents.orderenforcingwrapper import OrderEnforcingAgent
from citylearn.citylearn import CityLearnEnv

# Custom configure enviroment 
class Constants:
    episodes = 3
    schema_path = './data/citylearn_challenge_2022_phase_1/schema.json'

def action_space_to_dict(aspace):
    """ Only for box space """
    return { "high": aspace.high,
             "low": aspace.low,
             "shape": aspace.shape,
             "dtype": str(aspace.dtype)
    }

def env_reset(env):
    observations = env.reset()
    action_space = env.action_space
    observation_space = env.observation_space
    building_info = env.get_building_information()
    building_info = list(building_info.values())
    action_space_dicts = [action_space_to_dict(asp) for asp in action_space]
    observation_space_dicts = [action_space_to_dict(osp) for osp in observation_space]
    obs_dict = {"action_space": action_space_dicts,
                "observation_space": observation_space_dicts,
                "building_info": building_info,
                "observation": observations }
    return obs_dict

# Understand CityLearn Environment

env = CityLearnEnv(schema=Constants.schema_path)

# There is an action space for every building
# print(f' OBSERVATION SPACES {env.observation_space}')
# print(f' OBSERVATION SPACE for Builiding ONE is {env.observation_space[0]}')

# sample some actions
for building in range(5):
    print(f' SAMPLE OBSERVATION SPACE for Builiding ONE >>> {len(env.observation_space[building].sample()), env.observation_space[building].sample()}')

# we can see the observations are a 28 1D numpy array with every dimension defined by the range given in the spaces BOX

# There is an action space for every building
print(f' ACTION SPACES {env.action_space}')
print(f' ACTION SPACE for Builiding ONE is {env.action_space[0]}')

# sample some actions
for action in range(5):
    print(f' SAMPLE ACTION SPACE for Builiding ONE >>> {env.action_space[1].sample()}')

# we can observe the actions are continous in the range [-1,1]

from citylearn.agents.sac import SAC

# SAC??

# print(env_reset(env)["action_space"])
# env_reset(env)["observation_space"]
# env.reset()[0]

import random
Actions = [([random.uniform(-1,1)]) for _ in range(5)]
print(f' WE are about to take {Actions} \n')
next_state, reward, terminal, info = env.step(Actions)

print(f' NEXT STATE \n {next_state} \n')
print(f' REWARDS {reward} \n')
print(f' TERMINAL OR NOT >> {terminal} \n')
print(f' INFO {info}')


# obs_dict = env_reset(env)
# agent = OrderEnforcingAgent()
# print(agent.register_reset(obs_dict))
# env.step(agent.register_reset(obs_dict))

env.evaluate()

import numpy as np
import time

"""
Please do not make changes to this file. 
This is only a reference script provided to allow you 
to do local evaluation. The evaluator **DOES NOT** 
use this script for orchestrating the evaluations. 
"""

from agents.orderenforcingwrapper import OrderEnforcingAgent
from citylearn.citylearn import CityLearnEnv

class Constants:
    episodes = 5
    schema_path = './data/citylearn_challenge_2022_phase_1/schema.json'

def action_space_to_dict(aspace):
    """ Only for box space """
    return { "high": aspace.high,
             "low": aspace.low,
             "shape": aspace.shape,
             "dtype": str(aspace.dtype)
    }

def env_reset(env):
    observations = env.reset()
    action_space = env.action_space
    observation_space = env.observation_space
    building_info = env.get_building_information()
    building_info = list(building_info.values())
    action_space_dicts = [action_space_to_dict(asp) for asp in action_space]
    observation_space_dicts = [action_space_to_dict(osp) for osp in observation_space]
    obs_dict = {"action_space": action_space_dicts,
                "observation_space": observation_space_dicts,
                "building_info": building_info,
                "observation": observations }
    return obs_dict


def evaluate():
    print("Starting local evaluation")
    
    env = CityLearnEnv(schema=Constants.schema_path)
    agent = OrderEnforcingAgent()

    obs_dict = env_reset(env)

    agent_time_elapsed = 0

    step_start = time.perf_counter()
    actions = agent.register_reset(obs_dict)
    agent_time_elapsed += time.perf_counter()- step_start

    episodes_completed = 0
    num_steps = 0
    interrupted = False
    episode_metrics = []
    
    # Skipping to shorten training time
    days = 30*5
    training_steps = 24*days
    skipping = False
    
    try:
        while True:
            
            ### This is only a reference script provided to allow you 
            ### to do local evaluation. The evaluator **DOES NOT** 
            ### use this script for orchestrating the evaluations. 

            observations, _, done, _ = env.step(actions)
            if done or skipping:
                episodes_completed += 1
                metrics_t = env.evaluate()
                metrics = {"price_cost": metrics_t[0], "emmision_cost": metrics_t[1]}
                if np.any(np.isnan(metrics_t)):
                    raise ValueError("Episode metrics are nan, please contant organizers")
                episode_metrics.append(metrics)
                print(f"Episode complete: {episodes_completed} | Latest episode metrics: {metrics}", )

                obs_dict = env_reset(env)

                step_start = time.perf_counter()
                actions = agent.register_reset(obs_dict)
                agent_time_elapsed += time.perf_counter()- step_start
            else:
                step_start = time.perf_counter()
                actions = agent.compute_action(observations)
                agent_time_elapsed += time.perf_counter()- step_start
            
            num_steps += 1
            if num_steps % 1000 == 0:
                print(f"Num Steps: {num_steps}, Num episodes: {episodes_completed}")
            
            ### End training in set time
            if num_steps % training_steps == 0:
                print(f"Num Steps: {num_steps}, Num episodes: {episodes_completed}")
                if num_steps == training_steps:
                    print(f'ENDING TRAINING AFTER {training_steps} STEPS')
                    skipping = True

            if episodes_completed >= Constants.episodes:
                break
    except KeyboardInterrupt:
        print("========================= Stopping Evaluation =========================")
        interrupted = True
    
    if not interrupted:
        print("=========================Completed=========================")

    if len(episode_metrics) > 0:
        print("Average Price Cost:", np.mean([e['price_cost'] for e in episode_metrics]))
        print("Average Emmision Cost:", np.mean([e['emmision_cost'] for e in episode_metrics]))
    print(f"Total time taken by agent: {agent_time_elapsed}s")

if __name__ == '__main__':
    evaluate()