main_train.py

# import os
# import numpy as np
# import matplotlib.pyplot as plt
# from tqdm import tqdm
# import torch
# from smart_select_device import smart_select_device
# import csv
# import pandas as pd
# from plotManager import PlotManager


# from router_env.router_scheduler_env import RouterEnv
# from rl_agents.dqn_agent import DQNAgent


# plotter = PlotManager(bin_threshold=50, bin_size=10)
# plotter2 = PlotManager(bin_threshold=50, bin_size=5)
# agent_label = "dqn"

# device = smart_select_device(batch_size=128)

# variant = "tuned_v3.1.3"

# def evaluate_dqn_model(agent, suffix, scenario=1, episodes=100, max_timesteps=1000):
#     env = RouterEnv(scenario=scenario)
#     rewards, qos_rates, delays, switches = [], [], [], []

#     labels = ["Video", "Voice", "BestEffort"]

#     output_file = f"results/plots/dqn/csv/dqn_{variant}_{episodes}_episodes/episode_stats_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.csv"
#     header = [
#         "Episode", "Type", "Arrived", "Served", "QoS_Success", "QoS_Failure", "Dropped",
#         "Avg_Delay", "Success_Rate", "Failure_Rate", "Total_Reward", "QoS_Rate", "Switch_Count"
#     ]
#     plotter.save_csv_with_header(output_file, header)

#     for ep in range(episodes):
#         state = env.reset()
#         total_reward = 0
#         successful_qos = 0
#         total_served = 0
#         delay_sum = 0
#         delay_count = 0
#         switch_count = 0
#         prev_queue = env.current_queue

#         for _ in range(max_timesteps):
#             action = agent.act(state)
#             next_state, reward, done, info = env.step(action)
#             state = next_state
#             total_reward += reward

#             total_served += sum(info['total_served'])
#             successful_qos += sum(info['successful_qos'])
#             delay_sum += sum(info['average_delay'])
#             delay_count += len(info['average_delay'])

#             if action != 0 and env.current_queue != prev_queue:
#                 switch_count += 1
#             prev_queue = env.current_queue

#             if done:
#                 break

#         avg_qos = successful_qos / total_served if total_served > 0 else 0
#         avg_delay = delay_sum / delay_count if delay_count > 0 else 0
#          # Print per-queue summary
#         print(f"\n📋 Episode {ep + 1}")
#         print(f"{'Type':<12}{'Arrived':<10}{'Served':<10}{'QoS_Success':<12}{'Avg_Delay':<10}")
#         for i, label in enumerate(labels):
#             print(f"{label:<12}{info['total_arrived'][i]:<10}{info['total_served'][i]:<10}"
#                   f"{info['successful_qos'][i]:<12}{info['average_delay'][i]:.2f}")

#         # Save to CSV
#         with open(output_file, mode='a', newline='') as f:
#             writer = csv.writer(f)
#             for i in range(3):
#                 arrived = info['total_arrived'][i]
#                 served = info['total_served'][i]
#                 qos_success = info['successful_qos'][i]
#                 failed_qos = info['failed_qos'][i]  # Optional — add to env stats if not present
#                 dropped = arrived - served
#                 avg_delay = info['average_delay'][i]
#                 success_rate = qos_success / arrived if arrived > 0 else 0
#                 failure_rate = dropped / arrived if arrived > 0 else 0

#                 writer.writerow([
#                     ep + 1,
#                     labels[i],
#                     arrived,
#                     served,
#                     qos_success,
#                     failed_qos if 'failed_qos' in info else served - qos_success,  # fallback
#                     dropped,
#                     round(avg_delay, 2),
#                     round(success_rate, 3),
#                     round(failure_rate, 3),
#                     round(total_reward, 2),
#                     round(avg_qos, 3),
#                     switch_count
#                 ])
#         rewards.append(total_reward)
#         qos_rates.append(avg_qos)
#         delays.append(avg_delay)
#         switches.append(switch_count)

#     return rewards, qos_rates, delays, switches

# def train_dqn():
#     # variant = "tuned_v1"
#     for scenario in [1, 2]:

