#!/usr/bin/env python3
##import dependency libraries
import sys,time,glob,os,pickle,fileinput, random, h5py
from subprocess import Popen
from operator import itemgetter
import gc, fileinput
import numpy as np
from statsmodels.nonparametric.smoothers_lowess import lowess
from sklearn.neighbors import RadiusNeighborsRegressor as neighbors
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
import argparse

try:
    from pyminer.common_functions import *
except:
    from common_functions import *

##############################################################

def dispersion_main(args):
    if args.out_dir[-1]!='/':
    	args.out_dir+='/'
    if not os.path.isdir(args.out_dir):
        process_dir(args.out_dir)
    ##############################################################
    gc.enable()
    ## read in the input file
    if not args.hdf5:
        in_mat_str = np.array(read_table(args.infile))
        in_mat = np.array(in_mat_str[1:,1:],dtype=float)
        ID_list = in_mat_str[1:,0].tolist()
        del in_mat_str
    else:
        ## read in the hdf5 file
        print('reading hdf5 file:\n\t',args.infile)
        h5f = h5py.File(args.infile, 'r')
        in_mat=h5f["infile"]
        ID_list = read_file(args.ids,'lines')




    if args.log:
        if np.min(in_mat)<0.0:
            print("can't log transform with negative values... Bringing the matrix up to min = 1")
            in_mat+=min(in_mat)
        in_mat = np.log2(1 + in_mat)



    ## calculate the means
    means = np.mean(in_mat,axis=1)
    ## calculate the coefficient of variance
    variance = np.std(in_mat,axis=1)#/(means+1)
    ## square the coefficient of variance
    variance*=variance

    ## remove the nans

    print(np.sum(np.isnan(variance)))
    variance[np.isnan(variance)]=0
    print(np.sum(np.isnan(variance)))


    linear_dim = np.shape(means)


    ## plot (a random sample of) the relationship between the coefficient of variance and means
    max_sample_size = 50000
    sample_size = min([max_sample_size,linear_dim[0]])
    if sample_size < max_sample_size:
    	print('sampling the distribution for',sample_size,'points')
    	index_vect  = np.arange(linear_dim[0])
    	np.random.shuffle(index_vect)
    	sample = index_vect[:sample_size].tolist()
    	## always include the min and max values for interpolation
    	x_max_index = np.where(np.array(means) == np.max(means))[0][0]
    	x_min_index = np.where(np.array(means) == np.min(means))[0][0]
    	#print(sample)
    	#print(type(sample))
    	#sample = sample.tolist()
    	if x_max_index not in sample:
    		sample.append(x_max_index)
    	if x_min_index not in sample:
    		sample.append(x_min_index)
    	#print(sample)
    	print(type(sample))
    else:
        sample = np.arange(linear_dim[0], dtype =  int).tolist()
        #print(sample)
        #sys.exit()

    sample = np.sort(sample)
    ## remove the indices where the variance is zero
    keep_sample = np.where(variance[sample]>0)[0]
    sample = sample[keep_sample]

    plt.scatter(means[sample], variance[sample], c = 'black', s = 0.5)#,xlab = 'mean',ylab = 'CV')
    #plt.show()
    #sys.exit()

    ## do the lowess fit. This returns the expected values of the variance
    print('calculating the best fit curve')

    ## 

    neigh = neighbors(radius = 1.0, weights = 'uniform', leaf_size = 30)
    mean_max = max(means)
    mean_min = min(means)
    sds_min = min(variance)

    ## first calculate the lowess curve on the sample
    lowess_estimates_sample = lowess(variance[sample]-sds_min+1,means[sample]-mean_min+1, delta = 0.01*mean_max)-1
    # lowess_estimates_sample[:,0] = lowess_estimates_sample[:,0] + mean_min
    # lowess_estimates_sample[:,1] = lowess_estimates_sample[:,1] + sds_min

    ## remove the nans 
    lowess_estimates_sample[np.isnan(lowess_estimates_sample)]=0

    print(lowess_estimates_sample)

    ## remove x ties from the lowess as this messes up the interpolation
    values, indices = np.unique(lowess_estimates_sample[:,0], return_index=True)
    lowess_estimates_sample = np.array(lowess_estimates_sample[indices,:])
    print(lowess_estimates_sample)
    print(type(lowess_estimates_sample))
    #sys.exit()


    # ## positive control
    # x = np.linspace(0, 10, num=11, endpoint=True)
    # print(x)
    # y = np.cos(-x**2/9.0)
    # print(y)
    # f2 = interp1d(x, y, kind='cubic')
    # print(f2)
    # xnew = np.linspace(0, 10, num=41, endpoint=True)
    # print(xnew)
    # print(f2(xnew))

    print(np.shape(lowess_estimates_sample[:,0]))
    print(np.shape(lowess_estimates_sample[:,1]))

    ## cross the spread
    print("original_min_max:",min(lowess_estimates_sample[:,0]),max(lowess_estimates_sample[:,0]))


    epsilon = (max(lowess_estimates_sample[:,0]) - min(lowess_estimates_sample[:,0]))*0.01
    print('epsilon:',epsilon)

    #make the new x values evenly spaced a
    temp_min = min(lowess_estimates_sample[:,0])# - epsilon
    temp_max = max(lowess_estimates_sample[:,0])# + epsilon
    new_x = np.linspace(temp_min,temp_max, num=sample_size, endpoint = True)
    #new_x = np.linspace(temp_min,temp_max, num=sample_size, endpoint = True)
    #new_x[0] = temp_min
    #new_x[-1] = temp_max

    interpolation = interp1d(lowess_estimates_sample[:,0], lowess_estimates_sample[:,1], kind='cubic')


    print("new_min_max:",min(new_x),max(new_x))
    print(new_x)

    ## interpolate the lowess function
    new_y = interpolation(new_x)
    print(new_x,new_y)

    ## get the interplated lowess ready for nearest neighbor regression
    train_x = np.array([[x] for x in lowess_estimates_sample[:,0].tolist()])
    train_y = lowess_estimates_sample[:,1]

    full_x = np.array([[x] for x in means.tolist()])
    #train_x = full_x[sample]
    #train_y = variance[sample]

    print(train_x)
    print(train_y)
    #neigh.fit(train_x,train_y)
    print(new_x)
    print(new_y)
    new_x = np.array([[x] for x in new_x.tolist()])

    #plt.clf()
    plt.scatter(new_x,new_y, c = 'red', s = 0.75)
    plt.savefig(args.out_dir+'local_fit.png',dpi=600,bbox_inches='tight')
    #plt.show()
    #sys.exit()
    neigh.fit(new_x,new_y)

    ## now calculate the values for everything else
    bin_size = 100000
    total_vars = len(means)
    bins = []
    cur_bin = 0
    while cur_bin<total_vars:
    	bins.append(min(cur_bin, total_vars))
    	cur_bin+=bin_size

    bins.append(total_vars)
    print(bins)

    lowess_estimates = np.zeros((len(means),))
    # for i in range(1,len(bins)):
    # 	print("working on",bins[i-1],bins[i])
    # 	# if i%50000 == 0:
    # 	# 	print('\t',bins[i]/len(means))
    # 	lowess_estimates[bins[i-1]:bins[i]] += neigh.predict(full_x[bins[i-1]:bins[i]])

    for i in range(0,total_vars):
        x_dif_vect = np.abs(full_x[i]-new_x)
        closest_idx = int(np.argmin(x_dif_vect))
        lowess_estimates[i]=new_y[closest_idx]
        



    #lowess_estimates = neigh.predict(full_x)

    #lowess_estimates = lowess(variance[sample]+1,means[sample]+1)-1#, return_sorted = False)#, frac = 1e-3, delta = 1e-10*max(means))
    #lowess_estimates = [x[0] for x in lowess_estimates]


    #plt.scatter(means[sample],lowess_estimates[sample],c='red',s=0.5)
    #plt.savefig(args.out_dir+'local_fit_sample.png',dpi=600,bbox_inches='tight')
    #plt.show()
    #sys.exit()
    ## calculate the residuals
    print('getting the residuals')
    residuals = variance - lowess_estimates


    # print(lowess_estimates)
    # # sys.exit()
    # print(residuals - variance)
    # sys.exit()
    print('sample_means',means[:10])
    print('sample_estimates',lowess_estimates[:10])
    print('sample_CVs',variance[:10])
    print('sample_residuals',residuals[:10])


    ## reshape it back into the table
    resid_table = residuals

    print(np.sum(np.isnan(residuals)))
    nan_idx = np.where(np.isnan(residuals))[0]

    print('nan means',means[nan_idx])
    print('nan sds',variance[nan_idx])
    print('nan estimates',lowess_estimates[nan_idx])
    print('nan resid',residuals[nan_idx])


    #print(np.nan_std(residuals))
    st_dev_resid = np.std(residuals)
    print('sd residuals:',st_dev_resid)

    st_dev_cutoff = st_dev_resid*args.z_cutoff
    print('residual cutoff for overdispersion:',st_dev_cutoff)

    ## plot the residuals
    plt.clf()
    plt.scatter(means[sample],residuals[sample], c = 'black', s = 0.5)#,xlab = 'mean',ylab = 'CV')
    plt.plot([min(means),max(means)],[0,0],c='red')
    ## find the significant residuals
    resid_sample = residuals[sample]
    mean_sample = means[sample]
    sig_indices = np.where(resid_sample >= st_dev_cutoff)[0]
    plt.scatter(mean_sample[sig_indices], resid_sample[sig_indices], c = 'blue', s = 2)
    plt.savefig(args.out_dir+'residuals.png',dpi=600,bbox_inches='tight')
    #plt.show()


    ## determine which genes are locally overdispersed
    print('calculating local overdispersion')
    local_overdispersion_bool_table = resid_table >= st_dev_cutoff

    print('number of overdispersed genes:')
    number_overdispersed_per_group = np.sum(local_overdispersion_bool_table, axis = 0)
    print(number_overdispersed_per_group)



    ## get the boolean overdispersed table ready for writing to file
    #print(np.shape(local_overdispersion_bool_table))
    #print(local_overdispersion_bool_table)
    local_overdispersion_bool_table = np.array(local_overdispersion_bool_table,dtype = str)
    #print(local_overdispersion_bool_table)
    local_overdispersion_bool_table=local_overdispersion_bool_table.tolist()
    print(local_overdispersion_bool_table[:5],len(local_overdispersion_bool_table))
    print(np.shape(means)[0])
    ## add the gene names
    gene_names = ID_list
    for i in range(0,(np.shape(means)[0])):
        #print(i)
        local_overdispersion_bool_table[i] = [gene_names[i]] + [local_overdispersion_bool_table[i]]


    local_overdispersion_bool_table=np.array(local_overdispersion_bool_table, dtype = str)
    print(local_overdispersion_bool_table)
    print(np.shape(local_overdispersion_bool_table))

    write_table(local_overdispersion_bool_table,args.out_dir+'/locally_overdispersed_boolean_table.txt')




if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-infile","-i", 
        help="the input expression matrix",
        type = str)
    parser.add_argument("-out_dir", '-out','-o',
        help="the directory for output files",
        type = str)
    parser.add_argument("-log", 
        help="log transform the means",
        action = 'store_true',
        default = False)
    parser.add_argument("-hdf5", 
        help="the input file is an hdf5 file",
        action = 'store_true',
        default = False)
    parser.add_argument("-ids", '-ID_list','-id_list',
        help="if an hdf5 file is used, we'll need the list of IDs",
        type = str)
    parser.add_argument("-z_cutoff", "-z",
        help="the cutoff for how many standard deviations over the mean residual should be considered overdispersed. Default = 0.5",
        type = float,
        default = .5)
    args = parser.parse_args()
    dispersion_main(args)