from roukf.core import ROUKF
import argparse
import importlib
from dolfin import *
import dolfin
import logging

logging.getLogger().setLevel(logging.INFO)

parameters['form_compiler']['optimize'] = True
parameters['form_compiler']['cpp_optimize'] = True
parameters['form_compiler']['cpp_optimize_flags'] = ('-O3 -ffast-math '
                                                     '-march=native')


def print_timing():
    if dolfin.__version__ >= '2018':
        list_timings(TimingClear.clear, [TimingType.wall])
    else:
        list_timings(TimingClear_clear, [TimingType_wall])


def main(fwd_solver_module, inputfile):
    ''' Run ROUKF parameter estimation.

    Imports the given forward solver module and initializes the forward solver
    with the specified input file.
    Creates a ROUKF solver from the same input file and the instantiated
    forward solver object and solves the optimization problem specified in the
    input file.

    Args:
        fwd_solver_module (str): forward solver module to be imported and
            passed to the ROUKF solver
        inputfile (str): YAML input file with configuration of both
            forward and ROUKF solver
    '''

    fwd_solver = importlib.import_module(fwd_solver_module)
    fwd_solver = fwd_solver.init(inputfile)
    fwd_solver.logger.setLevel(logging.WARNING)

    kf = ROUKF(inputfile, fwd_solver)
    kf.logger.setLevel(logging.INFO)
    kf.solve()

    print_timing()


def get_parser():
    parser = argparse.ArgumentParser(
        description='''\
Run ROUKF parameter estimation.

1. Imports the given forward solver module and initializes the forward solver
   with the specified input file. The solver module is required to have a
   method `init()` which handles the complete setup and returns self.
2. Creates a ROUKF solver from the same input file and the instantiated
   forward solver object and solves the optimization problem specified in the
   input file.''',
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('fwd_solver', type=str,
                        help='Name of the forward solver module (see full'
                        ' documentation), such that it can be imported.\n'
                        'For example:\n navierstokes.solver or '
                        'hyperelasticity.solver')
    parser.add_argument('inputfile', type=str, help='path to YAML input file')
    return parser


if __name__ == '__main__':

    args = get_parser().parse_args()

    main(args.fwd_solver, args.inputfile)