import sys
from enum import Enum

class NoiseType(Enum):
    NONE = "none"
    NORMAL = "normal"
    UNIFORM = "uniform"


class ModelConfig:
    num_agents: int
    embedding_dim: int = 16
    input_dim: int = 1
    output_dim: int = 1



class TrainConfig:
    """Configuration for the training process."""
    learning_rate: float = 1e-3
    epochs: int = 100
    batch_size: int = 32
    verbose: bool = False  # Set to True to see epoch loss during training
    log: bool = True
    log_epoch_interval: int = 10
    noise_type: NoiseType = NoiseType.NONE
    noise_level: float = 0.01
    
    # --- New parameters for this script ---
    data_directory:str 
    # Threshold for converting attention scores to a binary graph
    f1_threshold: float = -0.5