import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
# from DataGen_DUCB import DataGen
from DataGen import DataGen
from StoExp import StoExp
from DPFit import DPFit
from CausalBound import CausalBound
from D_UCB import DUCB
from JZ_bound import JZ_bound

<<<<<<< HEAD
# Parameter configuration
D = 3
N = 3000
T = 5000
seed_num = np.random.randint(10000000)
=======

D = 100
N = 20000
Ns = 20
T = 5000

Mode = 'crazy'
>>>>>>> origin/master

seed_num = np.random.randint(10000000)
# seed_num = 160702
# seed_num = 12345678
# seed_num = 2748152
# seed_num = 110628


# seed_num = 7584732 # Case 2, good seednum
# seed_num = 5647597 # Case 2, good seednum
# seed_num = 1551761 # for N = 10000, T = 3000
# seed_num = 7576682 # for N = 20000, T = 5000
# seed_num = 6597310 # for N = 20000, T = 5000 (Case 3)

# seed_num = 5391860 # Case 3



# Generating Observation data
datagen = DataGen(D,N,seed_num)
OBS = datagen.obs_data()

# Externally policy
stoexp = StoExp(D)
JZ = JZ_bound()
[X,Y,Z] = JZ.sepOBS(OBS,D)
obslogit = stoexp.Logit(X,Z)
obsxgb = stoexp.XGB(X,Z)

policy_list = [obslogit,obsxgb]

# Bound construction
JZ = JZ_bound()
[L_obslogit, H_obslogit] = JZ.JZ_bounds(obslogit,OBS,D,N)
Y_obslogit = np.mean(datagen.poly_intv_data(obslogit, Z)['Y'])

[L_obsxgb, H_obsxgb] = JZ.JZ_bounds(obsxgb,OBS,D,N)
Y_obsxgb= np.mean(datagen.poly_intv_data(obsxgb, Z)['Y'])

Y_pis = [Y_obslogit, Y_obsxgb]

Bdd = [[L_obslogit, H_obslogit],[L_obsxgb, H_obsxgb]]
pl_list = [obslogit, obsxgb]
opt_pl = np.argmax([Y_obslogit, Y_obsxgb])
subopt_pl = 1-opt_pl
opt_Ypi = Y_pis[opt_pl]
subopt_Ypi = Y_pis[subopt_pl]






#
# plt.figure('Intervention')
# Intv['Y'].plot(kind='density')
#
# plt.figure('Observation')
# Obs['Y'].plot(kind='density')


# Construct the list of policies


# # Generated by observational data
# Z_obs, X_obs = stoexp.Construct_XZ(Obs)
# obslogit = stoexp.Logit(X_obs,Z_obs)
# obssvmlin = stoexp.SVM_linear(X_obs, Z_obs)
# obssvmrbf = stoexp.SVM_rbf(X_obs, Z_obs)
# obslda = stoexp.LDA(X_obs, Z_obs)
# obsxgb = stoexp.XGB(X_obs, Z_obs)
#
# # Generated by 'truncated' observational data
# TrunObs = stoexp.Trun_by_Y(Obs, 0.8, 0.7)
# Z_tobs, X_tobs = stoexp.Construct_XZ(TrunObs)
# tobslogit = stoexp.Logit(X_tobs,Z_tobs)
# tobssvmlin = stoexp.SVM_linear(X_tobs, Z_tobs)
# tobssvmrbf = stoexp.SVM_rbf(X_tobs, Z_tobs)
# tobslda = stoexp.LDA(X_tobs, Z_tobs)
# tobsxgb = stoexp.XGB(X_tobs, Z_tobs)
#
# # TrunIntv
# TrunIntv = stoexp.Trun_by_Y(Intv,0.8, 0.5)
# Z_tintv, X_tintv = stoexp.Construct_XZ(TrunIntv)
# tintvlogit = stoexp.Logit(X_tintv,Z_tintv)
# tintvsvmlin = stoexp.SVM_linear(X_tintv, Z_tintv)
# tintvsvmrbf = stoexp.SVM_rbf(X_tintv, Z_tintv)
# tintvxgb = stoexp.XGB(X_tintv, Z_tintv)

