#!/usr/bin/env python3
"""
Градиентный спуск для одномерной функции с тремя методами выбора шага:
1. Константный шаг
2. Золотое сечение (одномерная оптимизация)
3. Правило Армихо

Функция: f(x) = sqrt(x^2 + 9) / 4 + (5 - x) / 5
Область: [-3, 8]
"""

import sys
from pathlib import Path

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

import matplotlib.pyplot as plt
import numpy as np
from common.functions import TaskFunction1D
from common.gradient_descent import (
    GradientDescentResult1D,
    gradient_descent_1d,
    heavy_ball_1d,
)

# ============================================================================
# НАСТРОЙКИ
# ============================================================================

# Стартовая точка
X0 = -1.0

# Параметры сходимости
EPS_X = 0.05
EPS_F = 0.001
MAX_ITERS = 100

# Шаг для константного метода (небольшой, чтобы было 3-4+ итерации)
CONSTANT_STEP = 0.5

# Параметры для правила Армихо
ARMIJO_PARAMS = {
    "d_init": 2.0,
    "epsilon": 0.1,
    "theta": 0.5,
}

# Границы для золотого сечения
GOLDEN_SECTION_BOUNDS = (0.0, 30.0)

# Параметры для метода тяжёлого шарика
HEAVY_BALL_ALPHA = 0.5
HEAVY_BALL_BETA = 0.8

# Папка для сохранения графиков
OUTPUT_DIR = Path(__file__).parent / "plots"


# ============================================================================
# ВИЗУАЛИЗАЦИЯ
# ============================================================================


def plot_iteration(
    func: TaskFunction1D,
    result: GradientDescentResult1D,
    iter_idx: int,
    output_path: Path,
):
    """Построить график для одной итерации."""
    info = result.iterations[iter_idx]

    # Диапазон для графика
    a, b = func.domain
    x_plot = np.linspace(a, b, 500)
    y_plot = [func(x) for x in x_plot]

    plt.figure(figsize=(10, 6))

    # Функция
    plt.plot(x_plot, y_plot, "b-", linewidth=2, label="f(x)")

    # Траектория до текущей точки
    trajectory_x = [result.iterations[i].x for i in range(iter_idx + 1)]
    trajectory_y = [result.iterations[i].f_x for i in range(iter_idx + 1)]

    if len(trajectory_x) > 1:
        plt.plot(
            trajectory_x,
            trajectory_y,
            "g--",
            linewidth=1.5,
            alpha=0.7,
            label="Траектория",
        )

    # Предыдущие точки
    for i, (x, y) in enumerate(zip(trajectory_x[:-1], trajectory_y[:-1])):
        plt.plot(x, y, "go", markersize=6, alpha=0.5)

    # Текущая точка
    plt.plot(
        info.x,
        info.f_x,
        "ro",
        markersize=12,
        label=f"x = {info.x:.4f}, f(x) = {info.f_x:.4f}",
    )

    # Направление градиента (касательная)
    grad_scale = 0.5
    x_grad = np.array([info.x - grad_scale, info.x + grad_scale])
    y_grad = info.f_x + info.grad * (x_grad - info.x)
    plt.plot(
        x_grad, y_grad, "m-", linewidth=2, alpha=0.6, label=f"f'(x) = {info.grad:.4f}"
    )

    plt.xlabel("x", fontsize=12)
    plt.ylabel("f(x)", fontsize=12)
    plt.title(
        f"{result.method} — Итерация {info.iteration}\nШаг: {info.step_size:.6f}",
        fontsize=14,
        fontweight="bold",
    )
    plt.legend(fontsize=10, loc="upper right")
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()


def plot_final_result(
    func: TaskFunction1D,
    result: GradientDescentResult1D,
    output_path: Path,
):
    """Построить итоговый график с полной траекторией."""
    a, b = func.domain
    x_plot = np.linspace(a, b, 500)
    y_plot = [func(x) for x in x_plot]

    plt.figure(figsize=(10, 6))

    # Функция
    plt.plot(x_plot, y_plot, "b-", linewidth=2, label="f(x)")

    # Траектория
    trajectory_x = [it.x for it in result.iterations]
    trajectory_y = [it.f_x for it in result.iterations]
    plt.plot(
        trajectory_x, trajectory_y, "g-", linewidth=2, alpha=0.7, label="Траектория"
    )

    # Все точки
    for i, (x, y) in enumerate(zip(trajectory_x[:-1], trajectory_y[:-1])):
        plt.plot(x, y, "go", markersize=8, alpha=0.6)

    # Финальная точка
    plt.plot(
        result.x_star,
        result.f_star,
        "r*",
        markersize=20,
        label=f"x* = {result.x_star:.6f}\nf(x*) = {result.f_star:.6f}",
    )

    plt.xlabel("x", fontsize=12)
    plt.ylabel("f(x)", fontsize=12)
    plt.title(
        f"{result.method} — Результат\nИтераций: {len(result.iterations) - 1}",
        fontsize=14,
        fontweight="bold",
    )
    plt.legend(fontsize=10, loc="upper right")
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()


