"""
src/evaluate.py

Evaluation script that generates Figures 3, 4a, 4b, 4c.
"""

import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import os
import random

from src.models.environment import EnvironmentConfig, ChannelSimulator
from src.models.semantic_model import SemanticSurrogate
from src.models.allocator import ResourceAllocator
from src.models.baselines import BaselineModels


def set_seed(seed=42):
    np.random.seed(seed)
    random.seed(seed)


def generate_figure3(reports_dir="reports/figures"):
    """
    Figure 3: S-SE of the semantic-aware network with different models
    Varying Transmit Power vs S-SE for Semantic (Proposed) vs Fixed k_n (2, 4, 8)
    """
    print("Generating Figure 3...")
    powers_dbm = np.arange(-30, 20, 5)

    num_users = 10
    num_channels = 20
    env_config = EnvironmentConfig(num_users=num_users, num_channels=num_channels)
    env = ChannelSimulator(env_config)
    surrogate = SemanticSurrogate(L=12)
    allocator = ResourceAllocator(surrogate, env_config, K_max=8)

    num_trials = 50

    prop_sse_avg = []
    fixed_k2_avg = []
    fixed_k4_avg = []
    fixed_k8_avg = []

    for pt in tqdm(powers_dbm, desc="Fig 3: Transmit Power"):
        prop_sse = 0
        k2_sse = 0
        k4_sse = 0
        k8_sse = 0

        for _ in range(num_trials):
            snr_db, snr_linear = env.generate_channels(pt)

            _, _, total_sse = allocator.optimize_semantic_aware(snr_linear)
            prop_sse += total_sse / num_users

            k2_sse += allocator.evaluate_fixed_k(snr_linear, 2) / num_users
            k4_sse += allocator.evaluate_fixed_k(snr_linear, 4) / num_users
            k8_sse += allocator.evaluate_fixed_k(snr_linear, 8) / num_users

        prop_sse_avg.append(prop_sse / num_trials)
        fixed_k2_avg.append(k2_sse / num_trials)
        fixed_k4_avg.append(k4_sse / num_trials)
        fixed_k8_avg.append(k8_sse / num_trials)

    plt.figure(figsize=(8, 6))
    plt.plot(powers_dbm, prop_sse_avg, "b-o", label="Proposed algorithm")
    plt.plot(powers_dbm, fixed_k2_avg, "m--s", label="DeepSC-network ($k_n=2$)")
    plt.plot(powers_dbm, fixed_k4_avg, "r--^", label="DeepSC-network ($k_n=4$)")
    plt.plot(powers_dbm, fixed_k8_avg, "k--x", label="DeepSC-network ($k_n=8$)")

    plt.xlabel("Transmit power (dBm)")
    plt.ylabel("S-SE (words/s/Hz)")
    plt.grid(True)
    plt.legend()
    plt.savefig(os.path.join(reports_dir, "fig3.png"), dpi=300, bbox_inches="tight")
    plt.close()


