НПтаМ_ПР№5_Мірошниченко

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import matplotlib.pyplot as plt
import numpy as np


@dataclass(frozen=True)
class SecondOrderParams:
    g: float = 0.5


@dataclass(frozen=True)
class FirstOrderParams:
    g_theta: float = 0.6
    theta: float = 0.1
    alpha: float = 0.75


def epsilon_of_sigma(sigma: np.ndarray, T_e: float) -> np.ndarray:
    return sigma * (T_e - 1.0 + sigma**2) / (1.0 + sigma**2)


def temperature_of_sigma(sigma: np.ndarray, T_e: float) -> np.ndarray:
    return (T_e + 2.0 * sigma**2) / (1.0 + sigma**2)


def potential_second_order(sigma: np.ndarray, g: float, T_e: float) -> np.ndarray:
    return 0.5 * (1.0 - g) * sigma**2 + g * (1.0 - T_e / 2.0) * np.log(1.0 + sigma**2)


def critical_temperature_second_order(g: float) -> float:
    return 1.0 + 1.0 / g


def stationary_sigma_second_order(T_e: np.ndarray, g: float) -> np.ndarray:
    T_e = np.asarray(T_e, dtype=float)
    sigma0 = np.zeros_like(T_e)
    mask = T_e > critical_temperature_second_order(g)
    val = (g * T_e[mask] - (g + 1.0)) / (1.0 - g)
    sigma0[mask] = np.sqrt(np.clip(val, 0.0, None))
    return sigma0


def stationary_epsilon_second_order(sigma0: np.ndarray, g: float) -> np.ndarray:
    return sigma0 / g


def stationary_temperature_second_order(sigma0: np.ndarray, T_e: np.ndarray, g: float) -> np.ndarray:
    T_e = np.asarray(T_e, dtype=float)
    T0 = np.array(T_e, copy=True)
    mask = sigma0 > 0
    T0[mask] = critical_temperature_second_order(g)
    return T0


def g_of_sigma_first_order(sigma: np.ndarray, p: FirstOrderParams) -> np.ndarray:
    return p.g_theta * (1.0 + (p.theta**-1 - 1.0) / (1.0 + (sigma / p.alpha) ** 2))


def potential_first_order(sigma: np.ndarray, T_e: float, p: FirstOrderParams) -> np.ndarray:
    term1 = 0.5 * (1.0 - p.g_theta) * sigma**2
    term2 = p.g_theta * (1.0 - T_e / 2.0) * np.log(1.0 + sigma**2)
    pref = -0.5 * p.g_theta * p.alpha**2 * (p.theta**-1 - 1.0) / (p.alpha**2 - 1.0)
    term3 = pref * (
        (T_e - 2.0) * np.log(1.0 + sigma**2)
        + (p.alpha**2 - T_e + 1.0) * np.log(1.0 + sigma**2 / p.alpha**2)
    )
    return term1 + term2 + term3


def D_plateau_term(p: FirstOrderParams) -> float:
    return (
        16.0
        * p.g_theta**2
        * p.alpha**-4
        * (p.g_theta - 1.0)
        * (p.theta**-1 - 1.0)
        * (p.alpha**-2 - p.theta**-1)
    )


def plateau_temperature_first_order(p: FirstOrderParams) -> float:
    D = D_plateau_term(p)
    return p.g_theta**-2 * (
        p.g_theta * (p.alpha**2 + p.g_theta + 1.0 - p.alpha**2 * p.g_theta * p.theta**-1)
        + 2.0 * p.alpha**2 * p.g_theta * p.theta**-1 * (p.g_theta - 1.0)
        + 0.5 * p.alpha**4 * np.sqrt(max(D, 0.0))
    )


def barrier_zero_temperature_first_order(p: FirstOrderParams) -> float:
    return 1.0 + p.theta / p.g_theta


def D0_stationary(T_e: np.ndarray, p: FirstOrderParams) -> np.ndarray:
    T_e = np.asarray(T_e, dtype=float)
    return (
        (p.g_theta * (T_e - 1.0) / p.alpha**2 + p.g_theta / p.theta - 1.0 - p.alpha**-2) ** 2
        - 4.0 * p.alpha**-2 * (p.g_theta - 1.0) * (p.g_theta * (T_e - 1.0) / p.theta - 1.0)
    )


