# core/statistics.py
"""
Статистические функции для регрессии, проверки значимости, дисперсии и градиента.
"""
import numpy as np
try:
    from scipy import stats
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False

def compute_regression(X, y):
    """МНК-регрессия, возвращает коэффициенты, ошибки, t-статистики."""
    N, p = X.shape
    try:
        XTX = X.T @ X
        XTX_inv = np.linalg.inv(XTX)
        beta = XTX_inv @ (X.T @ y)
        y_pred = X @ beta
        residuals = y - y_pred
        df_res = N - p
        if df_res > 0:
            s2_res = np.sum(residuals**2) / df_res
            se = np.sqrt(np.diag(XTX_inv) * s2_res)
        else:
            s2_res = np.nan
            se = [np.nan] * p
        t_stats = beta / se if not np.isnan(se).any() else [np.nan] * p
    except np.linalg.LinAlgError:
        beta = [np.nan] * p
        se = [np.nan] * p
        t_stats = [np.nan] * p
        df_res = 0
        s2_res = np.nan
        residuals = np.full(N, np.nan)
        y_pred = np.full(N, np.nan)
        XTX_inv = None
    return {
        "beta": beta,
        "se": se,
        "t_stats": t_stats,
        "df_res": df_res,
        "s2_res": s2_res,
        "residuals": residuals,
        "y_pred": y_pred,
        "XTX_inv": XTX_inv
    }

def regression_summary(y, y_pred, k, df_res):
    """Вычисляет R², скорректированный R², F-статистику и p-value."""
    n = len(y)
    ss_total = np.sum((y - np.mean(y))**2)
    ss_res = np.sum((y - y_pred)**2)
    r2 = 1 - ss_res / ss_total if ss_total > 0 else np.nan
    if df_res > 0:
        r2_adj = 1 - (1 - r2) * (n - 1) / df_res
    else:
        r2_adj = np.nan
    if df_res > 0 and k > 0:
        f_stat = (r2 / k) / ((1 - r2) / df_res) if r2 < 1 else np.inf
    else:
        f_stat = np.nan
    if SCIPY_AVAILABLE and not np.isnan(f_stat) and f_stat != np.inf:
        p_value = 1 - stats.f.cdf(f_stat, k, df_res)
    else:
        p_value = np.nan
    return {"r2": r2, "r2_adj": r2_adj, "f_stat": f_stat, "f_pvalue": p_value}

def get_center_indices(plan, active_indices):
    """Возвращает индексы строк плана, где активные факторы равны 0."""
    if not plan:
        return []
    plan_arr = np.array(plan)
    if active_indices:
        active_cols = plan_arr[:, active_indices]
        center_mask = np.all(active_cols == 0, axis=1)
        return np.where(center_mask)[0].tolist()
    return []

def compute_center_stats(y_center):
    """Среднее и несмещённая дисперсия для центральных точек."""
    if len(y_center) < 2:
        return None, None
    return np.mean(y_center), np.var(y_center, ddof=1)

def chi2_confidence_interval(var, n, alpha=0.05):
    """Доверительный интервал для дисперсии на основе χ²-распределения."""
    if not SCIPY_AVAILABLE or n < 2:
        return np.nan, np.nan
    df = n - 1
    chi2_lower = stats.chi2.ppf(alpha / 2, df)
    chi2_upper = stats.chi2.ppf(1 - alpha / 2, df)
    return df * var / chi2_upper, df * var / chi2_lower

def get_t_critical(df, alpha=0.05):
    """Критическое значение t (двустороннее)."""
    if not SCIPY_AVAILABLE or df <= 0:
        return np.nan
    return stats.t.ppf(1 - alpha/2, df)

