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

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 Function2D, HimmelblauFunction, RavineFunction
from common.gradient_descent import (
    GradientDescentResult2D,
    gradient_descent_2d,
    heavy_ball_2d,
)
from matplotlib import cm

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

# Выбор функции: "himmelblau" или "ravine"
FUNCTION_CHOICE = "himmelblau"

# Стартовые точки для разных функций
START_POINTS = {
    "himmelblau": np.array([0.0, 0.0]),
    "ravine": np.array([1.0, 0.3]),  # Стартуем в овраге
}

# Параметры сходимости
EPS_X = 1e-2
EPS_F = 1e-2
MAX_ITERS = 200

# Шаг для константного метода (для оврага нужен маленький шаг!)
CONSTANT_STEPS = {
    "himmelblau": 0.01,
    "ravine": 0.01,  # Маленький шаг из-за большого градиента по y
}

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

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

# Параметры для метода тяжёлого шарика
HEAVY_BALL_PARAMS = {
    "himmelblau": {"alpha": 0.01, "beta": 0.7},
    "ravine": {"alpha": 0.01, "beta": 0.8},
}

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


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


def create_contour_grid(func: Function2D, resolution: int = 200):
    """Создать сетку для контурного графика."""
    (x1_min, x1_max), (x2_min, x2_max) = func.plot_bounds
    x1 = np.linspace(x1_min, x1_max, resolution)
    x2 = np.linspace(x2_min, x2_max, resolution)
    X1, X2 = np.meshgrid(x1, x2)

    Z = np.zeros_like(X1)
    for i in range(X1.shape[0]):
        for j in range(X1.shape[1]):
            Z[i, j] = func(np.array([X1[i, j], X2[i, j]]))

    return X1, X2, Z


def plot_iteration_2d(
    func: Function2D,
    result: GradientDescentResult2D,
    iter_idx: int,
    output_path: Path,
    X1: np.ndarray,
    X2: np.ndarray,
    Z: np.ndarray,
    levels: np.ndarray,
):
    """Построить контурный график для одной итерации."""
    info = result.iterations[iter_idx]

    fig, ax = plt.subplots(figsize=(10, 8))

    # Контурные линии
    contour = ax.contour(
        X1, X2, Z, levels=levels, colors="gray", alpha=0.6, linewidths=0.8
    )
    ax.clabel(contour, inline=True, fontsize=8, fmt="%.1f")

    # Заливка
    ax.contourf(X1, X2, Z, levels=levels, cmap=cm.viridis, alpha=0.3)

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

    if len(trajectory) > 1:
        ax.plot(
            trajectory[:, 0],
            trajectory[:, 1],
            "b-",
            linewidth=2,
            alpha=0.7,
            label="Траектория",
            zorder=5,
        )

    # Предыдущие точки
    for i in range(len(trajectory) - 1):
        ax.plot(
            trajectory[i, 0], trajectory[i, 1], "bo", markersize=6, alpha=0.5, zorder=6
        )

    # Текущая точка
    ax.plot(
        info.x[0],
        info.x[1],
        "ro",
        markersize=12,
        label=f"x = ({info.x[0]:.4f}, {info.x[1]:.4f})\nf(x) = {info.f_x:.4f}",
        zorder=7,
    )

    # Направление антиградиента
    grad_norm = np.linalg.norm(info.grad)
    if grad_norm > 0:
        direction = -info.grad / grad_norm * 0.5  # Нормализуем и масштабируем
        ax.arrow(
            info.x[0],
            info.x[1],
            direction[0],
            direction[1],
            head_width=0.1,
            head_length=0.05,
            fc="magenta",
            ec="magenta",
            alpha=0.7,
            zorder=8,
        )

    ax.set_xlabel("x₁", fontsize=12)
    ax.set_ylabel("x₂", fontsize=12)
    ax.set_title(
        f"{result.method} — Итерация {info.iteration}\n"
        f"Шаг: {info.step_size:.6f}, ||∇f|| = {np.linalg.norm(info.grad):.6f}",
        fontsize=14,
        fontweight="bold",
    )
    ax.legend(fontsize=10, loc="upper right")
    ax.set_aspect("equal")
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()


