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

from __future__ import annotations

import math
from pathlib import Path
from typing import Dict, Tuple

import matplotlib.pyplot as plt
import numpy as np

from numba import njit


G = 0.8
I_EPS = 0.8
D_NOISE = 0.8
I_SIGMA = 0.1
TAU_SIGMA = 1.0

POINTS: Dict[str, Tuple[float, float, str, str]] = {
    "1": (1.0, 3.8, "SF", "Рідинне тертя"),
    "2": (6.0, 2.5, "SS", "Переривчасте тертя"),
    "3": (3.0, 0.5, "DF", "Сухе тертя"),
}


@njit(cache=True)
def drift(sigma: float, i_t: float, t_e: float, g: float = G) -> float:
    d = 1.0 / (1.0 + sigma * sigma)
    return -sigma + g * sigma * (1.0 - (2.0 - t_e) * d)


@njit(cache=True)
def intensity(sigma: float, i_t: float, i_sigma: float = I_SIGMA, i_eps: float = I_EPS, g: float = G) -> float:
    d = 1.0 / (1.0 + sigma * sigma)
    return i_sigma + (i_eps + i_t * sigma * sigma) * g * g * d * d


@njit(cache=True)
def simulate_path(noise: np.ndarray, delta_t: float, i_t: float, t_e: float, sigma_0: float = 0.0) -> np.ndarray:
    n = noise.size
    out = np.empty(n + 1, dtype=np.float64)
    out[0] = sigma_0
    sigma = sigma_0
    noise_scale = math.sqrt(2.0 * D_NOISE * delta_t) / TAU_SIGMA
    drift_scale = delta_t / TAU_SIGMA
    for i in range(n):
        ff = drift(sigma, i_t, t_e)
        ii = intensity(sigma, i_t)
        sigma = sigma + ff * drift_scale + math.sqrt(ii) * noise_scale * noise[i]
        out[i + 1] = sigma
    return out


@njit(cache=True)
def simulate_prob_samples(noise: np.ndarray, delta_t: float, i_t: float, t_e: float, burn_in: int, thin: int, sigma_0: float = 0.0) -> np.ndarray:
    n = noise.size
    kept = (n - burn_in + thin - 1) // thin
    out = np.empty(kept, dtype=np.float64)
    sigma = sigma_0
    noise_scale = math.sqrt(2.0 * D_NOISE * delta_t) / TAU_SIGMA
    drift_scale = delta_t / TAU_SIGMA
    j = 0
    for i in range(n):
        ff = drift(sigma, i_t, t_e)
        ii = intensity(sigma, i_t)
        sigma = sigma + ff * drift_scale + math.sqrt(ii) * noise_scale * noise[i]
        if i >= burn_in and ((i - burn_in) % thin == 0):
            out[j] = sigma
            j += 1
    return out[:j]



def line_i(i_t: np.ndarray) -> np.ndarray:
    return 1.0 + 1.0 / G + 2.0 * G * D_NOISE * (i_t - 2.0 * I_EPS)



def curve_ii_sigma_search(i_t_values: np.ndarray, sigma_max: float = 10.0, dsigma: float = 0.001) -> tuple[np.ndarray, np.ndarray]:
    sigma = np.arange(0.0, sigma_max + dsigma, dsigma)
    x = 1.0 + sigma * sigma
    i_vals = []
    te_vals = []
    for i_t in i_t_values:
        te = 2.0 + ((1.0 - G) * x**3 - 2.0 * G * G * D_NOISE * i_t * x + 4.0 * G * G * D_NOISE * (i_t - I_EPS)) / (G * x**2)
        mask = (te[1:-1] < te[:-2]) & (te[1:-1] < te[2:])
        if np.any(mask):
            idxs = np.where(mask)[0] + 1
            idx = idxs[0]
            i_vals.append(i_t)
            te_vals.append(te[idx])
    return np.asarray(i_vals), np.asarray(te_vals)



def tricritical_point() -> tuple[float, float]:
    te = (2.0 / 3.0) * (1.0 + 2.0 / G - 2.0 * D_NOISE * G * I_EPS)
    i_t = (1.0 / (6.0 * G * D_NOISE)) * (1.0 / G - 1.0 + 8.0 * D_NOISE * G * I_EPS)
    return i_t, te



def analytical_probability(sigma_grid: np.ndarray, i_t: float, t_e: float) -> np.ndarray:
    i_vals = intensity(sigma_grid, i_t)
    f_vals = drift(sigma_grid, i_t, t_e)
    integrand = f_vals / i_vals
    integral = np.zeros_like(sigma_grid)
    dx = np.diff(sigma_grid)
    integral[1:] = np.cumsum(0.5 * (integrand[:-1] + integrand[1:]) * dx)
    u = np.log(i_vals) - integral / D_NOISE
    u = u - np.min(u)
    p = np.exp(-u)
    z = np.trapezoid(p, sigma_grid)
    return p / z



def numerical_probability(samples: np.ndarray, sigma_grid: np.ndarray, symmetrize: bool = True) -> np.ndarray:
    bins = np.linspace(sigma_grid[0], sigma_grid[-1], sigma_grid.size)
    hist, edges = np.histogram(samples, bins=bins, density=True)
    centers = 0.5 * (edges[:-1] + edges[1:])
    kernel = np.array([1, 2, 3, 2, 1], dtype=np.float64)
    kernel /= kernel.sum()
    smooth = np.convolve(hist, kernel, mode="same")
    if symmetrize:
        smooth = 0.5 * (smooth + smooth[::-1])
    return centers, smooth



