import numpy as np
import pandas as pd
from env import Env
import datetime


idx_start = 10  # first index is 1 !

n_sample = 30

df_xv = pd.read_csv("xv.csv", header=None, index_col=None).transpose()
lst_xv_keys = [
    'alpha',
    'percent_search',
    'is_RH_ratio',
    'is_FH_ratio'
]
df_xv.columns = lst_xv_keys

df_oa = pd.read_fwf("oa25.txt", header=None, widths=[1]*6)
n_row, n_col = df_oa.shape
model_para = {
    "n_worker": 1000,
    "n_firm": 100
}

lst_op_key = ['out_w_avg_salary', 'out_w_gini_salary', 'out_f_avg_profit', 'out_f_gini_profit', 'out_w_percent_hired']
lst_2d_op_avg = []
lst_2d_xv = []
for idx_row in range(idx_start-1, n_row, 1):
    print(f"Running the {idx_row + 1}-th experiment at {datetime.datetime.now()}.")
    lst_value = []
    for idx_col, k in enumerate(lst_xv_keys):
        oa_idx_level = int(df_oa.iat[idx_row, idx_col])
        lst_value.append(df_xv.iloc[oa_idx_level][k])
    lst_2d_xv.append(lst_value)
    dct_merge = {**model_para, **dict(zip(lst_xv_keys, lst_value))}
    # pprint.pprint(dct_merge)
    lst_2d_op = []

    for i in range(n_sample):
        print(f"-- {i + 1}-th sample at {datetime.datetime.now()}.")
        the_model = Env(dct_merge)

        while 1:
            the_model.step()
            if not the_model.running:
                break

        lst_2d_op.append([the_model.out_w_avg_salary, the_model.out_w_gini_salary,
                          the_model.out_f_avg_profit, the_model.out_f_gini_profit,
                          the_model.out_w_percent_hired])

    arr_op = np.array(lst_2d_op)
    lst_2d_op_avg.append(arr_op.mean(axis=0).tolist())

    # these codes below should be outside of loop. but temply inside
    arr_op_avg = np.array(lst_2d_op_avg)

    df_final = pd.concat([pd.DataFrame(lst_2d_xv, columns=lst_xv_keys).reset_index(drop=True),
                          df_oa.iloc[:, len(lst_xv_keys):].reset_index(drop=True)], axis=1)

    df_final = pd.concat([df_final.reset_index(drop=True),
                          pd.DataFrame(lst_2d_op_avg, columns=lst_op_key).reset_index(drop=True)], axis=1)

    df_final.to_excel(f'result/experiment_result_{idx_row + 1}.xlsx')
    df_final.to_csv(f'result/experiment_result_{idx_row + 1}.csv')