# LINE 1-17: ARE USED TO SILENCE LOGS & ALIGN NAMESPACES
import os
import warnings
import tensorflow as tf

os.environ["TF_USE_LEGACY_KERAS"] = "1"       
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"      
os.environ["TF_FUNCTION_NUMERIC_CHECKS"] = "0"

tf.keras.backend.clear_session()
tf.config.threading.set_intra_op_parallelism_threads(0) 
tf.config.threading.set_inter_op_parallelism_threads(0) 
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

warnings.filterwarnings("ignore", category=UserWarning, module="gym")
warnings.filterwarnings("ignore", category=RuntimeWarning, module="tf_agents.specs.array_spec")
warnings.filterwarnings("ignore", category=DeprecationWarning)

from tf_agents.agents.ppo import ppo_agent
from tf_agents.networks import network
from tf_agents.networks import categorical_projection_network
from tf_agents.environments import tf_py_environment

# Custom Modular Imports
from config import Config
from env import CyberPettingZooEnv
from utilities import configure_gpu, PettingZooToTFAgentsWrapper
from trainer import MultiAgentTrainer

# BULLETPROOF CUSTOM GNN POLICY NETWORK FOR TF-AGENTS
class CyberGNNNetwork(network.Network):
    """
    A custom TF-Agents Network that binds GNN feature extraction with 
    native CategoricalProjectionNetworks, explicitly forwarding outer_rank.
    """
    def __init__(self, input_tensor_spec, output_spec, fc_layer_params=(128, 64), is_value_net=False, name="CyberGNNNetwork"):
        super().__init__(input_tensor_spec=input_tensor_spec, state_spec=(), name=name)
        self.is_value_net = is_value_net
        
        # 1. Instantiate the Graph Convolution Weights
        self.w1 = tf.Variable(tf.keras.initializers.GlorotUniform()(shape=(2, 32)), trainable=True, name="gcn_w1")
        self.w2 = tf.Variable(tf.keras.initializers.GlorotUniform()(shape=(32, 32)), trainable=True, name="gcn_w2")
        
        # 2. Instantiate standard Dense MLPs for decision making
        self.dense_layers = []
        for units in fc_layer_params:
            self.dense_layers.append(tf.keras.layers.Dense(units, activation=tf.nn.tanh))
            
        # 3. Instantiate native TF-Agents projection heads
        if self.is_value_net:
            self.projection_head = tf.keras.layers.Dense(1, kernel_initializer=tf.keras.initializers.Orthogonal(1.0))
        else:
            self.projection_head = categorical_projection_network.CategoricalProjectionNetwork(output_spec)

    def call(self, observations, step_type=(), network_state=(), training=False):
        A_norm = observations['adjacency_matrix'] 
        X = observations['node_features']       

        outer_rank = X.shape.rank - 2
        
        # GNN Message-Passing Block
        h1 = tf.nn.tanh(tf.matmul(tf.matmul(A_norm, X), self.w1))
        h2 = tf.nn.tanh(tf.matmul(tf.matmul(A_norm, h1), self.w2))
        
        # Dynamically calculate batch structure to prevent rank mismatch crashes 
        batch_shape = tf.shape(h2)[:-2] 
        flat_features = tf.reshape(h2, tf.concat([batch_shape, [-1]], axis=0))
        
        # Standard Policy Dense Layer Block
        x = flat_features
        for layer in self.dense_layers:
            x = layer(x)
            
        # Output Generation Pass
        if self.is_value_net:
            value_predictions = self.projection_head(x)
            value_predictions = tf.squeeze(value_predictions, axis=-1)
            return value_predictions, network_state
        else:
            action_distributions, _ = self.projection_head(x, outer_rank=outer_rank, training=training)
            return action_distributions, network_state

# AGENT BUILD ORCHESTRATION WITH EXPLICIT VARIABLE REGISTRATION
def create_ppo_agent(env_time_step_spec, env_action_spec, agent_name, lr_schedule):
    """Creates an independent PPO Agent with guaranteed trainable variable tracking."""
    agent_obs_spec = env_time_step_spec.observation[agent_name]
    agent_action_spec = env_action_spec[agent_name]

    actor_net = CyberGNNNetwork(
        input_tensor_spec=agent_obs_spec,
        output_spec=agent_action_spec,
        fc_layer_params=Config.ACTOR_LAYERS,
        is_value_net=False,
        name=f"actor_gnn_{agent_name}"
    )
    
    value_net = CyberGNNNetwork(
        input_tensor_spec=agent_obs_spec,
        output_spec=None,
        fc_layer_params=Config.CRITIC_LAYERS,
        is_value_net=True,
        name=f"value_gnn_{agent_name}"
    )

    optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)

    agent = ppo_agent.PPOAgent(
        time_step_spec=env_time_step_spec._replace(
            observation=agent_obs_spec,
            reward=env_time_step_spec.reward[agent_name]
        ),
        action_spec=agent_action_spec,
        optimizer=optimizer,
        actor_net=actor_net,
        value_net=value_net,
        num_epochs=10,                                              
        discount_factor=Config.GAMMA,
        use_gae=True,                                                
        use_td_lambda_return=True,
        normalize_observations=False, 
        normalize_rewards=True,  
        value_pred_loss_coef=0.5 
    )
                
    agent.initialize()
    print(f"🔒 Verified Trainable Variables for {agent_name} Actor: {len(actor_net.trainable_variables)}")
    return agent

# MAIN RUNTIME EXECUTION LOOP
def main():
    configure_gpu()
    print(f"🚀 Initializing Custom GNN MARL Cyber Range Framework in **{Config.MODE}** mode.")
    
    raw_pz_env = CyberPettingZooEnv()
    py_env = PettingZooToTFAgentsWrapper(raw_pz_env)
    tf_env = tf_py_environment.TFPyEnvironment(py_env)
    
    time_step_spec = tf_env.time_step_spec()
    action_spec = tf_env.action_spec()
    
    print("🧠 Constructing Custom Graph Policy Networks for Red Agent...")
    red_agent = create_ppo_agent(time_step_spec, action_spec, 'red', Config.get_red_lr_schedule())
    
    print("🧠 Constructing Custom Graph Policy Networks for Blue Agent...")
    blue_agent = create_ppo_agent(time_step_spec, action_spec, 'blue', Config.get_blue_lr_schedule())
    
    trainer = MultiAgentTrainer(config=Config, tf_env=tf_env, agent_1=red_agent, agent_2=blue_agent)
    
    print("✨ Graph Neural Network pipeline successfully compiled. Launching training loop...")
    for epoch in range(1, Config.TOTAL_EPISODES + 1):
        trainer.train_epoch(epoch_idx=epoch)

if __name__ == '__main__':
    main()