def make_phase_diagram(out_path: Path) -> None:
    plt.rcParams.update({
        "font.family": "DejaVu Serif",
        "font.size": 12,
        "mathtext.fontset": "dejavuserif",
    })
    i_t_values = np.arange(0.0, 18.01, 0.01)
    te_i = line_i(i_t_values)
    i_curve, te_curve = curve_ii_sigma_search(i_t_values)
    i_tri, te_tri = tricritical_point()

    fig, ax = plt.subplots(figsize=(7.1, 5.3), dpi=180)
    ax.plot(i_t_values, te_i, lw=1.6, label="I")
    ax.plot(i_curve, te_curve, lw=1.6, label="II")
    ax.plot(i_tri, te_tri, marker="o", ms=4)
    ax.text(i_tri + 0.25, te_tri - 0.1, "T", fontsize=11)

    for key, (i_t, t_e, region, _) in POINTS.items():
        ax.plot(i_t, t_e, marker="D", ms=5, mfc="white", mec="black", mew=0.9)
        ax.text(i_t + 0.2, t_e + 0.05, key, fontsize=11)

    ax.text(1.2, 4.35, "SF", fontsize=12)
    ax.text(10.5, 3.05, "SS", fontsize=12)
    ax.text(5.8, 0.55, "DF", fontsize=12)
    ax.text(2.15, 3.95, "I", fontsize=12)
    ax.text(10.7, 1.1, "II", fontsize=12)

    ax.set_xlim(0, 18)
    ax.set_ylim(0, 5)
    ax.set_xlabel(r"$I_T$")
    ax.set_ylabel(r"$T_e$")
    ax.set_title("Фазова діаграма при g = 0.8, Iε = 0.8, D = 0.8", fontsize=12, pad=10)
    ax.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)



def make_time_series(out_path: Path, seed: int = 20260409) -> None:
    plt.rcParams.update({
        "font.family": "DejaVu Serif",
        "font.size": 11,
        "mathtext.fontset": "dejavuserif",
    })
    rng = np.random.default_rng(seed)
    t_all = 5000.0
    n = 200_000
    delta_t = t_all / n
    t = np.linspace(0.0, t_all, n + 1)

    fig, axes = plt.subplots(3, 1, figsize=(7.2, 7.6), dpi=180, sharex=True)
    for ax, key in zip(axes, ["1", "2", "3"]):
        i_t, t_e, region, region_full = POINTS[key]
        noise = rng.normal(size=n)
        sigma = simulate_path(noise, delta_t, i_t, t_e)
        ax.plot(t, sigma, lw=0.5)
        ax.set_ylabel("σ")
        ax.text(0.02, 0.87, f"Точка {key}: $I_T$={i_t:g}, $T_e$={t_e:g} ({region})", transform=ax.transAxes, fontsize=10)
        ax.set_ylim(-6.0, 6.0)
        ax.grid(alpha=0.2, linewidth=0.4)
    axes[-1].set_xlabel("t")
    fig.suptitle("Часові ряди напружень σ(t)", fontsize=13, y=0.995)
    fig.tight_layout(rect=(0, 0, 1, 0.985))
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)



def make_probability_figure(out_path: Path, seed: int = 20260410) -> None:
    plt.rcParams.update({
        "font.family": "DejaVu Serif",
        "font.size": 11,
        "mathtext.fontset": "dejavuserif",
    })
    sigma_grid = np.linspace(-6.0, 6.0, 1201)
    rng = np.random.default_rng(seed)

    fig, axes = plt.subplots(2, 1, figsize=(7.2, 7.6), dpi=180, sharex=True)

    for key in ["1", "2", "3"]:
        i_t, t_e, region, _ = POINTS[key]
        p = analytical_probability(sigma_grid, i_t, t_e)
        axes[0].plot(sigma_grid, p, lw=1.25)
        peak_idx = int(np.argmax(p))
        axes[0].text(sigma_grid[peak_idx] + 0.12, p[peak_idx] + 0.01, key, fontsize=11)

        n_steps = 4_000_000 if key == "1" else 2_000_000
        burn_in = 100_000
        thin = 10
        noise = rng.normal(size=n_steps)
        samples = simulate_prob_samples(noise, 0.01, i_t, t_e, burn_in=burn_in, thin=thin)
        centers, p_num = numerical_probability(samples, sigma_grid)
        axes[1].plot(centers, p_num, lw=1.15)
        peak_idx_num = int(np.argmax(p_num))
        axes[1].text(centers[peak_idx_num] + 0.12, p_num[peak_idx_num] + 0.01, key, fontsize=11)

    axes[0].text(0.05, 0.88, "a", transform=axes[0].transAxes, fontsize=14, fontstyle="italic")
    axes[1].text(0.05, 0.88, "b", transform=axes[1].transAxes, fontsize=14, fontstyle="italic")
    axes[0].set_ylabel("P(σ)")
    axes[1].set_ylabel("P(σ)")
    axes[1].set_xlabel("σ")
    axes[0].set_title("Аналітичний розподіл")
    axes[1].set_title("Чисельний розподіл за часовими рядами")
    for ax in axes:
        ax.grid(alpha=0.18, linewidth=0.4)
        ax.set_xlim(-6.0, 6.0)
    fig.suptitle("Порівняння аналітичних і чисельних розподілів ймовірності", fontsize=13, y=0.995)
    fig.tight_layout(rect=(0, 0, 1, 0.985))
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)



def main() -> None:
    base_dir = Path(__file__).resolve().parent
    out_dir = base_dir / "pr6_outputs"
    out_dir.mkdir(exist_ok=True)

    make_phase_diagram(out_dir / "phase_diagram.png")
    make_time_series(out_dir / "time_series.png")
    make_probability_figure(out_dir / "probability_distributions.png")

    print(f"Saved outputs to: {out_dir}")


if __name__ == "__main__":
    main()