#!/usr/bin/env python

import sys, os, shutil, glob
from optparse import OptionParser, make_option

import ufl
import ufl.algorithms

import sfc as sfc
from sfc.common.options import default_options

### Define available commandline options
parser = OptionParser()

parser.add_option('--outputdir', '-o', action="store", default=".", type="str",
                  help="Output directory.")
parser.add_option('--prefix', '-p', action="store", default="", type="str",
                  help="Prefix string.")
parser.add_option('-O', action="store", default=0, type="int",
                  help="Compiler optimization level.")
parser.add_option('--cache', '-c', action="store", default=0, type="int", 
                  help="Use JIT compiler cache.")
parser.add_option('--integration-method', '-i', action="store", default="quadrature", type="str", 
                  help="Integration method: [quadrature | symbolic].") # TODO: more?
parser.add_option('--integration-order', '-d', action="store", default=None, type=int, 
                  help="Integration order.")
parser.add_option('--dolfin-wrappers', '-w', action="store", default=0, type="int", 
                  help="Generate DOLFIN C++ wrapper code.")
parser.add_option('--benchmark', '-b', action="store", default=0, type="int", 
                  help="Run benchmark of element tensor computation.")
parser.add_option('--safemode', '-s', action="store", default=0, type="int", 
                  help="Use safe mode when generating (quadrature) code, writing complete integrand expressions unmodified.")
parser.add_option('--debug', '-g', action="store", default=0, type="int", 
                  help="Enable generation of debugging code.")
parser.add_option('--timing', '-t', action="store", default=0, type="int", 
                  help="Enable timing of code generation process.")
parser.add_option('--use_expand_indices2', '-e', action="store", default=0, type="int", 
                  help="Use new experimental version of expand_indices.")


# TODO: Add more options here! Synchronize with default_options() ...


### Parse argv
argv = sys.argv[1:]
(options, args) = parser.parse_args(args=argv)

# Get form filenames from remaining arguments
if len(args) == 0:
    print "Expecting form filename(s)."
    print 
    parser.print_usage()
    sys.exit(-1)

filenames = args
for f in filenames:
    if not os.path.exists(f):
        print "No such file as '%s'" % f
        print 
        parser.print_usage()
        sys.exit(-1)


### Update sfc options from commandline arguments 
sfc_options = default_options()

# Options for this script only:
outputdir = options.outputdir
cache = options.cache

# Options for SFC module behaviour
sfc_options.output.enable_timing = options.timing
sfc_options.code.integral.enable_debug_code = options.debug
sfc_options.compilation.enable_debug_code = options.debug

# Options for code generation
sfc_options.code.integral.integration_method = options.integration_method
sfc_options.code.integral.integration_order = options.integration_order
sfc_options.code.integral.use_expand_indices2 = options.use_expand_indices2
sfc_options.code.integral.safemode = options.safemode
sfc_options.code.dolfin_wrappers = options.dolfin_wrappers

# Options for compilation
O = options.O

# Options for postprocessing
benchmark = options.benchmark
if benchmark:
    import ufc_benchmark

# TODO: Add more options here! Synchronize with default_options() ...


### Load form file
all_forms = []
num_forms = 0
for filename in filenames:
    print "Loading form file '%s'" % filename
    forms = ufl.algorithms.load_forms(filename) # FIXME: Loading elements from a .ufl file is not supported here.
    all_forms.append((filename, forms))
    num_forms += len(forms)


### Compile forms
try:
    os.makedirs(outputdir)
except:
    print "Failed to create output directory", outputdir
os.chdir(outputdir)

print "Compiling %d form files with total %d forms..." % (len(all_forms), num_forms)
for (filename, forms) in all_forms:
    print "Compiling %d forms from file '%s'" % (len(forms), filename)
    
    if options.prefix:
        sfc_options.code.prefix = options.prefix
    else:
        sfc_options.code.prefix = os.path.basename(filename.strip()).replace(".ufl", "")
    
    if cache or benchmark:
        compiled_forms, module, formdatas = sfc.jit(forms, sfc_options)
        
        for compiled_form in compiled_forms:
            print "    Compiled form info:"
            print "        signature: %s" % compiled_form.signature()
            print "        rank %d" % compiled_form.rank()
            print "        %d coefficients." % compiled_form.num_coefficients()
            print "        %d cell integrals" % compiled_form.num_cell_integrals()
            print "        %d exterior facet integrals" % compiled_form.num_exterior_facet_integrals()
            print "        %d interior facet integrals" % compiled_form.num_interior_facet_integrals()
            if benchmark:
                ufc_benchmark.bench(compiled_form) # TODO: Probably need more args here
    else:
        sfc.generate_code(forms, sfc_options)