def generate_figure4a(reports_dir="reports/figures"):
    """
    Figure 4(a): S-SE versus the number of channels
    """
    print("Generating Figure 4(a)...")
    channels_list = np.arange(2, 21, 2)
    pt_dbm = 10.0
    num_users = 10

    num_trials = 50

    sem_sse = []
    ideal_sse = []
    g5_sse = []
    g4_sse = []

    surrogate = SemanticSurrogate(L=12)
    baselines = BaselineModels(mu=19.0, L=12)

    for M in tqdm(channels_list, desc="Fig 4a: Channels"):
        env_config = EnvironmentConfig(num_users=num_users, num_channels=M)
        env = ChannelSimulator(env_config)
        allocator = ResourceAllocator(surrogate, env_config, K_max=8)

        s_val = 0
        i_val = 0
        g5_val = 0
        g4_val = 0

        for _ in range(num_trials):
            snr_db, snr_linear = env.generate_channels(pt_dbm)

            _, _, total_sse = allocator.optimize_semantic_aware(snr_linear)
            s_val += total_sse

            # For baselines, use greedy matching or just sum max per user if M >= N
            from scipy.optimize import linear_sum_assignment

            ideal_matrix = baselines.calculate_baseline_sse(snr_linear, "ideal")
            ri, ci = linear_sum_assignment(-ideal_matrix)
            i_val += np.sum(ideal_matrix[ri, ci])

            g5_matrix = baselines.calculate_baseline_sse(snr_linear, "5G")
            ri, ci = linear_sum_assignment(-g5_matrix)
            g5_val += np.sum(g5_matrix[ri, ci])

            g4_matrix = baselines.calculate_baseline_sse(snr_linear, "4G")
            ri, ci = linear_sum_assignment(-g4_matrix)
            g4_val += np.sum(g4_matrix[ri, ci])

        sem_sse.append(s_val / num_trials)
        ideal_sse.append(i_val / num_trials)
        g5_sse.append(g5_val / num_trials)
        g4_sse.append(g4_val / num_trials)

    plt.figure(figsize=(8, 6))
    plt.plot(channels_list, sem_sse, "b-o", label="Semantic-aware network")
    plt.plot(channels_list, ideal_sse, "r--s", label="Ideal Shannon limit")
    plt.plot(channels_list, g5_sse, "k-.^", label="5G communications")
    plt.plot(channels_list, g4_sse, "m:x", label="4G communications")

    plt.xlabel("The number of channels")
    plt.ylabel("Sum S-SE (words/s/Hz)")
    plt.grid(True)
    plt.legend()
    plt.savefig(os.path.join(reports_dir, "fig4a.png"), dpi=300, bbox_inches="tight")
    plt.close()


def generate_figure4b(reports_dir="reports/figures"):
    """
    Figure 4(b): S-SE versus the transmit power
    """
    print("Generating Figure 4(b)...")
    powers_dbm = np.arange(-40, 25, 5)

    num_users = 10
    num_channels = 20

    env_config = EnvironmentConfig(num_users=num_users, num_channels=num_channels)
    env = ChannelSimulator(env_config)
    surrogate = SemanticSurrogate(L=12)
    allocator = ResourceAllocator(surrogate, env_config, K_max=8)
    baselines = BaselineModels(mu=19.0, L=12)

    num_trials = 50

    sem_sse = []
    ideal_sse = []
    g5_sse = []
    g4_sse = []

    from scipy.optimize import linear_sum_assignment

    for pt in tqdm(powers_dbm, desc="Fig 4b: Power"):
        s_val, i_val, g5_val, g4_val = 0, 0, 0, 0

        for _ in range(num_trials):
            snr_db, snr_linear = env.generate_channels(pt)

            _, _, total_sse = allocator.optimize_semantic_aware(snr_linear)
            s_val += total_sse / num_users

            ideal_matrix = baselines.calculate_baseline_sse(snr_linear, "ideal")
            ri, ci = linear_sum_assignment(-ideal_matrix)
            i_val += np.sum(ideal_matrix[ri, ci]) / num_users

            g5_matrix = baselines.calculate_baseline_sse(snr_linear, "5G")
            ri, ci = linear_sum_assignment(-g5_matrix)
            g5_val += np.sum(g5_matrix[ri, ci]) / num_users

            g4_matrix = baselines.calculate_baseline_sse(snr_linear, "4G")
            ri, ci = linear_sum_assignment(-g4_matrix)
            g4_val += np.sum(g4_matrix[ri, ci]) / num_users

        sem_sse.append(s_val / num_trials)
        ideal_sse.append(i_val / num_trials)
        g5_sse.append(g5_val / num_trials)
        g4_sse.append(g4_val / num_trials)

    plt.figure(figsize=(8, 6))
    plt.plot(powers_dbm, sem_sse, "b-o", label="Semantic-aware network")
    plt.plot(powers_dbm, ideal_sse, "r--s", label="Ideal Shannon limit")
    plt.plot(powers_dbm, g5_sse, "k-.^", label="5G communications")
    plt.plot(powers_dbm, g4_sse, "m:x", label="4G communications")

    plt.xlabel("Transmit power (dBm)")
    plt.ylabel("S-SE (words/s/Hz)")
    plt.grid(True)
    plt.legend()
    plt.savefig(os.path.join(reports_dir, "fig4b.png"), dpi=300, bbox_inches="tight")
    plt.close()


