"""Домашнее задание.

Нужно написать функцию для определения продолжительности эксперимента.

В качестве метрики в эксперименте используется не среднее, а 90% квантиль.
Нулевая гипотеза - 90% квантили равны, альтернативная гипотеза - 90% квантили не равны.
Эксперимент будет оцениваться с помощью 1000 итераций бутстрепа и центрального доверительного интервала c уровнем
значимости alpha=0.05.
По историческим данным известно, что измеряемые значения распределены нормально с параметрами mu=100 и std=10.
Для синтетических АБ экспериментов значения измерений в экспериментальной группе следует умножать на (1 + effect).
Допустимая сумма вероятностей ошибок и ожидаемый размер эффекта передаются в функцию в качестве аргументов.

Эксперимент может проводиться целое число дней. За один день получаем по 10 измерений в каждой из групп.
Нужно определить минимальное кол-во дней, которое позволит собрать необходимое кол-во данных, чтобы тест контролировал
сумму ошибок I и II рода на заданном уровне.
Если для проведения эксперимента с заданными параметрами требуется более 30 дней, то функция должна вернуть 30.

ПРИМЕР
- Если для эксперимент нужно 33 наблюдения в каждой группе, то необходимо собрать данные за 4 дня. Ответ = 4.
- Если для эксперимент нужно 9999 наблюдений в каждой группе, то необходимо более 30 дней. Ответ = 30.

ОЦЕНКА КАЧЕСТВА
Будет произведено 3 запуска функции с некоторыми параметрами, по полученным результатам будет вычислено MAE.
Оценка будет посчитана так:
SCORE = np.clip(11 - int(MAE) - penalty, 0, 10)
где penalty - штраф за медленную скорость вычисления. Ожидается, что три запуска функции отработают за 5 минут.
Каждые 30 секунд сверх 5 минут приводят к штрафному баллу. За 05:29 - не штрафа, за 05:31 - 1 штрафной балл, и тд.
Решение будет запускаться в Colab.
"""

import time

import numpy as np
import tqdm.notebook


def test_two(X, Y, alpha=0.05):
    X_boot = X[np.random.randint(0, X.shape[0], size=(X.shape[0], 1000, X.shape[1])), np.arange(X.shape[1])]
    X_boot_q = np.quantile(X_boot, 0.9, axis=0)
    Y_boot = Y[np.random.randint(0, Y.shape[0], size=(Y.shape[0], 1000, Y.shape[1])), np.arange(Y.shape[1])]
    Y_boot_q = np.quantile(Y_boot, 0.9, axis=0)

    boot_metrics = X_boot_q - Y_boot_q
    pe_metric = np.quantile(X, 0.9, axis=0) - np.quantile(Y, 0.9, axis=0)
    left = 2 * pe_metric - np.quantile(boot_metrics, 1 - alpha / 2, axis=0)
    right = 2 * pe_metric - np.quantile(boot_metrics, alpha / 2, axis=0)

    return ((left <= 0) & (right >= 0)).mean()  # Доля не отвержений нулевой гипотезы


def get_experiment_duration(error: float, effect: float) -> int:
    """Возвращает продолжительность эксперимента в днях.

    error - допустимая сумма ошибок I и II рода, возможны значения из отрезка [0.05, 0.5].
    effect - размер эффекта в процентах, возможны любые значения строго больше 0.
        Пример, effect=0.1 означает, что ожидается увеличение метрики на 10%.
    """
    for days in tqdm.notebook.tqdm(range(1, 31, 2)):
        X, Y1, Y2 = np.random.normal(100, 10, size=(3, 10 * days, 1000))
        Y2 = Y2 * (1 + effect)

        alpha = 1 - test_two(X, Y1)  # Ошибки первого рода
        beta = test_two(X, Y2)  # Ошибки второго рода
        print(alpha, beta)

        if alpha + beta < error:
            return days
    return 30


if __name__ == "__main__":
    errors = []

    # при проверке тут будут другие значения
    experiments = [
        {
            'params': {'error': 0.4, 'effect': 10},
            'true_answer': 1
        },
        {
            'params': {'error': 0.05, 'effect': 100.0},
            'true_answer': 30
        },
        # {
        #     'params': {'error': 0.05, 'effect': 0.00000001},
        #     'true_answer': 30
        # },
    ]

    t1 = time.time()
    for experiment in experiments:
        answer = get_experiment_duration(**experiment['params'])
        errors.append(answer - experiment['true_answer'])
    t2 = time.time()

    penalty = int(max(0, (t2 - t1 - 5 * 60) / 30))
    MAE = np.mean(np.abs(errors))
    SCORE = np.clip(11 - int(MAE) - penalty, 0, 10)

    print(f'time = {t2 - t1}')
    print(f'penalty = {penalty}')
    print(f'errors = {errors}')
    print(f'MAE = {MAE}')
    print(f'SCORE = {SCORE}')