def run_and_visualize(
    func: TaskFunction1D,
    method: str,
    method_name_short: str,
    **kwargs,
):
    """Запустить метод и создать визуализации."""

    result = gradient_descent_1d(
        func=func,
        x0=X0,
        step_method=method,
        eps_x=EPS_X,
        eps_f=EPS_F,
        max_iters=MAX_ITERS,
        **kwargs,
    )

    # Создаём папку для этого метода
    method_dir = OUTPUT_DIR / method_name_short
    method_dir.mkdir(parents=True, exist_ok=True)

    # Печатаем информацию
    print(f"\n{'=' * 80}")
    print(f"{result.method}")
    print("=" * 80)

    for info in result.iterations[:-1]:  # Без финальной точки
        print(
            f"Итерация {info.iteration:3d}: "
            f"x = {info.x:10.6f}, f(x) = {info.f_x:10.6f}, "
            f"f'(x) = {info.grad:10.6f}, шаг = {info.step_size:.6f}"
        )

        # Строим график для каждой итерации
        plot_iteration(
            func,
            result,
            info.iteration - 1,
            method_dir / f"iteration_{info.iteration:02d}.png",
        )

    # Итоговый результат
    print("-" * 80)
    print(f"x* = {result.x_star:.6f}")
    print(f"f(x*) = {result.f_star:.6f}")
    print(f"Итераций: {len(result.iterations) - 1}")

    # Финальный график
    plot_final_result(func, result, method_dir / "final_result.png")
    print(f"Графики сохранены в: {method_dir}")

    return result


def run_and_visualize_heavy_ball(
    func: TaskFunction1D,
    method_name_short: str,
    alpha: float,
    beta: float,
):
    """Запустить метод тяжёлого шарика и создать визуализации."""

    result = heavy_ball_1d(
        func=func,
        x0=X0,
        alpha=alpha,
        beta=beta,
        eps_x=EPS_X,
        eps_f=EPS_F,
        max_iters=MAX_ITERS,
    )

    # Создаём папку для этого метода
    method_dir = OUTPUT_DIR / method_name_short
    method_dir.mkdir(parents=True, exist_ok=True)

    # Печатаем информацию
    print(f"\n{'=' * 80}")
    print(f"{result.method}")
    print("=" * 80)

    for info in result.iterations[:-1]:
        print(
            f"Итерация {info.iteration:3d}: "
            f"x = {info.x:10.6f}, f(x) = {info.f_x:10.6f}, "
            f"f'(x) = {info.grad:10.6f}, шаг = {info.step_size:.6f}"
        )

        plot_iteration(
            func,
            result,
            info.iteration - 1,
            method_dir / f"iteration_{info.iteration:02d}.png",
        )

    # Итоговый результат
    print("-" * 80)
    print(f"x* = {result.x_star:.6f}")
    print(f"f(x*) = {result.f_star:.6f}")
    print(f"Итераций: {len(result.iterations) - 1}")

    # Финальный график
    plot_final_result(func, result, method_dir / "final_result.png")
    print(f"Графики сохранены в: {method_dir}")

    return result


def main():
    """Главная функция."""

    func = TaskFunction1D()

    print("=" * 80)
    print("ГРАДИЕНТНЫЙ СПУСК ДЛЯ ОДНОМЕРНОЙ ФУНКЦИИ")
    print("=" * 80)
    print(f"Функция: {func.name}")
    print(f"Область: {func.domain}")
    print(f"Стартовая точка: x₀ = {X0}")
    print(f"Параметры: eps_x = {EPS_X}, eps_f = {EPS_F}, max_iters = {MAX_ITERS}")

    # Создаём папку для графиков
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # 1. Константный шаг
    run_and_visualize(
        func,
        method="constant",
        method_name_short="constant",
        step_size=CONSTANT_STEP,
    )

    # 2. Золотое сечение
    run_and_visualize(
        func,
        method="golden_section",
        method_name_short="golden_section",
        golden_section_bounds=GOLDEN_SECTION_BOUNDS,
    )

    # 3. Правило Армихо
    run_and_visualize(
        func,
        method="armijo",
        method_name_short="armijo",
        armijo_params=ARMIJO_PARAMS,
    )

    # 4. Метод тяжёлого шарика
    run_and_visualize_heavy_ball(
        func,
        method_name_short="heavy_ball",
        alpha=HEAVY_BALL_ALPHA,
        beta=HEAVY_BALL_BETA,
    )

    print("\n" + "=" * 80)
    print("ГОТОВО! Все графики сохранены.")
    print("=" * 80)


if __name__ == "__main__":
    main()