#!/usr/bin/env python

## Copyright 2015-2022 Ian Jauslin
## 
## Licensed under the Apache License, Version 2.0 (the "License");
## you may not use this file except in compliance with the License.
## You may obtain a copy of the License at
## 
##     http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.


##########################################################################
##                                                                      ##
##              convert flow equations to various formats               ##
##                                                                      ##
##########################################################################

import sys
import re

def print_usage():
    print("\nusage:",file=sys.stderr)
    print("  export_flow_equation [-i infile] <format> [options]",file=sys.stderr)
    print("\nwhere <format> [options] is one of",file=sys.stderr)
    print("  C [-l oldl_symbol] [-L newl_symbol] [-C Constant_symbol] [-S]",file=sys.stderr)
    print("  javascript [-l oldl_symbol] [-L newl_symbol] [-C Constant_symbol]",file=sys.stderr)
    print("  LaTeX [-l oldl_symbol] [-L newl_symbol] [-C Constant_symbol] [-s]",file=sys.stderr)
    print("",file=sys.stderr)

# to C
def C_engine(argv,text):
    argc=len(argv)

    i=1
    # defaults
    lsym="l"
    Lsym="newL"
    Csym="C"
    oneline=0
    while (i<argc):
        # arguments requiring a variable
        if (argv[i] in ["-l","-L","-C"]):
            if(i>=argc-1):
                print("error: '"+argv[i]+"' must be followed by a string",file=sys.stderr)
                exit(-1)

        # symbol for l
        if (argv[i]=="-l"):
            lsym=argv[i+1]
            i=i+2
        # symbol for new l
        elif (argv[i]=="-L"):
            Lsym=argv[i+1]
            i=i+2
        # symbol for C
        elif (argv[i]=="-C"):
            Csym=argv[i+1]
            i=i+2

        # write each equation on a single line
        elif (argv[i]=="-S"):
            oneline=1
            i=i+1

    return(convert_C(text,lsym,Lsym,Csym,oneline))
        
# to LaTeX
def latex_engine(argv,text):
    argc=len(argv)

    i=1
    # defaults
    lsym=r'\\ell'
    Lsym=r'\\ell'
    Csym="C"
    oneline=1
    while (i<argc):
        # arguments requiring a variable
        if (argv[i] in ["-l","-L","-C"]):
            if(i>=argc-1):
                print("error: '"+argv[i]+"' must be followed by a string",file=sys.stderr)
                exit(-1)

        # symbol for l
        if (argv[i]=="-l"):
            lsym=argv[i+1]
            i=i+2
        # symbol for new l
        elif (argv[i]=="-L"):
            Lsym=argv[i+1]
            i=i+2
        # symbol for C
        elif (argv[i]=="-C"):
            Csym=argv[i+1]
            i=i+2

        # break lines in equation
        elif (argv[i]=="-s"):
            oneline=0
            i=i+1

    return(convert_latex(text,lsym,Lsym,Csym,oneline))

# convert to C format
def convert_C(text, lsym, Lsym, Csym, oneline):
    # make sure quotients are floats
    text=re.sub(r'([0-9])/',r'\1./',text)

    # remove newlines
    if (oneline==0):
        text=text.replace('\n','\\\n')
        # end of lines
        text=text.replace(',\\',';')
    else:
        text=text.replace('\n','')
        # end of lines
        text=text.replace(',',';\n')

    # remove extra space
    text=re.sub(r' {2,}',' ',text)

    # replace left hand side variables
    text=re.sub(r'\[(d*)C *([0-9]*)\] =',r'\1'+Csym+r'\2 =', text)
    text=re.sub(r'\[(d*)% *([0-9]*)\] =',r'\1'+Lsym+r'[\2] =', text)
    # replace right hand side variables
    text=re.sub(r'\[(d*)% *([0-9]*)\]',r'*\1'+lsym+r'[\2]',text)
    text=re.sub(r'\[(d*)C *([0-9]*)\]',r'*\1'+Csym+r'[\2]',text)
    # replace constants in right hand side (and expand powers as products)
    match=re.search(r'\[/C([0-9]*)\^([0-9]*)\]',text)
    while(match):
        power=int(match.group(2))
        # product of constants
        prod=''
        for i in range(power):
            prod=prod+'/'+Csym+match.group(1)
        text=text.replace('[/C'+match.group(1)+'^'+match.group(2)+']',prod)
        match=re.search(r'\[/C([0-9]*)\^([0-9]*)\]',text)

    # replace square roots
    text=re.sub(r's\{([0-9]*)\}',r'sqrt(\1)',text)
    text=text.replace(')sqrt',')*sqrt')

    # re-number variables in a sequential fashion
    variables=[]
    for match in re.finditer(r'\[([0-9]*)\]',text):
        index=int(match.group(1))
        if index not in variables:
            variables.append(index)
            variables.sort()
    for index in range(len(variables)):
        text=text.replace('['+str(variables[index])+']','[new'+str(index)+']')
    text=text.replace('[new','[')

    return(text+';')

