import os
import tempfile
import pytest
import torch.nn as nn
from aiia import AIIAConfig

def test_aiia_config_initialization():
    config = AIIAConfig()
    assert config.model_type == "AIIA"
    assert config.kernel_size == 3
    assert config.activation_function == "GELU"
    assert config.hidden_size == 512
    assert config.num_hidden_layers == 12
    assert config.num_channels == 3
    assert config.learning_rate == 5e-5

def test_aiia_config_custom_initialization():
    config = AIIAConfig(
        model_type="CustomModel",
        kernel_size=5,
        activation_function="ReLU",
        hidden_size=1024,
        num_hidden_layers=8,
        num_channels=1,
        learning_rate=1e-4
    )
    assert config.model_type == "CustomModel"
    assert config.kernel_size == 5
    assert config.activation_function == "ReLU"
    assert config.hidden_size == 1024
    assert config.num_hidden_layers == 8
    assert config.num_channels == 1
    assert config.learning_rate == 1e-4

def test_aiia_config_invalid_activation_function():
    with pytest.raises(ValueError):
        AIIAConfig(activation_function="InvalidFunction")

def test_aiia_config_to_dict():
    config = AIIAConfig()
    config_dict = config.to_dict()
    assert isinstance(config_dict, dict)
    assert config_dict["model_type"] == "AIIA"
    assert config_dict["kernel_size"] == 3

def test_aiia_config_save_pretrained_and_from_pretrained():
    with tempfile.TemporaryDirectory() as tmpdir:
        config = AIIAConfig(model_type="TempModel")
        save_pretrained_path = os.path.join(tmpdir, "config")
        config.save_pretrained(save_pretrained_path)

        loaded_config = AIIAConfig.from_pretrained(save_pretrained_path)
        assert loaded_config.model_type == "TempModel"
        assert loaded_config.kernel_size == 3
        assert loaded_config.activation_function == "GELU"

def test_aiia_config_save_pretrained_and_load_with_custom_attributes():
    with tempfile.TemporaryDirectory() as tmpdir:
        config = AIIAConfig(model_type="TempModel", custom_attr="value")
        save_pretrained_path = os.path.join(tmpdir, "config")
        config.save_pretrained(save_pretrained_path)

        loaded_config = AIIAConfig.from_pretrained(save_pretrained_path)
        assert loaded_config.model_type == "TempModel"
        assert loaded_config.custom_attr == "value"

def test_aiia_config_save_pretrained_and_load_with_nested_attributes():
    with tempfile.TemporaryDirectory() as tmpdir:
        config = AIIAConfig(model_type="TempModel", nested={"key": "value"})
        save_pretrained_path = os.path.join(tmpdir, "config")
        config.save_pretrained(save_pretrained_path)

        loaded_config = AIIAConfig.from_pretrained(save_pretrained_path)
        assert loaded_config.model_type == "TempModel"
        assert loaded_config.nested == {"key": "value"}