from dataclasses import dataclass
from typing import Dict, List, Tuple
import numpy as np
import matplotlib.pyplot as plt



@dataclass(frozen=True)
class SimulationConfig:
    n_nodes: int = 140
    dt: float = 0.01
    dy: float = 1.0
    amplitude: float = 4.0
    omega: float = 1.0
    wave_speed: float = 2.0
    t_max: float = 100.0
    target_times: Tuple[float, ...] = (
        5.85, 7.05, 12.36, 19.15, 21.25, 24.25, 28.35, 100.00
    )
    front_threshold_ratio: float = 0.1
    ignore_edge_points: int = 5
    drive_periods: float = 2.0

    @property
    def drive_stop_time(self) -> float:
        return self.drive_periods * (2.0 * np.pi / self.omega)



def simulate(cfg: SimulationConfig) -> Dict[str, np.ndarray]:

    n = cfg.n_nodes
    dt, dy = cfg.dt, cfg.dy
    A, w, c = cfg.amplitude, cfg.omega, cfg.wave_speed


    y = np.arange(n, dtype=float) * dy


    x = np.zeros(n, dtype=float)
    v = np.zeros(n, dtype=float)


    frame_times = np.asarray(cfg.target_times, dtype=float)
    frame_steps = np.clip(np.rint(frame_times / dt).astype(int), 0, int(np.rint(cfg.t_max / dt)))
    frame_map = {int(s): i for i, s in enumerate(frame_steps)}
    snapshots = np.zeros((len(frame_times), n), dtype=float)


    threshold = cfg.front_threshold_ratio * A
    arrival_times = np.full(n, np.nan, dtype=float)

    n_steps = int(np.rint(cfg.t_max / dt))
    c2_dt = (c * c) * dt / (dy * dy)


    for step in range(n_steps + 1):
        t = step * dt


        if step in frame_map:
            snapshots[frame_map[step]] = x


        if t <= cfg.drive_stop_time:
            x[0] = A * np.sin(w * t)
        else:
            x[0] = 0.0
        v[0] = 0.0


        x[-1] = 0.0
        v[-1] = 0.0


        lap = x[2:] - 2.0 * x[1:-1] + x[:-2]
        v[1:-1] += c2_dt * lap
        x[1:-1] += dt * v[1:-1]


        need = np.isnan(arrival_times)
        crossed = np.abs(x) >= threshold
        newly = need & crossed
        if np.any(newly):
            arrival_times[newly] = t


    L = cfg.ignore_edge_points
    core = arrival_times[L:n - L]
    avg_arrival = float(np.nanmean(core)) if np.any(~np.isnan(core)) else np.nan

    return {
        "y": y,
        "snapshots": snapshots,
        "arrival_times": arrival_times,
        "avg_arrival": np.array(avg_arrival),
    }



def plot_results(cfg: SimulationConfig, data: Dict[str, np.ndarray]) -> None:
    y = data["y"]
    Xs = data["snapshots"]
    arrival_times = data["arrival_times"]
    avg_arrival = float(data["avg_arrival"]) if np.ndim(data["avg_arrival"]) else float(data["avg_arrival"])

    n_frames = Xs.shape[0]
    rows, cols = 2, 4

    fig, axes = plt.subplots(rows, cols, figsize=(14, 6), sharey=True)
    axes = np.asarray(axes).reshape(-1)


    y_min = -cfg.amplitude * 1.05
    y_max = cfg.amplitude * 1.05

    for i, ax in enumerate(axes):
        t_show = cfg.target_times[i]
        ax.plot(y, Xs[i], linewidth=1.5)
        ax.set_title(f"t = {t_show:.2f} с", fontsize=10)
        ax.set_xlim(y[0], y[-1])
        ax.set_ylim(y_min, y_max)
        ax.grid(True, linewidth=0.6, alpha=0.4)
        ax.tick_params(labelsize=9)
        for spine in ax.spines.values():
            spine.set_linewidth(0.8)


    ax_last = axes[-1]
    L = cfg.ignore_edge_points
    for i in range(L, cfg.n_nodes - L):
        t_a = arrival_times[i]
        if not np.isnan(t_a):
            ax_last.axvline(i * cfg.dy, ymin=0.02, ymax=0.12, linewidth=0.6, linestyle=":", alpha=0.5)

    fig.suptitle("Розповсюдження одновимірної хвилі (кадри) і час приходу фронту", fontsize=12)
    fig.tight_layout(rect=(0, 0, 1, 0.95))
    plt.show()

    print(f"З рисунка бачимо, що середнiй час поширення фронту хвилi становить t ≈ {avg_arrival:.0f}.")



if __name__ == "__main__":
    cfg = SimulationConfig()
    results = simulate(cfg)
    plot_results(cfg, results)