from __future__ import division, print_function
# -*- coding: utf-8 -*-
"""
Created on Mon Aug 22 13:56:32 2016

@author: Trent
"""

from random import choice
from subprocess import Popen
import numpy
import random
import time
from sys import stdout

class SGD_Static:
    """
    Test header

    Attributes:
        scored (bool): test score arg

    Returns:
        list: test list return

    Example:
        >>> data = SGD_Static()
        >>> output = data.do_something()
    """

    def __init__(self,pdbfile,sequence_energy_file,wt_seq,sequence_alignment_file=False,reduce_alignment=1,pair_select_list=False,pair_dist=4,gamma_multiplier=2,regular_linear_regression=True,lasso=False,lambda_lasso_coef=.01,ridge=False,ridge_coef=.01,custom_tag=False):
        self.pdbfile = pdbfile
        self.sequence_energy_file = sequence_energy_file
        self.sequence_alignment_file = sequence_alignment_file
        self.pair_select_list = pair_select_list
        self.gamma_multiplier = gamma_multiplier
        self.wt_seq = wt_seq
        self.regular_linear_regression = regular_linear_regression
        self.lasso = lasso
        self.ridge = ridge
        self.lambda_lasso_coef = lambda_lasso_coef
        self.ridge_coef = ridge_coef
        self.custom_tag = custom_tag
        self.overall_solution_space=None
        self.energies=None
        self.sequences=None
        self.CC=None
        self.MAD=None
        self.model_parameters=None
        self.pair_dist=pair_dist
        self.reduce_alignment=reduce_alignment

    def static_model(self,mad_thresh=.01,model_output_text_file=False,display_plots=True):
        if not self.regular_linear_regression and not self.lasso and not self.ridge:
            raise ValueError('No model type selected, one of regular_linear_regression, lasso, or ridge must be true')
        if self.regular_linear_regression and self.lasso:
            raise ValueError('Cannot use both regular and lasso regression, select one or the other')
        if self.regular_linear_regression and self.ridge:
            raise ValueError('Cannot use both regular and ridge regression, select one or the other')
        if self.ridge and self.lasso:
            raise ValueError('Cannot use both ridge and lasso regression, select one or the other')
        if self.regular_linear_regression and self.lasso and self.ridge:
            raise ValueError('Cannot use regular, ridge, and lasso regression at the same time, select one or the other')

        starttime = time.time()

        #max_time = 252000
        max_time = 5000000
        start_time = time.time()

        if not self.custom_tag:
            if '/' in self.pdbfile:
                pdbtag = self.pdbfile.split('/')[-1].split('.')[0]
            else:
                pdbtag = self.pdbfile
            if '\\' in self.pdbfile:
                pdbtag = self.pdbfile.split('\\')[-1].split('.')[0]
        else:
            pdbtag = self.pdbfile
        if self.custom_tag:
             pdbtag = self.custom_tag

        print('Filename Tag will be = '+str(pdbtag))

        L = len(self.wt_seq) # the length of the protein

        overall_solution_space = import_sequence_alignment(self.sequence_alignment_file,self.reduce_alignment,L)
        print(overall_solution_space)

        #wildtype = 'MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG'
        wildtype = self.wt_seq
        wtlist = [wildtype[i] for i in range(len(wildtype))]
        reference_seq = wtlist
        print(reference_seq)

        h = open(self.sequence_energy_file,'r+')
        text = h.read()
        text = text.split('\n')
        del text[-1]
        sequences = [t.split(',')[0] for t in text]
        Energies = [float(t.split(',')[1]) for t in text]

        self.sequences=sequences
        self.energies=Energies

        if self.pair_select_list:
            position_pairs = pair_select(self.pdbfile,self.pair_dist,self.wt_seq)
        else:
            position_pairs = []

        # Sample a random sequence, evaluate its energy according to our energy fucntion, and store the sequence and energy data
        trainingsets = len(Energies)
        print('Number of training sets = '+str(len(Energies)))

        E_list = []
        random_Seq_list = []
        CC = []
        MAD = []
        #w = []
        Philist = []
        trainE = []
        trainPhi = []
        print('Stochastic Gradient Descent Multivariate Linear Regression')
        print('# of pair interactions = %s, gamma = %s/((i+1)**0.5)' % (len(position_pairs), self.gamma_multiplier))
        print('\n')

        for n in range(trainingsets):
            if time.time()-start_time<max_time:
                random_Seq_list.append(sequences[n])
                E_list.append(Energies[n])

                #Convert Seq to Phi Here
                seq_Phi = [1]
                for i in range(L):
                    Phi_i = sigma2Phi(random_Seq_list[n][i],i,overall_solution_space,reference_seq)
                    seq_Phi.extend(Phi_i)
                if self.pair_select_list:
                    for pair in position_pairs:
                        firstpos = pair[0]
                        secondpos = pair[1]
                        pairs_j = pairs(random_Seq_list[n][firstpos],random_Seq_list[n][secondpos],firstpos,secondpos,overall_solution_space,reference_seq)
                        seq_Phi.extend(pairs_j)
                else:
                    pass
                if n<100:
                    Philist.append(seq_Phi)
                    if n==0:
                        d = len(seq_Phi)
                        w = numpy.zeros(shape=(d,1)) #Starting vector for regular regression not L1
                        if self.lasso:
                            for f in range(1,len(Phi_i*L)):
                                w[f][0] = .01
                        print('Number of Variables in Model = '+str(d))
                if n>=100:
                    Philist.append(seq_Phi)
                    trainPhi = Philist[0]
                    del Philist[0]
                    E_testinglist = [E_list[-xx] for xx in range(1,101)]
                    E_testinglist.reverse()
                    trainE = E_list[n-100]
                    weight = 1

                    gamma = self.gamma_multiplier/((n-100+1)**0.5)

                    x = numpy.array([trainPhi]).T # column vector
                    if self.regular_linear_regression:
                        w = w-gamma*x*weight*(numpy.dot(x.T,w)*weight - trainE*weight)
                    #Ridge
                    if self.ridge:
                        w = w-gamma*(x*weight*(numpy.dot(x.T,w)*weight - trainE) + self.ridge_coef*w)
                    #Lasso
                    if self.lasso:
                        w = w-gamma*x*weight*(numpy.dot(x.T,w)*weight - trainE*weight)
                        for j in range(d):
                            w[j][0]=STF(w[j][0],gamma*self.lambda_lasso_coef)

                    # analyze the fit of the current model w
                    #Regular Matrix Multiplication
                    E_estimate=numpy.empty([len(Philist),1])
                    for s in range(len(Philist)):
                        E_estimate[s]=numpy.dot(Philist[s],w[:,0])
                    E_estimate = E_estimate[:,0]
                    cc=numpy.corrcoef(E_testinglist,E_estimate)[0,1]
                    mad=numpy.mean(abs(E_estimate-E_testinglist))
                    CC.append(cc) # calculate the correlation coeffcient, append into list
                    MAD.append(float(mad))# calculate the mean absolute deviation, append into list
                    stdout.write('\rTraining set = %s' % str(n+1)) #This prints the current training set index on same line as an update
                    stdout.flush()
                    madlist100 = []
                    if len(MAD)>=100:
                        for g in range(1,101):
                            madlist100.append(MAD[-g])
                        if all(mm<float(mad_thresh) for mm in madlist100):
                            break

        if model_output_text_file:
            open('w'+str(pdbtag)+'.txt','w')
            for vv in range(len(w)):
                open('w'+str(pdbtag)+'.txt','a').write(str(w[vv][0])+'\n')

        self.CC=CC
        self.MAD=MAD
        self.model_parameters=w
        if display_plots:
            import matplotlib.pyplot as plt
            plt.plot(range(len(MAD)),MAD)
            plt.show()

            plt.plot(range(len(CC)),CC)
            plt.show()

            plt.scatter(E_testinglist,E_estimate)
            plt.show()

        print('energy function fit:')
        print('\tCC = %0.2f' % CC[(len(MAD))-1])
        print('\tMAD = %0.2f' % MAD[(len(MAD))-1])

        #Test to see how many features are zero...
        zerolist = []
        for z in range(len(w)):
            if w[z][0]==0:
                zerolist.append(z)
        print('Number of Features Equal to Zero in the Model = %s' % len(zerolist))

        sumlist = []
        for i in range(1000):
            sumlist.append(MAD[-i])
        print('Avg MAD of last 1000 terms = '+str(sum(sumlist)/1000))

        elapsedtime = time.time() - starttime
        print('time = %0.2f' % elapsedtime)
        return w