# convert to LaTeX format
def convert_latex(text, lsym, Lsym, Csym, oneline):
    # remove newlines
    if (oneline==0):
        text=text.replace('\n','\\\\\n')
    else:
        # end of lines
        text=text.replace(',',',\\\\')

    # remove extra space
    text=re.sub(r' {2,}',' ',text)

    # remove unnecessary 1's
    text=re.sub(r'\(-1\)\[',r'-[',text)
    text=re.sub(r'\(1\)\[',r'[',text)

    # replace left hand side variables
    text=re.sub(r'\[(d*)C *([0-9]*)\] =',r'@\1'+Csym+r'_{\2} =', text)
    text=re.sub(r'\[(d*)% *([0-9]*)\] =',r'@\1'+Lsym+r'_{\2} =', text)
    # replace right hand side variables
    text=re.sub(r'\[(d*)% *([0-9]*)\]',r'@\1'+lsym+r'_{\2}',text)
    text=re.sub(r'\[(d*)C *([0-9]*)\]',r'@\1'+Csym+r'_{\2}',text)

    # rewrite derivatives
    # find the highest order of derivatives
    i=0
    dstr="d"
    index=0
    while (index>=0):
        index=text.find(r'@'+dstr,index)
        i=i+1
        dstr=dstr+"d"
    #replace derivatives
    for j in range(i):
        # number of derivs
        nrd=i-1-j
        dstr=""
        for k in range(nrd):
            dstr=dstr+"d"
        if (nrd>1):
            text=text.replace(r'@'+dstr,"\\partial^{"+str(nrd)+"}")
        elif (nrd==1):
            text=text.replace(r'@'+dstr,"\\partial ")
        else:
            text=text.replace(r'@'+dstr,"")

    # replace constants in right hand side
    text=re.sub(r'\[/C([0-9]*)\^([0-9]*)\]',r'\\frac1{'+Csym+r'_{\1}^{\2}}',text)

    # remove "^{1}"
    text=text.replace("^{1}","")

    # replace square roots
    text=text.replace('s{','\\sqrt{')

    # fractions
    text=re.sub(r'\((-?)([0-9]*)/([0-9]*)\)',r'\1\\frac{\2}{\3}',text)
    # numbers
    #text=re.sub(r'\(([-0-9]*)\)',r'\1',text)

    # fix signs
    text=re.sub(r'\+-',r'-',text)



    return(text)

# read arguments
argc=len(sys.argv)
if (argc<=1):
    print_usage()
    exit(-1)
i=1
infile=""
while i<argc:
    if (sys.argv[i]=="-i"):
        if(i<argc-1):
            infile=sys.argv[i+1]
            i=i+2
        else:
            print("error: '"+sys.argv[i]+"' must be followed by a filename",file=sys.stderr)
            exit(-1)

    else:
        # read input
        if (infile==""):
            # read input from stdin
            text=sys.stdin.read()
        else:
            f=open(infile,'r')
            text=f.read()
            f.close()


        if (sys.argv[i]=="C"):
            print(C_engine(sys.argv[i:],text))
            
        elif (sys.argv[i]=="javascript"):
            # reuse C engine but ensure that the single line option is used
            print(C_engine(sys.argv[i:]+["-S"],text))

        elif (sys.argv[i]=="LaTeX"):
            print(latex_engine(sys.argv[i:],text))

        else:
            print("error: '"+sys.argv[i]+"' is not supported",file=sys.stderr)
            exit(-1)
        break