import pandas
import time
import numpy as np
import Bio.SeqUtils as SeqUtil
import Bio.Seq as Seq
import util
import sys
import Bio.SeqUtils.MeltingTemp as Tm
import pickle
import itertools
import sklearn

def featurize_data(data, learn_options, Y, gene_position):
    '''
    assumes that data contains the 30mer
    returns set of features from which one can make a kernel for each one
    '''

    #print "Constructing features..."
    t0 = time.time()

    feature_sets = {}

    if learn_options["nuc_features"]:
        # spectrum kernels (position-independent) and weighted degree kernels (position-dependent)
        get_all_order_nuc_features(data['30mer'], feature_sets, learn_options, learn_options["order"], max_index_to_use=30)

    if learn_options["gc_features"]:
        gc_above_10, gc_below_10, gc_count = gc_features(data)
        feature_sets['gc_above_10'] = pandas.DataFrame(gc_above_10)
        feature_sets['gc_below_10'] = pandas.DataFrame(gc_below_10)
        feature_sets['gc_count'] = pandas.DataFrame(gc_count)

    if learn_options["include_gene_position"]:
        # gene_position_columns = ["Amino Acid Cut position", "Percent Peptide", "Nucleotide cut position"]
        # gene_position_columns = ["Percent Peptide", "Nucleotide cut position"]

        for set in gene_position.columns:
            set_name = set
            feature_sets[set_name] = pandas.DataFrame(gene_position[set])
        feature_sets["Percent Peptide <50%"] = feature_sets["Percent Peptide"] < 50

    if learn_options["include_gene_effect"]:
        print "including gene effect"
        gene_names = Y['Target gene']
        enc = sklearn.preprocessing.OneHotEncoder()
        label_encoder = sklearn.preprocessing.LabelEncoder()
        label_encoder.fit(gene_names)
        one_hot_genes = np.array(enc.fit_transform(label_encoder.transform(gene_names)[:, None]).todense())
        feature_sets["gene effect"] = pandas.DataFrame(one_hot_genes,
                                                       columns=["gene_%d" % i for i in range(one_hot_genes.shape[1])], index=gene_names.index)

    if learn_options['include_known_pairs']:
        feature_sets['known pairs'] = pandas.DataFrame(Y['test'])

    if learn_options["include_NGGX_interaction"]:
        feature_sets["NGGX"] = NGGX_interaction_feature(data)

    if learn_options["include_Tm"]:
        feature_sets["Tm"] = Tm_feature(data)

    if learn_options["include_sgRNAscore"]:
        feature_sets["sgRNA Score"] = pandas.DataFrame(data["sgRNA Score"])

    if learn_options["include_drug"]:
        # feature_sets["drug"] = pandas.DataFrame(data["drug"])
        drug_names = Y.index.get_level_values('drug').tolist()
        enc = sklearn.preprocessing.OneHotEncoder()
        label_encoder = sklearn.preprocessing.LabelEncoder()
        label_encoder.fit(drug_names)
        one_hot_drugs = np.array(enc.fit_transform(label_encoder.transform(drug_names)[:, None]).todense())
        feature_sets["drug"] = pandas.DataFrame(one_hot_drugs, columns=["drug_%d" % i for i in range(one_hot_drugs.shape[1])], index=drug_names)

    if learn_options['include_strand']:
        feature_sets['Strand effect'] = (pandas.DataFrame(data['Strand']) == 'sense')*1

    if learn_options["include_gene_feature"]:
        feature_sets["gene features"] = gene_feature(Y, data, learn_options)

    if learn_options["include_gene_guide_feature"] > 0:
        tmp_feature_sets = gene_guide_feature(Y, data, learn_options)
        for key in tmp_feature_sets:
            feature_sets[key] = tmp_feature_sets[key]

    if learn_options["include_microhomology"]:
        feature_sets["microhomology"] = get_micro_homology_features(Y['Target gene'], learn_options, data)

    t1 = time.time()
    #print "\t\tElapsed time for constructing features is %.2f seconds" % (t1-t0)

    check_feature_set_dimensions(feature_sets)

    if learn_options['normalize_features']:
        feature_sets = normalize_feature_sets(feature_sets)
        check_feature_set_dimensions(feature_sets)

    return feature_sets


