
import matplotlib.pyplot as plt
import numpy as np
import cvxpy as cvx
from gurobipy import *
from scipy.optimize import minimize
import datetime as datetime
import pickle
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib import collections as matcoll
from sklearn import svm
from models.crlogit.methods import *

import os
import sys
module_path = os.path.abspath(os.path.join('/Users/az/Box Sync/unconfoundedness'))
if module_path not in sys.path:
    sys.path.append(module_path)
# from scripts_running import *
from models.crlogit.unconfoundedness_fns import *
# from evals import *
# from greedy_partitioning import *
from models.crlogit.subgrad import *
from scipy.stats import norm
from scipy.stats import multivariate_normal
from copy import deepcopy

np.random.seed(42)

d = 5  # dimension of x
n = 800;
# parameters
rho = np.asarray([1 / np.sqrt(2), -1 / np.sqrt(2), 0, 0, 1 / np.sqrt(3)])  # normalize to unit 0.5
rho = rho / (np.dot(rho, rho) * 2)

beta_cons = 2.5
beta_x = np.asarray([0, .5, -0.5, 0, 0, 0])
beta_x_T = np.asarray([-1.5, 1, -1.5, 1., 0.5, 0])
beta_T = np.asarray([0, .75, -.5, 0, -1, 0, 0])
beta_T_conf = np.asarray([0, .75, -.5, 0, -1, 0])
mu_x = np.asarray([-1, .5, -1, 0, -1]);

alpha = -2
w = 1.5

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


# generate propensity model
def REAL_PROP_LOG(x, u):
    Gamma = 1.5
    print(x.shape)
    nominal_ = logistic_pol_asgn(beta_T_conf, x)
    #     return nominal_
    # set u = I[ Y(1)\mid x > Y(0) \mid x ]
    a_bnd, b_bnd = get_bnds(nominal_, Gamma)
    q_lo = 1 / b_bnd;
    q_hi = 1 / a_bnd
    opt_T = return_CATE_optimal_assignment(x, u)
    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(T, beta_cons, beta_x, beta_x_T, x, u):
    '''
    takes as input integral T
    '''
    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 real_risk_(T, x, u):
    return real_risk(T, beta_cons, beta_x, beta_x_T, x, u)

def real_risk_T_integer(T, x, u):
    T = get_sgn_0_1(T)
    return real_risk(T, beta_cons, beta_x, beta_x_T, x, u)


'''
takes in prob_1, n x m array of treatment assignment probabilities 
compute 
\sum_t 1/n \sum_i Y(t)\pi_i(t) 
'''
def real_risk_prob_oracle(prob_t, x, u):
    n_ts = prob_t.shape[1]
    risk = 0
    for t in range(n_ts):
        risk += np.mean( np.multiply(prob_t[:,t], real_risk_(t*np.ones(x.shape[0]), x, u) )  )
    return risk

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 generate_log_data(mu_x, n, beta_cons, beta_x, beta_x_T):
    # generate n datapoints from the same multivariate normal distribution
    d = len(mu_x)
    u = (np.random.rand(n) > 0.5)
    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 = REAL_PROP_LOG(x_, u)
    T = np.array(np.random.uniform(size=n) < true_Q).astype(int).flatten()
    T = T.reshape([n, 1]);
    T_sgned = np.asarray([1 if T[i] == 1 else -1 for i in range(n)]).flatten()
    nominal_propensities_pos = logistic_pol_asgn(beta_T_conf, x_)

    q0 = np.asarray([nominal_propensities_pos[i] if T[i] == 1 else 1 - nominal_propensities_pos[i] for i in range(n)])
    true_Q_obs = np.asarray([true_Q[i] if T[i] == 1 else 1 - true_Q[i] for i in range(n)])
    Y = np.zeros(n)
    for i in range(n):
        Y[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]
    # add random noise
    T = T.flatten()
    Y += 1 * np.random.randn(n)
    return [x_, u, T, Y, true_Q_obs, q0]


# generate propensity model
def REAL_PROP_LOG(x, u):
    Gamma = 1.5
    print(x.shape)
    nominal_ = logistic_pol_asgn(beta_T_conf, x)
    #     return nominal_
    # set u = I[ Y(T)\mid x > Y(-T) \mid x ]
    a_bnd, b_bnd = get_bnds(nominal_, Gamma)
    q_lo = 1 / b_bnd;
    q_hi = 1 / a_bnd
    opt_T = return_CATE_optimal_assignment(x, u)
    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_mt(x_, u, T, alpha, beta_tilde, beta_x_T,  eta_tilde, eta):
    n = len(T); risk = np.zeros(n)
    for i in range(len(T)):
        risk[i] = alpha[T[i]] + np.dot(beta_tilde.T, x_[i, :]) + np.dot(beta_x_T[:,T[i]].T, x_[i, :]) + eta[T[i]]*(u[i]) + eta_tilde*u[i]
    return risk