def compute_gradient_points(beta, base_levels, step_factor, num_steps, factor_names=None, maximize=True):
    """
    Генерирует точки вдоль градиента (крутое восхождение или спуск).
    beta – коэффициенты модели: [b0, b1, b2, ...]
    base_levels – список базовых уровней для каждого фактора
    step_factor – шаг по первому фактору
    num_steps – количество шагов
    factor_names – список названий факторов
    maximize – True для максимизации, False для минимизации
    """
    if len(beta) - 1 != len(base_levels):
        raise ValueError("Число коэффициентов (без b0) должно совпадать с числом факторов")
    direction = 1 if maximize else -1
    coefs = np.array(beta[1:]) * direction
    max_abs = np.max(np.abs(coefs))
    if max_abs == 0:
        raise ValueError("Все коэффициенты равны нулю, направление не определено.")
    norm_coefs = coefs / max_abs
    steps = step_factor * norm_coefs

    points = []
    for step_num in range(1, num_steps + 1):
        point = {}
        for i, name in enumerate(factor_names or [f"x{i+1}" for i in range(len(base_levels))]):
            point[name] = base_levels[i] + steps[i] * step_num
        predicted = beta[0] + np.sum(beta[1:] * np.array([point[name] for name in factor_names or range(len(base_levels))]))
        point["Шаг"] = step_num
        point["Предсказанный отклик"] = predicted
        points.append(point)
    return points


def anova_regression(y, y_pred, X, center_indices=None):
    """
    Выполняет полный дисперсионный анализ (ANOVA) для регрессионной модели.

    Аргументы:
        y – фактические значения отклика
        y_pred – предсказанные моделью значения
        X – матрица регрессоров (с единичным столбцом)
        center_indices – индексы центральных точек (опционально)

    Возвращает словарь с компонентами ANOVA:
        SSR, SSE, SST, df_reg, df_res, df_pe, df_lof,
        MSR, MSE, MS_pe, MS_lof,
        F_regression, p_regression, F_lof, p_lof
    """
    n = len(y)
    p = X.shape[1]          # число параметров (включая свободный член)
    df_reg = p - 1          # степени свободы регрессии (без свободного члена)
    df_res = n - p          # степени свободы остатка

    # Общая сумма квадратов (SST)
    y_mean = np.mean(y)
    SST = np.sum((y - y_mean)**2)

    # Регрессионная сумма квадратов (SSR)
    SSR = np.sum((y_pred - y_mean)**2)

    # Остаточная сумма квадратов (SSE)
    SSE = np.sum((y - y_pred)**2)

    # Средние квадраты
    MSR = SSR / df_reg if df_reg > 0 else np.nan
    MSE = SSE / df_res if df_res > 0 else np.nan

    # F-статистика для регрессии (значимость всей модели)
    F_reg = MSR / MSE if MSE > 0 else np.nan
    p_reg = 1 - stats.f.cdf(F_reg, df_reg, df_res) if not np.isnan(F_reg) and SCIPY_AVAILABLE else np.nan

    # ----- Разложение остатка на чистую ошибку и неадекватность (если есть центр) -----
    if center_indices is not None and len(center_indices) >= 2:
        y_center = y[center_indices]
        y_pred_center = y_pred[center_indices]
        n_center = len(center_indices)
        # Среднее в центре
        y_center_mean = np.mean(y_center)
        # Сумма квадратов чистой ошибки (SSPE) = сумма квадратов отклонений от среднего в центре
        SSPE = np.sum((y_center - y_center_mean)**2)
        df_pe = n_center - 1
        MS_pe = SSPE / df_pe if df_pe > 0 else np.nan

        # Сумма квадратов неадекватности (SSLOF) = SSE - SSPE
        SSLOF = SSE - SSPE
        df_lof = df_res - df_pe
        MS_lof = SSLOF / df_lof if df_lof > 0 else np.nan

        # F-статистика для проверки адекватности (lack-of-fit)
        F_lof = MS_lof / MS_pe if MS_pe > 0 and not np.isnan(MS_lof) else np.nan
        p_lof = 1 - stats.f.cdf(F_lof, df_lof, df_pe) if not np.isnan(F_lof) and SCIPY_AVAILABLE else np.nan
    else:
        SSPE = np.nan
        df_pe = 0
        MS_pe = np.nan
        SSLOF = np.nan
        df_lof = 0
        MS_lof = np.nan
        F_lof = np.nan
        p_lof = np.nan

    return {
        "SSR": SSR, "SSE": SSE, "SST": SST,
        "df_reg": df_reg, "df_res": df_res,
        "MSR": MSR, "MSE": MSE,
        "F_reg": F_reg, "p_reg": p_reg,
        "SSPE": SSPE, "SSLOF": SSLOF,
        "df_pe": df_pe, "df_lof": df_lof,
        "MS_pe": MS_pe, "MS_lof": MS_lof,
        "F_lof": F_lof, "p_lof": p_lof
    }
