from __future__ import print_function
import pickle
import warnings
from os import listdir
from os.path import isfile, join

from .Classes import Interface, Vault, Module, DerivedType, ObtainSourcefileInfo, Counter
from .Functions import insert_variable_precision_specification_to_allocatable, \
    insert_variable_precision_specification_to_namelist_parameter, insert_variable_precision_to_routine_variable, \
    add_after_call_routines, load_source_files, split_vars_by_allocatable, split_vars_by_main_or_routine, \
    split_by_appearance_in_namelist, replace_workalloc_in_file, preprocess_file, postprocess_file


def parse_sources(path_to_model_sources, save_vault=False):
    # Find the nemo source-files in sources
    # (Using the sources that already modified replacing real declarations by type(rpe_var) )
    nemo_sources = [f for f in listdir(path_to_model_sources) if
                    isfile(join(path_to_model_sources, f)) and f.count(".f90")]

    # Initialize two empty dictionaries that will contain all the routines and all the functions
    object_oriented_dictionary = {"procedures": {}, "interfaces": {}, "variables": {}, "modules": {},
                                  "derivedtypes": {}}

    # Initialize an empty dictionary that will contain all the source-files.
    sourcefiles = {}
    vault = Vault()

    print("\n\nParsing sourcecode...\n")

    # Going over all the nemo source-files loading the files,
    # parsing the routines and storing both things in dictionaries.
    for index, filename in enumerate(nemo_sources):
        module_name = filename.replace(".f90", "")
        sourcefiles[module_name] = ObtainSourcefileInfo(join(path_to_model_sources, filename))
        module = Module(module_name)
        f = sourcefiles[module_name]
        f.module = module
        f.find_routines()
        f.find_functions()
        f.find_module_variables()
        for procedure in f.objects:
            object_oriented_dictionary["procedures"][procedure.name] = procedure
            module.add_procedure(procedure)
        vault.add_module(module)

    # Going over all the nemo source-files parsing the interfaces and storing them in dictionary.
    for index, filename in enumerate(nemo_sources):
        module_name = filename.replace(".f90", "")
        f = sourcefiles[module_name]
        f.find_interfaces(object_oriented_dictionary["procedures"])
        # Getting and saving modules used
        f.module = vault[module_name]
        f.find_used_modules(vault)
        for interface in f.objects:
            if isinstance(interface, Interface):
                object_oriented_dictionary["interfaces"][interface.name] = interface
                vault[module_name].procedures.append(interface)

    for filename in sourcefiles:
        module_name = filename.replace(".f90", "")
        f = sourcefiles[module_name]
        f.find_types()
        for derived_type in f.objects:
            if isinstance(derived_type, DerivedType):
                object_oriented_dictionary["derivedtypes"][derived_type.name] = derived_type
                vault[module_name].procedures.append(derived_type)

    vault.process_variables()

    # Find implicit function declarations
    for filename in sourcefiles:
        module_name = filename.replace(".f90", "")
        f = sourcefiles[module_name]
        f.find_implicit_functions(vault=vault)
        for procedure in f.objects:
            if procedure not in vault[module_name].procedures:
                vault[module_name].procedures.append(procedure)

    vault.process_procedures()
    # Search which modules does each module has access to
    vault.expand_used_modules()
    # Search which modules have access to this module
    vault.expand_used_by()

    if save_vault:
        # This workaround is needed to allow pickle to write the file, don't know yet about its risks
        import sys
        sys.setrecursionlimit(2000)
        pickle.dump(vault, open("vault.pkl", "wb"))
    return vault


def replace_wrkalloc(path_to_input_sources, path_to_output_sources):
    # Find files in sources
    all_sources = [f for f in listdir(path_to_input_sources) if
                   isfile(join(path_to_input_sources, f)) and f.count(".f90")]
    # Process Sources!
    total = len(all_sources)
    for index, filename in enumerate(all_sources):
        print("\rReplacing wrkallocs %i/%i  %20s" % (index+1, total, filename), end="")
        replace_workalloc_in_file(path_to_input_sources, path_to_output_sources, filename)


def preprocess_sources(path_to_input_sources, path_to_output_sources, blacklist=[]):
    # Find files in sources
    all_sources = [f for f in listdir(path_to_input_sources) if
                   isfile(join(path_to_input_sources, f)) and f.count(".f90") and f not in blacklist]
    # Process Sources!
    for index, filename in enumerate(all_sources):
        total = len(all_sources)
        print("\rReplacing reals %i/%i  %20s" % (index+1, total, filename), end="")
        preprocess_file(path_to_input_sources, path_to_output_sources, filename)


