#!/usr/bin/env python
"""lpa.py

Usage: lpa.py mechanism_file num_complexes [order of interest]"""

import sys

import networkx as nx

import os
import sys
import errno
import cPickle as pickle
import datetime
import functools
import itertools as it
from multiprocessing import Pool

import numpy as np

from gratelpy.fragments import get_valid_fragments, score_fragment, get_sensible_fragments, pretty_print
from gratelpy.subgraphs import get_all_valid_subgraphs
from gratelpy.stoich import get_graph_stoich
from gratelpy.parse_mechanism import get_network_from_mechanism
from gratelpy.gensg import gen_valid_subgraphs, gen_valid_subgraphs_mp, gen_valid_subgraphs_mps


def get_valid_frags_with_caching(base_name, G, stoich_rank):

    fn_vfrags = base_name + '.vfrags'

    try:
        valid_fragments = pickle.load(open(fn_vfrags))
    except IOError, e:
        if e.errno == errno.ENOENT:
            valid_fragments = get_valid_fragments(G, stoich_rank)
            pickle.dump(valid_fragments, open(fn_vfrags, 'wb'))
        else: raise

    return valid_fragments

# def gen_valid_subgraphs_mps_with_caching(base_name, G, valid_fragments, stoich_rank):
    
#     vsg_name = base_name + '.vsg'
    
#     try:
#         valid_subgraphs = pickle.load(open(vsg_name))
#     except IOError, e:
#         if e.errno == errno.ENOENT:
#             valid_subgraphs = gen_valid_subgraphs_mps(G, valid_fragments, stoich_rank)
#             pickle.dump(valid_subgraphs, open(vsg_name, 'wb'))
#         else: raise

#     return valid_subgraphs

def display_valid_fragments(valid_fragments):
    print 'valid fragments received from \'validate_fragments\': '+str(
        len(valid_fragments))
    for f in valid_fragments:
        print f

def score_valid_fragments_with_caching(base_name, valid_subgraphs_result):
    # valid_subgraphs_result: the list of lists (fragment, subgraph components, and valid subgraphs) generated by get_all_valid_subgraphs
    
    frag_score_name = base_name + '.fs'

    # structure of list that get_all_valid_subgraphs returns
    vs_frag_i = 0
    vs_sc_i = 1
    vs_vs_i = 2

    # structure of list that score_fragment returns
    sf_frag_i = 0
    sf_sc_i = 1
    sf_K_S_i = 2

    try:
        frag_scores_res = pickle.load(open(frag_score_name))
    except IOError, e:
        if e.errno == errno.ENOENT:
            frag_scores_res = []
            for el in valid_subgraphs_result:
                curr_res = score_fragment(el[vs_vs_i], el[vs_sc_i], el[vs_frag_i])

                frag_print = pretty_print(el[vs_vs_i])
                
                frag_scores_res.append([frag_print, curr_res[sf_sc_i], el[vs_vs_i], curr_res[sf_K_S_i]])

            pickle.dump(frag_scores_res, open(frag_score_name, 'wb'))
        else: raise

    return frag_scores_res

def get_valid_subgraphs_with_caching(base_name, G, valid_fragments, stoich_rank):

    fn_vsg = base_name + '.vsg'

    try:
        valid_subgraphs = pickle.load(open(fn_vsg))
    except IOError, e:
        if e.errno == errno.ENOENT:
            valid_subgraphs = gen_valid_subgraphs_mps(G, valid_fragments, 
                stoich_rank)
            pickle.dump(valid_subgraphs, open(fn_vsg, 'wb'))
        else: raise

    return valid_subgraphs

def main():

    try:
        mechanism_file = sys.argv[1]
        num_complexes = int(sys.argv[2])
    except IndexError, e:
        print __doc__
        sys.exit(2)

    try:
            stoich_rank_user = int(sys.argv[3])
            stoich_rank_set = True
    except IndexError, e:
            stoich_rank_set = False

    base_name, ext = os.path.splitext(os.path.basename(mechanism_file))

    alpha, beta, dict_complexes, dict_constants, dict_complexes_reverse, dict_constants_reverse = get_network_from_mechanism(mechanism_file, num_complexes)

    # save dictionaries for later changes in visualization / naming of complexes and reactions
    dict_name = base_name + '.dict'
    pickle.dump({'constants_dict': dict_constants, 'complexes_dict': dict_complexes, 'constants_dict_reverse': dict_constants_reverse, 'complexes_dict_reverse': dict_complexes_reverse}, open(dict_name, 'wb'))

    # presently, the subgraph scoring function (score_subgraph in subgraphs.py) assumes that all stoichiometric coefficients equal 1!
    if np.max(alpha) > 1 or np.max(beta) > 1:
        raise Exception('presently, the subgraph scoring function (score_subgraph in subgraphs.py) assumes that all stoichiometric coefficients equal 1!')
    
    G, stoich, stoich_rank = get_graph_stoich(alpha, beta)

    if stoich_rank_set:
            stoich_rank = stoich_rank_user

    # save more data about current system
    system_data = base_name + '.sys'
    pickle.dump({'alpha': alpha, 'beta': beta, 'stoich': stoich, 'graph': G}, open(system_data, 'wb'))

    print 'rank of stoichiometry matrix = ' + str(stoich_rank)
    
    #valid_fragments = get_valid_frags_with_caching(base_name, G, stoich_rank)
    valid_fragments = get_sensible_fragments(G, stoich_rank)
    print 'numer of valid fragments: '+str(len(valid_fragments))

    #display_valid_fragments(valid_fragments)

    valid_subgraphs_result = get_valid_subgraphs_with_caching(base_name, G, valid_fragments, stoich_rank)
    
    print 'number of fragments we found valid subgraphs for: '+str(len(valid_subgraphs_result))

#    frag_scores_result = score_valid_fragments_with_caching(base_name, valid_subgraphs_result)

    return alpha, beta, dict_complexes, dict_constants, dict_complexes_reverse, dict_constants_reverse

if __name__ == '__main__':
    alpha, beta, dict_complexes, dict_constants, dict_complexes_reverse, dict_constants_reverse = main()
