#!/usr/bin/env python3
#
# crnsimulator: simulate chemical reaction networks using ODEs
#
# Written by Stefan Badelt (badelt@caltech.edu).
#
# Use at your own risk.
#
#
from __future__ import absolute_import, division, print_function

import os
import sys
import argparse

from crnsimulator import ReactionGraph, get_integrator, __version__
from crnsimulator.odelib_template import add_integrator_args
from crnsimulator.crn_parser import parse_crn_string, ParseException

class SimulationSetupError(Exception):
    pass

def main(args):
    """Translate a CRN into an ODE system. 

    Optional: Simulate ODEs on-the-fly.
    
    """

    # ********************* #
    # ARGUMENT PROCESSING 1 #
    # ..................... #
    filename = args.output + \
        '.py' if args.output[-3:] != '.py' else args.output
    odename = 'odesystem'

    input_crn = sys.stdin.readlines()
    input_crn = "".join(input_crn)

    try:
        crn, species = parse_crn_string(input_crn)
    except ParseException as ex:
        print('CRN-format parsing error:')
        print('Cannot parse line {:5d}: "{}"'.format(ex.lineno, ex.line))
        print('                          {} '.format(' ' * (ex.col-1) + '^'))
        raise SystemExit


    V = [] # sorted species (vertices) vector
    C = [] # corresponding concentration vector
    seen = set() # keep track of what species are covered

    # Move interesting species to the front, in the given order.
    labels = args.pyplot_labels
    for s in labels:
        if s in seen :
            raise SimulationSetupError('Multiple occurances of {} in labels.'.format(s))
        V.append(s)

        if species[s][0][0] != 'i': 
            raise NotImplementedError('Concentrations must be given as "initial" concentrations.')
        C.append(species[s][1])
        seen.add(s)

    # Append the remaining specified species 
    for s in sorted(species):
        if s in seen : continue
        V.append(s)
        if species[s][0][0] != 'i':
            raise NotImplementedError('Concentrations must be given as "initial" concentrations.')
        C.append(species[s][1])
        seen.add(s)

    # Split CRN into irreversible reactions
    new = []
    for [r, p, k] in crn:
        if None in k:
            print('# Set missing rates to 1.')
            k[:] = [x if x is not None else 1 for x in k]

        if len(k) == 2:
            new.append([r, p, k[0]])
            new.append([p, r, k[1]])
        else:
            new.append([r, p, k[0]])
    crn = new

    # **************** #
    # WRITE ODE SYSTEM #
    # ................ #
    if filename != 'odesystem.py' and not args.force and os.path.exists(filename):
        print('# Reading ODE system from existing file:', filename)
    else:
        # ******************* #
        # BUILD REACTIONGRAPH #
        # ................... #
        RG = ReactionGraph(crn)
        if len(RG.species) != len(V):
            print(len(V), sorted(V))
            print(len(RG.species), sorted(RG.species))
            raise SimulationSetupError('Confusion about which species appear in the reaction network!')

        # ********************* #
        # PRINT ODE TO TEMPLATE #
        # ..................... #
        filename, odename = RG.write_ODE_lib(sorted_vars = V, concvect = C,
                                             jacobian=not args.no_jacobian, filename=filename,
                                             odename=odename)
        print('# Wrote ODE system:', filename)

    # ******************* #
    # SIMULATE ODE SYSTEM #
    # ................... #
    if args.dryrun:
        print('# Dryrun: Simulate the ODE system using:')
        print("#  python {} --help ".format(filename))
    else:
        print('# Simulating the ODE system, change parameters using:')
        print("#  python {} --help ".format(filename))

        integrate = get_integrator(odename, filename)

        # ********************* #
        # ARGUMENT PROCESSING 2 #
        # ..................... #
        integrate(args)

    return


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--version', action='version', version='%(prog)s ' + __version__)
    parser.add_argument("--force", action='store_true',
            help="Overwrite existing files")
    parser.add_argument("--dryrun", action='store_true',
            help="Do not run the simulation, only write the files.")
    parser.add_argument("-o", "--output", default='odesystem', metavar='<str>',
            help="Name of ODE library files.")
    parser.add_argument("--no-jacobian", action='store_true',
            help="Do not compute the Jacobi-Matrix.")

    add_integrator_args(parser)

    args = parser.parse_args()

    main(args)