def check_feature_set_dimensions(feature_sets):
    '''
    Ensure the # of people is the same in each feature set
    '''
    N = None
    for ft in feature_sets.keys():
        N2 = feature_sets[ft].shape[0]
        if N is None:
            N = N2
        else:
            assert N == N2, "# of individuals do not match up across feature sets"


def NGGX_interaction_feature(data):
    '''
    assuming 30-mer, grab the NGGX _ _ positions, and make a one-hot
    encoding of the NX nucleotides yielding 4x4=16 features
    '''
    sequence = data['30mer'].values
    feat_NX = pandas.DataFrame()
    # check that GG is where we think
    for seq in sequence:
        if seq[25:27] != "GG":
            raise Exception("expected GG but found %s" % seq[25:27])
        NX = seq[24]+seq[27]
        NX_onehot = nucleotide_features(NX,order=2, include_pos_independent=False, max_index_to_use=2, prefix="NGGX")
        # NX_onehot[:] = np.random.rand(NX_onehot.shape[0]) ##TESTING RANDOM FEATURE
        feat_NX = pandas.concat([feat_NX, NX_onehot], axis=1)
    return feat_NX.T


def get_all_order_nuc_features(data, feature_sets, learn_options, maxorder, max_index_to_use, prefix=""):
    for order in range(1, maxorder+1):
        #print "\t\tconstructing order %s features" % order
        nuc_features_pd, nuc_features_pi = apply_nucleotide_features(data, order, learn_options["num_proc"],
                                                                     include_pos_independent=True, max_index_to_use=max_index_to_use, prefix=prefix)
        feature_sets['%s_nuc_pd_Order%i' % (prefix, order)] = nuc_features_pd
        if learn_options['include_pi_nuc_feat']:
            feature_sets['%s_nuc_pi_Order%i' % (prefix, order)] = nuc_features_pi
        #print "\t\t\t\t\t\t\tdone"


def countGC(s):
    '''
    GC content for only the 20mer, as per the Doench paper/code
    '''
    assert len(s) == 30, "seems to assume 30mer"
    return len(s[5:25].replace('A', '').replace('T', ''))


def SeqUtilFeatures(data):
    '''
    assuming '30-mer'is a key
    get melting temperature features from:
        0-the 30-mer ("global Tm")
        1-the Tm (melting temperature) of the DNA:RNA hybrid from positions 16 - 20 of the sgRNA, i.e. the 5nts immediately proximal of the NGG PAM
        2-the Tm of the DNA:RNA hybrid from position 8 - 15 (i.e. 8 nt)
        3-the Tm of the DNA:RNA hybrid from position 3 - 7  (i.e. 5 nt)
    '''
    sequence = data['30mer'].values
    num_features = 1
    featarray = np.ones((sequence.shape[0], num_features))
    for i, seq in enumerate(sequence):
        assert len(seq) == 30, "seems to assume 30mer"
        featarray[i, 0] = SeqUtil.molecular_weight(str(seq))

    feat = pandas.DataFrame(pandas.DataFrame(featarray))
    return feat


def organism_feature(data):
    '''
    Human vs. mouse
    '''
    organism = np.array(data['Organism'].values)
    feat = pandas.DataFrame(pandas.DataFrame(featarray))
    import ipdb; ipdb.set_trace()
    return feat


