import numpy as np 
import pandas as pd 
import pdb 
import torch 
import random 
from scipy import stats 

from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier, GradientBoostingRegressor, RandomForestRegressor

from sklearn.metrics.pairwise import cosine_similarity

def seed_everything(seed):
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)

def logistic_pol_asgn(theta, x):
    n = x.shape[0]
    theta = theta.flatten()
    if len(theta) == 1:
        logit = np.multiply(x, theta).flatten()
    else:
        logit = np.dot(x, theta).flatten()
    LOGIT_TERM_POS = np.ones(n)*1.0 / ( np.ones(n) + np.exp( -logit ))
    return LOGIT_TERM_POS

def real_risk(T, beta_cons, beta_x, beta_x_T, x, u, alpha, w):
    n = len(T)
    risk = np.zeros(n)
    for i in range(len(T)):
        risk[i] = T[i] * beta_cons + np.dot(beta_x.T, x[i, :]) + np.dot(beta_x_T.T, x[i, :] * T[i]) + alpha * (u[i]) * ((2 * T[i] - 1)) + w * (u[i])
    return risk


def return_CATE_optimal_assignment(x, u, beta_cons, beta_x, beta_x_T, alpha,w):
    n = x.shape[0]
    risk_T_1 = real_risk(np.ones(n), beta_cons, beta_x, beta_x_T, x, u,alpha,w)
    risk_T_0 = real_risk(np.zeros(n), beta_cons, beta_x, beta_x_T,x, u,alpha,w)
    opt_T = [1 if risk_T_1[k] < risk_T_0[k] else 0 for k in range(n)]
    return opt_T

def REAL_PROP_LOG(x, u, beta_T_conf, beta_cons, beta_x, beta_x_T, Gamma = 3,alpha=-2,w=1.5):
    nominal_ = logistic_pol_asgn(beta_T_conf, x)
    a_bnd, b_bnd = get_bnds(nominal_, Gamma) # propensity's bound 
    q_lo = 1 / b_bnd;
    q_hi = 1 / a_bnd
    opt_T = return_CATE_optimal_assignment(x, u, beta_cons, beta_x, beta_x_T, alpha,w)
    q_real = np.asarray([q_hi[i] if opt_T[i] == 1 else q_lo[i] for i in range(len(u))])
    return q_real


def real_risk_prob(prob_1, x, u):
    n = len(u)
    prob_1 = np.asarray(prob_1)
    return prob_1 * real_risk_(np.ones(n), x, u) + (1 - prob_1) * real_risk_(np.zeros(n), x, u)

def get_bnds(est_Q,LogGamma):
    n = len(est_Q)
    if type(LogGamma) is np.array and len(LogGamma) > 1:
        p_hi = est_Q
        p_lo = est_Q
        for i in range(n):
            p_hi[i] = np.multiply(np.exp(LogGamma[i]), est_Q[i] ) / (1 - est_Q[i] + np.multiply(np.exp(LogGamma[i]), est_Q[i] ))
            p_lo[i] = np.multiply(np.exp(-LogGamma[i]), est_Q[i] ) / (1 - est_Q[i] + np.multiply(np.exp(-LogGamma[i]), est_Q[i] ))
    else:
        p_hi = np.multiply(np.exp(LogGamma), est_Q ) / (np.ones(n) - est_Q + np.multiply(np.exp(LogGamma), est_Q ))
        p_lo = np.multiply(np.exp(-LogGamma), est_Q ) / (np.ones(n) - est_Q + np.multiply(np.exp(-LogGamma), est_Q ))
    assert (p_lo <= p_hi).all()
    a_bnd = 1/p_hi;
    b_bnd = 1/p_lo
    return [ a_bnd, b_bnd ]

def get_bnds1(wtilde,LogGamma):
    Gamma = np.exp(LogGamma)
    a_bnd = 1 + (wtilde - 1) / Gamma
    b_bnd = 1 + (wtilde - 1) * Gamma
    return [ a_bnd, b_bnd ]


def generate_log_data_pl(mu_x, n, beta_cons, beta_x, beta_x_T, beta_T_conf, Gamma = [0,1,2], alpha = -2, w=1.5):
    # human behavior model 
    num_person = len(Gamma)
    # random select person 
    hid = np.array([np.random.choice(num_person) for i in range(n)])
    d = len(mu_x)
    u = (np.random.rand(n) > 0.5)  # noise described in the paper 
    x = np.zeros([n, d])
    for i in range(n):
        x[i, :] = np.random.multivariate_normal(mean=mu_x * (2 * u[i] - 1), cov=np.eye(d))
    x_ = np.hstack([x, np.ones([n, 1])])

    # generate propensities
    true_Q_h = np.zeros([n,len(Gamma)])
    T_h = np.zeros([n,len(Gamma)])
    for id, gam in enumerate(Gamma):
        true_Q_h[:,id] = REAL_PROP_LOG(x_, u, beta_T_conf, beta_cons, beta_x, beta_x_T, gam, alpha,w)
        T_h[:,id] = np.array(np.random.uniform(size=n) < true_Q_h[:,id]).astype(int).flatten()   

    true_Q = true_Q_h[range(true_Q_h.shape[0]),hid]

    T = T_h[range(T_h.shape[0]),hid]
    T = T.reshape([n, 1]).astype(int)
    T_sgned = np.asarray([1 if T[i] == 1 else -1 for i in range(n)]).flatten()
    clf = LogisticRegression();
    clf.fit(x, T)
    propensities = clf.predict_proba(x)
    nominal_propensities_pos = propensities[:, 1]
    nominal_propensities_pos = logistic_pol_asgn(beta_T_conf, x_)

    true_Q_obs = np.asarray([true_Q[i] if T[i] == 1 else 1 - true_Q[i] for i in range(n)])

    q0_all = np.zeros((n,2)) 
    q0_all[:,1] = nominal_propensities_pos
    q0_all[:,0] = 1 - q0_all[:,1]
    q0 = np.asarray([nominal_propensities_pos[i] if T[i] == 1 else 1 - nominal_propensities_pos[i] for i in range(n)])

    Y_all = np.zeros((n,2))
    for i in range(n):
        Y_all[i,0] = 0 * beta_cons + np.dot(beta_x.T, x_[i, :]) + np.dot(beta_x_T.T, x_[i, :] * 0) + alpha * (u[i]) * ((2 * 0 - 1)) + w * (u[i])
        Y_all[i,1] = 1 * beta_cons + np.dot(beta_x.T, x_[i, :]) + np.dot(beta_x_T.T, x_[i, :] * 1) + alpha * (u[i]) * ((2 * 1 - 1)) + w * (u[i])
     
    # add random noise
    T = T.flatten()
    Y_all += 2*np.random.normal(size=(n,2))  
    Y = Y_all[range(len(T)),T]
    return [x, u, T, Y, true_Q_obs, q0, Y_all, q0_all, hid, T_h]  


def sigmoid(x):
    return 1 / (1 + np.exp(-x))



