import numpy as np
from utils import add_retraining_cost
from retrain.src.models.models_helper.dp import dp_iterative, retrains_iterative

name = "Oracle"


def run(C: np.ndarray, retrain_cost: float = 0) -> dict:
    """The optimal solution to the strategy problem. Uses an iterative Dynamic programming approach.

    Args:
        C (np.ndarray): the upper-triangular cost matrix
        retrain_cost (float, optional): the cost of retraining to be added to diagonals of cost matrix. Defaults to 0.

    Returns:
        dict: the strategy result dictionary
    """
    # gather the DP table
    if retrain_cost == 0:
        _C = C
    else:
        _C = add_retraining_cost(C, retrain_cost)
    dp = dp_iterative(_C)
    retrains = retrains_iterative(dp)
    result = {"retrains": retrains, "num_retrains": len(
        retrains)-1, "parameters": {}}
    return result


def optimize(C: np.ndarray, retrain_cost: float = 0) -> dict:
    return run(C, retrain_cost)