def get_micro_homology_features(gene_names, learn_options, X):
    # originally was flipping the guide itself as necessary, but now flipping the gene instead

    print "building microhomology features"
    feat = pandas.DataFrame(index=X.index)
    feat["mh_score"] = ""
    feat["oof_score"] = ""

    #with open(r"tmp\V%s_gene_mismatches.csv" % learn_options["V"],'wb') as f:
    if True:
        # number of nulceotides to take to the left and right of the guide
        k_mer_length_left = 9
        k_mer_length_right = 21
        for gene in gene_names.unique():
            gene_seq = Seq.Seq(util.get_gene_sequence(gene)).reverse_complement()
            guide_inds = np.where(gene_names.values == gene)[0]
            print "getting microhomology for all %d guides in gene %s" % (len(guide_inds), gene)
            for j, ps in enumerate(guide_inds):
                guide_seq = Seq.Seq(X['30mer'][ps])
                strand = X['Strand'][ps]
                if strand=='sense':              
                    gene_seq = gene_seq.reverse_complement()
                # figure out the sequence to the left and right of this guide, in the gene
                ind = gene_seq.find(guide_seq)                        
                if ind==-1:
                    gene_seq = gene_seq.reverse_complement()
                    ind = gene_seq.find(guide_seq)                        
                    #assert ind != -1, "still didn't work"
                    #print "shouldn't get here"
                else:
                    #print "all good"
                    pass
                #assert ind != -1, "could not find guide in gene"
                if ind==-1:
                    #print "***could not find guide %s for gene %s" % (str(guide_seq), str(gene))
                    #if.write(str(gene) + "," + str(guide_seq))
                    mh_score = 0
                    oof_score = 0
                else:
                    #print "worked"

                    assert gene_seq[ind:(ind+len(guide_seq))]==guide_seq, "match not right"
                    left_win = gene_seq[(ind - k_mer_length_left):ind]
                    right_win = gene_seq[(ind + len(guide_seq)):(ind + len(guide_seq) + k_mer_length_right)]

                    #if strand=='antisense':
                    #    # it's arbitrary which of sense and anti-sense we flip, we just want
                    #    # to keep them in the same relative alphabet/direction
                    #    left_win = left_win.reverse_complement()
                    #    right_win = right_win.reverse_complement()
                    assert len(left_win.tostring())==k_mer_length_left
                    assert len(right_win.tostring())==k_mer_length_right
                                    
                    sixtymer = str(left_win) + str(guide_seq) + str(right_win)
                    assert len(sixtymer)==60, "should be of length 60"
                    mh_score, oof_score = microhomology.compute_score(sixtymer)

                feat.ix[ps,"mh_score"] = mh_score
                feat.ix[ps,"oof_score"] = oof_score
            print "computed microhomology of %s" % (str(gene))           
    
    return pandas.DataFrame(feat, dtype='float')


def local_gene_seq_features(gene_names, learn_options, X):
    
    print "building local gene sequence features"
    feat = pandas.DataFrame(index=X.index)
    feat["gene_left_win"] = ""
    feat["gene_right_win"] = ""

    # number of nulceotides to take to the left and right of the guide
    k_mer_length = learn_options['include_gene_guide_feature']
    for gene in gene_names.unique():
        gene_seq = Seq.Seq(util.get_gene_sequence(gene)).reverse_complement()
        for ps in np.where(gene_names.values==gene)[0]:
            guide_seq = Seq.Seq(X['30mer'][ps])
            strand = X['Strand'][ps]
            if strand=='sense':
                guide_seq = guide_seq.reverse_complement()
                #gene_seq = gene_seq.reverse_complement()
            # figure out the sequence to the left and right of this guide, in the gene
            ind = gene_seq.find(guide_seq)
            if ind ==-1:
                #gene_seq = gene_seq.reverse_complement()
                #ind = gene_seq.find(guide_seq)
                assert ind != -1, "could not find guide in gene"
            assert gene_seq[ind:(ind+len(guide_seq))]==guide_seq, "match not right"
            left_win = gene_seq[(ind - k_mer_length):ind]
            right_win = gene_seq[(ind + len(guide_seq)):(ind + len(guide_seq) + k_mer_length)]

            if strand=='antisense':
                # it's arbitrary which of sense and anti-sense we flip, we just want
                # to keep them in the same relative alphabet/direction
                left_win = left_win.reverse_complement()
                right_win = right_win.reverse_complement()
            assert not left_win.tostring()=="", "k_mer_context, %s, is too large" % k_mer_length
            assert not left_win.tostring()=="", "k_mer_context, %s, is too large" % k_mer_length
            assert len(left_win)==len(right_win), "k_mer_context, %s, is too large" % k_mer_length
            feat.ix[ps,"gene_left_win"] = left_win.tostring()
            feat.ix[ps,"gene_right_win"] = right_win.tostring()
        print "featurizing local context of %s" % (gene)

    feature_sets = {}
    get_all_order_nuc_features(feat["gene_left_win"], feature_sets, learn_options, learn_options["order"], max_index_to_use=sys.maxint, prefix="gene_left_win")
    get_all_order_nuc_features(feat["gene_right_win"], feature_sets, learn_options, learn_options["order"], max_index_to_use=sys.maxint, prefix="gene_right_win")
    return feature_sets

