#!/bin/python3

# SPDX-License-Identifier: GPL-2.0-only

import matplotlib.animation as animation
import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.lines import Line2D

#
# Plot the traces generated by three-gpp-channel-consistency-example
#

# -----------------------------------------------------------------------------
# Load data
# -----------------------------------------------------------------------------
df = pd.read_csv("3gpp-channel-consistency-output.txt", sep=r"\s+", comment="#")

# Column names from the ns-3 trace
TIME_COL = "Time[s]"
TX_X_COL = "TxPosX[m]"
TX_Y_COL = "TxPosY[m]"
RX_X_COL = "RxPosX[m]"
RX_Y_COL = "RxPosY[m]"
STATE_COL = "ChannelState"
SNR_COL = "SNR[dB]"

# Convenience series
time = df[TIME_COL]
tx_x = df[TX_X_COL]
tx_y = df[TX_Y_COL]
rx_x = df[RX_X_COL]
rx_y = df[RX_Y_COL]
snr = df[SNR_COL]

# Channel state (for LOS/NLOS vertical lines)
state = df[STATE_COL].astype(str)

# -----------------------------------------------------------------------------
# Varying rates
# -----------------------------------------------------------------------------
snr_diff = snr.diff().fillna(0.0)
df["SNRdiff"] = snr_diff
# -----------------------------------------------------------------------------
# Compute LOS/NLOS transition times and types
# -----------------------------------------------------------------------------
prev_state = state.shift(1, fill_value=state.iloc[0])

times_LN = []  # LOS -> NLOS(NLOSv)
times_NL = []  # NLOS(NLOSv) -> LOS

for t, s_prev, s_cur in zip(time, prev_state, state):
    if s_prev == s_cur:
        continue

    s_prev_is_nlos = s_prev.startswith("NLOS")
    s_cur_is_nlos = s_cur.startswith("NLOS")

    if s_prev == "LOS" and s_cur_is_nlos:
        times_LN.append(t)
    elif s_prev_is_nlos and s_cur == "LOS":
        times_NL.append(t)


def add_state_vlines(ax, add_legend=False):
    """Draw vertical lines for LOS/NLOS transitions, with different types."""
    # NLOS -> LOS (green, solid)
    for t in times_NL:
        ax.axvline(t, color="green", linestyle="-", linewidth=0.9, alpha=0.6)
    # LOS -> NLOS (red, dashed)
    for t in times_LN:
        ax.axvline(t, color="red", linestyle="--", linewidth=0.9, alpha=0.6)

    if add_legend:
        legend_lines = [
            Line2D([0], [0], color="green", linestyle="-", linewidth=1.2, label="NLOS → LOS"),
            Line2D([0], [0], color="red", linestyle="--", linewidth=1.2, label="LOS → NLOS"),
        ]
        ax.legend(handles=legend_lines, loc="upper right", fontsize=8)


# -----------------------------------------------------------------------------
# Global plot params
# -----------------------------------------------------------------------------
zoom_window_s = 0.05  # SNR zoom window (example: 0.05 = 50 ms)

t_min = time.min()
t_max = time.max()

# SNR limits
snr_min = snr.min()
snr_max = snr.max()
snr_margin = 3.0
# snr_ymin = snr_min - snr_margin
# snr_ymax = snr_max + snr_margin
snr_ymax = 30
snr_ymin = -60


# ΔSNR limits: symmetric around 0 with minimum span
max_abs_dsnr = float(np.max(np.abs(snr_diff)))
if max_abs_dsnr < 1.0:
    max_abs_dsnr = 1.0  # at least +/- 1 dB
# snr_diff_ymin = -1.2 * max_abs_dsnr
# snr_diff_ymax = +1.2 * max_abs_dsnr

snr_diff_ymin = -35
snr_diff_ymax = +35