def stationary_branches_first_order(T_e: np.ndarray, p: FirstOrderParams) -> tuple[np.ndarray, np.ndarray]:
    T_e = np.asarray(T_e, dtype=float)
    D0 = D0_stationary(T_e, p)
    sqrt_D0 = np.sqrt(np.clip(D0, 0.0, None))

    pref = 0.5 * p.alpha**2 / (p.g_theta - 1.0)
    base = 1.0 + p.alpha**-2 - p.g_theta * p.alpha**-2 * (T_e - 1.0) - p.g_theta * p.theta**-1

    sigma_max_sq = pref * (base + sqrt_D0)
    sigma_min_sq = pref * (base - sqrt_D0)

    sigma_unstable = np.sqrt(np.where(sigma_max_sq >= 0.0, sigma_max_sq, np.nan))
    sigma_stable = np.sqrt(np.where(sigma_min_sq >= 0.0, sigma_min_sq, np.nan))

    T_c0_lower = plateau_temperature_first_order(p)
    T_c0_upper = barrier_zero_temperature_first_order(p)
    mask = (T_e >= T_c0_lower) & (T_e <= T_c0_upper)

    sigma_stable = np.where(mask, sigma_stable, np.nan)
    sigma_unstable = np.where(mask, sigma_unstable, np.nan)
    return sigma_stable, sigma_unstable


def _save_figure(fig: plt.Figure, out_path: Path) -> None:
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=200, bbox_inches="tight")
    plt.close(fig)


def plot_part1_epsilon_sigma(out_dir: Path) -> Path:
    sigma = np.linspace(0.0, 4.0, 500)
    T_values = [1.0, 2.0, 4.2]

    fig, ax = plt.subplots(figsize=(8, 5))
    for T_e in T_values:
        ax.plot(sigma, epsilon_of_sigma(sigma, T_e), label=rf"$T_e = {T_e}$")
    ax.set_xlabel(r"$\sigma$")
    ax.set_ylabel(r"$\varepsilon$")
    ax.set_title(r"Частина 1. Залежність $\varepsilon(\sigma)$")
    ax.grid(True, alpha=0.3)
    ax.legend()

    path = out_dir / "figure_1_epsilon_sigma.png"
    _save_figure(fig, path)
    return path


def plot_part1_temperature_sigma(out_dir: Path) -> Path:
    sigma = np.linspace(0.0, 4.0, 500)
    T_values = [1.0, 2.0, 4.2]

    fig, ax = plt.subplots(figsize=(8, 5))
    for T_e in T_values:
        ax.plot(sigma, temperature_of_sigma(sigma, T_e), label=rf"$T_e = {T_e}$")
    ax.set_xlabel(r"$\sigma$")
    ax.set_ylabel(r"$T$")
    ax.set_title(r"Частина 1. Залежність $T(\sigma)$")
    ax.grid(True, alpha=0.3)
    ax.legend()

    path = out_dir / "figure_2_temperature_sigma.png"
    _save_figure(fig, path)
    return path


def plot_part1_potential(out_dir: Path, g: float = 0.5) -> Path:
    sigma = np.linspace(-3.5, 3.5, 800)
    T_values = [1.0, 4.2]
    T_c0 = critical_temperature_second_order(g)

    fig, ax = plt.subplots(figsize=(8, 5))
    for T_e in T_values:
        suffix = r"$< T_{c0}$" if T_e < T_c0 else r"$> T_{c0}$"
        ax.plot(sigma, potential_second_order(sigma, g, T_e), label=rf"$T_e = {T_e}$ {suffix}")
    ax.set_xlabel(r"$\sigma$")
    ax.set_ylabel(r"$V(\sigma)$")
    ax.set_title(rf"Частина 1. Потенціал $V(\sigma)$ для $g = {g}$, $T_{{c0}} = {T_c0:.1f}$")
    ax.grid(True, alpha=0.3)
    ax.legend()

    path = out_dir / "figure_3_potential_second_order.png"
    _save_figure(fig, path)
    return path


def plot_part1_stationary_sigma(out_dir: Path) -> Path:
    T_grid = np.linspace(1.0, 4.0, 500)
    g_values = [0.5, 0.6, 0.7, 0.8, 0.9]

    fig, ax = plt.subplots(figsize=(8, 5))
    for g in g_values:
        ax.plot(T_grid, stationary_sigma_second_order(T_grid, g), label=rf"$g = {g}$")
    ax.set_xlabel(r"$T_e$")
    ax.set_ylabel(r"$\sigma_0$")
    ax.set_title(r"Частина 1. Стаціонарні значення $\sigma_0(T_e)$")
    ax.grid(True, alpha=0.3)
    ax.legend()

    path = out_dir / "figure_4_stationary_sigma_second_order.png"
    _save_figure(fig, path)
    return path