class SGD_Online:

    def __init__(self,wt_seq,pdbfile,sequence_alignment_file=False,reduce_alignment=2,output_energy_file=False,
                 custom_tag=False,pair_select_list=False,pair_dist=4,gamma_multiplier=2,
                 regular_linear_regression=True,lasso=False,lambda_lasso_coef=.01,ridge=False,
                 ridge_coef=.01,output_cc=True,output_mad=True,
                 rosetta_score_path='/home/romeroroot/code/rosetta_src_2016.17.58663_bundle/main/source/bin/',
                 rosetta_database_path='/home/romeroroot/code/rosetta_src_2016.17.58663_bundle/main/database'):

        'Takes in a pdb file, sequence alignment file, wild type sequence, and energy scoring script'
        'We provide two scoring scripts score_sequence_amber.py and score_sequence_rosetta.py'
        'The rosetta script requires rosetta protein structure software'
        'The amber script requires openmm molecular modeling software as well as rosetta for energy minimization'
        'this class must run in command line to make use of energy scoring scripts which output energies to command line'
        self.pdbfile = pdbfile
        self.pair_select_list = pair_select_list
        self.gamma_multiplier = gamma_multiplier
        self.wt_seq = wt_seq
        self.regular_linear_regression = regular_linear_regression
        self.lasso = lasso
        self.ridge = ridge
        self.lambda_lasso_coef = lambda_lasso_coef
        self.ridge_coef = ridge_coef
        self.output_energy_file = output_energy_file
        self.output_cc = output_cc
        self.output_mad = output_mad
        self.custom_tag = custom_tag
        self.overall_solution_space=None
        self.CC=None
        self.MAD=None
        self.model_parameters=None
        self.reduce_alignment = reduce_alignment
        self.pair_dist=pair_dist
        self.sequence_alignment_file=sequence_alignment_file
        self.rosetta_score_path = rosetta_score_path
        self.rosetta_database_path = rosetta_database_path

    def online_model(self,energy_scoring_function='rosetta',num_mutations=4,max_computation_time=252000,max_training_sets=50000,mad_cutoff=1,w_start_file=False):
        #Scoring function is either 'rosetta' or 'amber'

        if not self.regular_linear_regression and not self.lasso and not self.ridge:
            raise ValueError('No model type selected, one of regular_linear_regression, lasso, or ridge must be true')
        if self.regular_linear_regression and self.lasso:
            raise ValueError('Cannot use both regular and lasso regression, select one or the other')
        if self.regular_linear_regression and self.ridge:
            raise ValueError('Cannot use both regular and ridge regression, select one or the other')
        if self.ridge and self.lasso:
            raise ValueError('Cannot use both ridge and lasso regression, select one or the other')
        if self.regular_linear_regression and self.lasso and self.ridge:
            raise ValueError('Cannot use regular, ridge, and lasso regression at the same time, select one or the other')

        starttime = time.time()

        max_time = max_computation_time

        if w_start_file:
            hh = open(w_start_file,'r+')
            newtext = hh.read()
            newtext = newtext.split('\n')
            del newtext[-1]
            w = []
            for numbers in newtext:
                w.append(float(numbers))

        wildtype = self.wt_seq
        wtlist = [wildtype[i] for i in range(len(wildtype))]

        L = len(wildtype)

        overall_solution_space = import_sequence_alignment(self.sequence_alignment_file,self.reduce_alignment,L)
        print(overall_solution_space)

        trainingsets = max_training_sets
        madcutoff = mad_cutoff

        #Here Creating the tag that follows to output files...
        #This is skipped if a custom_tag variable is defined
        if not self.custom_tag:
            if '/' in self.pdbfile:
                pdbtag = self.pdbfile.split('/')[-1].split('.')[0]
            else:
                pdbtag = self.pdbfile
            if '\\' in self.pdbfile:
                pdbtag = self.pdbfile.split('\\')[-1].split('.')[0]
        else:
            pdbtag = self.pdbfile
        if self.custom_tag:
             pdbtag = self.custom_tag

        print('Filename Tag will be = '+str(pdbtag))

        if self.output_energy_file:
            seqEfilename = 'ubiquitin_energies_'+str(pdbtag)+'.txt'

         # the length of the protein

        reference_seq = wtlist
        print(reference_seq)

        if self.pair_select_list:
            position_pairs = pair_select(self.pdbfile,self.pair_dist,self.wt_seq)
        else:
            position_pairs = []

        #Definitions

        # Sample a random sequence, evaluate its energy according to our energy fucntion, and store the sequence and energy data

        E_list = []
        CC = []
        MAD = []
        Philist = []
        print('Stochastic Gradient Descent')
        print('# of pair interactions = %s, %s maximum training sequences, gamma = %s/((i+1)**0.5)' % (len(position_pairs),trainingsets, self.gamma_multiplier))

        if self.output_energy_file:
            open(seqEfilename,'w')
        if self.output_cc:
            open('CC'+str(pdbtag)+'.txt','w')
        if self.output_mad:
            open('MAD'+str(pdbtag)+'.txt','w')

        for n in range(trainingsets):
            if time.time()-starttime<max_time:
                wildtype = self.wt_seq
                wtlist = [wildtype[i] for i in range(len(wildtype))]
                randAAlist = []
                for i in range(num_mutations):
                    randAAlist.append(random.randint(0,L-1))
                for mutantpositions in randAAlist:
                    wtlist[mutantpositions] = choice(overall_solution_space[mutantpositions])
                random_seq = ''.join(wtlist)

                #cmd = 'python '+str(self.energy_scoring_script)+' %s' % random_seq #1. generate a terminal command as a string
                #output = Popen(cmd,shell=True,stdout=PIPE).communicate() #2. send the command to the terminal shell
                #seq_E = float(output[0]) #3. read the result and convert to a float

                'Add a new energy scoring function here if wanted'
                'Input should be random_seq and an associated pdbfile'
                'Should create a definition similar to the two energy function definitions provided'
                'if there is a need or want for a different energy function'
                if energy_scoring_function=='amber':
                    seq_E = score_sequence_amber(random_seq,pdbfile=self.pdbfile,
                                                   rosetta_path = self.rosetta_score_path,
                                                   rosetta_db = self.rosetta_database_path)

                if energy_scoring_function=='rosetta':
                    seq_E = score_sequence_rosetta(random_seq,pdbfile=self.pdbfile,
                                                   rosetta_path = self.rosetta_score_path,
                                                   rosetta_db = self.rosetta_database_path)

                #This bypasses Amber not returning scores for certain sequences
                if seq_E=='Nan':
                    print('Sequence did not generate a score')
                    pass

                if seq_E!='Nan':
                    #print(seq_E)
                    seq_E = float(seq_E) #3. read the result and convert to a float

                    print(random_seq, seq_E)
                    if self.output_energy_file:
                        open(seqEfilename,'a').write(random_seq+','+str(seq_E)+'\n')
                    E_list.append(seq_E)
                    #Convert Seq to Phi Here
                    seq_Phi = [1]
                    for i in range(L):
                        Phi_i = sigma2Phi(random_seq[i],i,overall_solution_space,reference_seq)
                        seq_Phi.extend(Phi_i)
                    if self.pair_select_list:
                        for pair in position_pairs:
                            firstpos = pair[0]
                            secondpos = pair[1]
                            pairs_j = pairs(random_seq[firstpos],random_seq[secondpos],firstpos,secondpos,overall_solution_space,reference_seq)
                            seq_Phi.extend(pairs_j)
                    else:
                        pass
                    if len(E_list)-1<100:
                        Philist.append(seq_Phi)
                        if n==0:
                            d = len(seq_Phi)
                            if not w_start_file:
                                w = numpy.zeros(shape=(d,1)) #Starting vector for regular regression not L1
                                if self.lasso:
                                    for f in range(1,len(Phi_i*L)):
                                        w[f][0] = .01
                            print('Number of Variables in Model = '+str(d))
                    if len(E_list)-1>=100:
                        Philist.append(seq_Phi)
                        trainPhi = Philist[0]
                        del Philist[0]
                        E_testinglist = [E_list[-xx] for xx in range(1,101)]
                        E_testinglist.reverse()
                        trainE = E_list[len(E_list)-1-100]
                        weight = 1

                        gamma = self.gamma_multiplier/((n+1)**0.5) # have to do plus one cause first round i = 0

                        x = numpy.array([trainPhi]).T # column vector
                        if self.regular_linear_regression:
                            w = w-gamma*x*weight*(numpy.dot(x.T,w)*weight - trainE*weight)
                        #Ridge
                        if self.ridge:
                            w = w-gamma*(x*weight*(numpy.dot(x.T,w)*weight - trainE) + self.ridge_coef*w)
                        #Lasso
                        if self.lasso:
                            w = w-gamma*x*weight*(numpy.dot(x.T,w)*weight - trainE*weight)
                            for j in range(d):
                                w[j][0]=STF(w[j][0],gamma*self.lambda_lasso_coef)

                        E_estimate=numpy.empty([len(Philist),1])
                        for s in range(len(Philist)):
                            E_estimate[s]=numpy.dot(Philist[s],w[:,0])
                        E_estimate = E_estimate[:,0]
                        cc=numpy.corrcoef(E_testinglist,E_estimate)[0,1]
                        mad=numpy.mean(abs(E_estimate-E_testinglist))
                        print('cc = '+str(cc))
                        print('mad = '+str(mad))
                        CC.append(cc) # calculate the correlation coeffcient, append into list
                        MAD.append(mad)# calculate the mean absolute deviation, append into list
                        if self.output_cc:
                            open('CC'+str(pdbtag)+'.txt','a').write(str(cc)+'\n')
                        if self.output_mad:
                            open('MAD'+str(pdbtag)+'.txt','a').write(str(mad)+'\n')
                        madlist100 = []
                        if len(MAD)>=100:
                            for g in range(1,101):
                                madlist100.append(MAD[-g])
                            if all(mm<float(madcutoff) for mm in madlist100):
                                break

        open('w'+str(pdbtag)+'.txt','w')
        #open('wState0.txt','w')
        for vv in range(len(w)):
            open('w'+str(pdbtag)+'.txt','a').write(str(w[vv][0])+'\n')

        self.CC=CC
        self.MAD=MAD
        self.model_parameters=w

        print('energy function fit:')
        print('\tCC = %0.2f' % CC[(len(MAD))-1])
        print('\tMAD = %0.2f' % MAD[(len(MAD))-1])

        #plt.scatter(E_testinglist,E_estimate)
        #plt.show()

        #Test to see how many features are zero...
        zerolist = []
        for z in range(len(w)):
            if w[z][0]==0:
                zerolist.append(z)
        print('Number of Features Equal to Zero in the Model = %s' % len(zerolist))

        elapsedtime = time.time() - starttime
        print('time = %0.2f' % elapsedtime)
        return w

