import numpy as np
from scipy.stats import logistic
import scipy.optimize

from collections import Counter
from nneighbor import *
from ai_matching import *

import sys

def logLikelihoodLogit(vTheta, mX, vY):
	"""
	Copmutes Log Likelihood functino for logistic model

	Arguments:
		vTheta, 	1d array of size k, vector of estimates
		mX, 		2d array of size (N,k), covariates
		vY, 		1d array of size N, outcomes
	"""
    test = 	vY * ( np.dot(mX, vTheta) - np.log( 1.0 + np.exp(np.dot(mX, vTheta)) ) ) \
    		+ (1-vY) * ( -np.log( 1.0 + np.exp(np.dot(mX, vTheta)) ) )
    
    return -np.sum(test)

def index_wide(vW, index_set):
	"""
	For each index of a control subject, this function finds the index of its closest
	n treatment neighbors.

	Arguments:
		vW, 		1d array of size N, treatment indicator variable
		index_set, 	1d array of size N1, set of indices of control and matched treatment subject

	"""
	N1 = Counter(vW)[0]
	N = vW.shape[0]
	result = [ [i]+index_set[i].tolist() for i in range(N) if vW[i] == 0]
	result = [val for sublist in result for val in sublist]
	return result

# set seed
np.random.seed(6000)

# magic numbers for simulation
N = 20000
k = 10

# generate outcome
vY_control = np.random.rand(N)
vY_treatment = 1.3 * np.random.rand(N)
vY = variables_mix(vY_control,vY_treatment)

# prepare covariates
mX = np.concatenate((np.ones((N,1)), np.random.rand(N,k-1)), axis=1)

# true theta
vTheta_true = np.random.uniform(low=-1, high=1, size=(k,))

# treatment
p = logistic.cdf(np.dot(mX, vTheta_true), loc=0, scale=1)
vW = np.random.binomial(1, p, size=(N,))

# estimate theta
vTheta_init = np.ones(k)
vTheta_est = scipy.optimize.fmin_bfgs(logLikelihoodLogit, vTheta_init, args = (mX, vW), gtol = 1e-3)

print "True Theta: \n", vTheta_true, "\n"
print "Estimated Theta: \n", vTheta_est, "\n"

# matching
vPscore = logistic.cdf(np.dot(mX, vTheta_est), loc=0, scale=1)
vMatchIndex = index_set(vW, vPscore, 1, False, False)
index_ordering = index_wide(vW, vMatchIndex)

print "Number of treated municipalities: \n", Counter(vW)[0], "\n"
print "Number of both control and treatment municipalities: \n", len(index_ordering), "\n"

# update
vY = vY[index_ordering]
vW = vW[index_ordering]
mX = mX[index_ordering,:]
vPscore = vPscore[index_ordering]
index = range(len(index_ordering))

# abadie imbens
sigma_hat_sq, std_err = ai_matching(index, vY, vW, mX, vPscore, 2, 1)

print "Value of std_err: ", std_err
