from copy import deepcopy
from cPickle import load, dump
from sympy import Symbol
from numpy import linspace
from multiprocessing import Pool
from cobra.solvers.glpk_solver import solve as optimize

def write_lp_file(cobra_model, the_objective='G6931_No_Compartment_sink', objective_sense='maximize', lp_file_prefix='e_coli_me_model'):
    '''
    required input:
    cobra_model: the COBRApy model structure provided
    
    output:
    An LP file to be used with soplex to solve the model
    
    '''
    
    for rxn in cobra_model.reactions:
        for met in rxn._metabolites:
            met._reaction.add(rxn)

    basis_compatible_encode_dict = {}

    counter = 0
    for rxn in cobra_model.reactions:
        the_id = 'r_'+str(counter)
        basis_compatible_encode_dict[rxn.id] = the_id
        rxn.id= the_id
        counter +=1

    counter = 0
    for met in cobra_model.metabolites:
        the_id = 'm_'+str(counter)
        basis_compatible_encode_dict[met.id] = the_id
        met.id= the_id
        counter +=1

    constructed_problem = optimize(cobra_model,the_objective=the_objective, the_problem='setup',objective_sense=objective_sense)

    constructed_problem.write(cpxlp='' + lp_file_prefix+'.lp')
    
    fout = open('basis_compatible_encode_dict.pickle','w')
    dump(basis_compatible_encode_dict, fout)
    fout.close()

def parameterize_model(cobra_model, substitution_dict, with_capacity_constraint=True):
    count = 0;
    number_of_sig_digits=4
    from decimal import Decimal, getcontext
    getcontext().prec=number_of_sig_digits

    for reaction in cobra_model.reactions:
        count+=1;

        if 'sympy' in str(type(reaction.lower_bound)):
            reaction.lower_bound = Decimal(str(reaction.lower_bound.evalf(subs=substitution_dict))).normalize()
        
        if 'sympy' in str(type(reaction.upper_bound)):
            reaction.upper_bound = Decimal(str(reaction.upper_bound.evalf(subs=substitution_dict))).normalize()

        for k, v in reaction._metabolites.items():
            if  'sympy' in str(type(v)):
                reaction._metabolites[cobra_model.metabolites.get_by_id(k.id)] = Decimal(str(v.evalf(subs=substitution_dict))).normalize() # is this consistent with before?

        if count % 1000 == 0:
            print '%1.2f percent done substituting...'%(100.*float(count)/len(cobra_model.reactions))

    for metabolite in cobra_model.metabolites:
        if 'sympy' in str(type(metabolite._bound)):
            metabolite._bound = Decimal(str(metabolite._bound.evalf(subs=substitution_dict))).normalize()

    if with_capacity_constraint==False:
        cobra_model.metabolites.get_by_id('biomass_capacity_lower')._bound = 1000.0 # not constraining at all  (x+y+z <= 1000)
        cobra_model.metabolites.get_by_id('biomass_capacity_upper')._bound = 1000.0 # not constraining at all  (-x -y -z <= 1000)

if __name__ == "__main__":
    # set parameters
    substitution_dict = {
                        "growth_rate_in_per_hour": 1.0,
                        "mRNA_mean_lifetime_minutes": 5.0,
                        "enzyme_efficiency_in_per_second": 65.0,                   
                        "proportion_of_rna_that_is_mrna" : 0.02, 
                        "proportion_of_rna_that_is_trna" : 0.12,
                        "proportion_of_rna_that_is_rrna" : 0.86, 
                        "average_rna_nt_mw" : 324, 
                        "average_aa_mw" : 109, 
                        "mass_rrna_per_ribosome" : 1700.*1000, 
                        "average_trna_mw" :25000, 
                        "percent_nadh_dehydrogenase_that_ndh1": 0.5,
                        "gam_value" : 35.,
                        "ngam_value" : 1.,
                        "unmodeled_protein_proportion_of_proteome": 0.45,
                        "enzyme_efficiency_scaling": 1.0 # ignore
                        }
    
    print 'Loading symbolic model pickle...'
    f = open('e_coli_me_model_symbolic_cobra_model.pickle','r')
    cobra_model = load(f)
    f.close()

    # providing unlimited D-Glucose, a 'batch' growth simulation. The default model provides 10 mmol /gDW h D-Glucose.
    cobra_model.reactions.get_by_id('D-Glucose_source').upper_bound = 1000; 
    
    # set a dummy objective
    cobra_model.reactions.get_by_id('G6931_No_Compartment_sink').objective_coefficient = 1; 
    
    
    print 'Substituting parameters (this may take awhile)...'
    parameterize_model(cobra_model, substitution_dict, with_capacity_constraint=True)

    print 'Writing LP file...'
    write_lp_file(cobra_model)