def score_sequence_amber(newseq=None,pdbfile=None,rosetta_path = '/home/romeroroot/code/rosetta_src_2016.17.58663_bundle/main/source/bin/',rosetta_db = '/home/romeroroot/code/rosetta_src_2016.17.58663_bundle/main/database'):
    'Provide rosetta path to fixbb.linuxgccrelease'
    from random import choice
    from os import remove
    import simtk.openmm.app as app
    import simtk.openmm as op
    import simtk.unit as unit

    def rand_tag():
        '''creates a unique tag that can be used to identify the current files a script is working with.  necessary for running the same program multiple times in the same directory'''
        alpha = 'abcdefghijklmnopqrstuvwxyz0123456789'
        tag = ''.join([choice(alpha) for i in range(15)])
        return tag

    def thread_repack_score_amber(pdbfile,newseq):

        tag = rand_tag() # generate a random tag to associate with this run

        # generate a resfile to mutate the pdb
        resfile = """
        #header
        NATAA # keep this so all sites aren't designed
        USE_INPUT_SC # this to also consider the WT rotamer, if we're mutating WT to WT

        start

        #body\n"""

        for i in range(len(newseq)):
            resfile += '%i    A    PIKAA %s \n'%(i+1,newseq[i])

        open('resfile_'+tag+'.res','w').write(resfile)

        # the options for the run
        options = ['nice',
                   rosetta_path+'fixbb.linuxgccrelease', # fixbb is the program used for threading a repacking
                   '-s '+pdbfile,
                   '-database '+rosetta_db, # good to explicitly define location of rosetta DB
                   '-resfile resfile_'+tag+'.res', # a resfile to specify mutations
                   '-out:suffix _'+tag] # this adds a suffix to the score output file: score_[RNDTAG].sc

        # run fixbb
        Popen(options,stdout=open('/dev/null','w')).wait() # send stdout to the trash, report stderr

        # read the output
        scorefile = open('score_%s.sc'%tag).read().split('\n')
        names = scorefile[1].split(':')[1].split()
        values = scorefile[2].split(':')[1].split()
        score = dict((names[i],float(values[i])) for i in range(len(names)-1)) # use i-1 because the last entry is the filename
        pdbname = (pdbfile[:-4]+'_'+tag+'_0001.pdb').split('/')[-1]

        # now load into openMM and score with amber
        pdb = app.PDBFile(pdbname)
        forcefield = app.ForceField('amber03.xml','amber03_obc.xml')
        system = forcefield.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=None)
        integrator = op.LangevinIntegrator(300*unit.kelvin, 1/unit.picosecond, 1e-9*unit.picoseconds)
        simulation = app.Simulation(pdb.topology, system, integrator)
        simulation.context.setPositions(pdb.positions)

        #Must do try/except here because Amber occasionally throws an Exception error and does not return a score
        try:
            simulation.minimizeEnergy()
            state = simulation.context.getState(getPositions=True, getEnergy=True)
            score = state.getPotentialEnergy()/unit.kilojoule_per_mole

        except Exception:
            score = 'Nan'
            pass

        #delete the evidence
        remove('score_%s.sc'%tag)
        remove('resfile_'+tag+'.res')
        remove(pdbname)

        return score

    score = thread_repack_score_amber(pdbfile,newseq)
    return score

def score_sequence_rosetta(newseq,pdbfile,rosetta_path = '/home/romeroroot/code/rosetta_src_2016.17.58663_bundle/main/source/bin/',rosetta_db = '/home/romeroroot/code/rosetta_src_2016.17.58663_bundle/main/database'):
    from random import choice
    from subprocess import Popen
    from os import remove

    def rand_tag():
        '''creates a unique tag that can be used to identify the current files a script is working with.  necessary for running the same program multiple times in the same directory'''
        alpha = 'abcdefghijklmnopqrstuvwxyz0123456789'
        tag = ''.join([choice(alpha) for i in range(15)])
        return tag

    def fast_thread_repack_score(pdbfile,newseq,savepdb=False):
        """this function inputs a pdbfile name and a sequence, mutates the pdb file, repacks, and returns the scores as a dict
        REQUIRES: pdbfile to be consecutivly numbered from 1 to N, and len(newseq) = N. Used when we want to score seqence variants many times
        """

        tag = rand_tag() # generate a random tag to associate with this run

        # generate a resfile to mutate the pdb
        resfile = """
        #header
        NATAA # keep this so all sites aren't designed
        USE_INPUT_SC # this to also consider the WT rotamer, if we're mutating WT to WT

        start

        #body\n"""

        for i in range(len(newseq)):
            resfile += '%i    A    PIKAA %s \n'%(i+1,newseq[i])

        open('resfile_'+tag+'.res','w').write(resfile)

        # the options for the run
        options = ['nice',
                   rosetta_path+'fixbb.linuxgccrelease', # fixbb is the program used for threading a repacking
                   '-s '+pdbfile,
                   '-database '+rosetta_db, # good to explicitly define location of rosetta DB
                   '-resfile resfile_'+tag+'.res', # a resfile to specify mutations
                   '-out:suffix _'+tag] # this adds a suffix to the score output file: score_[RNDTAG].sc

        # run fixbb
        Popen(options,stdout=open('/dev/null','w')).wait() # send stdout to the trash, report stderr

        # read the output
        scorefile = open('score_%s.sc'%tag).read().split('\n')
        names = scorefile[1].split(':')[1].split()
        values = scorefile[2].split(':')[1].split()
        score = dict((names[i],float(values[i])) for i in range(len(names)-1)) # use i-1 because the last entry is the filename

        #delete the evidence
        remove('score_%s.sc'%tag)
        remove('resfile_'+tag+'.res')

        pdbname = (pdbfile[:-4]+'_'+tag+'_0001.pdb').split('/')[-1]
        if savepdb:
            return score,pdbname
        else:
            remove(pdbname)
            return score

    score = fast_thread_repack_score(pdbfile,newseq)
    return score['total_score']

def pair_select(pdbfile,dist,wt_seq):
    "Here we select the maximum distance for a pair of AAs to be considered important"
    angstrom_distance = dist
    #########################################################################
    #this should really use the 'state' pdb files....
    L = len(wt_seq)

    #AAs = ['A', 'C', 'D', 'E', 'F', 'G', 'H','I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W','Y'] #all 20 amino acids

    pdbdata = open(pdbfile).read()

    pdblines = pdbdata.split('\n')

    coordinates = []
    for line in pdblines:
        if line[:4]=='ATOM':
            coordinates.append(line)

    #"Why is there a lot of extra sequences in the pdb file"
    #for i in range(4123):
    #    coordinates.pop()

    "Here we make lists of all atoms x,y,z coordinates and pair them by which AA they are in into lists, each list corresponds to the atoms in each AA"
    xcoors = {}
    for q in range(1,L+1):
        xcoors['AAxcoor'+str(q)] = []

    for j in range(L+1):
        for line in coordinates:
            if int(line[22:26])==j:
                x = line[30:38]
                xnum = float(x)
                xcoors['AAxcoor'+str(j)].append(xnum)

    ycoors = {}
    for q in range(1,L+1):
        ycoors['AAycoor'+str(q)] = []

    for k in range(L+1):
        for line in coordinates:
            if int(line[22:26])==k:
                y = line[38:46]
                ynum = float(y)
                ycoors['AAycoor'+str(k)].append(ynum)

    zcoors = {}
    for q in range(1,L+1):
        zcoors['AAzcoor'+str(q)] = []

    for s in range(L+1):
        for line in coordinates:
            if int(line[22:26])==s:
                z = line[46:54]
                znum = float(z)
                zcoors['AAzcoor'+str(s)].append(znum)

    "Pair all xyz together"
    listofxyz = {}
    for q in range(1,L+1):
        listofxyz['listofxyz'+str(q)] = []

    "XYZ points of all atoms sorted into lists of corresponding AAs"
    for q in range(1,L+1):
        for j in range(len(xcoors['AAxcoor'+str(q)])):
            listofxyz['listofxyz'+str(q)].append((xcoors['AAxcoor'+str(q)][j],ycoors['AAycoor'+str(q)][j],zcoors['AAzcoor'+str(q)][j]))

    "Here we find the distance between every atom in the protein separated by the atoms that belong to specific AAs"
    atomtoatom = {}
    for i in range(1,L+1):
        for j in range(1,L+1):
            if i<j:
                atomtoatom['atomtoatom'+str(i)+'-'+str(j)] = []
                for k in range(len(listofxyz['listofxyz'+str(i)])):
                    for f in range(len(listofxyz['listofxyz'+str(j)])):
                        atomtoatom['atomtoatom'+str(i)+'-'+str(j)].append(numpy.linalg.norm(numpy.array(listofxyz['listofxyz'+str(i)])[k]-numpy.array(listofxyz['listofxyz'+str(j)])[f]))

    "Here we take the minimum distance between atoms between all AAs"
    minimumdistances = {}
    for i in range(1,L+1):
        for j in range(1,L+1):
            if i<j and atomtoatom['atomtoatom'+str(i)+'-'+str(j)] != []: #New shit here...
                minimumdistances['minimumdistances'+str(i)+'-'+str(j)] = []
                minimumdistances['minimumdistances'+str(i)+'-'+str(j)] = min(atomtoatom['atomtoatom'+str(i)+'-'+str(j)])

    "Here we determine which pairs are close enough by setting an angstrom_distance constraint"
    closeatoms = []
    for i in range(1,L+1):
        for j in range(1,L+1):
            if all([i<j and minimumdistances['minimumdistances'+str(i)+'-'+str(j)]<float(angstrom_distance)]):
            #if all([i<j and minimumdistances['minimumdistances'+str(i)+'-'+str(j)]<float(angstrom_distance) and abs(i-j)!=1]):
                closeatoms.append((i,j))

    "Here we subtract of (1,1) from each pair to match the indecies of the regression format"

    newcloseatoms = numpy.array(closeatoms) - (1,1)
    newcloseatoms = map(tuple, newcloseatoms)

    print(len(newcloseatoms))
    return newcloseatoms

