"""
Parser for formal Chemical Reaction Networks.

Note: This file was adapted from the python package: `Nuskell`, which is
    developed at the *DNA and Natural Algorithms Group*, Caltech.  
    Original code written by Seung Woo Shin (seungwoo.theory@gmail.com).

Test using tests/test_crn_parser.py.
"""

from pyparsing import (Word, Literal, Group, Suppress, Combine, Optional, ParseException,
                       alphas, nums, alphanums, delimitedList, StringStart, StringEnd, LineEnd,
                       ZeroOrMore, OneOrMore, pythonStyleComment, ParseElementEnhance)

class CRNParseError(Exception):
    pass

def crn_document_setup(modular = False):
    """Parse a formal chemical reaction network.

    Args:
      modular <optional:bool>: Adds an additional nesting for modules within a
        CRN. Use one line per module (';' separates reactions).

    Format:
      # A list of reactions, optionally with reaction rates:
      # <- this is a comment!
      B + B -> C    # [k = 1]
      C + A <=> D   # [kf = 1, kr = 1]
      <=> A  [kf = 15, kr = 6]

      # Note that you can write multiple reactions in one line:
      A + 2C -> E [k = 13.78]; E + F <=> 2A  [kf = 13, kr = 14]

    Returns:

    """
    # NOTE: If you want to add support for multiple modules per line, you can use
    # the '|' character.

    W = Word
    G = Group
    S = Suppress
    O = Optional
    C = Combine
    L = Literal

    def T(x, tag):
        """ Return a *Tag* to distinguish (ir)reversible reactions """
        def TPA(tag):
            return lambda s, l, t: [tag] + t.asList()
        return x.setParseAction(TPA(tag))

    crn_DWC = "".join(
        [x for x in ParseElementEnhance.DEFAULT_WHITE_CHARS if x != "\n"])
    ParseElementEnhance.setDefaultWhitespaceChars(crn_DWC)

    identifier = W(alphas, alphanums + "_")

    multiplier = W(nums)
    species = G(O(multiplier) + identifier)

    number = W(nums, nums)
    num_flt = C(number + O(L('.') + number))
    num_sci = C(number + O(L('.') + number) + L('e') + O(L('-') | L('+')) + W(nums))
    gorf = num_sci | num_flt

    # Make specification of forward, backward, reverse more flexible
    kf = S('kf') | S('fw')
    kr = S('kr') | S('bw') | S('rv')

    k = G(S('[') + O(S('k') + S('=')) + gorf + S(']'))
    rev_k = G(S('[') + kf + S('=') + gorf + S(',') + kr + S('=') + gorf + S(']')) | G(S('[') + gorf + S(',') + gorf + S(']'))

    concentration = T(species + S('@') + G(L("initial") | L("i") | L("constant") | L("c")) + G(gorf), 'concentration')

    reaction = T(G(O(delimitedList(species, "+"))) +
                 S("->") +
                 G(O(delimitedList(species, "+"))) + O(k), 'irreversible')

    rev_reaction = T(G(O(delimitedList(species, "+"))) +
                     S("<=>") +
                     G(O(delimitedList(species, "+"))) + O(rev_k), 'reversible')

    expr = G(reaction | rev_reaction | concentration)

    if modular:
        module = G(expr + ZeroOrMore(S(";") + expr))
    else:
        module = expr + ZeroOrMore(S(";") + expr)

    crn = OneOrMore(module + ZeroOrMore(S(LineEnd())))

    document = StringStart() + ZeroOrMore(S(LineEnd())) + crn + StringEnd()
    document.ignore(pythonStyleComment)
    return document

def flint(inp):
    try:
        return int(inp) if float(inp) == int(float(inp)) else float(inp)
    except OverflowError:
        return inp

def post_process(crn, modular = False, defaultrate = 1, defaultmode = 'initial', defaultconc = 0):
    """Process a parsed CRN.

    Does:
     - remove 'reversible' and 'irreversible' tags.
        (this information is stored implicitly with one or two rates).
     - Translates species multipliers to multisets
     - Extracts the set of all species in the CRN.
     - Adds 'None' as reaction rate if no rate was specified.

    Does not:
     - Check if reactions are specified twice
     - Check stochiometry

    Returns:
        new: the CRN in format [[r],[p],k]
        species: dictionary species[name]=('initial', float)

    """
    def remove_multipliers(species):
        flat = []
        for s in species:
            if len(s) == 1:
                flat.append(s[0])
            elif len(s) == 2:
                ss = [s[1]] * int(s[0])
                flat.extend(ss)
        return flat

    default = dict()
    species = dict()
    def get_module(crn):
        new = []
        for line in crn:
            if line[0] == 'concentration':
                spe = line[1][0]
                ini = 'initial' if line[2][0][0] == 'i' else 'constant'
                num = line[3][0]
                if spe in species:
                    raise CRNParseError(f'Duplicate concentration specification for {spe}!')
                species[spe] = (ini, flint(num))
                continue
            elif len(line) == 3:
                # No rate specified
                t, r, p = line
                r = remove_multipliers(r)
                p = remove_multipliers(p)
                if t == 'reversible':
                    new.append([r, p, [defaultrate, defaultrate]])
                elif t == 'irreversible':
                    new.append([r, p, [defaultrate]])
                else:
                    raise CRNParseError('Wrong CRN format!')
            elif len(line) == 4:
                t, r, p, k = line
                r = remove_multipliers(r)
                p = remove_multipliers(p)
                k = [flint(v) for v in k]
                if t == 'reversible':
                    assert len(k) == 2
                    new.append([r, p, k])
                elif t == 'irreversible':
                    assert len(k) == 1
                    new.append([r, p, k])
                else:
                    raise CRNParseError('Wrong CRN format!')
            else:
                raise CRNParseError('Wrong CRN format!')
            for s in r + p:
                if s not in default:
                    default[s] = (defaultmode, defaultconc)
        return new

    if modular:
        new = []
        for module in crn:
            mcrn = get_module(module)
            if mcrn:
                new.append(mcrn)
    else:
        new = get_module(crn)

    default.update(species)
    return new, default

def parse_crn_file(filename, process = True, modular = False, **kwargs):
    """Parses a CRN from a file.

    Args:
      filename (<str>): The name of the file for parsing.
      process (optional <bool>): Process the format of the returned CRN.

    Returns:
      crn (<lol>): List of list representation of a CRN
      species (<set()>): A set of all involved species (only when process=True)

    """
    crn_document = crn_document_setup(modular)
    if process:
        return post_process(crn_document.parseFile(filename, parseAll=True).asList(), modular, **kwargs)
    else:
        return crn_document.parseFile(filename, parseAll=True).asList()

def parse_crn_string(data, process = True, modular = False, **kwargs):
    """Parses a CRN from a string.

    Args:
      filename (<str>): The name of the file for parsing.
      process (optional <bool>): Process the format of the returned CRN.

    Returns:
      crn (<lol>): List of list representation of a CRN
      species (<set()>): A set of all involved species (only when process=True)

    """
    crn_document = crn_document_setup(modular)
    if process:
        return post_process(crn_document.parseString(data).asList(), modular, **kwargs)
    else:
        return crn_document.parseString(data).asList()


