import sys
import numpy as np
from scipy.stats import chi2
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import log_loss
from scipy.io import loadmat
import pickle
import sklearn
import scipy.io as sio
import math
import os

def save_object(obj, filename):
    with open(filename, 'wb') as output:  # Overwrites any existing file.
        pickle.dump(obj, output, pickle.HIGHEST_PROTOCOL)
        
def load_object(filename):
    with open(filename, 'rb') as input:  # Overwrites any existing file.
        obj = pickle.load(input)
    return obj


def likelihood_ratio_test(features_alternate, labels, lr_model, features_null=None):
    """
    Compute the likelihood ratio test for a model trained on the set of features in
    `features_alternate` vs a null model.  If `features_null` is not defined, then
    the null model simply uses the intercept (class probabilities).  Note that
    `features_null` must be a subset of `features_alternative` -- it can not contain
    features that are not in `features_alternate`.
    Returns the p-value, which can be used to accept or reject the null hypothesis.
    """
    labels = np.array(labels)
    features_alternate = np.array(features_alternate)
    
    if features_null:
        features_null = np.array(features_null)
        
        if features_null.shape[1] >= features_alternate.shape[1]:
            raise(ValueError, "Alternate features must have more features than null features")
        
        lr_model.fit(features_null, labels)
        null_prob = lr_model.predict_proba(features_null)[:, 1]
        df = features_alternate.shape[1] - features_null.shape[1]
    else:
        null_prob = sum(labels) / float(labels.shape[0]) * \
                    np.ones(labels.shape)
        df = features_alternate.shape[1]
    
    lr_model.fit(features_alternate, np.ravel(labels))
    alt_prob = lr_model.predict_proba(features_alternate)

    alt_log_likelihood = -log_loss(labels,
                                   alt_prob,
                                   normalize=False)
    null_log_likelihood = -log_loss(labels,
                                    null_prob,
                                    normalize=False)

    G = 2 * (alt_log_likelihood - null_log_likelihood)
    p_value = chi2.sf(G, df)

    return p_value


id_fold    = math.floor(int(sys.argv[1])%10)

if( not os.path.exists('interactions/pretrain'+str(id_fold+1)+'/model')):
    os.mkdir('interactions/pretrain'+str(id_fold+1)+'/model')

for i in range(7):
    vis = loadmat('interactions/pretrain'+str(id_fold+1)+'/batch'+str(i)+'.mat')
    if i==0:
        X = vis['map']
        Y = vis['class']
    else:
        X = np.concatenate((X,vis['map']))
        Y = np.concatenate((Y,vis['class']))

      
loc = vis['indx'] 
       
r_path = load_object('remaining_paths3.pkl')

pid = 0

for p in range(0, len(r_path)):
    
    for path_n in r_path[p]:
        
        pp = np.load('../paths/path'+str(path_n)+'.npy')
        
        # Create feaure vector using the importance scores along a path.
        x = np.zeros((np.shape(X)[0],len(pp)-1))
        y = np.zeros((np.shape(X)[0],1))
        
        for pp_ind in range(np.shape(pp)[0]-1):
            idd = np.where(np.logical_and(loc[0]==pp[pp_ind+1], loc[1]==pp[pp_ind]))[0][0]
            x[:,pp_ind] = X[:,idd]
            
        y = Y.copy()
        
        # initialize model.
        model = sklearn.linear_model.LogisticRegression(penalty='none', max_iter=500)   
        
        # Perform likelihood test.
        pval = likelihood_ratio_test(x, y, model)
                
        # Save results.        
        n_vis = {'pval':pval, 'path':pp}
        sio.savemat('interactions/pretrain'+str(id_fold+1)+'/model/vis'+str(path_n)+'.mat',n_vis)
        
        
        
        
        
        
        
        
        