def import_sequence_alignment(seq_align_file,reduce_alignment,seq_len):
    if seq_align_file:
        h = open(seq_align_file,'r+')
        text = h.read()
        text = text.split('\n')
        del text[-1]

        overall_solution_space = []
        for topaminoacids in text:
            tempsolspace = []
            for i in range(len(topaminoacids)):
                tempsolspace.append(topaminoacids[i])
            overall_solution_space.append(tempsolspace)

        #This part deletes the last AA at each position to have one less AA per position
        #Take this out Eventually
        if reduce_alignment:
            for i in range(reduce_alignment):
                 for top5 in overall_solution_space:
                     del top5[-1]
            return overall_solution_space
        else:
            return overall_solution_space
    else:
        AAs = ['A', 'C', 'D', 'E', 'F', 'G', 'H','I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W','Y'] #all 20 amino acids
        overall_solution_space = []
        for i in range(len(seq_len)):
            overall_solution_space.append(AAs)
        return overall_solution_space

#Definitions
def sigma2Phi(sigma,i,overall_solution_space,reference_seq):
    """This takes in an amino acid and a position and returns the 1x19 binary indicator vector Phi"""
    AAchanges = [aa for aa in overall_solution_space[i] if aa!=reference_seq[i]] # these are all 19 possible AA changes (i.e. AAs that are not the reference sequence AA)
    #AAchanges = [aa for aa in AAs if aa!=reference_seq[i]]
    Phi = []
    for aa in AAchanges:
        if sigma==aa:
            Phi.append(1)
        else:
            Phi.append(0)
    return Phi

def pairs(aa1,aa2,i,h,overall_solution_space,reference_seq):
    #This generates the pairs associated with currently solution space
    #Similar to sigma2Phi but for the pairs
    AAchanges1 = [aa for aa in overall_solution_space[i] if aa!=reference_seq[i]]
    AAchanges2 = [bb for bb in overall_solution_space[h] if bb!=reference_seq[h]]
    Phi = []
    for aa in AAchanges1:
        for bb in AAchanges2:
            if (aa1==aa and aa2==bb):
                Phi.append(1)
            else:
                Phi.append(0)
    return Phi

#Used in the lasso function
def STF(a,z):
    if a>z:
        return a-z
    if a<-z:
        return a+z
    if -z<=a<=z:
        return 0


#!/usr/bin/env python3

import argparse
from collections import Counter
import pdb
import pickle
import re
import sys
import time

import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch import optim
import torch.nn.functional as F
import torch.multiprocessing as mp

import data_producer

from multiprocessing import set_start_method

parser = argparse.ArgumentParser()
parser.add_argument("--train", type=str, default="", help="training file")
parser.add_argument("--vocab", type=str, default="", help="vocab pickle file")
parser.add_argument("--save", type=str, default="csv.pth.tar", help="saved model filename")
parser.add_argument("--size", type=int, default=300, help="word embedding dimension")
parser.add_argument("--window", type=int, default=5, help="context window size")
parser.add_argument("--sample", type=float, default=1e-5, help="subsample threshold")
parser.add_argument("--negative", type=int, default=10, help="number of negative samples")
parser.add_argument("--delta", type=float, default=0.15, help="create new sense for a type if similarity lower than this value.")
parser.add_argument("--min_count", type=int, default=5, help="minimum frequency of a word")
parser.add_argument("--processes", type=int, default=4, help="number of processes")
parser.add_argument("--num_workers", type=int, default=6, help="number of workers for data processsing")
parser.add_argument("--iter", type=int, default=3, help="number of iterations")
parser.add_argument("--lr", type=float, default=-1.0, help="initial learning rate")
parser.add_argument("--batch_size", type=int, default=100, help="(max) batch size")
parser.add_argument("--cuda", action='store_true', default=False, help="enable cuda")
parser.add_argument("--multi_proto", action='store_true', default=False, help="True: multi-prototype, False:single-prototype")

MAX_SENT_LEN = 1000

# Build the vocabulary.
def file_split(f, delim=' \t\n', bufsize=1024):
    prev = ''
    while True:
        s = f.read(bufsize)
        if not s:
            break
        tokens = re.split('['+delim+']{1,}', s)
        if len(tokens) > 1:
            yield prev + tokens[0]
            prev = tokens[-1]
            for x in tokens[1:-1]:
                yield x
        else:
            prev += s
    if prev:
        yield prev

def build_vocab(args):
    vocab = Counter()
    word_count = 0
    for word in file_split(open(args.train)):
        vocab[word] += 1
        word_count += 1
        if word_count % 10000 == 0:
            sys.stdout.write('%d\r' % len(vocab))
    freq = {k:v for k,v in vocab.items() if v >= args.min_count}
    word_count = sum([freq[k] for k in freq])
    word_list = sorted(freq, key=freq.get, reverse=True)
    word2idx = {}
    for i,w in enumerate(word_list):
        word2idx[w] = i

    print("Vocab size: %ld" % len(word2idx))
    print("Words in train file: %ld" % word_count)
    vars(args)['vocab_size'] = len(word2idx)
    vars(args)['train_words'] = word_count

    return word2idx, word_list, freq

class CSV(nn.Module):
    def __init__(self, args):
        super(CSV, self).__init__()
        self.global_embs = nn.Embedding(args.vocab_size+1, args.size, padding_idx=args.vocab_size, sparse=True)
        self.sense_embs = nn.Embedding(args.vocab_size*5, args.size, sparse=True)
        self.ctx_weight = torch.nn.Parameter(torch.ones(2*args.window, args.size))
        self.word2sense = [ [i] for i in range(args.vocab_size) ]
        '''
        word2sense = np.zeros((args.vocab_size, 5), dtype='int32')
        for i in range(args.vocab_size):
            word2sense[i, 0] = i
        self.word2sense = torch.nn.Parameter(torch.from_numpy(word2sense).int())
        self.word_sense_cnts = torch.nn.Parameter(torch.ones((args.vocab_size,)).int())
        '''

        self.global_embs.weight.data.uniform_(-0.5/args.size, 0.5/args.size)
        self.sense_embs.weight.data.uniform_(-0.5/args.size, 0.5/args.size)

        self.n_senses = args.vocab_size
        self.sense_capacity = args.vocab_size*5
        self.batch_size = args.batch_size
        self.size = args.size
        self.window = args.window
        self.negative = args.negative
        self.pad_idx = args.vocab_size

    def get_context_feats(self, ctx_type_indices):
        ctx_type_embs = self.global_embs(ctx_type_indices)
        return torch.sum(ctx_type_embs * self.ctx_weight, 1).cpu().data.numpy()

    def get_possible_sense_embs(self, type_indices, cuda=True):
        sense_indices = []
        sense2idx = {}
        for type_id in type_indices:
            for s_id in self.word2sense[type_id]:
                if s_id not in sense2idx:
                    sense2idx[s_id] = len(sense_indices)
                    sense_indices.append( s_id )
        sense_indices = np.array(sense_indices)

        if cuda:
            sense_embs = self.sense_embs(Variable(torch.LongTensor(sense_indices).cuda()))
            return sense2idx, sense_embs.cpu().data.numpy()
        else:
            sense_embs = self.sense_embs(Variable(torch.LongTensor(sense_indices)))
            return sense2idx, sense_embs.data.numpy()

    def forward(self, data):
        ctx_type_indices = data[:, 0:2*self.window]
        pos_sense_idx = data[:, 2*self.window+1]
        neg_sense_indices = data[:, 2*self.window+2:2*self.window+2+self.negative]
        neg_mask = data[:, 2*self.window+2+self.negative:].float()

        ctx_type_embs = self.global_embs(ctx_type_indices)
        pos_sense_embs = self.sense_embs(pos_sense_idx)
        neg_sense_embs = self.sense_embs(neg_sense_indices)

        ctx_feats = torch.sum(ctx_type_embs * self.ctx_weight, 1, keepdim=True)

        # Neg Log Likelihood
        pos_ips = torch.sum(ctx_feats[:,0,:] * pos_sense_embs, 1)
        pos_loss = torch.sum( -F.logsigmoid(torch.clamp(pos_ips,max=10,min=-10)))
        neg_ips = torch.bmm(neg_sense_embs, ctx_feats.permute(0,2,1))[:,:,0]
        neg_loss = torch.sum( -F.logsigmoid(torch.clamp(-neg_ips,max=10,min=-10)) * neg_mask )

        return pos_loss + neg_loss

# Initialize model.
def init_net(args):
    if args.lr == -1.0:
        vars(args)['lr'] = 0.05
    return CSV(args)

def save_model(filename, model, args, word2idx):
    torch.save({
        'word2idx':word2idx,
        'args':args,
        #'word2sense': model.word2sense,
        'n_senses': model.n_senses,
        'params': model.state_dict()
    }, filename)

def load_model(filename):
    checkpoint = torch.load(filename)
    word2idx = checkpoint['word2idx']
    args = checkpoint['args']
    model = CSV(args)
    if args.cuda:
        model.cuda()

    model.global_embs.weight.data = checkpoint['params']['global_embs.weight']
    model.sense_embs.weight.data = checkpoint['params']['sense_embs.weight']
    model.ctx_weight.data = checkpoint['params']['ctx_weight']
    model.word2sense = checkpoint['word2sense']
    #model.word2sense.data = checkpoint['params']['word2sense']
    #model.word_sense_cnts.data = checkpoint['params']['word_sense_cnts']
    model.n_senses = checkpoint['n_senses']

    return model, word2idx

