import re


# Function only used in the case of NEMO to replace the wrkalloc calls with actual allocations.
def replace_wrkalloc(text):
    def remove_comments(_text):
        # Function to remove the comments
        _lines = _text.split("\n")
        for _line in _lines:
            comments = _line.find("!")
            if comments > -1:
                new_line = _line[:comments]
                _text = _text.replace(_line, new_line)
        return _text

    def argument_type(_argument, _text):
        # Find if the argument is a variable or a dimension
        try:
            int(_argument)
        except ValueError:
            if _argument in ["jpi", "jpj", "jpk"]:
                return "u",
            else:
                for line in _text.split("\n"):
                    match = re.search("[^\!]*real.*::.*" + _argument, line, flags=re.I)
                    if match:
                        dimensions = line.count(":") - 2
                        return "r", dimensions
                    match = re.search("[^\!]*integer.*allocatable.*::.*" + _argument, line, flags=re.I)
                    if match:
                        dimensions = line.count(":") - 2
                        return "i", dimensions

        return "u",

    def strip_arguments(_line, _text):
        # Get the dimensions and the arguments of the wrk_alloc call"
        comments = _line.find("!")
        if comments > -1:
            _line = _line[:comments]
        opening = _line.find("oc(")
        closing_indices = [i for i, x in enumerate(_line) if x == ")"]
        new_line = _line[opening + 3:closing_indices[-1]]
        arguments = [x.strip() for x in new_line.split(",")]
        dimensions = []
        variables = []
        _text = remove_comments(_text)

        for argument in arguments:
            type = argument_type(argument, _text)
            if type[0] == 'r' or type[0] == "i":
                variables.append(argument)
            if type[0] == 'u':
                if argument.count("kistart"):
                    displacement = argument.split("=")[-1]
                    dimensions[0] = displacement + " : " + dimensions[0] + " + " + displacement
                elif argument.count("kjstart"):
                    displacement = argument.split("=")[-1]
                    dimensions[1] = displacement + " : " + dimensions[1] + " + " + displacement
                elif argument.count("kkstart"):
                    displacement = argument.split("=")[-1]
                    dimensions[2] = displacement + " : " + dimensions[2] + " + " + displacement
                elif argument.count("klstart"):
                    displacement = argument.split("=")[-1]
                    dimensions[3] = displacement + " : " + dimensions[3] + " + " + displacement
                else:
                    dimensions.append(argument)
        return variables, dimensions

    list_of_things_to_replace = [["dimension", "DIMENSION"], ["pointer", "POINTER"], ["wrk_alloc (", "wrk_alloc("],
                                 ["wrk_dealloc (", "wrk_dealloc("]]

    match = re.search('\n[^\!]*call *wrk_alloc *\(', text, flags=re.I)
    if match:
        # Replace some strings defined above to simplify the process
        for element_to_replace in list_of_things_to_replace:
            text = text.replace(element_to_replace[0], element_to_replace[1])

        # Divide the text by lines
        lines = text.split("\n")

        # Replace pointer for allocate when needed
        type_dec_end = [i for i, l in enumerate(lines) if l.count("END TYPE") and not l.count("(")]
        type_dec_start = [i for i, l in enumerate(lines) if l.count("TYPE") and not l.count("(") and i not in type_dec_end]

        for i, line in enumerate(lines):
            if line.count("POINTER"):
                if type_dec_start:
                    substitute = True
                    for _j in range(len(type_dec_start)):
                        if type_dec_start[_j] < i < type_dec_end[_j]:
                            type_name = lines[type_dec_start[_j]].split()[-1]
                            if type_name == "arrayptr":
                                substitute = False
                    if substitute:
                        lines[i] = line.replace("POINTER", "ALLOCATABLE")
                else:
                    lines[i] = line.replace("POINTER", "ALLOCATABLE")
        # Rebuild text
        text = "".join([l + "\n" for l in lines[:]])

        # Look for lines containing wrk_alloc
        lines_containing_wrk_alloc = [i for i, m in enumerate(lines) if
                                      re.search('^[^\!]*call *wrk_alloc *\(', m, flags=re.I)]
        # Look for lines containing wrk_dealloc
        lines_containing_wrk_dealloc = [i for i, m in enumerate(lines) if
                                        re.search('^[^\!]*call *wrk_dealloc *\(', m, flags=re.I)]
        for n in lines_containing_wrk_alloc:
            # Find if there is a condition for the allocation
            if re.search('if.*call_wrkalloc *\(', lines[n], flags=re.I):
                allocation_has_condition = True
            else:
                allocation_has_condition = False

            if allocation_has_condition:
                closing_parentesis = lines[n].find(")")
                condition = lines[n][:closing_parentesis + 1] + "   "

            new_line = ""
            variables, dimensions = strip_arguments(lines[n], text)
            for variable in variables:
                alloc = "      ALLOCATE( " + variable + " ( "
                if allocation_has_condition:
                    alloc = condition + alloc
                for num, dimension in enumerate(dimensions):
                    alloc += dimension
                    if num < len(dimensions) - 1:
                        alloc += " , "
                alloc += " ) )"
                new_line += alloc + "\n"

            text = text.replace(lines[n], new_line)
        for n in lines_containing_wrk_dealloc:
            if re.search('if.*call.*wrk_dealloc *\(', lines[n], flags=re.I):
                deallocation_has_condition = True
            else:
                deallocation_has_condition = False

            if deallocation_has_condition:
                closing_parentesis = lines[n].find(")")
                condition = lines[n][:closing_parentesis + 1] + "   "
            new_line = ""
            variables, dimensions = strip_arguments(lines[n], text)
            for variable in variables:
                dealloc = "      DEALLOCATE( " + variable
                if deallocation_has_condition:
                    dealloc = condition + dealloc
                dealloc += " )"
                new_line += dealloc + "\n"

            text = text.replace(lines[n], new_line)
        # Fix substituting pointers inside TYPE definition
        # Divide the text by lines
        lines = text.split("\n")

    return text