def generate_figure4c(reports_dir="reports/figures"):
    """
    Figure 4(c): S-SE versus the transforming factor
    """
    print("Generating Figure 4(c)...")
    mu_list = np.arange(18, 41, 2)
    pt_dbm = 10.0
    num_users = 10
    num_channels = 20

    env_config = EnvironmentConfig(num_users=num_users, num_channels=num_channels)
    env = ChannelSimulator(env_config)
    surrogate = SemanticSurrogate(L=12)
    allocator = ResourceAllocator(surrogate, env_config, K_max=8)

    num_trials = 50

    sem_sse_avg = 0  # Semantic-aware is independent of mu in terms of bits, but evaluated in words
    # Actually Semantic is purely dependent on DeepSC params. So it stays flat across mu changes!

    sem_sse_vals = []

    for _ in range(num_trials):
        snr_db, snr_linear = env.generate_channels(pt_dbm)
        _, _, total_sse = allocator.optimize_semantic_aware(snr_linear)
        sem_sse_avg += total_sse / num_users

    sem_sse_avg /= num_trials

    sem_sse = [sem_sse_avg] * len(mu_list)

    ideal_sse = []
    g5_sse = []
    g4_sse = []

    from scipy.optimize import linear_sum_assignment

    for mu in tqdm(mu_list, desc="Fig 4c: Transforming Factor"):
        baselines = BaselineModels(mu=mu, L=12)
        i_val, g5_val, g4_val = 0, 0, 0

        for _ in range(num_trials):
            snr_db, snr_linear = env.generate_channels(pt_dbm)

            ideal_matrix = baselines.calculate_baseline_sse(snr_linear, "ideal")
            ri, ci = linear_sum_assignment(-ideal_matrix)
            i_val += np.sum(ideal_matrix[ri, ci]) / num_users

            g5_matrix = baselines.calculate_baseline_sse(snr_linear, "5G")
            ri, ci = linear_sum_assignment(-g5_matrix)
            g5_val += np.sum(g5_matrix[ri, ci]) / num_users

            g4_matrix = baselines.calculate_baseline_sse(snr_linear, "4G")
            ri, ci = linear_sum_assignment(-g4_matrix)
            g4_val += np.sum(g4_matrix[ri, ci]) / num_users

        ideal_sse.append(i_val / num_trials)
        g5_sse.append(g5_val / num_trials)
        g4_sse.append(g4_val / num_trials)

    plt.figure(figsize=(8, 6))
    plt.plot(mu_list, sem_sse, "b-o", label="Semantic-aware network")
    plt.plot(mu_list, ideal_sse, "r--s", label="Ideal Shannon limit")
    plt.plot(mu_list, g5_sse, "k-.^", label="5G communications")
    plt.plot(mu_list, g4_sse, "m:x", label="4G communications")

    plt.xlabel("Transforming factor (bits/word)")
    plt.ylabel("S-SE (words/s/Hz)")
    plt.grid(True)
    plt.legend()
    plt.savefig(os.path.join(reports_dir, "fig4c.png"), dpi=300, bbox_inches="tight")
    plt.close()


if __name__ == "__main__":
    set_seed(42)
    os.makedirs("reports/figures", exist_ok=True)
    generate_figure3()
    generate_figure4a()
    generate_figure4b()
    generate_figure4c()
    print("Evaluation completed! Figures saved in reports/figures/")
