from typing import Any, List
from pytlex_core.TimeX import TimeX
from pytlex_core import Link
import datetime

"""
Main function of the Connectivity_Increaser script. 
@:arg output_as_dict, result from the partitioner in 
                    in dictionary form
@:arg number_of_links in the original TimeML graph
"""

def connect_partitions(output_as_dict, number_of_links) -> Link:
    """
    create a dictionary to store the disconnected partitions
    anchor is defined as the partition that holds the DCT
    connectivity candidates are other partitions that have
    time expressions that could be connected with the dct
    """
    disconnected_partitions = {"anchor": [], "connectivity_candidates": []}

    """
    if a partition is the main partition and has the DCT, then it is an anchor
    if the partition has the DCT but it is not a main partition, it is also the anchor
    other partitions with time expressions are candidates for connectivity
    """
    for partition in (output_as_dict["main_graphs"]+ output_as_dict["subordination_graphs"]):
        if partition.type == "main_graph" and get_dct(partition):
            disconnected_partitions["anchor"] = partition
        elif get_dct(partition):
            disconnected_partitions["anchor"] = partition
        elif has_time_expressions(partition.nodes):
            disconnected_partitions["connectivity_candidates"].append(partition)

    #add to the dictionary an entry "suggested links", where
    output_as_dict["suggested_links"] = try_suggest_link(disconnected_partitions, number_of_links)

"""
@:arg set of nodes in a partition
@:returns true if one of the nodes is a time expression, false otherwise
"""
def has_time_expressions(nodes):
    for node in nodes:
        try:
            if node.tID is not None:
                return True
        except AttributeError:
            pass
    return False


"""
@:arg the dictionary of disconnected partitions
@:arg the number of links in the original timeML graph
@:returns the list of suggested links for this partition
"""
def try_suggest_link(disconnected_partitions, number_of_links):
    #get the dct from the anchor
    dct = get_dct(disconnected_partitions["anchor"])
    suggested_links = []

    #if the dct exists
    if dct:
        try:
            #go through the candidates
            for partition in disconnected_partitions["connectivity_candidates"]:
                #ge the time expressions for this partition
                time_expressions = get_timexs(partition.nodes)
                #suggest a link between the dct and any time expressions
                suggestion = suggest_link(dct, time_expressions, number_of_links)
                #if there was a suggestion, add it to the list
                if suggestion is not None:
                    suggested_links.append(suggestion)
            #return the list
            return suggested_links
        except KeyError:
            #if there are no connectivity candidates, return an empty list
            return suggested_links

"""
@:arg the partition that contains the dct
@:returns the DCT
"""
def get_dct(partition_with_dct) -> TimeX:
    try:
        for node in partition_with_dct.nodes:
            try:
                #a DCT is a node that contains "CREATION_TIME" or "PUBLICATION_TIME" as their function value
                if node.documentFunction == "CREATION_TIME" or  node.documentFunction == "PUBLICATION_TIME":
                    return node
            except AttributeError:
                pass
    except AttributeError:
        return None

"""
@:arg a set of nodes
@:returns time expressions that fulfill a number of conditions
"""
def get_timexs(nodes: Any):
    timexs = []
    for node in nodes:
        try:
            if node.tID:
                if node.type not in {"DURATION", "SET"} and \
                        node.value not in {"PAST_REF", "PRESENT_REF", "FUTURE_REF", "PXY"}:
                    timexs.append(node)
        except AttributeError:
            pass
    return timexs