def gene_feature(Y, X, learn_options):
    '''
    Things like the sequence of the gene, the DNA Tm of the gene, etc.
    '''

    gene_names = Y['Target gene']

    gene_length = np.zeros((gene_names.values.shape[0], 1))
    gc_content = np.zeros((gene_names.shape[0], 1))
    temperature = np.zeros((gene_names.shape[0], 1))
    molecular_weight = np.zeros((gene_names.shape[0], 1))

    for gene in gene_names.unique():
        seq = util.get_gene_sequence(gene)
        gene_length[gene_names.values==gene] = len(seq)
        gc_content[gene_names.values==gene] = SeqUtil.GC(seq)
        temperature[gene_names.values==gene] = Tm.Tm_staluc(seq, rna=False)
        molecular_weight[gene_names.values==gene] = SeqUtil.molecular_weight(seq, 'DNA')

    all = np.concatenate((gene_length, gc_content, temperature, molecular_weight), axis=1)
    df = pandas.DataFrame(data=all, index=gene_names.index, columns=['gene length',
                                                                     'gene GC content',
                                                                     'gene temperature',
                                                                     'gene molecular weight'])
    return df

def gene_guide_feature(Y, X, learn_options):
    #features, which are related to parts of the gene-local to the guide, and
    #possibly incorporating the guide or interactions with it

    #expensive, so pickle if necessary
    gene_file = r"..\data\gene_seq_feat_V%s_km%s.ord%s.pickle" % (learn_options['V'], learn_options['include_gene_guide_feature'], learn_options['order'])

    if False: #os.path.isfile(gene_file): #while debugging, comment out
        print "loading local gene seq feats from file %s" % gene_file
        with open(gene_file, "rb") as f: feature_sets = pickle.load(f)
    else:
        feature_sets = local_gene_seq_features(Y['Target gene'], learn_options, X)
        print "writing local gene seq feats to file %s" % gene_file
        with open(gene_file, "wb") as f: pickle.dump(feature_sets, f)

    return feature_sets


def gc_cont(seq):
    return (seq.count('G') + seq.count('C'))/float(len(seq))



def Tm_feature(data):
    '''
    assuming '30-mer'is a key
    get melting temperature features from:
        0-the 30-mer ("global Tm")
        1-the Tm (melting temperature) of the DNA:RNA hybrid from positions 16 - 20 of the sgRNA, i.e. the 5nts immediately proximal of the NGG PAM
        2-the Tm of the DNA:RNA hybrid from position 8 - 15 (i.e. 8 nt)
        3-the Tm of the DNA:RNA hybrid from position 3 - 7  (i.e. 5 nt)
    '''
    sequence = data['30mer'].values
    featarray = np.ones((sequence.shape[0],4))
    for i, seq in enumerate(sequence):
        if seq[25:27]!="GG":
            raise Exception("expected GG but found %s" % seq[25:27])
        rna = False
        featarray[i,0] = Tm.Tm_staluc(seq, rna=rna)        #30mer Tm
        featarray[i,1] = Tm.Tm_staluc(seq[20:25], rna=rna) #5nts immediately proximal of the NGG PAM
        featarray[i,2] = Tm.Tm_staluc(seq[12:20], rna=rna)   #8-mer
        featarray[i,3] = Tm.Tm_staluc(seq[7:12], rna=rna)      #5-mer

    feat = pandas.DataFrame(featarray, index=data.index, columns=["Tm global_%s" % rna, "5mer_end_%s" %rna, "8mer_middle_%s" %rna, "5mer_start_%s" %rna])

    return feat

def gc_features(data):
    gc_count = data['30mer'].apply(countGC)
    gc_count.name = 'GC count'
    gc_above_10 = (gc_count > 10)*1
    gc_above_10.name = 'GC > 10'
    gc_below_10 = (gc_count < 10)*1
    gc_below_10.name = 'GC < 10'
    return gc_above_10, gc_below_10, gc_count



