# from pynars.NARS.DataStructures._py.Link import TermLink
from math import log2
from ..Budget import Budget
from ..Truth import Truth
from .ExtendedBooleanFunctions import *
from copy import deepcopy
from .UncertaintyMappingFunctions import w_to_c
from ..Config import Config
from .Tools import truth_to_quality

def Budget_revision(budget_task: Budget, truth_task: Truth, truth_belief: Truth, truth_derived: Truth, budget_tasklink: Budget=None, budget_termlink: Budget=None, replace=True, replace_tasklink=True, replace_termlink=True):
    '''
    budget_task (Budget):
        budget, of a task, to be revised.
    truth_task (Truth):
        truth of the task.
    truth_belief (Truth):
        truth of the belief.
    truth_derived (Truth):
        truth of a task derived from the task and the belief.
    replace (bool): 
        whether to revise the `budget_task` without a copy.
        default: True
    budget_tasklink (Budget): default: None
        budget, of the tasklink whose post-link task is the one with the `budget`, to be revised.
    replace_tasklink (bool):  
        whether to revise the `budget_tasklink` without a copy.
        default:True
    budget_termlink (Budget): default: None
        budget, of the termlink whose post-link task is the one with the `budget`, to be revised.
    replace_termlink (bool):  
        whether to revise the `replace_termlink` without a copy.
        default:True

    Ref: OpenNARS 3.1.0 BudgetFunctions.java line 72~118
        Evaluate the quality of a revision, then de-prioritize the premises
    '''
    if not replace: budget_task = deepcopy(budget_task)
    diff_task = abs(truth_task.e - truth_derived.e)
    budget_task.priority = And(budget_task.priority, 1-diff_task)
    budget_task.durability = And(budget_task.durability, 1-diff_task)

    if budget_tasklink is not None:
        if not replace_tasklink: budget_tasklink = deepcopy(budget_tasklink)
        budget_tasklink.priority = And(budget_task.priority, diff_task)
        budget_tasklink.durability = And(budget_task.durability, diff_task)
    if budget_termlink is not None:
        if not replace_termlink: budget_termlink = deepcopy(budget_termlink)
        diff_belief = abs(truth_belief.e - truth_derived.e)
        budget_termlink.priority = And(budget_termlink.priority, 1-diff_belief)
        budget_termlink.durability = And(budget_termlink.durability, 1-diff_belief)
    diff = truth_derived.c - max(truth_task.c, truth_belief.c)
    priority = Or(diff, budget_task.priority)
    durability = Average(diff, budget_task.durability)
    quality = truth_to_quality(truth_derived)
    return Budget(priority, durability, quality), budget_task, budget_tasklink, budget_termlink

def Budget_inference(quality: float, budget_tasklink: Budget, budget_termlink: Budget=None, complexity: float=1.0):
    '''
    Ref. OpenNARS 3.1.0 BudgetFunctions.java line 292~317.
    '''
    complexity = 1 + log2(complexity)
    p = budget_tasklink.priority
    d = budget_tasklink.durability/complexity
    q = quality/complexity
    if budget_termlink is not None:
        p = Or(p, budget_termlink.priority)
        d = And(d, budget_termlink.durability)
        # budget_termlink.priority = min(1.0, Or(budget_termlink.priority, Or(q, budget_belief.priority)))
        # budget_termlink.durability = min(1.0-Config.budget_epsilon, Or(budget_termlink.durability, q))
    # d = Scalar(d)
    # q = Scalar(q)
    return Budget(p, d, q) # , budget_termlink


def Budget_forward(truth_new_task: Truth, budget_tasklink: Budget, budget_termlink: Budget=None):
    return Budget_inference(truth_to_quality(truth_new_task), budget_tasklink, budget_termlink, 1.0)

def Budget_backward(truth_new_task: Truth, budget_tasklink: Budget, budget_termlink: Budget=None):
    return Budget_inference(truth_to_quality(truth_new_task), budget_tasklink, budget_termlink, 1.0)

