"""IRIS model configurations. Each config is tested and has known parameter count."""

CONFIGS = {
    "iris-tiny": {"latent_channels": 32, "dim": 256, "patch_size": 4, "num_blocks": 6, "num_heads": 4, "max_iterations": 8, "ffn_expansion": 2, "params_M": 10.3, "tokens": 16, "description": "Ultra-mobile, 10M params, 16 tokens, trains on Colab free tier"},
    "iris-small": {"latent_channels": 32, "dim": 512, "patch_size": 4, "num_blocks": 6, "num_heads": 8, "max_iterations": 8, "ffn_expansion": 2, "params_M": 40.0, "tokens": 16, "description": "Mobile, 40M params, 16 tokens, trains on 16GB GPU"},
    "iris-base": {"latent_channels": 32, "dim": 512, "patch_size": 2, "num_blocks": 8, "num_heads": 8, "max_iterations": 8, "ffn_expansion": 2, "params_M": 53.4, "tokens": 64, "description": "Base quality, 53M params, 64 tokens, trains on 16GB GPU"},
    "iris-medium": {"latent_channels": 32, "dim": 768, "patch_size": 2, "num_blocks": 12, "num_heads": 12, "max_iterations": 8, "ffn_expansion": 2, "params_M": 181.2, "tokens": 64, "description": "Full quality, 181M params, 64 tokens, needs 24GB GPU"},
    "iris-large": {"latent_channels": 32, "dim": 1024, "patch_size": 2, "num_blocks": 16, "num_heads": 16, "max_iterations": 8, "ffn_expansion": 2, "params_M": 430.9, "tokens": 64, "description": "Maximum quality, 431M params, 64 tokens, needs 40GB+ GPU"},
}


def get_model_config(name: str) -> dict:
    if name not in CONFIGS:
        raise ValueError(f"Unknown config: {name}. Available: {list(CONFIGS.keys())}")
    cfg = CONFIGS[name].copy()
    for k in ["params_M", "tokens", "description"]:
        cfg.pop(k, None)
    return cfg