def plot_part2_potential(out_dir: Path, p: FirstOrderParams) -> Path:
    sigma = np.linspace(-5.0, 5.0, 1200)
    T_c0_lower = plateau_temperature_first_order(p)
    T_c0_upper = barrier_zero_temperature_first_order(p)
    T_values = [0.50, T_c0_lower, 0.90, 1.30]

    fig, ax = plt.subplots(figsize=(8, 5))
    for T_e in T_values:
        if np.isclose(T_e, T_c0_lower, atol=1e-6):
            label = rf"$T_e = T_c^0 \approx {T_e:.3f}$"
        elif T_e < T_c0_lower:
            label = rf"$T_e = {T_e:.2f} < T_c^0$"
        elif T_e < T_c0_upper:
            label = rf"$T_c^0 < T_e = {T_e:.2f} < T_{{c0}}$"
        else:
            label = rf"$T_e = {T_e:.2f} > T_{{c0}}$"
        ax.plot(sigma, potential_first_order(sigma, T_e, p), label=label)
    ax.set_xlabel(r"$\sigma$")
    ax.set_ylabel(r"$V(\sigma)$")
    ax.set_title(r"Частина 2. Потенціал $V(\sigma)$ для переходу першого роду")
    ax.grid(True, alpha=0.3)
    ax.legend()

    path = out_dir / "figure_5_potential_first_order.png"
    _save_figure(fig, path)
    return path


def plot_part2_stationary_branches(out_dir: Path, p: FirstOrderParams) -> Path:
    T_c0_lower = plateau_temperature_first_order(p)
    T_c0_upper = barrier_zero_temperature_first_order(p)
    T_grid = np.linspace(T_c0_lower, T_c0_upper + 0.18, 600)
    sigma_stable, sigma_unstable = stationary_branches_first_order(T_grid, p)

    fig, ax = plt.subplots(figsize=(8, 5))
    ax.plot(T_grid, sigma_stable, label=r"стійка гілка $\sigma_0$")
    ax.plot(T_grid, sigma_unstable, linestyle="--", label=r"нестійка гілка $\sigma_m$")
    ax.axvline(T_c0_lower, linestyle=":")
    ax.axvline(T_c0_upper, linestyle=":")
    ax.text(T_c0_lower + 0.01, 2.2, r"$T_c^0$")
    ax.text(T_c0_upper + 0.01, 2.2, r"$T_{c0}$")
    ax.set_xlabel(r"$T_e$")
    ax.set_ylabel(r"$\sigma$")
    ax.set_title(r"Частина 2. Стаціонарні гілки $\sigma(T_e)$")
    ax.grid(True, alpha=0.3)
    ax.legend()

    path = out_dir / "figure_6_stationary_branches_first_order.png"
    _save_figure(fig, path)
    return path


def build_all_figures(out_dir: Path) -> list[Path]:
    p2 = FirstOrderParams()
    saved = [
        plot_part1_epsilon_sigma(out_dir),
        plot_part1_temperature_sigma(out_dir),
        plot_part1_potential(out_dir, g=0.5),
        plot_part1_stationary_sigma(out_dir),
        plot_part2_potential(out_dir, p2),
        plot_part2_stationary_branches(out_dir, p2),
    ]
    return saved


def print_summary(saved_paths: Iterable[Path]) -> None:
    p2 = FirstOrderParams()
    g = 0.5
    T_c0_2nd = critical_temperature_second_order(g)
    T_c0_1st_lower = plateau_temperature_first_order(p2)
    T_c0_1st_upper = barrier_zero_temperature_first_order(p2)

    print("Згенеровано файли:")
    for path in saved_paths:
        print(f" - {path}")

    print("\nКлючові параметри:")
    print(f"Частина 1: для g = {g} маємо T_c0 = 1 + 1/g = {T_c0_2nd:.3f}")
    print(
        "Частина 2: "
        f"T_c^0 ≈ {T_c0_1st_lower:.3f}, "
        f"T_c0 = 1 + θ/g_θ ≈ {T_c0_1st_upper:.3f}"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Побудова всіх рисунків для практичної роботи №5."
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=Path("pr5_output"),
        help="Папка, куди буде збережено рисунки.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    saved = build_all_figures(args.out_dir)
    print_summary(saved)


if __name__ == "__main__":
    main()