"""
Reference plots for Resource Allocation for Text Semantic Communications
Generated from paper images for verification purposes.

Run: python reference_plots.py
Output: workspace/resource_allocation/analysis/reference_images/
"""

import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
from scipy.interpolate import PchipInterpolator

OUTPUT_DIR = Path("workspace/resource_allocation/analysis/reference_images")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


def plot_figure_2():
    """
    Figure 2: The semantic similarity for DeepSC
    """
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection="3d")

    # Generate data
    snr = np.linspace(-10, 20, 30)
    k_n = np.linspace(0, 20, 30)
    SNR, KN = np.meshgrid(snr, k_n)

    # Approximate function for similarity
    # Logistic-like function depending on SNR and k_n
    z = 0.4 + 0.6 / (1 + np.exp(-0.3 * (SNR + 5)) * np.exp(-0.2 * (KN - 5)))
    z = np.clip(z, 0.4, 1.0)

    surf = ax.plot_surface(SNR, KN, z, cmap="viridis", edgecolor="none", alpha=0.9)

    ax.set_xlabel("SNR, $\gamma_{n,m}$ (dB)")
    ax.set_ylabel("$k_n$ (symbols/word)")
    ax.set_zlabel(r"$\xi_{n,m}$")
    ax.set_zlim(0.4, 1.0)

    ax.view_init(elev=30, azim=225)
    plt.savefig(OUTPUT_DIR / "fig2.png", dpi=150)
    plt.close()
    print("Generated: fig2.png")


def plot_figure_3():
    """
    Figure 3: The S-SE of the semantic-aware network with different models
    """
    M = np.arange(1, 11)

    # Approximate values from visual inspection
    proposed = np.array([0.24, 0.48, 0.72, 0.96, 1.18, 1.20, 1.20, 1.20, 1.20, 1.20])
    conv_k3 = np.zeros(10)
    conv_k5 = np.array([0.20, 0.39, 0.58, 0.77, 0.94, 0.96, 0.97, 0.97, 0.97, 0.97])
    conv_k7 = np.array([0.14, 0.28, 0.42, 0.56, 0.70, 0.70, 0.70, 0.70, 0.70, 0.70])
    conv_k9 = np.array([0.11, 0.22, 0.33, 0.44, 0.54, 0.54, 0.54, 0.54, 0.54, 0.54])

    plt.figure(figsize=(8, 6))
    plt.plot(M, proposed, "rd-", label="Proposed model")
    plt.plot(M, conv_k3, "ko-", label="Conventional model, $k_n = 3$", fillstyle="none")
    plt.plot(M, conv_k5, "k+-", label="Conventional model, $k_n = 5$")
    plt.plot(M, conv_k7, "k*-", label="Conventional model, $k_n = 7$")
    plt.plot(M, conv_k9, "kx-", label="Conventional model, $k_n = 9$")

    plt.xlabel("Number of channels, $M$")
    plt.ylabel("S-SE, $\Phi$ (suts/s/Hz) $\\times (I/L)$")
    plt.xticks(np.arange(1, 11))
    plt.yticks(np.arange(0, 1.5, 0.2))
    plt.grid(True)
    plt.legend(loc="lower right")
    plt.xlim(1, 10)
    plt.ylim(0, 1.3)

    plt.savefig(OUTPUT_DIR / "fig3.png", dpi=150)
    plt.close()
    print("Generated: fig3.png")


def plot_figure_4a():
    """
    Figure 4(a): The S-SE versus the number of channels
    """
    M = np.arange(1, 11)

    # Approximate values from visual inspection
    semantic = np.array([0.24, 0.48, 0.72, 0.96, 1.18, 1.20, 1.20, 1.20, 1.20, 1.20])
    ideal = np.array([0.21, 0.40, 0.55, 0.68, 0.79, 0.82, 0.84, 0.85, 0.86, 0.87])
    g5 = np.array([0.13, 0.26, 0.37, 0.47, 0.56, 0.58, 0.59, 0.60, 0.60, 0.60])
    g4 = np.array([0.12, 0.23, 0.31, 0.39, 0.46, 0.48, 0.49, 0.49, 0.50, 0.50])

    plt.figure(figsize=(8, 6))
    plt.plot(M, semantic, "rd-", label="Semantic")
    plt.plot(M, ideal, "ko-", label="Ideal", fillstyle="none")
    plt.plot(M, g5, "k*-.", label="5G")
    plt.plot(M, g4, "ks--", label="4G", fillstyle="none")

    plt.xlabel("Number of channels, $M$")
    plt.ylabel("S-SE, $\Phi$ (suts/s/Hz) $\\times (I/L)$")
    plt.xticks(np.arange(1, 11))
    plt.yticks(np.arange(0, 1.6, 0.2))
    plt.grid(True)
    plt.legend(loc="upper left")
    plt.xlim(1, 10)
    plt.ylim(0, 1.4)

    plt.savefig(OUTPUT_DIR / "fig4a.png", dpi=150)
    plt.close()
    print("Generated: fig4a.png")


