interface gui agent enviroment interaction 
import random
import math
from typing import List, Tuple, Dict, Any
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import gymnasium as gym
from gymnasium import spaces
import plotly.graph_objects as go
from collections import deque

# Set random seeds for reproducibility
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)


class AttackDefenseEnv(gym.Env):
    """
    Custom Environment for Attack-Defense interactions.
    """

    metadata = {"render.modes": ["human"]}

    def __init__(self) -> None:
        super(AttackDefenseEnv, self).__init__()

        # Define the number of attack and defense types
        self.num_attacks: int = 10
        self.num_defenses: int = 10

        # Define attack and defense emojis with strengths
        self.attack_emojis: List[str] = [
            "🔥", "💧", "⚡", "🌪️", "🌋",
            "🦠", "☢️", "🔫", "🪖", "💣"
        ]
        self.defense_emojis: List[str] = [
            "🛡️", "🧱", "🏰", "🛠️", "🪦",
            "🥽", "🪖", "🚧", "🪤", "🧨"
        ]

        # Randomly assign strengths to attacks and defenses
        self.attack_strengths: List[float] = [
            random.uniform(0.5, 1.5) for _ in range(self.num_attacks)
        ]
        self.defense_strengths: List[float] = [
            random.uniform(0.5, 1.5) for _ in range(self.num_defenses)
        ]

        # Action space: Choose a defense (with repetition allowed)
        self.action_space = spaces.Discrete(self.num_defenses)

        # Observation space: Current attack
        self.observation_space = spaces.Discrete(self.num_attacks)

        # Initialize state
        self.current_attack: int = 0
        self.done: bool = False

    def reset(self) -> int:
        """
        Resets the environment to an initial state and returns the initial observation.
        """
        self.current_attack = random.randint(0, self.num_attacks - 1)
        self.done = False
        return self.current_attack

    def step(self, action: int) -> Tuple[int, float, bool, Dict[str, Any]]:
        """
        Executes one time step within the environment.

        Args:
            action (int): The index of the defense chosen.

        Returns:
            observation (int): The next attack.
            reward (float): The reward for the action taken.
            done (bool): Whether the episode has ended.
            info (dict): Additional information.
        """
        if not self.action_space.contains(action):
            raise ValueError(f"Invalid action: {action}")

        # Calculate the effectiveness of the defense against the attack
        attack_strength = self.attack_strengths[self.current_attack]
        defense_strength = self.defense_strengths[action]

        # Simple reward function: Higher defense strength against attack strength yields higher reward
        reward = defense_strength - attack_strength

        # For simplicity, let's say each episode consists of 10 attacks
        self.done = True  # Single step per episode

        # Sample next attack
        next_attack = random.randint(0, self.num_attacks - 1)

        self.current_attack = next_attack

        info = {
            "attack": self.attack_emojis[self.current_attack],
            "defense": self.defense_emojis[action],
            "attack_strength": attack_strength,
            "defense_strength": defense_strength,
            "next_attack": self.attack_emojis[next_attack]
        }

        return next_attack, reward, self.done, info

    def render(self, mode: str = "human") -> None:
        """
        Renders the environment.

        Args:
            mode (str): The mode to render with.
        """
        if mode == "human":
            print(f"Current Attack: {self.attack_emojis[self.current_attack]} "
                  f"(Strength: {self.attack_strengths[self.current_attack]})")