def real_risk_mt_(x_, u, T):
    return real_risk_mt(x_, u, T,*outcome_dgp_params)

'''
takes in prob_1, n x m array of treatment assignment probabilities 
compute 
\sum_t 1/n \sum_i Y(t)\pi_i(t) 
'''
def real_risk_prob_oracle_mt(prob_t, x, u, *outcome_dgp_params):
    n_ts = prob_t.shape[1]
    risk = 0
    for t in range(n_ts):
        risk += np.mean( np.multiply(prob_t[:,t], real_risk_mt_(x, u, t*np.ones(n).astype(int)) )  )
    return risk
#     return prob_1 * real_risk_T_integer(np.ones(n), x, u) + (1 - prob_1) * real_risk_(np.zeros(n), x, u)


def return_CATE_optimal_assignment_mt(x, u, n_ts, *outcome_dgp_params):
    n = x.shape[0]
    risk = np.zeros([n,n_ts])
    for k in range(n_ts):
        risk[:,k] = real_risk_mt(x, u, k*np.ones(n).astype(int), *outcome_dgp_params)
    opt_T = [np.argmin(risk[i,:]) for i in range(n)]
    return opt_T

''' real propensity that generates treatment assignments for multiple treatments 
'''
def real_propensity_mt(x, u, beta_T_conf_prop, log_gamma, *outcome_dgp_params ):
    nominal_ = logistic_pol_asgn_mt(beta_T_conf_prop, x)
    # set u = I[ Y(T)\mid x > Y(-T) \mid x ]
    opt_T = return_CATE_optimal_assignment_mt(x, u, beta_T_conf_prop.shape[1], *outcome_dgp_params)
    n_ts = beta_T_conf_prop.shape[1]
    q_lo = np.zeros([x.shape[0], n_ts]); q_hi = np.zeros([x.shape[0], n_ts])
    for k in range(n_ts):
        a_bnd, b_bnd = get_bnds(nominal_[:,k], log_gamma)
        q_lo[:,k] = 1 / b_bnd;
        q_hi[:,k] = 1 / a_bnd
    q_real = deepcopy(nominal_)
    for i in range(x.shape[0]):
        if opt_T[i] == 1:
            q_real[i,1] = q_hi[i, 1 ]
            q_real[i,2] = nominal_[i,2]
            q_real[i,0] = 1 - q_real[i,1] - q_real[i,2]
        else:
            q_real[i,:] = nominal_[i,:]
    # clip and renormalize
    q_real = np.clip(q_real, 0.01, 0.99)
    for i in range(x.shape[0]):
        q_real[i,:] = q_real[i,:] / sum(q_real[i,:])
    return q_real


'''
# take in a vector of
beta_cons: 
beta_tilde: \tilde{\beta} (confounding interaction with x) 
beta_x_T: [d x k] array of coefficient vectors, one for each treatment

'''
def generate_log_data_mt(mu_x, n, beta_T_conf_prop, n_ts, log_gamma, alpha, beta_tilde, beta_x_T,  eta_tilde, eta):
    # generate n datapoints from the same multivariate normal distribution
    outcome_dgp_params = [alpha, beta_tilde, beta_x_T,  eta_tilde, eta ]
    d = len(mu_x);
    u = (np.random.rand(n) > 0.5)  # Bernoulli noise realization;
    x = np.zeros([n, d])
    for i in range(n):
        x[i, :] = np.random.uniform(low=np.ones(d)*-3,high = np.ones(d)*3)

        # 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 = real_propensity_mt(x_, u, beta_T_conf_prop, log_gamma, *outcome_dgp_params ) # in n x k for multiple treatments
    T = np.asarray([np.random.choice( list(range(n_ts)), p = true_Q[i,:] ) for i in range(n) ]).astype(int).flatten()
    q0 = logistic_pol_asgn_mt_obsA(beta_T_conf_prop, x_, T)
    true_Q_obs = np.asarray([true_Q[i,T[i]] for i in range(n)])
    Y = np.zeros(n)
    for i in range(n):
        Y[i] = alpha[T[i]] + np.dot(beta_tilde.T, x_[i, :]) + np.dot(beta_x_T[:,T[i]].T, x_[i, :]) + eta[T[i]]*(u[i]) + eta_tilde*u[i]
    # add random noise
    Y += np.random.randn(n)  # + np.random.randn()*np.asarray(T)*2 ;
    return [x_, u, T, Y, true_Q_obs, q0]
