import numpy as np
import random as rand
from scipy.stats import chi2


def isVer(p):
    return 0 <= p <= 1


def ver(p1, p2, error=0.01):
    result = {}
    for i in range(100):
        if i == 0:
            result[i] = p1
        else:
            ver_p = (
                p2 * (1 - p1) ** i * (1 - p2) ** (i - 1)
                + p1 * (1 - p1) ** i * (1 - p2) ** i
            )
            result[i] = ver_p
            if ver_p < error:
                break
    return result


def getVer(first_p, second_p, count):
    result = {}
    if isVer(first_p) and isVer(second_p):
        for i in range(count):
            local_p_1 = rand.random()
            local_p_2 = rand.random()
            sum = 0
            while local_p_1 > first_p and local_p_2 > second_p:
                sum += 1
                local_p_1 = rand.random()
                local_p_2 = rand.random()

            if local_p_1 < first_p:
                result[sum] = result.get(sum, 0) + 1
            else:
                result[sum + 1] = result.get(sum + 1, 0) + 1
    else:
        print("p not ver")
    return result


number_intervals = int(input("Введите количество интервалов: "))
intervals = []
for i in range(number_intervals):
    interval = float(input(f"Введите верхнюю границу интервала {i+1}: "))
    intervals.append(interval - 1)

first_p = float(input("first_p = "))
second_p = float(input("second_p = "))
error = 0.0001
count = int(input(f"count = "))
theoretical_prob = ver(first_p, second_p, error)

theoretical_frequencies = [0] * (len(intervals) + 1)
for k, freq in theoretical_prob.items():
    for i, upper_bound in enumerate(intervals):
        if k <= upper_bound:
            theoretical_frequencies[i] += freq * count
            break
    else:
        theoretical_frequencies[-1] += freq * count

print("Теоретические вероятности:")
for i in range(len(theoretical_frequencies)):
    print(f"q{i} = {theoretical_frequencies[i]/count}")
alpha = float(input("Введите уровень значимости (например, 0.05): "))
number_exp = int(input("number_exp = "))
accept_hyp = 0
reject_hyp = 0
for i in range(number_exp):
    empirical_prob = getVer(first_p, second_p, count)
    empirical_frequencies = [0] * (len(intervals) + 1)
    for k, freq in empirical_prob.items():
        for i, upper_bound in enumerate(intervals):
            if k <= upper_bound:
                empirical_frequencies[i] += freq
                break
        else:
            empirical_frequencies[-1] += freq
    chi_squared = sum(
        (empirical_frequencies[i] - theoretical_frequencies[i]) ** 2
        / theoretical_frequencies[i]
        for i in range(len(theoretical_frequencies))
    )
    critical_value = chi2.ppf(1 - alpha, df=number_intervals)

    if chi_squared < critical_value:
        accept_hyp += 1
    else:
        reject_hyp += 1
print(f"{accept_hyp = }")
print(f"{reject_hyp = }")