from typing import Tuple
import math


def linear_warmup_and_cosine_protocol(
        f_values: Tuple[float, float, float],
        x_milestones: Tuple[int, int, int, int]):
    """
    There are 5 regions:
    1. constant at f0 for x < x0
    2. linear increase from f0 to f1 for x0 < x < x1
    3. constant at f1 for x1 < x < x2
    4. cosine protocol from f1 to f2 for x2 < x < x3
    5. constant at f2 for x > x3

    If you want a linear_ramp followed by a cosine_decay only simply set:
    1. x0=0 (to eliminate the first constant piece)
    2. x2=x1 (to eliminate the second constant piece)
    3. max_epochs=x3 (to make the simulation stop after the linear or cosine decay)
    """
    assert x_milestones[0] <= x_milestones[1] <= x_milestones[2] <= x_milestones[3]

    def fn(step):
        if step <= x_milestones[0]:
            return float(f_values[0])
        elif (step > x_milestones[0]) and (step <= x_milestones[1]):
            m = float(f_values[1] - f_values[0]) / float(max(1, x_milestones[1] - x_milestones[0]))
            return float(f_values[0]) + m * float(step - x_milestones[0])
        elif (step > x_milestones[1]) and (step <= x_milestones[2]):
            return float(f_values[1])
        elif (step > x_milestones[2]) and (step <= x_milestones[3]):
            progress = float(step - x_milestones[2]) / float(max(1, x_milestones[3] - x_milestones[2]))  # in (0,1)
            tmp = 0.5 * (1.0 + math.cos(math.pi * progress))  # in (1,0)
            return float(f_values[2]) + tmp * float(f_values[1] - f_values[2])
        else:
            return float(f_values[2])

    return fn