# -----------------------------------------------------------------------------
# Frame downsampling for animation (lightweight)
# Simulation: 40 s @ 1 ms => ~40 000 samples
# We cap GIF to ~800 frames => ~1 frame each 50 ms
# -----------------------------------------------------------------------------
n_samples = len(df)
max_frames = 800  # target max frames in the GIF
frame_step = max(1, n_samples // max_frames)
frame_indices = range(0, n_samples, frame_step)

# -----------------------------------------------------------------------------
# 2×2 Animated Figure:
# row 0: [map, SNR global]
# row 1: [SNR zoom, ΔSNR]
# -----------------------------------------------------------------------------
fig, axes = plt.subplots(2, 2, figsize=(24, 9), dpi=150)

ax_map = axes[0, 0]
ax_snr = axes[0, 1]
ax_snr_zoom = axes[1, 0]
ax_snr_diff = axes[1, 1]

# -----------------------------------------------------------------------------
# MAP
# -----------------------------------------------------------------------------
ax_map.set_xlabel("X [m]")
ax_map.set_ylabel("Y [m]")
ax_map.set_aspect("equal")

ax_map.set_xlim(-25, 600)
ax_map.set_ylim(-25, 1000)

tx_circle = patches.Circle((0.0, 0.0), 5.0, color="blue", alpha=0.35)
rx_circle = patches.Circle((0.0, 0.0), 5.0, color="red", alpha=0.35)
ax_map.add_patch(tx_circle)
ax_map.add_patch(rx_circle)

buildings = pd.read_csv(
    "3gpp-channel-consistency-buildings.txt", sep=r"\s+", comment="#", header=None
)

for _, b in buildings.iterrows():
    x0, y0, x1, y1 = b
    rect = patches.Rectangle((x0, y0), x1 - x0, y1 - y0, color="gray", alpha=0.5)
    ax_map.add_patch(rect)

# -----------------------------------------------------------------------------
# SNR Global
# -----------------------------------------------------------------------------
ax_snr.set_xlabel("Time [s]")
ax_snr.set_ylabel("SNR [dB]")
ax_snr.grid(True)
ax_snr.set_xlim(t_min, t_max)
ax_snr.set_ylim(snr_ymin, snr_ymax)
(snr_line,) = ax_snr.plot([], [], "k-", linewidth=1.5)

add_state_vlines(ax_snr, add_legend=True)

# -----------------------------------------------------------------------------
# SNR Zoom
# -----------------------------------------------------------------------------
ax_snr_zoom.set_xlabel("Time [s]")
ax_snr_zoom.set_ylabel("SNR [dB] (zoom)")
ax_snr_zoom.grid(True)
ax_snr_zoom.set_ylim(snr_ymin, snr_ymax)
(snr_zoom_line,) = ax_snr_zoom.plot([], [], "k-", linewidth=1.5)

add_state_vlines(ax_snr_zoom)

# -----------------------------------------------------------------------------
# ΔSNR (global)
# -----------------------------------------------------------------------------
ax_snr_diff.set_xlabel("Time [s]")
ax_snr_diff.set_ylabel("ΔSNR [dB]")
ax_snr_diff.grid(True)
ax_snr_diff.set_xlim(t_min, t_max)
ax_snr_diff.set_ylim(snr_diff_ymin, snr_diff_ymax)
(snr_diff_line,) = ax_snr_diff.plot([], [], "k-", linewidth=1.5)

add_state_vlines(ax_snr_diff)

# -----------------------------------------------------------------------------
# Static SNR + ΔSNR figure (2×1)
# -----------------------------------------------------------------------------
fig_snr, axes_snr = plt.subplots(2, 1, figsize=(10, 6), sharex=True)

axes_snr[0].plot(time, snr, "k-", linewidth=1.2)
axes_snr[0].set_ylabel("SNR [dB]")
axes_snr[0].grid(True)
add_state_vlines(axes_snr[0], add_legend=True)

axes_snr[1].plot(time, snr_diff, "k-", linewidth=1.2)
axes_snr[1].set_ylabel("ΔSNR [dB]")
axes_snr[1].set_xlabel("Time [s]")
axes_snr[1].grid(True)
add_state_vlines(axes_snr[1])

axes_snr[0].set_xlim(t_min, t_max)
axes_snr[0].set_ylim(snr_ymin, snr_ymax)
axes_snr[1].set_ylim(snr_diff_ymin, snr_diff_ymax)
fig_snr.tight_layout()
fig_snr.savefig("channel_consistency_snr_and_dsnr.png", dpi=300, bbox_inches="tight")
plt.close(fig_snr)


# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------
# Animation update function (uses downsampled frame_indices)
# -----------------------------------------------------------------------------
def update(frame_idx):
    # frame_idx is an actual index into df (from frame_indices)
    t_now = time.iloc[frame_idx]

    # Map
    tx_circle.set_center((tx_x.iloc[frame_idx], tx_y.iloc[frame_idx]))
    rx_circle.set_center((rx_x.iloc[frame_idx], rx_y.iloc[frame_idx]))
    ax_map.set_title(f"Time = {t_now:.3f} s")

    # Global SNR
    snr_line.set_data(time.iloc[: frame_idx + 1], snr.iloc[: frame_idx + 1])

    # SNR zoom
    t0 = max(t_min, t_now - zoom_window_s / 2.0)
    t1 = min(t_max, t_now + zoom_window_s / 2.0)
    mask = (time >= t0) & (time <= t1)

    ax_snr_zoom.set_xlim(t0, t1)
    snr_zoom_line.set_data(time[mask], snr[mask])

    # ΔSNR (global)
    snr_diff_line.set_data(time.iloc[: frame_idx + 1], snr_diff.iloc[: frame_idx + 1])

    return (tx_circle, rx_circle, snr_line, snr_zoom_line, snr_diff_line)


# -----------------------------------------------------------------------------
# Run animation (downsampled frames, lighter GIF)
# -----------------------------------------------------------------------------
ani = animation.FuncAnimation(
    fig,
    update,
    frames=frame_indices,
    interval=0,
    blit=False,
)

ani.save("3gpp-channel-consistency.gif", writer="pillow", fps=25, dpi=90)