# Training
def train_process_sent_producer(p_id, data_queue, word_count_actual, word_list, word2idx, freq, args):
    n_proc = 1 if args.stage == 2 else args.processes
    N = 1 if args.stage == 2 else args.iter
    neg = 0 if args.stage == 2 else args.negative

    if args.negative > 0:
        table_ptr_val = data_producer.init_unigram_table(word_list, freq, args.train_words)

    train_file = open(args.train)
    file_pos = args.file_size * p_id // n_proc
    train_file.seek(file_pos, 0)
    while True:
        try:
            train_file.read(1)
        except UnicodeDecodeError:
            file_pos -= 1
            train_file.seek(file_pos, 0)
        else:
            train_file.seek(file_pos, 0)
            break

    batch_count = 0
    batch_placeholder = np.zeros((args.batch_size, 2*args.window+2+2*neg), 'int64')

    for it in range(N):
        train_file.seek(file_pos, 0)

        last_word_cnt = 0
        word_cnt = 0
        sentence = []
        prev = ''
        eof = False
        while True:
            if eof or train_file.tell() > file_pos + args.file_size / n_proc:
                break

            while True:
                s = train_file.read(1)
                if not s:
                    eof = True
                    break
                elif s == ' ' or s == '\t':
                    if prev in word2idx:
                        sentence.append(prev)
                    prev = ''
                    if len(sentence) >= MAX_SENT_LEN:
                        break
                elif s == '\n':
                    if prev in word2idx:
                        sentence.append(prev)
                    prev = ''
                    break
                else:
                    prev += s

            if len(sentence) > 0:
                # subsampling
                sent_id = []
                if args.sample != 0:
                    sent_len = len(sentence)
                    i = 0
                    while i < sent_len:
                        word = sentence[i]
                        f = freq[word] / args.train_words
                        pb = (np.sqrt(f / args.sample) + 1) * args.sample / f;

                        if pb > np.random.random_sample():
                            sent_id.append( word2idx[word] )
                        i += 1

                if len(sent_id) < 2:
                    word_cnt += len(sentence)
                    sentence.clear()
                    continue

                next_random = (2**24) * np.random.randint(0, 2**24) + np.random.randint(0, 2**24)
                chunk = data_producer.cbow_producer(sent_id, len(sent_id), table_ptr_val, args.window,
                            neg, args.vocab_size, args.batch_size, next_random)

                chunk_pos = 0
                while chunk_pos < chunk.shape[0]:
                    remain_space = args.batch_size - batch_count
                    remain_chunk = chunk.shape[0] - chunk_pos

                    if remain_chunk < remain_space:
                        take_from_chunk = remain_chunk
                    else:
                        take_from_chunk = remain_space

                    batch_placeholder[batch_count:batch_count+take_from_chunk, :] = chunk[chunk_pos:chunk_pos+take_from_chunk, :]
                    batch_count += take_from_chunk

                    if batch_count == args.batch_size:
                        data_queue.put(batch_placeholder)
                        batch_count = 0

                    chunk_pos += take_from_chunk

                word_cnt += len(sentence)
                if word_cnt - last_word_cnt > 10000:
                    with word_count_actual.get_lock():
                        word_count_actual.value += word_cnt - last_word_cnt
                    last_word_cnt = word_cnt
                sentence.clear()

        with word_count_actual.get_lock():
            word_count_actual.value += word_cnt - last_word_cnt
        print(p_id, it, file_pos, train_file.tell(), args.file_size)
    if batch_count > 0:
        data_queue.put(batch_placeholder[:batch_count,:])
    data_queue.put(None)
    print(p_id, file_pos, train_file.tell(), args.file_size)

def train_process(p_id, word_count_actual, word2idx, word_list, freq, args, model):
    data_queue = mp.SimpleQueue()

    lr = args.lr
    #optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)
    optimizer = optim.Adagrad(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)

    t = mp.Process(target=train_process_sent_producer, args=(p_id, data_queue, word_count_actual, word_list, word2idx, freq, args))
    t.start()

    #n_iter = 1 if args.stage == 2 else args.iter
    n_iter = args.iter
    # get from data_queue and feed to model
    prev_word_cnt = 0
    while True:
        chunk = data_queue.get()
        if chunk is None:
            break
        else:
            # lr anneal & output
            if word_count_actual.value - prev_word_cnt > 10000:
                #if args.lr_anneal:
                #    lr = args.lr * (1 - word_count_actual.value / (n_iter * args.train_words))
                #    if lr < 0.0001 * args.lr:
                #        lr = 0.0001 * args.lr
                #    for param_group in optimizer.param_groups:
                #        param_group['lr'] = lr

                #sys.stdout.write("\rAlpha: %0.8f, Progess: %0.2f, Words/sec: %f, word_cnt: %d" % (lr, word_count_actual.value / (n_iter * args.train_words) * 100, word_count_actual.value / (time.monotonic() - args.t_start), word_count_actual.value))
                sys.stdout.write("\rProgess: %0.2f, Words/sec: %f, word_cnt: %d" % (word_count_actual.value / (n_iter * args.train_words) * 100, word_count_actual.value / (time.monotonic() - args.t_start), word_count_actual.value))
                sys.stdout.flush()
                prev_word_cnt = word_count_actual.value

            if args.stage == 1:
                if args.cuda:
                    data = Variable(torch.LongTensor(chunk).cuda(), requires_grad=False)
                else:
                    data = Variable(torch.LongTensor(chunk), requires_grad=False)

                optimizer.zero_grad()
                loss = model(data)
                loss.backward()
                optimizer.step()
                model.global_embs.weight.data[args.vocab_size].fill_(0)

            elif args.stage == 3:
                if args.cuda:
                    data = Variable(torch.LongTensor(chunk).cuda(), requires_grad=False)
                else:
                    data = Variable(torch.LongTensor(chunk), requires_grad=False)

                #type_ids = chunk[:, 2*args.window+1:2*args.window+2+2*args.negative]
                type_ids = chunk[:, 2*args.window+1:2*args.window+2+args.negative]
                type_ids = np.reshape(type_ids, (type_ids.shape[0] * type_ids.shape[1]))
                sense2idx, sense_embs = model.get_possible_sense_embs(type_ids.tolist())

                # get type_idx from chunk, and do sense selection here.
                context_feats = model.get_context_feats(data[:, :2*args.window])

                chunk = data_producer.select_sense(chunk, context_feats, sense2idx, sense_embs,
                            model.word2sense, chunk.shape[0], args.size, args.window, args.negative)

                if args.cuda:
                    data = Variable(torch.LongTensor(chunk).cuda(), requires_grad=False)
                else:
                    data = Variable(torch.LongTensor(chunk), requires_grad=False)

                optimizer.zero_grad()
                loss = model(data)
                loss.backward()
                optimizer.step()
                model.global_embs.weight.data[args.vocab_size].fill_(0)
    t.join()

def train_process_stage2(p_id, word_count_actual, word2idx, word_list, freq, args, model):
    data_queue = mp.SimpleQueue()

    sense_embs = model.sense_embs.weight.data.numpy()
    counter_list = np.zeros((model.sense_capacity), dtype='float32')

    t = mp.Process(target=train_process_sent_producer, args=(p_id, data_queue, word_count_actual, word_list, word2idx, freq, args))
    t.start()

    n_iter = 1
    # get from data_queue and feed to model
    prev_word_cnt = 0
    while True:
        chunk = data_queue.get()
        if chunk is None:
            break
        else:
            if word_count_actual.value - prev_word_cnt > 10000:
                sys.stdout.write("\rProgess: %0.2f, Words/sec: %f, word_cnt: %d" % (word_count_actual.value / (n_iter * args.train_words) * 100, word_count_actual.value / (time.monotonic() - args.t_start), word_count_actual.value))
                sys.stdout.flush()
                prev_word_cnt = word_count_actual.value

            if args.cuda:
                data = Variable(torch.LongTensor(chunk).cuda(), requires_grad=False)
            else:
                data = Variable(torch.LongTensor(chunk), requires_grad=False)

            context_feats = model.get_context_feats(data[:, :2*args.window])

            # update sense_embs
            create_cnt = data_producer.create_n_update_sense(chunk[:, 2*args.window+1], context_feats, sense_embs, model.word2sense, counter_list, chunk.shape[0], args.size, args.delta, model.n_senses)
            model.n_senses += create_cnt

            #if model.n_senses + args.batch_size > model.sense_capacity:
            #    new_capacity = model.sense_capacity * 3 // 2
            #    counter_list = np.concatenate( (counter_list, np.ones((new_capacity - model.sense_capacity),dtype='float32')), axis=0)
            #    zero = np.zeros((new_capacity - model.sense_capacity, args.size), 'float32')
            #    sense_embs = np.concatenate((sense_embs, zero), 0)
            #    model.sense_capacity = new_capacity
            #    print("\nexapnded sense_embs: %d" % model.n_senses)
    t.join()

    sense_embs[:model.n_senses, :] = sense_embs[:model.n_senses, :] / counter_list[:model.n_senses, None]

if __name__ == '__main__':
    set_start_method('forkserver')

    args = parser.parse_args()
    print("Starting training using file %s" % args.train)
    train_file = open(args.train)
    train_file.seek(0, 2)
    vars(args)['file_size'] = train_file.tell()

    word_count_actual = mp.Value('L', 0)

    if args.vocab == '':
        word2idx, word_list, freq = build_vocab(args)
    else:
        with open(args.vocab, 'rb') as f:
            word2idx, word_list, freq, pos2idx, dep2id = pickle.load(f)
            word_count = sum([freq[k] for k in freq])
            vars(args)['vocab_size'] = len(word2idx)
            vars(args)['train_words'] = word_count
            print("Vocab size: %ld" % len(word2idx))
            print("Words in train file: %ld" % word_count)

    model = init_net(args)
    model.share_memory()
    if args.cuda:
        model.cuda()

    # stage 1, learn robust context representation.
    vars(args)['stage'] = 1
    print("Stage 1")
    vars(args)['lr_anneal'] = True
    vars(args)['t_start'] = time.monotonic()
    processes = []
    for p_id in range(args.processes):
        p = mp.Process(target=train_process, args=(p_id, word_count_actual, word2idx, word_list, freq, args, model))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
    del processes
    print("\nStage 1, ", time.monotonic() - args.t_start, " secs ", word_count_actual.value)
    filename = args.save
    if not filename.endswith('.pth.tar'):
        filename += '.stage1.pth.tar'
    save_model(filename, model, args, word2idx)

    if args.multi_proto:
        # stage 2, create new sense in a non-parametric way.
        # Freeze model paramters except sense_embs, and use only 1 process to prevent race condition
        old_batch_size = vars(args)['batch_size']
        model.global_embs.requires_grad = False
        model.ctx_weight.requires_grad = False
        model.sense_embs = model.sense_embs.cpu()
        vars(args)['stage'] = 2
        vars(args)['batch_size'] = 5000
        print("\nStage 2")
        word_count_actual.value = 0
        vars(args)['t_start'] = time.monotonic()
        train_process_stage2(0, word_count_actual, word2idx, word_list, freq, args, model)

        if args.cuda:
            model.cuda()
        print("\nStage 2, ", time.monotonic() - args.t_start, " secs")
        print("Current # of senses: %d" % model.n_senses)
        pdb.set_trace()
        filename = args.save
        if not filename.endswith('.pth.tar'):
            filename += '.stage2.pth.tar'
        save_model(filename, model, args, word2idx)

        # stage 3, no more sense creation.
        vars(args)['lr'] = args.lr * 0.01
        vars(args)['batch_size'] = old_batch_size
        model.global_embs.requires_grad = True
        model.ctx_weight.requires_grad = True
        vars(args)['stage'] = 3
        print("\nBegin stage 3")
        word_count_actual.value = 0
        vars(args)['t_start'] = time.monotonic()
        processes = []
        for p_id in range(args.processes):
            p = mp.Process(target=train_process, args=(p_id, word_count_actual, word2idx, word_list, freq, args, model))
            p.start()
            processes.append(p)

        for p in processes:
            p.join()

        print("\nStage 3, ", time.monotonic() - args.t_start, " secs")

    # save model
    filename = args.save
    if not filename.endswith('.pth.tar'):
        filename += '.stage3.pth.tar'
    save_model(filename, model, args, word2idx)
    print("")