def plot_final_result_2d(
    func: Function2D,
    result: GradientDescentResult2D,
    output_path: Path,
    X1: np.ndarray,
    X2: np.ndarray,
    Z: np.ndarray,
    levels: np.ndarray,
):
    """Построить итоговый контурный график с полной траекторией."""
    fig, ax = plt.subplots(figsize=(10, 8))

    # Контурные линии
    contour = ax.contour(
        X1, X2, Z, levels=levels, colors="gray", alpha=0.6, linewidths=0.8
    )
    ax.clabel(contour, inline=True, fontsize=8, fmt="%.1f")

    # Заливка
    ax.contourf(X1, X2, Z, levels=levels, cmap=cm.viridis, alpha=0.3)

    # Траектория
    trajectory = np.array([it.x for it in result.iterations])
    ax.plot(
        trajectory[:, 0],
        trajectory[:, 1],
        "b-",
        linewidth=2,
        alpha=0.8,
        label="Траектория",
        zorder=5,
    )

    # Все точки
    for i, point in enumerate(trajectory[:-1]):
        ax.plot(point[0], point[1], "bo", markersize=6, alpha=0.5, zorder=6)

    # Стартовая точка
    ax.plot(
        trajectory[0, 0],
        trajectory[0, 1],
        "go",
        markersize=12,
        label=f"Старт: ({trajectory[0, 0]:.2f}, {trajectory[0, 1]:.2f})",
        zorder=7,
    )

    # Финальная точка
    ax.plot(
        result.x_star[0],
        result.x_star[1],
        "r*",
        markersize=20,
        label=f"x* = ({result.x_star[0]:.4f}, {result.x_star[1]:.4f})\n"
        f"f(x*) = {result.f_star:.6f}",
        zorder=8,
    )

    ax.set_xlabel("x₁", fontsize=12)
    ax.set_ylabel("x₂", fontsize=12)
    ax.set_title(
        f"{result.method} — Результат\nИтераций: {len(result.iterations) - 1}",
        fontsize=14,
        fontweight="bold",
    )
    ax.legend(fontsize=10, loc="upper right")
    ax.set_aspect("equal")
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()


def run_and_visualize_2d(
    func: Function2D,
    x0: np.ndarray,
    method: str,
    method_name_short: str,
    X1: np.ndarray,
    X2: np.ndarray,
    Z: np.ndarray,
    levels: np.ndarray,
    max_plot_iters: int = 20,
    **kwargs,
):
    """Запустить метод и создать визуализации."""

    result = gradient_descent_2d(
        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)

    # Определяем какие итерации визуализировать
    total_iters = len(result.iterations) - 1
    if total_iters <= max_plot_iters:
        plot_indices = list(range(total_iters))
    else:
        # Выбираем равномерно распределённые итерации
        step = total_iters / max_plot_iters
        plot_indices = [int(i * step) for i in range(max_plot_iters)]
        if total_iters - 1 not in plot_indices:
            plot_indices.append(total_iters - 1)

    for idx, info in enumerate(result.iterations[:-1]):
        print(
            f"Итерация {info.iteration:3d}: "
            f"x = ({info.x[0]:10.6f}, {info.x[1]:10.6f}), "
            f"f(x) = {info.f_x:12.6f}, ||∇f|| = {np.linalg.norm(info.grad):10.6f}, "
            f"шаг = {info.step_size:.6f}"
        )

        # Строим график только для выбранных итераций
        if idx in plot_indices:
            plot_iteration_2d(
                func,
                result,
                idx,
                method_dir / f"iteration_{info.iteration:03d}.png",
                X1,
                X2,
                Z,
                levels,
            )

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

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

    return result


