import numpy as np
import pandas as pd


import sys
import os
import yaml

# Temporary fix to to make agents module available, will probably need to turn src into a python package
# so that it can be imported from anywhere
sys.path.insert(0, os.path.abspath('src'))


from agents import two_stage_thompson_sampling #, thomspon_sampling
from Environment import Environment
from simulator import BaseSimulator, TwoStageSimulator

try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader

yaml_parms = {}
context = open("./configs/poster_context.yaml", 'r')
yaml_parms["context"] = yaml.load(context, Loader=Loader)
outcome_model = open("./configs/poster_outcome_model.yaml", 'r')
yaml_parms["outcome_model"] = yaml.load(outcome_model, Loader=Loader)


config_file = open("configs/poster_agent.yaml", 'r')
config = yaml.load(config_file, Loader=Loader)

# Simulate a bunch of prompts


# curr_loc = ["home", "work", "other"]
# prev_steps = np.arange(1,5)
# prompts = np.arange(1,4)
# prompt_df_base = pd.DataFrame([(loc, steps, prompt) for loc in curr_loc for steps in prev_steps for prompt in prompts])
# prompt_df_base.columns = ["curr_loc", "prev_steps", "prompt"]

# dimensions = ["optimism", "gender", "formality", "clarity", "urgency"]
# prompt_centers = {
#     1 : {
#         "optimism" : 3,
#         "gender" : 3,
#         "formality" : 3,
#         "clarity": 3,
#         "urgency": 3
#     },
#     2 : {
#         "optimism" : 3.5,
#         "gender" : 3.25,
#         "formality" : 1,
#         "clarity": 1,
#         "urgency": 1
#     },
#     3 : {
#         "optimism" : .5,
#         "gender" : .5,
#         "formality" : 1,
#         "clarity": 3.5,
#         "urgency": 3.25
#     }
# }

#Rudimentary way of just simulating out toy text dimensions
# def get_dim_mean(row, dim, prompt_centers):
#     return(prompt_centers[row["prompt"]][dim] + .1*row["prev_steps"])

# def sim_prompts(base_df, dims, dim_center_by_prompt, noise_sd):
#     output_df = base_df
#     prompts = np.unique(base_df["prompt"])
#     for dim in dims:
#         output_df.loc[:, dim] = output_df.apply(get_dim_mean, axis=1, args=(dim, dim_center_by_prompt))
#     output_df.loc[:, dims] = output_df.loc[:, dims] + np.random.normal(0, noise_sd, output_df.loc[:, dims].shape)
#     return(output_df)

# prompts_to_sim = 1000
# prompt_db = sim_prompts(prompt_df_base, dimensions, prompt_centers, .25)

# for i in range(1, prompts_to_sim):
#     prompt_db = pd.concat([prompt_db, sim_prompts(prompt_df_base, dimensions, prompt_centers, .25)])

prompt_db = pd.read_csv("data/sim_params/prompt_db_w_vae.csv")


env = Environment(prompt_db, yaml_parms)

K = prompt_db.prompt.nunique()
# Remove context / prompt columns to get number of dimensions in db
d = len(set(prompt_db.columns) - set(['prompt', 'prev_steps', 'curr_loc']))
agent = two_stage_thompson_sampling.TwoStageThompsonSampling(prompt_db, config)

simulator = TwoStageSimulator(agent =  agent, env  = env, n_steps = 5, unique_id = 1)

simulator.run_experiment()

print(simulator.results)