class DQNAgent:
    """
    Deep Q-Network Agent
    """

    def __init__(
        self,
        state_size: int,
        action_size: int,
        replay_buffer_size: int = 10000,
        batch_size: int = 64,
        gamma: float = 0.99,
        lr: float = 1e-3,
        epsilon_start: float = 1.0,
        epsilon_end: float = 0.01,
        epsilon_decay: int = 500
    ) -> None:
        self.state_size: int = state_size
        self.action_size: int = action_size
        self.memory: deque = deque(maxlen=replay_buffer_size)
        self.batch_size: int = batch_size
        self.gamma: float = gamma

        self.epsilon: float = epsilon_start
        self.epsilon_end: float = epsilon_end
        self.epsilon_decay: int = epsilon_decay
        self.steps_done: int = 0

        # Define the policy network
        self.policy_net: nn.Module = self.build_network()
        self.target_net: nn.Module = self.build_network()
        self.update_target_network()

        self.optimizer: optim.Optimizer = optim.Adam(self.policy_net.parameters(), lr=lr)
        self.loss_fn: nn.Module = nn.MSELoss()

    def build_network(self) -> nn.Module:
        """
        Builds the neural network model.

        Returns:
            nn.Module: The neural network model.
        """
        model = nn.Sequential(
            nn.Embedding(self.state_size, 64),
            nn.Flatten(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, self.action_size)
        )
        return model

    def update_target_network(self) -> None:
        """
        Updates the target network to match the policy network.
        """
        self.target_net.load_state_dict(self.policy_net.state_dict())

    def remember(self, state: int, action: int, reward: float, next_state: int, done: bool) -> None:
        """
        Stores experience in the replay buffer.

        Args:
            state (int): Current state.
            action (int): Action taken.
            reward (float): Reward received.
            next_state (int): Next state.
            done (bool): Whether the episode has ended.
        """
        self.memory.append((state, action, reward, next_state, done))

    def act(self, state: int) -> int:
        """
        Selects an action using an epsilon-greedy policy.

        Args:
            state (int): The current state.

        Returns:
            int: The action to take.
        """
        self.steps_done += 1
        self.epsilon = self.epsilon_end + (1.0 - self.epsilon_end) * math.exp(-1.0 * self.steps_done / self.epsilon_decay)

        if random.random() < self.epsilon:
            return random.randint(0, self.action_size - 1)
        else:
            state_tensor = torch.LongTensor([state])
            with torch.no_grad():
                q_values = self.policy_net(state_tensor)
                return q_values.argmax().item()

    def replay(self) -> float:
        """
        Samples a batch from memory and performs a training step.

        Returns:
            float: The loss value.
        """
        if len(self.memory) < self.batch_size:
            return 0.0

        batch = random.sample(self.memory, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states_tensor = torch.LongTensor(states)
        actions_tensor = torch.LongTensor(actions)
        rewards_tensor = torch.FloatTensor(rewards)
        next_states_tensor = torch.LongTensor(next_states)
        dones_tensor = torch.FloatTensor(dones)

        # Compute Q(s, a)
        q_values = self.policy_net(states_tensor).gather(1, actions_tensor.unsqueeze(1)).squeeze()

        # Compute target Q values
        with torch.no_grad():
            next_q_values = self.target_net(next_states_tensor).max(1)[0]
            target_q_values = rewards_tensor + self.gamma * next_q_values * (1 - dones_tensor)

        # Compute loss
        loss = self.loss_fn(q_values, target_q_values)

        # Optimize the model
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return loss.item()

    def save(self, path: str) -> None:
        """
        Saves the policy network.

        Args:
            path (str): The file path to save the model.
        """
        torch.save(self.policy_net.state_dict(), path)

    def load(self, path: str) -> None:
        """
        Loads the policy network.

        Args:
            path (str): The file path to load the model from.
        """
        self.policy_net.load_state_dict(torch.load(path))
        self.update_target_network()


def plot_metrics(
    rewards: List[float],
    losses: List[float],
    episode: int
) -> None:
    """
    Plots the training metrics using Plotly.

    Args:
        rewards (List[float]): List of rewards per episode.
        losses (List[float]): List of losses per training step.
        episode (int): Current episode number.
    """
    fig = go.Figure()

    # Reward Plot
    fig.add_trace(go.Scatter(
        y=rewards,
        mode='lines+markers',
        name='Reward'
    ))

    # Loss Plot (scaled to fit in the same plot)
    if losses:
        scaled_losses = [loss / max(losses) for loss in losses]
        fig.add_trace(go.Scatter(
            y=scaled_losses,
            mode='lines+markers',
            name='Scaled Loss',
            yaxis='y2'
        ))

    # Layout
    fig.update_layout(
        title=f"Training Metrics - Episode {episode}",
        xaxis_title="Episode",
        yaxis_title="Reward",
        yaxis2=dict(
            title="Scaled Loss",
            overlaying='y',
            side='right'
        ),
        legend=dict(x=0.01, y=0.99),
        template="plotly_dark"
    )

    fig.show()


def train_agent(
    env: AttackDefenseEnv,
    agent: DQNAgent,
    episodes: int = 1000,
    target_update: int = 10
) -> None:
    """
    Trains the DQN agent in the given environment.

    Args:
        env (AttackDefenseEnv): The environment.
        agent (DQNAgent): The RL agent.
        episodes (int): Number of training episodes.
        target_update (int): How often to update the target network.
    """
    rewards_history: List[float] = []
    losses_history: List[float] = []

    for episode in range(1, episodes + 1):
        state = env.reset()
        total_reward = 0.0
        done = False

        while not done:
            action = agent.act(state)
            next_state, reward, done, info = env.step(action)
            agent.remember(state, action, reward, next_state, done)
            loss = agent.replay()
            if loss != 0.0:
                losses_history.append(loss)
            state = next_state
            total_reward += reward

        rewards_history.append(total_reward)

        # Update the target network
        if episode % target_update == 0:
            agent.update_target_network()

        # Visualization
        if episode % 100 == 0 or episode == 1:
            print(f"Episode {episode}, Total Reward: {total_reward:.2f}, Epsilon: {agent.epsilon:.2f}")
            plot_metrics(rewards_history, losses_history, episode)

    # Plot final metrics
    plot_metrics(rewards_history, losses_history, episodes)


def main() -> None:
    """
    The main function to set up the environment and agent, and start training.
    """
    try:
        env = AttackDefenseEnv()
        agent = DQNAgent(state_size=env.num_attacks, action_size=env.num_defenses)
        train_agent(env, agent, episodes=1000, target_update=50)
    except Exception as e:
        print(f"An error occurred during training: {e}")


if __name__ == "__main__":
    main()