"""
@:arg dct, the DCT of the graph
@:arg timexs, the list of time expressions from the candidate partition
@:arg number_of_links, the number of links in the graph
@:returns a link that will connect the DCT with another time expression
"""
def suggest_link(dct: TimeX, timexs: List[TimeX], number_of_links):
    """
    In each step of the algorithm, take the value for month, date and day
    of one of the time expressions and compare with the respective values
    of the dct. The link tempral relation will be BEFORE, AFTER or INCLUDES
    accordingly
    """
    if timexs:
        dct_year = int(dct.value[0:4])
        dct_month = int(dct.value[5:7])
        dct_day = int(dct.value[8:10])
        for time_expression in timexs:
            time_value = time_expression.value
            if 'X' not in time_value[0: 4]:
                time_expression_year = int(time_value[0: 4])
            else:
                continue

            if time_expression_year < dct_year:
                number_of_links += 1
                return Link.Link(number_of_links, "TLINK", "BEFORE", time_expression, dct)
            elif time_expression_year > dct_year:
                number_of_links += 1
                return Link.Link(number_of_links, "TLINK", "AFTER", time_expression, dct)
            else:
                time_expression_month = time_value[5:7]
                if len({"X", "S", "U", "F", "A", "P", "E"}.intersection(set(time_value[5:]))) != 0:
                    number_of_links += 1
                    return Link.Link(number_of_links, "TLINK", "INCLUDES", time_expression, dct)
                elif time_expression_month == "":
                    number_of_links += 1
                    return Link.Link(number_of_links, "TLINK", "INCLUDES", time_expression, dct)
                elif "Q" in time_expression_month:
                    dct_month_as_quarter = convert_month_to_quarter(dct_month)
                    if time_expression_month < dct_month_as_quarter:
                        return Link.Link(number_of_links + 1, "TLINK", "BEFORE", time_expression, dct)
                    elif time_expression_month > dct_month_as_quarter:
                        number_of_links += 1
                        return Link.Link(number_of_links, "TLINK", "AFTER", time_expression, dct)
                    else:
                        number_of_links += 1
                        return Link.Link(number_of_links, "TLINK", "INCLUDES", time_expression, dct)
                elif "H" in time_expression_month:
                    dct_month_as_half = convert_month_to_half(dct_month)
                    if time_expression_month < dct_month_as_half:
                        return Link.Link(number_of_links + 1, "TLINK", "BEFORE", time_expression, dct)
                    elif time_expression_month > dct_month_as_half:
                        number_of_links += 1
                        return Link.Link(number_of_links, "TLINK", "AFTER", time_expression, dct)
                    else:
                        number_of_links += 1
                        return Link.Link(number_of_links, "TLINK", "INCLUDES", time_expression, dct)
                elif "W" in time_expression_month:
                    if time_value[5:8] != time_expression_month:
                        time_expression_month = time_value[5:8]
                    _, week, _ = datetime.date(dct_year,dct_month,dct_day).isocalendar()
                    dct_as_week = "W" +str(week)
                    if time_expression_month < dct_as_week:
                        return Link.Link(number_of_links + 1, "TLINK", "BEFORE", time_expression, dct)
                    elif time_expression_month > dct_as_week:
                        number_of_links += 1
                        return Link.Link(number_of_links, "TLINK", "AFTER", time_expression, dct)
                    else:
                        number_of_links += 1
                        return Link.Link(number_of_links, "TLINK", "INCLUDES", time_expression, dct)
                else:
                    if int(time_expression_month) < dct_month:
                        return Link.Link(number_of_links + 1, "TLINK", "BEFORE", time_expression, dct)
                    elif int(time_expression_month) > dct_month:
                        number_of_links += 1
                        return Link.Link(number_of_links, "TLINK", "AFTER", time_expression, dct)
                    else:
                        time_expression_day = time_value[8:10]
                        if time_expression_day == "" or len({"X", "S", "U", "F", "A", "P", "E"}.intersection(set(time_value[8:]))) != 0:
                            number_of_links += 1
                            return Link.Link(number_of_links, "TLINK", "INCLUDES", time_expression, dct)
                        else:
                            if int(time_expression_day) < dct_day:
                                return Link.Link(number_of_links + 1, "TLINK", "BEFORE", time_expression, dct)
                            elif int(time_expression_day) > dct_day:
                                number_of_links += 1
                                return Link.Link(number_of_links, "TLINK", "AFTER", time_expression, dct)
                            else:
                                number_of_links += 1
                                return Link.Link(number_of_links, "TLINK", "INCLUDES", time_expression, dct)
    return None

"""
@:arg an integer value representing a month
@returns the string representation of that integer as a quarter of the year 
"""
def convert_month_to_quarter(month):
    if month in {1,2,3}:
        return "Q1"
    elif month in {4,5,6}:
        return "Q2"
    elif month in {7,8,9}:
        return "Q3"
    return "Q4"

"""
@:arg an integer value representing a month
@returns the string representation of that integer as a half of the year 
"""
def convert_month_to_half(month):
    if month in {1,2,3,4,5,6}:
        return "H1"
    return "H2"