def Budget_backward_weak(truth_new_task: Truth, budget_tasklink: Budget, budget_termlink: Budget=None):
    return Budget_inference(w_to_c(1, Config.k)*truth_to_quality(truth_new_task), budget_tasklink, budget_termlink, 1.0)

def Budget_forward_compound(content, truth_new_task: Truth, budget_tasklink: Budget, budget_termlink: Budget=None):
    '''Ref. OpenNARS 3.1.0 BudgetFunctions.java line 254~257.'''
    return Budget_inference(truth_to_quality(truth_new_task), budget_tasklink, budget_termlink, Config.complexity_unit if content is None else Config.complexity_unit*content.complexity)

def Budget_backward_compound(content: Term, budget_tasklink: Budget, budget_termlink: Budget=None):
    return Budget_inference(1.0, budget_tasklink, budget_termlink, Config.complexity_unit*content.complexity)

def Budget_backward_weak_compound(content: Term, budget_tasklink: Budget, budget_termlink: Budget=None):
    return Budget_inference(w_to_c(1, Config.k), budget_tasklink, budget_termlink, Config.complexity_unit*content.complexity)


'''Bag'''
def Budget_decay(budget: Budget, replace=True):
    '''
    Ref: The Conceptual Design of OpenNARS 3.1.0
    Ref: OpenNARS 3.1.0 BudgetFunctions.java line 176~196
        Decrease Priority after an item is used, called in Bag. After a constant time, p should become d*p. Since in this period, the item is accessed c*p times, each time p-q should multiple d^(1/(c*p)). The intuitive meaning of the parameter "forgetRate" is: after this number of times of access, priority 1 will become d, it is a system parameter
        adjustable in run time.
    '''
    if not replace: budget = deepcopy(budget)
    Q = Config.quality_min
    C = Config.cycles_forget
    p = budget.priority
    q = budget.quality * Q
    d = budget.durability
    budget.priority = q + (p-q)*pow(d, 1.0/(p*C))
    # budget.priority = q + (p-q)*pow(d, 1.0/((p-q)*C)) # the implementation in OpenNARS 3.0.4
    return budget

def Budget_merge(budget_base: Budget, budget_merged: Budget, replace=True):
    '''
    Ref: The Conceptual Design of OpenNARS 3.1.0
        When an item is added into a bag where there is another one with the same key, the two will be merged, with their budget accumulated. In this process, the two quality values should be the same, and if not, the max operator is used. The two priority values are combined using or, so that the result will be no smaller than either of the two, while still remains in the [0, 1] range. For the same reason, or is used to combine the two durability values. Consequently, repeatedly appeared items will get more resources, both at the moment and in the near future.
    Ref: OpenNARS 3.1.0 BudgetFunctions.java line 161~176, 206~209
        There are two options.
        1) use the `max` function to all three values;
        2) use the `or` function to `priority`, use `arithmetic average` funciton to `durability`, and keep the original value of `quality` of the concept unchanged.
    Here the implementation is accordant to the description in the Conceptual Design.
    '''
    if not replace: budget_base = deepcopy(budget_base)
    # # implementation in the Conceptual Design.
    # budget_base.priority = max(budget_base.priority, budget_merged.priority)
    # budget_base.durability = Or(budget_base.durability, budget_merged.durability)
    # budget_base.quality = Or(budget_base.quality, budget_merged.quality)

    # implementation in OpenNARS 3.1.0 or 3.0.4:
    budget_base.priority = Or(budget_base.priority, budget_merged.priority)
    budget_base.durability = Average(budget_base.durability, budget_merged.durability)
    budget_base.quality = budget_base.quality

    # TODO: which implementation is better?
    # Note (2021.1.25):
    # I find that if I adopt the former one, which is consistent with the conceptual design, the budget of a term will be very large, because when build term links, there will be several times to call this function. Hence, if we choose the first one, `Concept._build_term_links(...)` and `Concept._build_task_links(...)` should be modified.
    
    return budget_base


'''Task'''
'''Concept'''
'''Task-Link'''
'''Term-Link'''
'''Belief and Desire'''
'''Processing Units'''
'''Goal Evaluations'''