def plot_figure_4b():
    """
    Figure 4(b): The S-SE versus the transmit power
    """
    p_n = np.arange(-40, 25, 5)

    # Approximate function to match shapes
    # Semantic: logistic curve
    semantic = 1.21 / (1 + np.exp(-0.25 * (p_n + 5)))

    # Ideal: mostly linear in higher dBm, slower in lower
    # Use Shannon approx log2(1 + SNR)
    snr_linear_ideal = 10 ** ((p_n - 10) / 10)  # arbitrary scaling to match
    ideal = 0.15 * np.log2(1 + 10 ** ((p_n + 15) / 10))
    ideal = np.clip(ideal, 0, 1.35)

    # 5G and 4G: similar to semantic but lower cap and shifted
    g5 = 0.7 / (1 + np.exp(-0.15 * (p_n - 5)))
    g4 = 0.68 / (1 + np.exp(-0.15 * (p_n - 8)))

    # Slight manual adjustments to match visual points
    ideal = np.interp(
        p_n,
        [-40, -30, -20, -10, 0, 10, 20, 23],
        [0.0, 0.01, 0.05, 0.15, 0.38, 0.72, 1.15, 1.32],
    )

    plt.figure(figsize=(8, 6))
    plt.plot(p_n, semantic, "rd-", label="Semantic")
    plt.plot(p_n, ideal, "ko-", label="Ideal", fillstyle="none")
    plt.plot(p_n, g5, "k*-.", label="5G")
    plt.plot(p_n, g4, "ks--", label="4G", fillstyle="none")

    plt.xlabel("Transmit power, $p_n$ (dBm)")
    plt.ylabel("S-SE, $\Phi$ (suts/s/Hz) $\\times (I/L)$")
    plt.xticks([-40, -30, -20, -10, 0, 10, 23])
    plt.yticks(np.arange(0, 1.6, 0.2))
    plt.grid(True)
    plt.legend(loc="upper left")
    plt.xlim(-40, 23)
    plt.ylim(0, 1.4)

    plt.savefig(OUTPUT_DIR / "fig4b.png", dpi=150)
    plt.close()
    print("Generated: fig4b.png")


def plot_figure_4c():
    """
    Figure 4(c): The S-SE versus the transforming factor
    """
    mu = np.arange(18, 42, 2)

    # Values extracted from plot visually
    semantic = np.ones_like(mu) * 1.18

    # Conventional decrease roughly as 1/mu
    # Ideal at mu=18 is ~1.78. 1.78 * 18 = 32.04
    ideal = 32.04 / mu

    # 5G at mu=18 is ~1.25. 1.25 * 18 = 22.5
    g5 = 22.5 / mu

    # 4G at mu=18 is ~1.02. 1.02 * 18 = 18.36
    g4 = 18.36 / mu

    plt.figure(figsize=(8, 6))
    plt.plot(mu, semantic, "rd-", label="Semantic")
    plt.plot(mu, ideal, "ko-", label="Ideal", fillstyle="none")
    plt.plot(mu, g5, "k*-.", label="5G")
    plt.plot(mu, g4, "ks--", label="4G", fillstyle="none")

    plt.xlabel("Transforming factor, $\mu$ (bits/word)")
    plt.ylabel("S-SE, $\Phi$ (suts/s/Hz) $\\times (I/L)$")
    plt.xticks(np.arange(18, 42, 2))
    plt.yticks(np.arange(0.4, 2.0, 0.2))
    plt.grid(True)
    plt.legend(loc="upper right")
    plt.xlim(18, 40)
    plt.ylim(0.4, 1.8)

    plt.savefig(OUTPUT_DIR / "fig4c.png", dpi=150)
    plt.close()
    print("Generated: fig4c.png")


def main():
    """Generate all reference plots."""
    print("Generating reference plots...")
    plot_figure_2()
    plot_figure_3()
    plot_figure_4a()
    plot_figure_4b()
    plot_figure_4c()
    print(f"\nAll plots saved to: {OUTPUT_DIR}")


if __name__ == "__main__":
    main()