# policy_list = [obslogit, obssvmlin, obssvmrbf, obslda,
#                tobslogit, tobssvmlin, tobssvmrbf, tobsxgb, tobslda,
#                tintvlogit, tintvsvmlin, tintvsvmrbf, tintvxgb
#                ]

# List of policies to be considered


# Number of policies
N_poly = len(policy_list)

# List of List of each policy's arm pulling
X_pl_list = []

# List of List of each policy's arm pulling result
Y_pl_list = []
for pl in policy_list: # For each policy pl
    # For all context Z, compute pl's arm choice
    X_pl = np.round(pl.predict(Z_obs),0)

    # For all context Z,  compute the result of pl's arm pull.
    Y_pl = datagen.gen_policy_Y(np.matrix(X_pl))

    # Store pl's arm pulling and result in the List.
    X_pl_list.append(X_pl)
    Y_pl_list.append(Y_pl)
    print(np.mean(Y_pl))

# Compute the upper bound
C_list = []
for pl in policy_list: # For each policy pl,
    # Consider obsxgb is a P(x|z) from observational data
    # Compute the cross-entropy between P(x|z) and pl.
    c_pl = stoexp.Compute_C(obsxgb,pl,Z_obs)
    C_list.append(c_pl)

# Store the result of all Pl's arm pulling in the DataFrame.
Y_policy = pd.DataFrame(np.matrix(np.array(Y_pl_list).T))

# Initial true DPMM for obs
Yobs = Obs['Y']
init_compo = 50
DPFit_obs = DPFit(Yobs, init_compo)
DPFit_obs.Conduct()
dpobs = DPFit_obs.dpgmm
init_compo = sum(1 for x in dpobs.weights_ if x > 1e-2)
DPFit_obs = DPFit(Yobs, init_compo)
DPFit_obs.Conduct()
dpobs = DPFit_obs.dpgmm

iter_opt = 100
Bound_list = []
Model_list = []
Weight_list = []
True_Mu = []

for pl in range(len(policy_list)): # For each policy pl,
    # P(Y)|pl // Distribution of Pl's Y
    Yintv_pl = Y_policy[pl]

    # Cross entropy
    C_pl = C_list[pl]

    # Min/Max of mu estimate
    CB = CausalBound(dpobs, C_pl)

    # Lower and Upper bound of estimate from the observation
    LB, UB, min_mu_list, max_mu_list, min_weight_list, max_weight_list = CB.Conduct_Optimization(C_pl, init_compo, iter_opt)

    # Collection of [LB,UB] of policy pl
    Bound_list.append([LB,UB])
    # Model_list.append([lower,upper])
    # Weight_list.append([lower_weight, upper_weight])

    # Unknown true mu estimate
    True_Mu.append(np.mean(Yintv_pl))

print(Bound_list)
print(True_Mu)


''' From UCB '''
K = 1
ducb = DUCB(Bound_list,policy_list,Z_obs,Y_pl_list,X_pl_list,K,T)
[UCB_result, BUCB_result] = ducb.Bandit_Run()
prob_opt_list, cum_regret_list, Sto_pick, mu_k_dict = UCB_result
prob_opt_list_B, cum_regret_list_B, Sto_pick_B, mu_k_dict_B = BUCB_result

print(prob_opt_list)
print(prob_opt_list_B)

print(cum_regret_list)
print(cum_regret_list_B)

# font = {'family' : 'normal',
#         'weight' : 'bold',
#         'size'   : 22}

#
plt.figure()
plt.rc('font', size=40)          # controls default text sizes
plt.rc('xtick', labelsize=25)    # fontsize of the tick labels
plt.rc('ytick', labelsize=25)    # fontsize of the tick labels
plt.rc('legend', fontsize=30)    # legend fontsize

plt.title('Case 2')
plt.ylabel('Cumulative regret')
plt.xlabel('Trials')
cum_UCB = plt.plot(cum_regret_list,label='D-UCB')
cum_BUCB = plt.plot(cum_regret_list_B,label='B-D-UCB')
plt.setp(cum_UCB, linewidth=5)       # set both to dashed
plt.setp(cum_BUCB, linewidth=5)       # set both to dashed
plt.legend()

