import numpy as np
from tqdm import tqdm

def gen_array(n_candidates=10):
    """
    Randomly generate `n_candidates` candidate rank from 1->n_candidates
        Assumption: P(rank of ith candidate = k) is random uniform (1 <= i, k <= n_candidates)
        (meaning each position has equally chance to be the best, second best, worst, etc.)
        
    Example:
    >> gen_array(10)
    6,  5,  1,  4,  3,  7,  8,  9,  2, 10
    """
    assert n_candidates > 1
    
    return np.random.choice(range(1, n_candidates+1),
                            size=n_candidates,
                            replace=False)

def simulate_stop_after_r(r, n_candidates, tries=10_000):
    """
    Estimate the winning chance of "stopping after r" strategy
    (Winning means the best candidate overall has been picked)
    
    The stopping after r strategy:
        - Ignore the first r candidates
        - Only consider candidates from r+1 to n_candidates:
            - if candidate is the best so far, pick
            - if candidate is not the best so far, move on 
                until a best so far is found or ran out of candidates
            - pray
    """
    assert r <= n_candidates
    
    win_count = 0
    
    for _ in range(tries):
        arr = gen_array(n_candidates)
        
        # ignore the first r-1 candidates, but keeping tab of the best ignored
        best_ignored = max(arr[:r+1])
        
        found = False
        
        # consider each candidate from r+1 onwards
        for current in range(r+1, n_candidates):
        
            
            if arr[current] > best_ignored:
                # pick if current is better than best ignored
                pick = arr[current]
                found = True
                break
            
            
        if found:
            # is pick the best?
            if pick == n_candidates:
                win_count += 1
             
    
    return round(win_count / tries, 4)
        
    
n_candidates = 10
p_win = dict()

for r in tqdm(range(n_candidates)):
    p_win[r] = simulate_stop_after_r(r, n_candidates)

print(p_win)