import numpy as np
from scipy.stats import norm

def generate_data(sample_size, corr, mean=2000, sigma=300):
    """Генерируем коррелированные данные исходной метрики и ковариаты.
    
    sample_size - размер выборки
    corr - корреляция исходной метрики с ковариатой
    mean - среднее значение исходной метрики
    sigma - стандартное отклонение исходной метрики

    Возвращает:
        'metric' - значения исходной метрики,
        'covariate' - значения ковариаты.
    """
    np.random.seed(1337)
    means = np.array([mean, mean])
    cov = sigma ** 2 * np.array([[1, corr], [corr, 1]])
    data = np.random.multivariate_normal(means, cov, sample_size).astype(int)
    metric, covariate = data[:, 0], data[:, 1]
    return metric, covariate

def get_theta_for_cuped(values, covariate):
    '''Функция получает значение theta
    values - значения
    covariate - значение ковариаты
    Возвращает:
    theta - значение theta'''
    import numpy as np

    cov = np.cov(values, covariate)[0][1]
    var = np.var(covariate)
    theta = cov / var
    return theta

def get_mde(values, alpha = 0.05, beta = 0.2):
    '''На вход получает значения метрики
    На выходе 
    mde_abs - абсолютный MDE
    mde_rel - относительный MDE'''

    std = np.std(values)
    mu = np.mean(values)
    sample_size = len(values)

    f = norm.ppf(1 - alpha / 2) + norm.ppf(1 - beta)

    mde_abs = 2 * std * f / np.sqrt(sample_size)
    mde_rel = mde_abs / mu * 100

    return mde_abs, mde_rel

corr = 0.6
values, covariate = generate_data(20000, corr)
print(f' Коэффициент корреляции - {np.corrcoef(values, covariate)[0][1]}')
theta = get_theta_for_cuped(values, covariate)

print('')

cuped_metric = values - theta * covariate
cuped_metric_with_mu = values - theta * (covariate - np.mean(covariate))

print(f' Стандартное отклонение метрики - {np.std(values)}')
print(f' Стандартное отклонение CUPED метрики - {np.std(cuped_metric)}')
print(f' Старндартное отклонение CUPED метрики с поправкой - {np.std(cuped_metric_with_mu)}')
print('')
print(f' Среднее метрики - {np.mean(values)}')
print(f' Стандартное отклонение CUPED метрики - {np.mean(cuped_metric)}')
print(f' Старндартное отклонение CUPED метрики с поправкой - {np.mean(cuped_metric_with_mu)}')


mde_standard_abs, mde_standard_rel = get_mde(values)
mde_cuped_abs, mde_cuped_rel = get_mde(cuped_metric)
mde_cuped_with_mu_abs, mde_cuped_with_mu_rel = get_mde(cuped_metric_with_mu)

print('')

print(f' Абслютный и относительный MDE метрики {float(mde_standard_abs), float(mde_standard_rel)}')
print(f' Абслютный и относительный MDE CUPED метрики {float(mde_cuped_abs), float(mde_cuped_rel)}')
print(f' Абслютный и относительный MDE CUPED метрики с поправкой {float(mde_cuped_with_mu_abs), float(mde_cuped_with_mu_rel)}')