def normalize_features(data,axis):
    '''
    input: Pandas.DataFrame of dtype=np.float64 array, of dimensions
    mean-center, and unit variance each feature
    '''
    data -= data.mean(axis)
    data /= data.std(axis)
    # remove rows with NaNs
    data = data.dropna(1)
    if np.any(np.isnan(data.values)): raise Exception("found NaN in normalized features")
    return data

def apply_nucleotide_features(seq_data_frame, order, num_proc, include_pos_independent, max_index_to_use, prefix=""):

    fast = True

    if fast:
        feat_pd = seq_data_frame.apply(nucleotide_features, args=(order, include_pos_independent, max_index_to_use, prefix, 'pos_dependent'))
        feat_pi = seq_data_frame.apply(nucleotide_features, args=(order, include_pos_independent, max_index_to_use, prefix, 'pos_independent'))
    else:
        # old, much slower code
        feat_pd = pandas.DataFrame()
        feat_pi = pandas.DataFrame()
        for s in seq_data_frame.values:
            assert not (s==''), "string is empty"
            feat_pd_tmp, feat_pi_tmp = nucleotide_features(s, order, include_pos_independent, max_index_to_use, prefix=prefix)
            feat_pd = pandas.concat([feat_pd,feat_pd_tmp], axis=1)
            feat_pi = pandas.concat([feat_pi,feat_pi_tmp], axis=1)
        
        feat_pd = feat_pd.T
        feat_pi = feat_pi.T  
    
    return feat_pd, feat_pi



def nucleotide_features(s, order, include_pos_independent, max_index_to_use, prefix="", feature_type='all'):
    '''
    compute position-specific order-mer features for the 4-letter alphabet
    (e.g. for a sequence of length 30, there are 30*4 single nucleotide features
          and (30-1)*4^2=464 double nucleotide features
    '''
    if max_index_to_use is not None:
        s = s[:max_index_to_use]
    #assert(len(s)==30, "length not 30")
    #s = s[:30] #cut-off at thirty to clean up extra data that they accidentally left in, and were instructed to ignore in this way
    raw_alphabet = ['A', 'T', 'C', 'G']
    alphabet = ["".join(i) for i in itertools.product(raw_alphabet, repeat=order)]
    features_pos_dependent = np.zeros(len(alphabet)*(len(s)-(order-1)))
    features_pos_independent = np.zeros(np.power(len(raw_alphabet),order))


    #for position in range(0, len(s)-order, 1): JENN 9/4/2014 failing when len(s)=2
    for position in range(0, len(s)-order+1, 1):
        nucl = s[position:position+order]
        features_pos_dependent[alphabet.index(nucl)+(position*len(alphabet))] = 1.0
        features_pos_independent[alphabet.index(nucl)] += 1.0
    index_dependent = ['%s_pd.Order%d_P%d' % (prefix, order,i) for i in range(len(features_pos_dependent))]

    if feature_type == 'all' or feature_type == 'pos_independent':
        if include_pos_independent:
            index_independent = ['%s_pi.Order%d_P%d' % (prefix, order,i) for i in range(len(features_pos_independent))]
            if feature_type == 'all':
                return pandas.Series(features_pos_dependent,index=index_dependent), pandas.Series(features_pos_independent,index=index_independent)
            else:
                return pandas.Series(features_pos_independent, index=index_independent)
    
    if np.any(np.isnan(features_pos_dependent)): raise Exception("found nan features in features_pos_dependent")
    if np.any(np.isnan(features_pos_independent)): raise Exception("found nan features in features_pos_independent")
    
    return pandas.Series(features_pos_dependent, index=index_dependent)



def normalize_feature_sets(feature_sets):
    '''
    zero-mean, unit-variance each feature within each set
    '''

    print "Normalizing features..."
    t1 = time.time()

    new_feature_sets = {}
    for set in feature_sets:
         new_feature_sets[set] = normalize_features(feature_sets[set],axis=0)
         if np.any(np.isnan(new_feature_sets[set].values)):
             raise Exception("found Nan feature values in set=%s" % set)

    t2 = time.time()
    print "\t\tElapsed time for normalizing features is %.2f seconds" % (t2-t1)

    return new_feature_sets