def run_and_visualize_heavy_ball(
    func: Function2D,
    x0: np.ndarray,
    method_name_short: str,
    X1: np.ndarray,
    X2: np.ndarray,
    Z: np.ndarray,
    levels: np.ndarray,
    alpha: float,
    beta: float,
    max_plot_iters: int = 20,
):
    """Запустить метод тяжёлого шарика и создать визуализации."""

    result = heavy_ball_2d(
        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)

    # Определяем какие итерации визуализировать
    total_iters = len(result.iterations) - 1
    if total_iters <= max_plot_iters:
        plot_indices = list(range(total_iters))
    else:
        step = total_iters / max_plot_iters
        plot_indices = [int(i * step) for i in range(max_plot_iters)]
        if total_iters - 1 not in plot_indices:
            plot_indices.append(total_iters - 1)

    for idx, info in enumerate(result.iterations[:-1]):
        print(
            f"Итерация {info.iteration:3d}: "
            f"x = ({info.x[0]:10.6f}, {info.x[1]:10.6f}), "
            f"f(x) = {info.f_x:12.6f}, ||∇f|| = {np.linalg.norm(info.grad):10.6f}, "
            f"шаг = {info.step_size:.6f}"
        )

        if idx in plot_indices:
            plot_iteration_2d(
                func,
                result,
                idx,
                method_dir / f"iteration_{info.iteration:03d}.png",
                X1,
                X2,
                Z,
                levels,
            )

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

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

    return result


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

    # Выбираем функцию
    if FUNCTION_CHOICE == "himmelblau":
        func = HimmelblauFunction()
    elif FUNCTION_CHOICE == "ravine":
        func = RavineFunction()
    else:
        raise ValueError(f"Unknown function: {FUNCTION_CHOICE}")

    x0 = START_POINTS[FUNCTION_CHOICE]
    constant_step = CONSTANT_STEPS[FUNCTION_CHOICE]

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

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

    # Создаём сетку для контурных графиков (один раз)
    print("\nСоздание сетки для контурных графиков...")
    X1, X2, Z = create_contour_grid(func, resolution=200)

    # Уровни для контурных линий
    z_min, z_max = Z.min(), Z.max()
    if FUNCTION_CHOICE == "himmelblau":
        # Логарифмические уровни для лучшей визуализации
        levels = np.array([0.5, 1, 2, 5, 10, 20, 40, 80, 150, 300, 500])
        levels = levels[levels < z_max]
    elif FUNCTION_CHOICE == "ravine":
        # Уровни для овражной функции - эллипсы
        levels = np.array([0.01, 0.05, 0.1, 0.2, 0.5, 1, 2, 3, 5, 7, 10])
        levels = levels[levels < z_max]
    else:
        levels = np.linspace(z_min, min(z_max, 100), 20)

    # Убедимся, что уровни уникальны и отсортированы
    levels = np.unique(levels)

    # 1. Константный шаг
    run_and_visualize_2d(
        func,
        x0,
        method="constant",
        method_name_short="constant",
        X1=X1,
        X2=X2,
        Z=Z,
        levels=levels,
        step_size=constant_step,
    )

    # 2. Золотое сечение
    run_and_visualize_2d(
        func,
        x0,
        method="golden_section",
        method_name_short="golden_section",
        X1=X1,
        X2=X2,
        Z=Z,
        levels=levels,
        golden_section_bounds=GOLDEN_SECTION_BOUNDS,
    )

    # 3. Правило Армихо
    run_and_visualize_2d(
        func,
        x0,
        method="armijo",
        method_name_short="armijo",
        X1=X1,
        X2=X2,
        Z=Z,
        levels=levels,
        armijo_params=ARMIJO_PARAMS,
    )

    # 4. Метод тяжёлого шарика
    hb_params = HEAVY_BALL_PARAMS[FUNCTION_CHOICE]
    run_and_visualize_heavy_ball(
        func,
        x0,
        method_name_short="heavy_ball",
        X1=X1,
        X2=X2,
        Z=Z,
        levels=levels,
        alpha=hb_params["alpha"],
        beta=hb_params["beta"],
    )

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


if __name__ == "__main__":
    main()