import numpy as np
import matplotlib.pyplot as plt

MAX_LATENCY = 500

LATENCIES = [60, 65, 70, 80, 90, None, 79, None, 60, 70, 80, 65, 45, 50, 60, 65,  60, 70, 80, 65, 45, 50, 60, 65,  60, 70, 80, 65, 45, None, 50, 60, 65, 76, 60, 70, 80, 65, 45, 50, 60, 65,  60, 70, 80, 65, 45,]

def compute_scores(
    latencies,
    availability_window_size: int,
    latency_window_size: int,
    weight_avail: float,
    is_median,
    pow_factor_lat,
    pow_factor_avail,
    is_recency,
    lambda_decay,
):
    availability_arr = []
    latency_arr = []
    
    scores_avail = []
    scores_lat = []
    scores_combined = []

    for point in latencies:
        availability = 1 if point is not None else 0
        
        availability_arr.append(availability)
        availability_arr = availability_arr[-availability_window_size:]
        
        if is_recency:
            score_avail = 0
            weight_sum = 0
            for (j, avail) in enumerate(availability_arr):
                delta = availability_window_size - j
                wi = np.exp(-lambda_decay*delta)
                weight_sum += wi
                score_avail += wi * avail
            score_avail /= weight_sum
        else:
            score_avail = (sum(availability_arr) / len(availability_arr)) ** pow_factor_avail
        
        scores_avail.append(score_avail)

        if point is not None:
            latency_arr.append(point)
            latency_arr = latency_arr[-latency_window_size:]

        if is_median:
            metric = np.median(latency_arr)
        else:
            metric = np.mean(latency_arr)

        score_lat = (1 - (metric / MAX_LATENCY)) ** pow_factor_lat
        scores_lat.append(score_lat)

        score = weight_avail * score_avail + (1 - weight_avail) * score_lat
        
        scores_combined.append(score)

    return scores_avail, scores_lat, scores_combined


if __name__ == "__main__":    
    WEIGHT_AVAIL = 0.5 # WEIGHT_LATENCY = 1.0 - WEIGHT_AVAIL
    AVAIL_WINDOW_SIZE = 30
    LAT_WINDOW_SIZE = 15
    POW_FACTOR_LAT = 1
    POW_FACTOR_AVAIL = 1
    LAMBDA_DECAY = 0.1
    IS_MEDIAN = False
    IS_RECENCY = True
    
    scores_a, scores_l, combined_scores = compute_scores(
        latencies=LATENCIES,
        availability_window_size=AVAIL_WINDOW_SIZE,
        is_median=IS_MEDIAN,
        is_recency=IS_RECENCY,
        lambda_decay=LAMBDA_DECAY,
        latency_window_size=LAT_WINDOW_SIZE,
        pow_factor_avail=POW_FACTOR_AVAIL,
        pow_factor_lat=POW_FACTOR_LAT,
        weight_avail=WEIGHT_AVAIL,
    )
    
    x = list(range(1, len(LATENCIES) + 1))
    
    plt.plot(x, scores_a, "--o", markersize=2,  label="availability scores")
    plt.plot(x, scores_l, markersize=2, linewidth=2, label="latency scores")
    plt.plot(x, combined_scores, markersize=2, linewidth=2, label="combined scores")
  
    plt.legend()
    plt.show()