from typing import Callable, Dict, Optional, List
import numpy as np
from models.utils import add_retraining_cost, compute_strategy_cost

name = "Periodic"


def run(C: np.ndarray, period: int, offset: Optional[int] = -1, online_offset: Optional[int] = 0) -> dict:
    T = C.shape[0]-1
    assert period <= T*2-1 and period > 0
    retrains = list(range(offset+online_offset, T+1, period))
    # if 0 not in retrains:
    #     retrains = [0]+retrains
    result = {"retrains": retrains, "num_retrains": len(
        retrains)-1, "parameters": {"period": period, "offset": offset, "online_offset": online_offset}}
    return result


def run_cost(period: int, C: np.ndarray):
    return compute_strategy_cost(run(C, period), C)


def optimize(C: np.ndarray, retrain_cost):
    T = C.shape[0]-1
    _C = add_retraining_cost(C, retrain_cost)

    costs = np.array([run_cost(p, _C) for p in range(1, T*2)])
    best_period = costs.argmin()+1
    # Find offset when needed to run online
    online_offset = (best_period - T % best_period) - 1
    result = run(C, best_period)
    result["parameters"]["online_offset"] = online_offset
    return result
