import matplotlib.pyplot as plt
import matplotlib
import numpy as np
matplotlib.use('TkAgg')

# 1. Функция для генерации линейной зависимости
def F(b0, b1, N):
    # Генерация случайных данных для X
    x = np.random.uniform(0, 10, size=N)
    x.sort()  # Упорядочиваем массив для наглядности
    err = np.random.normal(0, 2, N)  # Добавляем случайный шум
    Y = b0 + b1 * x + err  # Вычисляем Y с линейной зависимостью и шумом
    # Построение графика рассеяния и идеальной линейной функции
    plt.scatter(x, Y, facecolors='none', edgecolor='blue', s=40)
    plt.plot(x, b0 + b1 * x, color='red')
    plt.title("Линейная зависимость с шумом")
    plt.xlabel('X')
    plt.ylabel('Data Y')
    plt.show()
    return Y, x  # Возвращаем массивы X и Y

# 2. Функция для генерации квадратичной зависимости
def F2(a, b, c, N):
    xp = np.random.uniform(0, 10, size=N)
    xp.sort()
    err = np.random.normal(0, 5, N)  # Шум для квадратичной зависимости
    Yp = a * xp ** 2 + b * xp + c + err
    # Построение графика рассеяния и идеальной квадратичной функции
    plt.scatter(xp, Yp, facecolors='none', edgecolor='blue', s=40)
    plt.plot(xp, a * xp ** 2 + b * xp + c, color='red')
    plt.title("Квадратичная зависимость с шумом")
    plt.xlabel('X')
    plt.ylabel('Data Yp')
    plt.show()
    return Yp, xp

# Генерация данных
Y, x = F(6, 9, 500)  # Линейная зависимость
Yp, xp = F2(3, 1, -2, 500)  # Квадратичная зависимость

# 3. Метод наименьших квадратов (МНК) для расчета коэффициентов линейной регрессии
def MHk(Y, x):
    # Вычисление коэффициента наклона (b1)
    b1 = sum((x - np.mean(x)) * (Y - np.mean(Y))) / sum((x - np.mean(x)) ** 2)
    # Вычисление свободного члена (b0)
    b0 = np.mean(Y) - b1 * np.mean(x)
    return b1, b0

b1, b0 = MHk(Y, x)  # Коэффициенты для линейной зависимости
print(f"Линейные коэффициенты: b1={b1:.2f}, b0={b0:.2f}")

# Аналогичная функция для квадратичных данных
def MHkp(Yp, xp):
    b1p = sum((xp - np.mean(xp)) * (Yp - np.mean(Yp))) / sum((xp - np.mean(xp)) ** 2)
    b0p = np.mean(Yp) - b1p * np.mean(xp)
    return b1p, b0p

b1p, b0p = MHkp(Yp, xp)
print(f"Коэффициенты квадратичной регрессии: b1={b1p:.2f}, b0={b0p:.2f}")

# 4. Построение регрессионных линий
Y_l = b0 + b1 * x  # Линейная регрессия
Y_s = b0p + b1p * xp  # Квадратичная регрессия

# 5. Вычисление SSE и коэффициента детерминации (R²)
SSE = sum((Y - (b0 + b1 * x)) ** 2)  # Ошибка для линейной регрессии
SST = sum((Y - np.mean(Y)) ** 2)  # Общая дисперсия данных
R2_lin = 1 - (SSE / SST)  # Коэффициент детерминации для линейной
print(f"SSE линейной: {SSE:.2f}, R^2 линейной функции: {R2_lin:.3f}")

SSEp = sum((Yp - Y_s) ** 2)  # Ошибка для квадратичной регрессии
SSTp = sum((Yp - np.mean(Yp)) ** 2)
R2_quad = 1 - (SSEp / SSTp)  # Коэффициент детерминации для квадратичной
print(f"SSE квадратичной: {SSEp:.2f}, R^2 квадратичной функции: {R2_quad:.3f}")

# 6. Построение графиков
fig, axs = plt.subplots(2, 2, figsize=(14, 10))

# Линейная регрессия
axs[0, 0].scatter(x, Y, facecolors='none', edgecolor='blue', s=40)
axs[0, 0].plot(x, Y_l, color='red', label=f'Линейная регрессия, R^2 = {R2_lin:.3f}')
axs[0, 0].set_xlabel('X')
axs[0, 0].set_ylabel('Y')
axs[0, 0].set_title('Линейная регрессия')
axs[0, 0].legend()

# Квадратичная регрессия
axs[0, 1].scatter(xp, Yp, facecolors='none', edgecolor='blue', s=40)
axs[0, 1].plot(xp, Y_s, color='red', label=f'Квадратичная регрессия, R^2 = {R2_quad:.3f}')
axs[0, 1].set_xlabel('X')
axs[0, 1].set_ylabel('Yp')
axs[0, 1].set_title('Квадратичная регрессия')
axs[0, 1].legend()

# Остатки для линейной регрессии
residuals_lin = Y - (b0 + b1 * x)
axs[1, 0].scatter(x, residuals_lin, color='green')
axs[1, 0].axhline(y=0, color='black', linestyle='--')
axs[1, 0].set_xlabel('X')
axs[1, 0].set_ylabel('Остатки')
axs[1, 0].set_title('Остатки для линейной регрессии')

# Остатки для квадратичной регрессии
residuals_quad = Yp - Y_s
axs[1, 1].scatter(xp, residuals_quad, color='green')
axs[1, 1].axhline(y=0, color='black', linestyle='--')
axs[1, 1].set_xlabel('X')
axs[1, 1].set_ylabel('Остатки')
axs[1, 1].set_title('Остатки для квадратичной регрессии')

plt.show()

# 7. Гистограммы остатков
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
ax1.hist(residuals_lin, bins=20, color='blue', alpha=0.7, edgecolor='black')
ax1.set_xlabel('Остатки')
ax1.set_ylabel('Количество')
ax1.set_title('Гистограмма остатков для линейной регрессии')

ax2.hist(residuals_quad, bins=20, color='blue', alpha=0.7, edgecolor='black')
ax2.set_xlabel('Остатки')
ax2.set_ylabel('Количество')
ax2.set_title('Гистограмма остатков для квадратичной регрессии')

plt.show()

# 8. Тест Голдфельда-Квандта
def goldfeld_quandt_test(e):
    n = 500
    k = 166
    e1 = e[:k]
    e3 = e[(n-k):]
    s1 = np.sum(e1**2)
    s3 = np.sum(e3**2)
    f = s3 / s1
    return f

f_lin = goldfeld_quandt_test(residuals_lin)
f_sqr = goldfeld_quandt_test(residuals_quad)

print(f"Статистика F для линейной функции: {f_lin:.3f}")
print(f"Статистика F для квадратичной функции: {f_sqr:.3f}")