#         episodes = 100
#         max_timesteps = 1000
#         print(f"\n🔧 Training DQN on Scenario {scenario}...")
#         env = RouterEnv(scenario=scenario)
#         state_size = env.observation_space.shape[0]
#         action_size = env.action_space.n
#         agent = DQNAgent(state_size, action_size, device=device, variant = variant, episodes = episodes, scenario = scenario)
#         # # === PATCH: Load pretrained Scenario 1 weights if training Scenario 2 ===
#         transfer_from_s1 = False
#         # if scenario == 2:
#         #     pretrained_model_path = f"results/models/dqn/dqn_router_env_tuned_{variant}_scenario_1_{episodes}_episodes.pth"
#         #     if os.path.exists(pretrained_model_path):
#         #         print(f"🚀 Loading pretrained model from Scenario 1: {pretrained_model_path}")
#         #         agent.q_network.load_state_dict(torch.load(pretrained_model_path))
#         #         agent.target_network.load_state_dict(torch.load(pretrained_model_path))
#         #         transfer_from_s1 = True
#         #         for param_group in agent.optimizer.param_groups:
#         #             param_group['lr'] = 1e-5
#         #     else:
#         #         print("⚠️ Pretrained model not found. Training from scratch.")
#         if scenario == 2:
#             base_dir = os.path.dirname(os.path.abspath(__file__))
#             pretrained_model_path = os.path.join(base_dir, "results/models/dqn", f"dqn_router_env_{variant}_scenario_1_{episodes}_episodes.pth")
            
#             print(f"🔍 Checking model path: {pretrained_model_path}")
#             if os.path.exists(pretrained_model_path):
#                 print(f"🚀 Loading pretrained model from Scenario 1: {pretrained_model_path}")
#                 agent.load(pretrained_model_path)
#                 transfer_from_s1 = True
#                 for param_group in agent.optimizer.param_groups:
#                     param_group['lr'] = 1e-5
#             else:
#                 print("❌ File not found at path:", pretrained_model_path)

#         suffix = "_tansfer_from_s1" if transfer_from_s1 else ""
        
#         training_log_file = f"results/plots/dqn/csv/dqn_{variant}_{episodes}_episodes/training_dqn_log_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.csv"
#         header = ["Episode", "Total_Reward", "Epsilon", "Memory_Size"]
#         plotter.save_csv_with_header(training_log_file, header)
        

#         for ep in tqdm(range(episodes), desc=f"Scenario {scenario} Training"):
#             state = env.reset()
#             total_reward = 0
            
#             for t in range(max_timesteps):
#                 action = agent.act(state)
#                 next_state, reward, done, info = env.step(action)
#                 agent.remember(state, action, reward, next_state, done)
#                 agent.replay()
#                 agent.decay_epsilon()
#                 state = next_state
#                 total_reward += reward
#                 if done:
#                     break
#             # Log training stats
#             with open(training_log_file, 'a', newline='') as f:
#                 writer = csv.writer(f)
#                 writer.writerow([ep + 1, round(total_reward, 2), round(agent.epsilon, 3), len(agent.memory)])

#         agent.save(f"results/models/dqn/dqn_router_env_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.pth")
#         print(f"✅ DQN model saved to: results/models/dqn_router_env_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.pth")

#         print(f"\n📊 Evaluating DQN for {variant} version Scenario {scenario}... episode {episodes}:")
#         rewards, qos_rates, delays, switches = evaluate_dqn_model(agent, suffix, scenario=scenario)

#         # After evaluation
#         plot_dir = f"results/plots/dqn/dqn_{variant}_{episodes}_episodes{suffix}"
#         os.makedirs(plot_dir, exist_ok=True)

#         plotter2.plot_metric(rewards, "Reward", f"Total Reward (Scenario {scenario})", f"{plot_dir}/dqn_reward_scenario_{scenario}_{episodes}_episodes{suffix}.png", agent_label=agent_label)
#         plotter2.plot_metric(qos_rates, "QoS Rate", f"QoS Success Rate (Scenario {scenario})", f"{plot_dir}/dqn_qos_scenario_{scenario}_{episodes}_episodes{suffix}.png", agent_label=agent_label)
#         plotter2.plot_metric(delays, "Delay", f"Average Delay (Scenario {scenario})", f"{plot_dir}/dqn_delay_scenario_{scenario}_{episodes}_episodes{suffix}.png", agent_label=agent_label)
#         plotter2.plot_metric(switches, "Switches", f"Switch Count (Scenario {scenario})", f"{plot_dir}/dqn_switches_scenario_{scenario}_{episodes}_episodes{suffix}.png", agent_label=agent_label)
        
#         print(f"📈 Plots saved for DQN Scenario {variant} version {scenario} for {episodes}")
#         plotter.generate_csv_plots(episodes,scenario, variant,suffix, agent_label = agent_label)

#         # Load the newly saved evaluation CSV and plot advanced traffic metrics
#         output_file = f"results/plots/dqn/csv/dqn_{variant}_{episodes}_episodes/episode_stats_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.csv"
#         df = pd.read_csv(output_file)
#         plotter.plot_traffic_stats(df, base_path=f"{plot_dir}/traffic_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}")
#         plotter.plot_success_failure_rates(df, base_path=f"{plot_dir}/traffic_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}")
#         plotter.plot_combined_metrics(df, base_path=f"{plot_dir}/traffic_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}")

