"""
Hyperparameters for the DualEdit baseline.

Based on DualEdit paper (COLM 2025), Appendix D, LLaVA-1.5 configuration.
"""

from dataclasses import dataclass, field
from typing import List
import yaml


@dataclass
class DualEditHyperParams:
    # Model
    model_name: str = "llava-hf/llava-1.5-7b-hf"
    device: int = 0

    # Architecture
    hidden_size: int = 4096
    adapter_mid_dim: int = 1024
    cross_att_head_n: int = 8
    img_tok_n: int = 576  # (336/14)^2 for LLaVA-1.5

    # Layer configuration (from paper: vision=19, text=16 for LLaVA-1.5)
    vision_adapter_layer: int = 19
    text_adapter_layer: int = 16
    llm_layer_tmp: str = "language_model.model.layers.{}"

    # Gating
    gating_threshold: float = 0.6

    # Training
    edit_lr: float = 1e-4
    n_iterations: int = 5000
    batch_size: int = 4
    reliability_weight: float = 1.0
    generality_weight: float = 1.0
    locality_weight: float = 1.0
    checkpoint_every: int = 500

    # Algorithm name (for validation)
    alg_name: str = "DualEdit"

    @classmethod
    def from_hparams(cls, hparams_path: str) -> "DualEditHyperParams":
        """Load from YAML file."""
        if not hparams_path.endswith(".yaml"):
            hparams_path += ".yaml"
        with open(hparams_path, "r") as f:
            data = yaml.safe_load(f)
        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})