def postprocess_sources(path_to_input_sources, path_to_output_sources, vault):
    # Find files in sources
    all_sources = [f for f in listdir(path_to_input_sources) if
                   isfile(join(path_to_input_sources, f)) and f.count(".f90")]
    # all_sources = ['limcons.f90']
    # all_sources = ['sbcrnf.f90']
    total = len(all_sources)
    for index, filename in enumerate(all_sources):
        print("\rFixing issues %i/%i  %20s" % (index + 1, total, filename), end="")
        postprocess_file(path_to_input_sources, path_to_output_sources, filename=filename, vault=vault)


def assign_sbits_to_rpe_variables(path_to_sources, vault=None):
    print("\n\nAssigning significant bits...")

    def condition(v):
        return v.type == "rpe_var" and not v.is_parameter and not v.name.count("tmp_var_")

    source_files = load_source_files(path_to_sources)
    rpe_variables = [var for var in vault.variables if condition(var)]
    allocatable, non_allocatable = split_vars_by_allocatable(rpe_variables)
    var_counter = Counter()
    # Allocatables:
    for variable in allocatable:
        insert_variable_precision_specification_to_allocatable(variable, source_files, var_counter)

    # Non allocatables that appear in a namelist:
    na_in_namelist, na_not_in_namelist = split_by_appearance_in_namelist(non_allocatable, source_files)
    for variable in na_in_namelist:
        insert_variable_precision_specification_to_namelist_parameter(variable, source_files, var_counter)
    in_main, not_in_main = split_vars_by_main_or_routine(na_not_in_namelist)

    # Variables inside routines
    for variable in not_in_main:
        if variable.intent != "in":
            insert_variable_precision_to_routine_variable(variable, source_files, var_counter)

    # Adding sbits assignation and apply_truncation after routine calls
    insert_variable_precision_after_routine_call(source_files, vault=vault, counter=var_counter)

    # To save the variable id into the vault.
    save = True
    if save:
        # Saving modified files
        for filename, source_file in source_files.items():
            source_file.save_modified_file(join(path_to_sources, filename))
        # This workaround is needed to allow pickle to write the file, don't know yet about its risks
        import sys
        sys.setrecursionlimit(2000)
        pickle.dump(vault, open("vault.pkl", "wb"))
    return


def insert_variable_precision_after_routine_call(source_files, vault=None, counter=None):
    for name, source_file in source_files.items():
        # print name
        add_after_call_routines(source_file, vault, counter)
        source_file.rebuild_text()
        # source_file.save_modified_file(join(path_to_sources, source_file.filename))
    return


def create_read_precisions_module_and_add_to_file(path_to_sources, target_module=None, vault=None):
    print("\n\nCreating module to read precision namelist...")
    # Adding the use statement in all the files.
    line_to_add = "   USE read_precisions\n"
    source_files = load_source_files(path_to_sources)

    for name, source_file in source_files.items():
        source_file.lines.insert(1, line_to_add)
        source_file.rebuild_text()
        source_file.save_modified_file(join(path_to_sources, source_file.filename))

    # Create module
    rpe_vars = [var for var in vault.variables if var.type == "rpe_var"]
    assigned_rpe_vars = [var for var in rpe_vars if var.id is not None]
    ids = [var.id for var in assigned_rpe_vars]
    # Path to template is hardcoded
    with open("AdditionalFiles/read_precisions_template", "r") as input_file:
        template = input_file.read()
    module = template.replace("%NUMBER_OF_RPE_VARS%", "%i" % max(ids))
    with open(join(path_to_sources, "read_precisions.f90"), "w") as output_file:
        output_file.write(module)
    # Put call into target module

    if target_module:
        if isfile(join(path_to_sources, target_module)):
            with open(join(path_to_sources, target_module), "r") as input_file:
                original_source = input_file.read()

            lines = original_source.split("\n")
            lines.insert(1, "use read_precisions")
            for index, line in enumerate(lines):
                if line.lower().strip() == "private":
                    lines.pop(index)
            original_source = "\n".join(lines)

            with open(join(path_to_sources, target_module), "w") as output_file:
                output_file.write(original_source)
        else:
            warnings.warn("File %s not found in sources.\nUnable to insert read_precisions module")