# if __name__ == "__main__":
#     train_dqn()

import os
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import torch
from smart_select_device import smart_select_device
import csv
import pandas as pd
from plotManager import PlotManager


from router_env.tabularStyleRouterEnv import TabularStyleRouterEnv

from rl_agents.dqn_agent import DQNAgent


plotter = PlotManager(bin_threshold=50, bin_size=35)
plotter2 = PlotManager(bin_threshold=50, bin_size=17)
agent_label = "dqn"

device = smart_select_device(batch_size=128)

variant = "tuned_v3.1.4"

def evaluate_dqn_model(agent, suffix, scenario=1, episodes=350, max_timesteps=1000):
    env = TabularStyleRouterEnv(scenario=scenario)
    rewards, qos_rates, delays, switches = [], [], [], []

    labels = ["Video", "Voice", "BestEffort"]

    output_file = f"results/plots/dqn/csv/dqn_{variant}_{episodes}_episodes/episode_stats_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.csv"
    header = [
        "Episode", "Type", "Arrived", "Served", "QoS_Success", "QoS_Failure", "Dropped",
        "Avg_Delay", "Success_Rate", "Failure_Rate", "Total_Reward", "QoS_Rate", "Switch_Count"
    ]
    plotter.save_csv_with_header(output_file, header)

    for ep in range(episodes):
        state = env.reset().astype(np.float32)
        total_reward = 0
        successful_qos = 0
        total_served = 0
        delay_sum = 0
        delay_count = 0
        switch_count = 0
        prev_queue = env.current_queue

        for _ in range(max_timesteps):
            action = agent.act(state)
            next_state, reward, done, info = env.step(action)
            next_state = next_state.astype(np.float32)
            state = next_state
            total_reward += reward

            total_served += sum(info['total_served'])
            successful_qos += sum(info['successful_qos'])
            delay_sum += sum(info['average_delay'])
            delay_count += len(info['average_delay'])

            if action != 0 and env.current_queue != prev_queue:
                switch_count += 1
            prev_queue = env.current_queue

            if done:
                break

        avg_qos = successful_qos / total_served if total_served > 0 else 0
        avg_delay = delay_sum / delay_count if delay_count > 0 else 0
         # Print per-queue summary
        print(f"\n📋 Episode {ep + 1}")
        print(f"{'Type':<12}{'Arrived':<10}{'Served':<10}{'QoS_Success':<12}{'Avg_Delay':<10}")
        for i, label in enumerate(labels):
            print(f"{label:<12}{info['total_arrived'][i]:<10}{info['total_served'][i]:<10}"
                  f"{info['successful_qos'][i]:<12}{info['average_delay'][i]:.2f}")

        # Save to CSV
        with open(output_file, mode='a', newline='') as f:
            writer = csv.writer(f)
            for i in range(3):
                arrived = info['total_arrived'][i]
                served = info['total_served'][i]
                qos_success = info['successful_qos'][i]
                failed_qos = info['failed_qos'][i]  # Optional — add to env stats if not present
                dropped = arrived - served
                avg_delay = info['average_delay'][i]
                success_rate = qos_success / arrived if arrived > 0 else 0
                failure_rate = dropped / arrived if arrived > 0 else 0

                writer.writerow([
                    ep + 1,
                    labels[i],
                    arrived,
                    served,
                    qos_success,
                    failed_qos if 'failed_qos' in info else served - qos_success,  # fallback
                    dropped,
                    round(avg_delay, 2),
                    round(success_rate, 3),
                    round(failure_rate, 3),
                    round(total_reward, 2),
                    round(avg_qos, 3),
                    switch_count
                ])
        rewards.append(total_reward)
        qos_rates.append(avg_qos)
        delays.append(avg_delay)
        switches.append(switch_count)

    return rewards, qos_rates, delays, switches