# -*- coding: utf-8 -*-
import numpy as np
import astropy.units as u
import astropy.constants as const
from astropy.time import Time
import os,inspect
from EXOSIMS.util.eccanom import eccanom

class Observatory(object):
    """Observatory class template

    This class contains all variables and methods necessary to perform
    Observatory Definition Module calculations in exoplanet mission simulation.

    Args:
        \*\*specs:
            user specified values
        forceStaticEphem (bool):
            If True, forces use of stored ephemerides for solar system objects
            and uses eccanom utility to propagate.  Defaults to False.
        spkpath (str):
            Path to SPK file on disk (Defaults to de432s.bsp).

    Attributes:
        settlingTime (Quantity):
            instrument settling time after repoint (default with units of day)
        thrust (Quantity):
            occulter slew thrust (default with units of mN)
        slewIsp (Quantity):
            occulter slew specific impulse (default with units of s)
        scMass (Quantity):
            occulter (maneuvering sc) wet mass (default with units of kg)
        dryMass (Quantity):
            occulter (maneuvering sc) dry mass (default with units of kg)
        coMass (Quantity):
            telescope (or non-maneuvering sc) mass (default with units of kg)
        occulterSep (Quantity):
            occulter-telescope distance (default with units of km)
        skIsp (Quantity):
            station-keeping specific impulse (default with units of s)
        kogood (ndarray):
            1D numpy ndarray of booleans where True is observable target star
            in the target list
        r_sc (Quantity):
            1D numpy ndarray of observatory postion vector (default with units
            of km)
        currentSep (Quantity):
            current occulter separation (default with units of km)
        flowRate (Quantity):
            slew flow rate (default with units of kg/day)

    Notes:
        For finding positions of solar system bodies, this routine will attempt to
        use the jplephem module and a local SPK file on disk.  The module can be
        installed via pip or from source.  The default SPK file can be downloaded from
        here: http://naif.jpl.nasa.gov/pub/naif/generic_kernels/spk/planets/de430.bsp
        and should be placed in the Observatory subdirectory of EXOSIMS.

    """

    _modtype = 'Observatory'
    _outspec = {}

    def __init__(self, settlingTime=1., thrust=450., slewIsp=4160.,\
                 scMass=6000., dryMass=3400., coMass=5800., skIsp=220.,\
                 defburnPortion=0.05, spkpath=None, forceStaticEphem=False,\
                 **specs):

        # default Observatory values
        # instrument settling time after repoint (days)
        self.settlingTime = float(settlingTime)*u.day
        # occulter slew thrust (mN)
        self.thrust = float(thrust)*u.mN
        # occulter slew specific impulse (s)
        self.slewIsp = float(slewIsp)*u.s
        # occulter (maneuvering sc) initial (wet) mass (kg)
        self.scMass = float(scMass)*u.kg
        # occulter (maneuvering sc) dry mass (kg)
        self.dryMass = float(dryMass)*u.kg
        # telescope (or non-maneuvering sc) mass (kg)
        self.coMass = float(coMass)*u.kg
        # station-keeping Isp (s)
        self.skIsp = float(skIsp)*u.s
        # default burn portion
        self.defburnPortion = float(defburnPortion)

        # occulter-telescope distance (km)
        self.occulterSep = 55000.*u.km

        #if jplephem is available, we'll use that for propagating solar system bodies
        #otherwise, use static ephemeris
        if not forceStaticEphem:
            try:
                from jplephem.spk import SPK
                self.havejplephem = True
            except ImportError:
                print "WARNING: Module jplephem not found, using static solar system ephemeris."
                self.havejplephem = False
        else:
            self.havejplephem = False
            print "Using static solar system ephemeris."

        # populate outspec
        for att in self.__dict__.keys():
            dat = self.__dict__[att]
            self._outspec[att] = dat.value if isinstance(dat,u.Quantity) else dat

        # initialize values updated by functions
        # observatory keepout booleans
        self.kogood = np.array([])
        # observatory orbit position vector (km)
        self.r_sc = np.zeros(3)*u.km
        # current occulter separation
        self.currentSep = self.occulterSep

        # set values derived from quantities above
        # slew flow rate (kg/day)
        self.flowRate = (self.thrust/const.g0/self.slewIsp).to('kg/day')

        #define function for calculating obliquity of the ecliptic
        #(arg Julian centuries from J2000)
        self.obe = lambda TDB: 23.439279 - 0.0130102*TDB - 5.086e-8*(TDB**2) + \
                5.565e-7*(TDB**3) + 1.6e-10*(TDB**4) + 1.21e-11*(TDB**5)

        # if you have jplephem, load spice file. otherwise, load static ephem.
        if self.havejplephem:
            if (spkpath is None) or not(os.path.exists(spkpath)):
                # if the path does not exist, load the default de432s.bsp
                    classpath = os.path.split(inspect.getfile(self.__class__))[0]
                    classpath = os.path.normpath(os.path.join(classpath,'..','Observatory'))
                    filename = 'de432s.bsp'
                    spkpath = os.path.join(classpath, filename)
            self.kernel = SPK.open(spkpath)
        else:
            """All ephemeride data from Vallado Appendix D.4
            Values are:
            a     e             i               O                       w                   lM
            sma   eccentricity  inclination     long. ascending node    long. perihelion    mean longitude
            AU    N/A           deg             deg                     deg                 deg
            """

            # Store Mercury ephemerides data (ecliptic)
            Mercurya = 0.387098310
            Mercurye = [0.20563175, 0.000020406, -0.0000000284, -0.00000000017]
            Mercuryi = [7.004986, -0.0059516, 0.00000081, 0.000000041]
            MercuryO = [48.330893, -0.1254229, -0.00008833, -0.000000196]
            Mercuryw = [77.456119, 0.1588643, -0.00001343, 0.000000039]
            MercurylM = [252.250906, 149472.6746358, -0.00000535, 0.000000002]
            Mercury = self.SolarEph(Mercurya, Mercurye, Mercuryi, MercuryO, Mercuryw, MercurylM)

            # Store Venus epemerides data (ecliptic)
            Venusa = 0.723329820
            Venuse = [0.00677188, -0.000047766, 0.0000000975, 0.00000000044]
            Venusi = [3.394662, -0.0008568, -0.00003244, 0.000000010]
            VenusO = [76.679920, -0.2780080, -0.00014256, -0.000000198]
            Venusw = [131.563707, 0.0048646, -0.00138232, -0.000005332]
            VenuslM = [181.979801, 58517.8156760, 0.00000165, -0.000000002]
            Venus = self.SolarEph(Venusa, Venuse, Venusi, VenusO, Venusw, VenuslM)

            # Store Earth ephemerides data (ecliptic)
            Eartha = 1.000001018
            Earthe = [0.01670862, -0.000042037, -0.0000001236, 0.00000000004]
            Earthi = [0., 0.0130546, -0.00000931, -0.000000034]
            EarthO = [174.873174, -0.2410908, 0.00004067, -0.000001327]
            Earthw = [102.937348, 0.3225557, 0.00015026, 0.000000478]
            EarthlM = [100.466449, 35999.3728519, -0.00000568, 0.]
            Earth = self.SolarEph(Eartha, Earthe, Earthi, EarthO, Earthw, EarthlM)

            # Store Mars ephemerides data (ecliptic)
            Marsa = 1.523679342
            Marse = [0.09340062, 0.000090483, -0.0000000806, -0.00000000035]
            Marsi = [1.849726, -0.0081479, -0.00002255, -0.000000027]
            MarsO = [49.558093, -0.2949846, -0.00063993, -0.000002143]
            Marsw = [336.060234, 0.4438898, -0.00017321, 0.000000300]
            MarslM = [355.433275, 19140.2993313, 0.00000261, -0.000000003]
            Mars = self.SolarEph(Marsa, Marse, Marsi, MarsO, Marsw, MarslM)

            # Store Jupiter ephemerides data (ecliptic)
            Jupitera = [5.202603191, 0.0000001913]
            Jupitere = [0.04849485, 0.000163244, -0.0000004719, -0.00000000197]
            Jupiteri = [1.303270, -0.0019872, 0.00003318, 0.000000092]
            JupiterO = [100.464441, 0.1766828, 0.00090387, -0.000007032]
            Jupiterw = [14.331309, 0.2155525, 0.00072252, -0.000004590]
            JupiterlM = [34.351484, 3034.9056746, -0.00008501, 0.000000004]
            Jupiter = self.SolarEph(Jupitera, Jupitere, Jupiteri, JupiterO, Jupiterw, JupiterlM)

            # Store Saturn ephemerides data (ecliptic)
            Saturna = [9.554909596, -0.0000021389]
            Saturne = [0.05550862, -0.000346818, -0.0000006456, 0.00000000338]
            Saturni = [2.488878, 0.0025515, -0.00004903, 0.000000018]
            SaturnO = [113.665524, -0.2566649, -0.00018345, 0.000000357]
            Saturnw = [93.056787, 0.5665496, 0.00052809, 0.000004882]
            SaturnlM = [50.077471, 1222.1137943, 0.00021004, -0.000000019]
            Saturn = self.SolarEph(Saturna, Saturne, Saturni, SaturnO, Saturnw, SaturnlM)

            # Store Uranus ephemerides data (ecliptic)
            Uranusa = [19.218446062, -0.0000000372, 0.00000000098]
            Uranuse = [0.04629590, -0.000027337, 0.0000000790, 0.00000000025]
            Uranusi = [0.773196, -0.0016869, 0.00000349, 0.000000016]
            UranusO = [74.005947, 0.0741461, 0.00040540, 0.000000104]
            Uranusw = [173.005159, 0.0893206, -0.00009470, 0.000000413]
            UranuslM = [314.055005, 428.4669983, -0.00000486, 0.000000006]
            Uranus = self.SolarEph(Uranusa, Uranuse, Uranusi, UranusO, Uranusw, UranuslM)

            # Store Neptune ephemerides data (ecliptic)
            Neptunea = [30.110386869, -0.0000001663, 0.00000000069]
            Neptunee = [0.00898809, 0.000006408, -0.0000000008]
            Neptunei = [1.769952, 0.0002257, 0.00000023, -0.000000000]
            NeptuneO = [131.784057, -0.0061651, -0.00000219, -0.000000078]
            Neptunew = [48.123691, 0.0291587, 0.00007051, 0.]
            NeptunelM = [304.348665, 218.4862002, 0.00000059, -0.000000002]
            Neptune = self.SolarEph(Neptunea, Neptunee, Neptunei, NeptuneO, Neptunew, NeptunelM)

            # Store Pluto ephemerides data (ecliptic)
            Plutoa = [39.48168677, -0.00076912]
            Plutoe = [0.24880766, 0.00006465]
            Plutoi = [17.14175, 0.003075]
            PlutoO = [110.30347, -0.01036944]
            Plutow = [224.06676, -0.03673611]
            PlutolM = [238.92881, 145.2078]
            Pluto = self.SolarEph(Plutoa, Plutoe, Plutoi, PlutoO, Plutow, PlutolM)

            #store all as dictionary:
            self.planets = {'Mercury': Mercury,
                            'Venus': Venus,
                            'Earth': Earth,
                            'Mars': Mars,
                            'Jupiter': Jupiter,
                            'Saturn': Saturn,
                            'Uranus': Uranus,
                            'Neptune': Neptune,
                            'Pluto': Pluto}

    def __str__(self):
        """String representation of the Observatory object

        When the command 'print' is used on the Observatory object, this method
        will print the attribute values contained in the object"""

        for att in self.__dict__.keys():
            print '%s: %r' % (att, getattr(self, att))

        return 'Observatory class object attributes'

    def orbit(self, time):
        """Finds observatory orbit position vector, returns True if successful

        This method finds the observatory position vector (heliocentric
        equatorial frame) as a 1D numpy array with astropy Quantity units of km
        and stores it in self.r_sc.

        This defines the data type expected, orbits are determined by specific
        instances of Observatory classes.

        Args:
            time (Time):
                absolute time

        Returns:
            success (bool):
                True if successful, False if not

        """

        self.r_sc = np.array([float(time.mjd),float(time.mjd),float(time.mjd)])*u.km
        b = np.isfinite(self.r_sc) # finds if all values are finite
        success = all(b) # returns True if all values of self.r_sc are finite

        return success

    def keepout(self, currentTime, TL, koangle):
        """Finds keepout Boolean values, returns True if successful

        This method finds the keepout Boolean values for each target star where
        True is an observable star and stores the 1D numpy array in self.kogood.

        Args:
            currentTime (Time):
                absolute time
            TL (TargetList or StarCatalog):
                TargetList or StarCatalog class object
            koangle (float):
                telescope keepout angle in degrees

        Returns:
            success (bool):
                True if successful, False if not

        """

        # update spacecraft orbital position
        a = self.orbit(currentTime)

        self.kogood = np.array([True for row in TL.Name])

        # check to make sure all elements in self.kogood are Boolean
        b = [isinstance(element, np.bool_) for element in self.kogood]
        c = [a, b]
        # return True if orbital position is successful and all elements of
        # self.kogood are Boolean
        success = all(c)

        return success

    def solarSystem_body_position(self, time, bodyname):
        """Finds position vector for solar system objects

        This passes all arguments to one of spk_body or keplerplanet, depending
        on the value of self.havejplephem.

        Args:
            time (Time):
                absolute time
            bodyname (str):
                solar system object name

        Returns:
            r_body (Quantity):
                heliocentric equatorial position vector in 1D numpy ndarray
                (units of km)

        """

        if self.havejplephem:
            return self.spk_body(time,bodyname)
        else:
            return self.keplerplanet(time,bodyname)

    def spk_body(self, time, bodyname):
        """Finds position vector for solar system objects

        This method uses spice kernel from NAIF to find heliocentric
        equatorial position vectors (astropy Quantity in km) for solar system
        objects.

        Args:
            time (Time):
                absolute time
            bodyname (str):
                solar system object name

        Returns:
            r_body (Quantity):
                heliocentric equatorial position vector in 1D numpy ndarray
                (units of km)

        """

        # dictionary of solar system bodies available in spice kernel
        bodies = {'Mercury':199,
                  'Venus':299,
                  'Earth':399,
                  'Mars':4,
                  'Jupiter':5,
                  'Saturn':6,
                  'Uranus':7,
                  'Neptune':8,
                  'Pluto':9,
                  'Sun':10,
                  'Moon':301}

        assert bodies.has_key(bodyname),\
                 "%s is not a recognized body name."%(bodyname)

        if bodies[bodyname] == 199:
            r_body = (self.kernel[0,1].compute(time.jd) + \
                    self.kernel[1,199].compute(time.jd) - \
                    self.kernel[0,10].compute(time.jd))*u.km
        elif bodies[bodyname] == 299:
            r_body = (self.kernel[0,2].compute(time.jd) + \
                    self.kernel[2,299].compute(time.jd) - \
                    self.kernel[0,10].compute(time.jd))*u.km
        elif bodies[bodyname] == 399:
            r_body = (self.kernel[0,3].compute(time.jd) + \
                    self.kernel[3,399].compute(time.jd) - \
                    self.kernel[0,10].compute(time.jd))*u.km
        elif bodies[bodyname] == 301:
            r_body = (self.kernel[0,3].compute(time.jd) + \
                    self.kernel[3,301].compute(time.jd) - \
                    self.kernel[0,10].compute(time.jd))*u.km
        else:
            r_body = (self.kernel[0,bodies[bodyname]].compute(time.jd) - \
                    self.kernel[0,10].compute(time.jd))*u.km

        return r_body

    def keplerplanet(self, time, bodyname):
        """Finds position vector for solar system objects

        This method uses algorithms 2 and 10 from Vallado 2013 to find
        heliocentric equatorial position vectors (astropy Quantity in km) for
        solar system objects.

        Args:
            time (Time):
                absolute time
            bodyname (str):
                solar system object name

        Returns:
            r_body (Quantity):
                heliocentric equatorial position vector in 1D numpy ndarray
                (units of km)

        """

        if bodyname == 'Moon':
            r_Earth = self.keplerplanet(time, 'Earth')
            return r_Earth + self.moon_earth(time)

        assert self.planets.has_key(bodyname),\
                "%s is not a recognized body name."%(bodyname)

        planet = self.planets[bodyname]
        # find Julian centuries from J2000
        TDB = self.cent(time)
        # update ephemeride data
        a = self.propeph(planet.a, TDB)
        e = self.propeph(planet.e, TDB)
        i = np.radians(self.propeph(planet.i, TDB))
        O = np.radians(self.propeph(planet.O, TDB))
        w = np.radians(self.propeph(planet.w, TDB))
        lM = np.radians(self.propeph(planet.lM, TDB))
        # Find mean anomaly and argument of perigee
        M = np.mod(lM - w,2*np.pi)
        wp = np.mod(w - O,2*np.pi)
        # Find eccentric anomaly
        E = eccanom(M,e)[0]
        # Find true anomaly
        nu = np.arctan2(np.sin(E) * np.sqrt(1 - e**2), np.cos(E) - e)
        # Find semiparameter
        p = a*(1 - e**2)
        # position vector (km) in orbital plane
        r_planet = np.array([(p*np.cos(nu)/(1 + e*np.cos(nu))), (p*np.sin(nu))/(1 + e*np.cos(nu)), 0.])
        # position vector (km) in ecliptic plane
        r_planet = np.dot(np.dot(self.rot(-O,3),self.rot(-i,1)),np.dot(self.rot(-wp,3),r_planet))
        # find obliquity of the ecliptic
        obe = self.obe(TDB)
        # position vector (km) in heliocentric equatorial frame
        r_planet = np.dot(self.rot(np.radians(-obe),1),r_planet)*u.km

        return r_planet

    def moon_earth(self, time):
        """Finds geocentric equatorial position vector (km) for Earth's moon

        This method uses Algorithm 31 from Vallado 2013 to find the geocentric
        equatorial position vector for Earth's moon.

        Args:
            time (Time):
                absolute time

        Returns:
            r_moon (Quantity):
                geocentric equatorial position vector in 1D numpy array (units
                of km)

        """

        TDB = self.cent(time)
        la = np.radians(218.32 + 481267.8813*TDB + \
            6.29*np.sin(np.radians(134.9 + 477198.85*TDB)) -
            1.27*np.sin(np.radians(259.2 - 413335.38*TDB)) +
            0.66*np.sin(np.radians(235.7 + 890534.23*TDB)) +
            0.21*np.sin(np.radians(269.9 + 954397.70*TDB)) -
            0.19*np.sin(np.radians(357.5 + 35999.05*TDB)) -
            0.11*np.sin(np.radians(186.6 + 966404.05*TDB)))

        phi = np.radians(5.13*np.sin(np.radians(93.3 + 483202.03*TDB)) +
            0.28*np.sin(np.radians(228.2 + 960400.87*TDB)) -
            0.28*np.sin(np.radians(318.3 + 6003.18*TDB)) -
            0.17*np.sin(np.radians(217.6 - 407332.20*TDB)))

        P = np.radians(0.9508 + 0.0518*np.cos(np.radians(134.9 + 477198.85*TDB)) +
            0.0095*np.cos(np.radians(259.2 - 413335.38*TDB)) +
            0.0078*np.cos(np.radians(235.7 + 890534.23*TDB)) +
            0.0028*np.cos(np.radians(269.9 + 954397.70*TDB)))

        e = np.radians(23.439291 - 0.0130042*TDB - 1.64e-7*TDB**2 + 5.04e-7*TDB**3)

        r = 1./np.sin(P)*6378.137 # km

        r_moon = r*np.array([np.cos(phi)*np.cos(la),
            np.cos(e)*np.cos(phi)*np.sin(la) - np.sin(e)*np.sin(phi),
            np.sin(e)*np.cos(phi)*np.sin(la) + np.cos(e)*np.sin(phi)])*u.km

        return r_moon

    def starprop(self, currentTime, TL, sInd):
        """Finds target star position vector (km) for current time (MJD)

        Args:
            currentTime (Time):
                absolute time
            TL (TargetList or StarCatalog):
                TargetList or StarCatalog class object
            sInd (int):
                index for star catalog information

        Returns:
            r_star (Quantity): star position vector in heliocentric equatorial
                frame in 1D numpy ndarray (units of km)

        """

        # right ascension and declination
        ra = TL.coords.ra[sInd]
        dec = TL.coords.dec[sInd]

        # set J2000 epoch
        j2000 = Time(2000., format='jyear')

        # directions
        p0 = np.array([-np.sin(ra), np.cos(ra), 0.])
        q0 = np.array([-np.sin(dec)*np.cos(ra), -np.sin(dec)*np.sin(ra), np.cos(dec)])
        r0 = TL.coords[sInd].cartesian.xyz/TL.coords[sInd].distance

        # proper motion vector
        mu0 = p0*TL.pmra[sInd] + q0*TL.pmdec[sInd]

        # space velocity vector
        v = mu0/TL.parx[sInd]*u.AU + r0*TL.rv[0]

        # stellar position vector
        r_star = TL.coords[sInd].cartesian.xyz + v*(currentTime.mjd - j2000.mjd)*u.day

        return r_star.to('km')

    def cent(self, currentTime):
        """Finds time in Julian centuries since J2000 epoch

        This quantity is needed for many algorithms from Vallado 2013.

        Args:
            currentTime (Time):
                absolute time

        Returns:
            TDB (float):
                time in Julian centuries since the J2000 epoch

        """

        j2000 = Time(2000., format='jyear')
        TDB = (currentTime.jd - j2000.jd)/36525.

        return TDB

    def propeph(self, x, TDB):
        """Propagates ephemeride to current time and returns this value

        This method propagates the ephemerides from Vallado 2013 to the current
        time.

        Args:
            x (list):
                ephemeride list (maximum of 4 elements)
            TDB (float):
                time in Julian centuries since the J2000 epoch

        Returns:
            y (float):
                ephemeride value at current time

        """

        if isinstance(x, list):
            if len(x) < 4:
                q = 4 - len(x)
                i = 0
                while i < q:
                    x.append(0.)
                    i += 1
        elif (isinstance(x, float) or isinstance(x, int)):
            x = [float(x)]
            i = 0
            while i < 3:
                x.append(0.)
                i += 1

        y = x[0] + x[1]*TDB + x[2]*(TDB**2) + x[3]*(TDB**3)

        return y

    def rot(self, th, axis):
        """Finds the rotation matrix of angle th about the axis value

        Args:
            th (float):
                rotation angle in radians
            axis (int):
                integer value denoting rotation axis (1,2, or 3)

        Returns:
            matrix (ndarray):
                rotation matrix defined as numpy ndarray

        """

        if axis == 1:
            return np.array([[1., 0., 0.],
                       [0., np.cos(th), np.sin(th)],
                       [0., -np.sin(th), np.cos(th)]])
        elif axis == 2:
            return np.array([[np.cos(th), 0., -np.sin(th)],
                       [0., 1., 0.],
                       [np.sin(th), 0., np.cos(th)]])
        elif axis == 3:
            return np.array([[np.cos(th), np.sin(th), 0.],
                       [-np.sin(th), np.cos(th), 0.],
                       [0., 0., 1.]])

    def distForces(self, TK, TL, sInd):
        """Finds lateral and axial disturbance forces on an occulter

        Args:
            TK (TimeKeeping):
                TimeKeeping class object
            TL (TargetList):
                TargetList class object
            sInd (int):
                index of target star

        Returns:
            dF_lateral, dF_axial (Quantity, Quantity):
                lateral and axial disturbance forces (units of N)

        """

        # occulter separation distance
        occulterSep = self.occulterSep

        # get spacecraft position vector
        self.orbit(TK.currentTimeAbs)
        r_Ts = self.r_sc
        # sun -> earth position vector
        r_Es = self.solarSystem_body_position(TK.currentTimeAbs, 'Earth')
        # sun -> target star vector
        r_ts = self.starprop(TK.currentTimeAbs, TL, sInd)
        # Telescope -> target vector and unit vector
        r_tT = r_ts - r_Ts
        u_tT = r_tT/np.sqrt(np.sum(r_tT**2))
        # sun -> occulter vector
        r_Os = r_Ts + occulterSep*u_tT
        # Earth-Moon barycenter -> spacecraft vectors
        r_TE = r_Ts - r_Es
        r_OE = r_Os - r_Es

        # force on occulter
        F_sO = (-const.G*const.M_sun*self.scMass*r_Os/np.sqrt(np.sum(r_Os**2)**3)).to('N')
        mEMB = const.M_sun/328900.56
        F_EO = (-const.G*mEMB*self.scMass*r_OE/np.sqrt(np.sum(r_OE**2)**3)).to('N')

        F_O = F_sO + F_EO

        # force on telescope
        F_sT = (-const.G*const.M_sun*self.coMass*r_Ts/np.sqrt(np.sum(r_Ts**2))**3).to('N')
        F_ET = (-const.G*mEMB*self.coMass*r_TE/np.sqrt(np.sum(r_TE**2))**3).to('N')
        F_T = F_sT + F_ET

        # differential force
        dF = ((F_O/self.scMass - F_T/self.coMass)*self.scMass).to('N')

        dF_axial = np.dot(dF.to('N'), u_tT)*u.N
        dF_lateral = np.sqrt(np.sum((dF - dF_axial*u_tT)**2))
        dF_axial = np.abs(dF_axial)

        return dF_lateral, dF_axial

    def mass_dec(self, dF_lateral, t_int):
        """Returns mass_used and deltaV

        The values returned by this method are used to decrement spacecraft
        mass for station-keeping.

        Args:
            dF_lateral (Quantity):
                lateral force on occulter (units of force)
            t_int (Quantity):
                integration time (units of time)

        Returns:
            intMdot, mass_used, deltaV (Quantity, Quantity, Quantity):
                mass flow rate (units like kg/day),
                mass used in station-keeping (units of mass),
                change in velocity required for station-keeping (velocity units
                like km/s)

        """

        intMdot = (1./np.cos(np.radians(45.))*np.cos(np.radians(5.))*dF_lateral/const.g0/self.skIsp).to('kg/s')
        mass_used = (intMdot*t_int).to('kg')
        deltaV = (dF_lateral/self.scMass*t_int).to('km/s')

        return intMdot, mass_used, deltaV

    class SolarEph:
        """Solar system ephemerides class

        This class takes the constants in Appendix D.4 of Vallado as inputs
        and stores them for use in defining solar system ephemerides at a
        given time.

        Args:
            a (list):
                semimajor axis list (in AU)
            e (list):
                eccentricity list
            i (list):
                inclination list
            O (list):
                right ascension of the ascending node list
            w (list):
                longitude of periapsis list
            lM (list):
                mean longitude list

        Each of these lists has a maximum of 4 elements. The values in
        these lists are used to propagate the solar system planetary
        ephemerides for a specific solar system planet.

        Attributes:
            a (list):
                list of semimajor axis (in AU)
            e (list):
                list of eccentricity
            i (list):
                list of inclination
            O (list):
                list of right ascension of the ascending node
            w (list):
                list of longitude of periapsis
            lM (list):
                list of mean longitude values

        Each of these lists has a maximum of 4 elements. The values in
        these lists are used to propagate the solar system planetary
        ephemerides for a specific solar system planet."""

        def __init__(self, a, e, i, O, w, lM):

            # attach units of AU
            self.a = a*u.AU
            # change to km
            self.a = self.a.to('km')
            # strip dimensions
            self.a = self.a.value
            if not isinstance(self.a, float):
                self.a = self.a.tolist()
            # store list of dimensionless eccentricity values
            self.e = e
            # store list of inclination values (degrees)
            self.i = i
            # store list of right ascension of ascending node values (degrees)
            self.O = O
            # store list of longitude of periapsis values (degrees)
            self.w = w
            # store list of mean longitude values (degrees)
            self.lM = lM

        def __str__(self):
            """String representation of the SolarEph object

            When the command 'print' is used on the SolarEph object, this
            method will print the attribute values contained in the object"""

            for att in self.__dict__.keys():
                print '%s: %r' % (att, getattr(self, att))

            return 'SolarEph class object attributes'

import numpy,import astropy,from astropy