__version__ = "0.2"
__author__ = "Reidmen Arostica"
__date__ = "18/07/2020"

# Libraries 
from solvers import iNSE_Monolithic_Solver, iNSE_CT_Solver
from solvers import testing_comparative_plots

# Basic info to user
# print("SOLVER VERSION: {}".format(__version__))
# print("__author__: {}".format(__author__))


def testing_defs_ALE(solver_to_use='LU', gcl_stab=True, jac_at_n=False):
    """
    Computes the iNSE problem under several deformations (including the Identity).
    - Identity -> ('x[0]', 'x[1]')
    - High Oscillation (init with expansion  ) -> ('x[0]*(1+0.9*sin(16*pi*t/2))', 'x[1]')
    - High Oscillation (init with contraction) -> ('x[0]*(1+0.9*sin(16*pi*t/2))', 'x[1]')
    """
    print("""Initializing testing of two deformation cases with input parameters:
    Solver to use -solver_to_use- : {}
    GCL-stabilization term -gcl_stab- : {}
    Evaluation of jacobian at n in time-derivate term -jac_at_n- : {}
    """.format(solver_to_use, gcl_stab, jac_at_n))
    # Define dict of deformation used and output names
    dict_def = {
    #    'Identity': ('x[0] + 0*t', 'x[1]'),
        'High_Osc_exp': ('x[0]*(1+0.9*sin(16*pi*t/2))', 'x[1]'),
        'High_Osc_con': ('x[0]*(1-0.9*sin(16*pi*t/2))', 'x[1]')
    }
    dict_def_names = {
    #    'Identity': ('Idx', 'Idy'),
        'High_Osc_exp': ('High_Osc_exp', 'Idy'),
        'High_Osc_con': ('High_Osc_con', 'Idy')
    }
    dict_defs = {def_type: None for def_type in dict_def.keys()}
    dict_vels = {def_type: None for def_type in dict_def.keys()}
    dict_pres = {def_type: None for def_type in dict_def.keys()}
    for def_type in dict_defs.keys():
        print("\nInitializing iNSE problem with Deformation type: {}\n".format(def_type))
        parameters = {
            'gcl_stab': gcl_stab, # True, False whenever we want to impose CGL in the scheme
            'jac_at_n': jac_at_n, # Choose True, False for taking the Jacobian evaluated at time k+1 or k
            'init_cond': True,
            'Nx': 6*16,
            'Ny': 16,
            'T': 2.0, #0.5, 1.0, 2.0
            'mu': 0.035,
            'rho_f': 1.2,
            'tau': 0.01, # 0.01
            'deg_U': 2,
            'deg_P': 1
        }
        deformation = { 
            'expression': dict_def[def_type],
            'output_name': dict_def_names[def_type]
        }
        stabilization = {
            'type': None
        }
        print("\nSetup with initial condition: {}".format(True))
        dict_defs[def_type], dict_vels[def_type], dict_pres[def_type] = iNSE_Monolithic_Solver(parameters, deformation, stabilization, solver_to_use=solver_to_use)
        print("\n----------------Computation Completed!----------------\n")
    print("\n----------------Simulations Done!----------------.\n")

    # Return dictionaries with values for further processing...
    return dict_defs, dict_vels, dict_pres


def testing_defs_ALE_CT(solver_to_use='LU', gcl_stab=True, jac_at_n=False):
    """
    Computes the iNSE problem under several deformations (including the Identity).
    - Identity -> ('x[0]', 'x[1]')
    - High Oscillation (init with expansion  ) -> ('x[0]*(1+0.9*sin(16*pi*t/2))', 'x[1]')
    - High Oscillation (init with contraction) -> ('x[0]*(1+0.9*sin(16*pi*t/2))', 'x[1]')
    """
    print("""Initializing testing of two deformation cases with input parameters:
    Solver to use -solver_to_use- : {}
    """.format(solver_to_use))
    # Define dict of deformation used and output names
    dict_def = {
    #    'Identity': ('x[0] + 0*t', 'x[1]'),
        'High_Osc_exp': ('x[0]*(1+0.9*sin(16*pi*t/2))', 'x[1]'),
        'High_Osc_con': ('x[0]*(1-0.9*sin(16*pi*t/2))', 'x[1]')
    }
    dict_def_names = {
    #    'Identity': ('Idx', 'Idy'),
        'High_Osc_exp': ('High_Osc_exp', 'Idy'),
        'High_Osc_con': ('High_Osc_con', 'Idy')
    }
    dict_defs = {def_type: None for def_type in dict_def.keys()}
    dict_vels = {def_type: None for def_type in dict_def.keys()}
    dict_pres = {def_type: None for def_type in dict_def.keys()}
    for def_type in dict_defs.keys():
        print("\nInitializing iNSE problem with Deformation type: {}\n".format(def_type))
        parameters = {
            'gcl_stab': gcl_stab,
            'jac_at_n': jac_at_n,
            'init_cond': True,
            'Nx': 6*16,
            'Ny': 16,
            'T': 2.0, #0.5, 1.0, 2.0
            'mu': 0.035,
            'rho_f': 1.2,
            'tau': 0.01, # 0.01
            'deg_U': 1,
            'deg_P': 1
        }
        deformation = { 
            'expression': dict_def[def_type],
            'output_name': dict_def_names[def_type]
        }
        stabilization = {
            'type': None
        }
        print("\nSetup with initial condition: {}".format(True))
        dict_defs[def_type], dict_vels[def_type], dict_pres[def_type] = iNSE_CT_Solver(parameters, deformation, stabilization, solver_to_use=solver_to_use)
        print("\n----------------Computation Completed!----------------\n")
    print("\n----------------Simulations Done!----------------.\n")

    # Return dictionaries with values for further processing...
    return dict_defs, dict_vels, dict_pres


def testing_ALE(solver_to_use='LU', gcl_stab=True):
    """
    Computes testing_defs_ALE (each deformation case) for two different monolithic schemes, i.e.
    using evaluation of J at {n, n+1}, both with GCL + Temam stabilizations.
    """
    print("The solution of the linear system is done with built-in PETS Krylov solver.")
    dict_fields_mt = {jac_name: None for jac_name in ['j_implicit', 'j_explicit']}
    dict_fields_ct = {jac_name: None for jac_name in ['j_implicit', 'j_explicit']}
    for jac_name, jac_at_n in zip(['j_implicit', 'j_explicit'], [False, True]):
        dict_fields_mt[jac_name] = testing_defs_ALE(solver_to_use=solver_to_use, gcl_stab=gcl_stab, jac_at_n=jac_at_n)
        dict_fields_ct[jac_name] = testing_defs_ALE_CT(solver_to_use=solver_to_use, gcl_stab=gcl_stab, jac_at_n=jac_at_n)
    # Generate comparative figures for cases 'j_implicit', 'j_explicit'
    testing_comparative_plots(dict_fields_mt, dict_fields_ct, gcl_stab, solver_to_use)
    return dict_fields_mt, dict_fields_ct

if __name__ == "__main__":
    #testing_defs_ALE_CT(solver_to_use='Krylov')
    #testing_defs_ALE(solver_to_use='LU')
    testing_ALE(solver_to_use='LU') # 'LU', 'Krylov'