def train_dqn():
    # variant = "tuned_v1"
    for scenario in [1, 2]:

        episodes = 350
        max_timesteps = 1000
        print(f"\n🔧 Training DQN on Scenario {scenario}...")
        env = TabularStyleRouterEnv(scenario=scenario)
        state_size = env.observation_space.shape[0]
        action_size = env.action_space.n
        agent = DQNAgent(state_size, action_size, device=device, variant = variant, episodes = episodes, scenario = scenario)
        # # === PATCH: Load pretrained Scenario 1 weights if training Scenario 2 ===
        transfer_from_s1 = False
        # if scenario == 2:
        #     pretrained_model_path = f"results/models/dqn/dqn_router_env_tuned_{variant}_scenario_1_{episodes}_episodes.pth"
        #     if os.path.exists(pretrained_model_path):
        #         print(f"🚀 Loading pretrained model from Scenario 1: {pretrained_model_path}")
        #         agent.q_network.load_state_dict(torch.load(pretrained_model_path))
        #         agent.target_network.load_state_dict(torch.load(pretrained_model_path))
        #         transfer_from_s1 = True
        #         for param_group in agent.optimizer.param_groups:
        #             param_group['lr'] = 1e-5
        #     else:
        #         print("⚠️ Pretrained model not found. Training from scratch.")
        # if scenario == 2:
        #     base_dir = os.path.dirname(os.path.abspath(__file__))
        #     pretrained_model_path = os.path.join(base_dir, "results/models/dqn", f"dqn_router_env_{variant}_scenario_1_{episodes}_episodes.pth")
            
        #     print(f"🔍 Checking model path: {pretrained_model_path}")
        #     if os.path.exists(pretrained_model_path):
        #         print(f"🚀 Loading pretrained model from Scenario 1: {pretrained_model_path}")
        #         agent.load(pretrained_model_path)
        #         transfer_from_s1 = True
        #         for param_group in agent.optimizer.param_groups:
        #             param_group['lr'] = 1e-5
        #     else:
        #         print("❌ File not found at path:", pretrained_model_path)

        suffix = "_tansfer_from_s1" if transfer_from_s1 else ""
        
        training_log_file = f"results/plots/dqn/csv/dqn_{variant}_{episodes}_episodes/training_dqn_log_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.csv"
        header = ["Episode", "Total_Reward", "Epsilon", "Memory_Size"]
        plotter.save_csv_with_header(training_log_file, header)
        

        for ep in tqdm(range(episodes), desc=f"Scenario {scenario} Training"):
            state = env.reset().astype(np.float32)
            total_reward = 0
            
            for t in range(max_timesteps):
                action = agent.act(state)
                next_state, reward, done, info = env.step(action)
                agent.remember(state, action, reward, next_state, done)
                agent.replay()
                agent.decay_epsilon()
                state = next_state
                total_reward += reward
                if done:
                    break
            # Log training stats
            with open(training_log_file, 'a', newline='') as f:
                writer = csv.writer(f)
                writer.writerow([ep + 1, round(total_reward, 2), round(agent.epsilon, 3), len(agent.memory)])

        agent.save(f"results/models/dqn/dqn_router_env_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.pth")
        print(f"✅ DQN model saved to: results/models/dqn_router_env_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.pth")

        print(f"\n📊 Evaluating DQN for {variant} version Scenario {scenario}... episode {episodes}:")
        rewards, qos_rates, delays, switches = evaluate_dqn_model(agent, suffix, scenario=scenario)

        # After evaluation
        plot_dir = f"results/plots/dqn/dqn_{variant}_{episodes}_episodes{suffix}"
        os.makedirs(plot_dir, exist_ok=True)

        plotter2.plot_metric(rewards, "Reward", f"Total Reward (Scenario {scenario})", f"{plot_dir}/dqn_reward_scenario_{scenario}_{episodes}_episodes{suffix}.png", agent_label=agent_label)
        plotter2.plot_metric(qos_rates, "QoS Rate", f"QoS Success Rate (Scenario {scenario})", f"{plot_dir}/dqn_qos_scenario_{scenario}_{episodes}_episodes{suffix}.png", agent_label=agent_label)
        plotter2.plot_metric(delays, "Delay", f"Average Delay (Scenario {scenario})", f"{plot_dir}/dqn_delay_scenario_{scenario}_{episodes}_episodes{suffix}.png", agent_label=agent_label)
        plotter2.plot_metric(switches, "Switches", f"Switch Count (Scenario {scenario})", f"{plot_dir}/dqn_switches_scenario_{scenario}_{episodes}_episodes{suffix}.png", agent_label=agent_label)
        
        print(f"📈 Plots saved for DQN Scenario {variant} version {scenario} for {episodes}")
        plotter.generate_csv_plots(episodes,scenario, variant,suffix, agent_label = agent_label)

        # Load the newly saved evaluation CSV and plot advanced traffic metrics
        output_file = f"results/plots/dqn/csv/dqn_{variant}_{episodes}_episodes/episode_stats_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}.csv"
        df = pd.read_csv(output_file)
        plotter.plot_traffic_stats(df, base_path=f"{plot_dir}/traffic_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}")
        plotter.plot_success_failure_rates(df, base_path=f"{plot_dir}/traffic_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}")
        plotter.plot_combined_metrics(df, base_path=f"{plot_dir}/traffic_{variant}_scenario_{scenario}_{episodes}_episodes{suffix}")

if __name__ == "__main__":
    train_dqn()