import os
import sys
import argparse
import datetime
import subprocess as sp
from lark import Lark, Token
from jinja2 import Template
from post import Stage1, Stage2, Stage3, Stage4
from resources.m_template import mod_form

parser = argparse.ArgumentParser()
parser.add_argument("term_file", help="name of file containing postprocessing term spec")
parser.add_argument("-b", "--build-info", action="store_true", help="print build_info instead of code")
parser.add_argument("-x", "--print-latex", action="store_true", help="print latex equation instead of code")
parser.add_argument("-p", "--print-report", action="store_true", help="print code generation report instead of code")
args = parser.parse_args()

calc_grammar = """
    ?varlist: "[" [NAME ("," NAME)*] "]"

    ?start: statement*

    ?statement: NAME [ attr_list ] "=" sum -> assign_var
              | avg "{" [NAME ("," NAME)*] "}" -> assign_avg_var
              | varlist

    attr_list: "(" [attr_pair ("," attr_pair)*] ")"

    ?attr_pair: NAME "=" BOOL
              | NAME "=" INT
              | NAME "=" ESCAPED_STRING

    ?sum: product
        | sum "+" product   -> add
        | sum "-" product   -> sub

    ?product: atom
        | product "*" atom  -> mul
        | product "/" atom  -> div

    ?atom: NUMBER           -> number
         | "-" atom         -> neg
         | NAME             -> var
         | NAME "'"         -> fluc
         | "$" NAME         -> env
         | "(" sum ")"      -> paren
         | inlinefunc "(" sum ")" -> icall
         | mathfunc "(" sum ("," sum)* ")" -> fcall
         | derivative "(" NAME ")" -> dnx

    avg: "avg" [NAME]

    ?inlinefunc: "sqr" -> sqr
               | "pow3" -> pow3

    ?mathfunc: "log" -> log
             | "exp" -> exp
             | "sqrt" -> sqrt
             | "abs" -> abs
             | "rxn_rate" -> rxn_rate
             | "$" NAME   -> udf

    ?derivative: "ddx"  -> ddx
               | "d2dx" -> d2dx
               | "ddy" -> ddy
               | "d2dy" -> d2dy
               | "ddz" -> ddz
               | "d2dz" -> d2dz

    %import common.CNAME -> NAME
    %import common.NUMBER
    %import common.ESCAPED_STRING
    %import common.INT
    %import common.WS

    BOOL: "true" | "false"

    COMMENT: /#.*/

    %ignore COMMENT
    %ignore WS
"""

def tok_to_bool(tok):
    "Convert the value of `tok` from string to bool, while maintaining line number & column."
    # tok.type == 'BOOL'
    return Token.new_borrow_pos(tok.type, tok.value == "true", tok)

def tok_to_int(tok):
    "Convert the value of `tok` from string to int, while maintaining line number & column."
    # tok.type == 'INT'
    return Token.new_borrow_pos(tok.type, int(tok), tok)

def tok_to_str(tok):
    "Convert the value of `tok` from string to string, while maintaining line number & column."
    # tok.type == 'ESCAPED_STRING'
    return Token.new_borrow_pos(tok.type, tok.value.strip('"'), tok)

class VersionInfo(object):

    def __init__ (self, input_raw):

        bd = 'build date: {}'
        bb = 'build base: {}'

        self.build_date = bd.format(datetime.datetime.today().ctime())

        try:
            self.build_base = bb.format(sp.check_output(["git", "rev-parse", "HEAD"]).strip())
        except sp.CalledProcessError:
            self.build_base = "None"

        try:
            sp.check_output("git diff --exit-code".split())
            sp.check_output("git diff --cached --exit-code".split())
            self.off_base = ""
        except sp.CalledProcessError:
            self.off_base = "built with changes not commited"

        tinfo = '''
================================================================================
************************************ terms *************************************
--------------------------------------------------------------------------------
{}
================================================================================
'''

        self.term_info = tinfo.format(input_raw)

    def fparams (self):
        return "\n".join([self.build_date, self.build_base, self.off_base, self.term_info])



def test(terms_raw, report=False, latex=False):
    """Compiles the post-processing term specifications from DSL into targets.

    This compiles the Lark AST through all 4 pipeline stages (Stage 1 to 4) 
    and outputs the resulting Fortran source code via Jinja2, a LaTeX equation 
    dictionary, or a code generation IR report.

    Args:
        terms_raw (str): The raw string contents of the DSL terms specification input file.
        report (bool): If True, generates and saves the intermediate code generation report.
        latex (bool): If True, prints LaTeX equations of the averaged terms to stdout.
    """

    parser = Lark(calc_grammar,
            parser='lalr',
            lexer_callbacks =
            {
            'ESCAPED_STRING': tok_to_str,
            'INT': tok_to_int,
            'BOOL': tok_to_bool
            }
            )

    tree = parser.parse(terms_raw)

    ir1 = Stage1(tree)

    ir2 = Stage2(ir1)

    ir3 = Stage3(ir2)

    ir4 = Stage4(ir3)

    cdict = ir4.print_program()

    if report:
        ir4.save_ir()
    elif latex:
        print("{")
        for avg in ir4.averaged.values():
            print('"{}" : '.format(avg.name))
            # print r"\begin{equation}"
            print('r"${}$" ,'.format(avg.latex))
            # print r"\end{equation}"
        print("}")
    else:
        print(Template(mod_form).render(**cdict))



def build_info(terms_raw):
    """Generates and prints build information for the current code generation run.

    Args:
        terms_raw (str): The raw string contents of the DSL terms specification input file.
    """

    vinfo = VersionInfo(terms_raw)

    print(vinfo.fparams())



if __name__ == '__main__':

    with open(args.term_file) as inputfile:
        terms_raw = ((inputfile.read()))

    if args.build_info:
        build_info(terms_raw)
    else:
        test(terms_raw, report=args.print_report, latex=args.print_latex)