# core/statistics.py (дополнение)

def build_quadratic_X(plan):
    """
    Строит расширенную матрицу регрессоров для квадратичной модели.
    plan: numpy array (N x k) – матрица плана в кодированных переменных.
    Возвращает: X (N x p), где p = 1 + 2k + k*(k-1)//2
    (свободный член, линейные, квадраты, попарные произведения)
    """
    N, k = plan.shape
    X = np.ones((N, 1))               # свободный член
    X = np.hstack([X, plan])          # линейные (k столбцов)
    X = np.hstack([X, plan**2])       # квадраты (k столбцов)
    # взаимодействия (k*(k-1)//2 столбцов)
    for i in range(k):
        for j in range(i+1, k):
            X = np.hstack([X, (plan[:, i] * plan[:, j]).reshape(-1, 1)])
    return X

def fit_quadratic_model(plan, y):
    """
    Выполняет регрессию для квадратичной модели по матрице плана.
    Возвращает словарь с результатами (аналогично compute_regression, но с добавлением структуры).
    """
    X = build_quadratic_X(plan)
    reg = compute_regression(X, y)   # используем существующую функцию
    # Добавляем информацию о структуре модели
    k = plan.shape[1]
    p_total = X.shape[1]
    reg["k"] = k
    reg["X"] = X
    reg["plan"] = plan
    reg["p_total"] = p_total
    # Можно добавить имена коэффициентов (для вывода)
    coef_names = ["Свободный член"]
    coef_names += [f"x{i+1}" for i in range(k)]
    coef_names += [f"x{i+1}²" for i in range(k)]
    for i in range(k):
        for j in range(i+1, k):
            coef_names.append(f"x{i+1}·x{j+1}")
    reg["coef_names"] = coef_names
    return reg

def analyze_quadratic_model(beta, k, X, y):
    """
    Анализирует квадратичную модель: находит стационарную точку и тип экстремума.
    beta – коэффициенты (включая все члены модели)
    k – число факторов
    X – матрица регрессоров (для расчёта предсказанных значений)
    y – отклик
    Возвращает словарь с результатами.
    """
    # Извлекаем коэффициенты
    b0 = beta[0]
    b_lin = beta[1:1+k]                # линейные
    b_quad = beta[1+k:1+2*k]           # квадраты
    b_inter = beta[1+2*k:]             # взаимодействия

    # Строим матрицу B (гессиан) для квадратичной формы
    B = np.diag(2 * b_quad)
    idx = 0
    for i in range(k):
        for j in range(i+1, k):
            B[i, j] = b_inter[idx]
            B[j, i] = b_inter[idx]
            idx += 1

    # Стационарная точка: x* = -0.5 * inv(B) * b_lin
    try:
        x_star = -0.5 * np.linalg.solve(B, b_lin)
    except np.linalg.LinAlgError:
        x_star = None

    # Предсказанное значение в стационарной точке
    if x_star is not None:
        y_star = b0 + np.dot(b_lin, x_star) + np.dot(x_star, B @ x_star) / 2
    else:
        y_star = None

    # Тип экстремума
    if x_star is not None:
        eigvals = np.linalg.eigvalsh(B)
        if np.all(eigvals < 0):
            extremum_type = "Максимум"
        elif np.all(eigvals > 0):
            extremum_type = "Минимум"
        elif np.any(eigvals < 0) and np.any(eigvals > 0):
            extremum_type = "Седловая точка"
        else:
            extremum_type = "Неопределён (есть нулевые собственные значения)"
    else:
        extremum_type = "Не найден (матрица B вырождена)"

    return {
        "x_star": x_star,
        "y_star": y_star,
        "extremum_type": extremum_type,
        "B": B,
        "eigenvalues": np.linalg.eigvalsh(B) if x_star is not None else None
    }