from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import yaml


@dataclass
class ModelConfig:
    model_type: str = "mamba"
    embed_dim: int = 512
    depth: int = 24
    predictor_depth: int = 2
    drop_path_rate: float = 0.1
    rms_norm: bool = False
    fused_add_norm: bool = True
    residual_in_fp32: bool = True
    bimamba_type: str = "v2"
    if_bimamba: bool = False
    mixer_type: str = "mamba"
    if_devide_out: bool = True
    momentum: float = 0.996
    final_momentum: float = 0.9999
    norm_target: bool = True
    num_heads: int = 8
    mlp_ratio: float = 4.0
    head_type: str = "transformer"
    num_classes: int = 3
    head_depth: int = 2
    head_num_heads: int = 8
    head_mlp_ratio: float = 4.0
    head_proj_drop: float = 0.1
    head_drop_path: float = 0.1
    mlp_hidden: int = 512
    mlp_depth: int = 4
    mlp_dropout: float = 0.1
    freeze_backbone: bool = False


@dataclass
class DataConfig:
    train_list: str = ""
    val_list: str = ""
    test_list: Optional[str] = None
    csv: Optional[str] = None
    id_column: str = "Subject"
    label_column: str = "Group_idx"
    label_mode: str = "multiclass"
    path_id_mode: str = "auto"
    normal_label: int = 2
    batch_size: int = 8
    num_workers: int = 8
    memory_map: bool = True
    T_prime: int = 30
    tau_seconds: float = 6.0
    default_tr: Optional[float] = None


@dataclass
class TrainingConfig:
    epochs: int = 30
    lr: float = 5e-4
    lr_backbone: Optional[float] = None
    lr_head: Optional[float] = None
    weight_decay: float = 0.05
    warmup_epochs: int = 2
    mask_ratio: float = 0.65
    grad_clip: float = 1.0
    grad_accumulation_steps: int = 1
    seed: int = 42
    use_amp: bool = False
    local_rank: int = 0
    world_size: int = 1


@dataclass
class LoggingConfig:
    log_interval: int = 20
    checkpoint_dir: str = "./checkpoints"
    log_dir: str = "./logs"
    resume: Optional[str] = None


@dataclass
class RunConfig:
    model: ModelConfig = field(default_factory=ModelConfig)
    data: DataConfig = field(default_factory=DataConfig)
    training: TrainingConfig = field(default_factory=TrainingConfig)
    logging: LoggingConfig = field(default_factory=LoggingConfig)
    pretrain_checkpoint: Optional[str] = None
    from_scratch: bool = False
    use_checkpoint_config: bool = True


def _update_dataclass(obj, values: dict):
    for key, value in values.items():
        if hasattr(obj, key):
            setattr(obj, key, value)


def load_config(path: Optional[str]) -> RunConfig:
    cfg = RunConfig()
    if not path:
        return cfg
    data = yaml.safe_load(Path(path).read_text()) or {}
    if "model" in data:
        _update_dataclass(cfg.model, data["model"] or {})
    if "data" in data:
        _update_dataclass(cfg.data, data["data"] or {})
    if "training" in data:
        _update_dataclass(cfg.training, data["training"] or {})
    if "logging" in data:
        _update_dataclass(cfg.logging, data["logging"] or {})
    for key in ["pretrain_checkpoint", "from_scratch", "use_checkpoint_config"]:
        if key in data:
            setattr(cfg, key, data[key])
    return cfg


def apply_checkpoint_config(model_cfg: ModelConfig, checkpoint_config: dict) -> None:
    keys = [
        "model_type", "embed_dim", "depth", "predictor_depth", "drop_path_rate",
        "rms_norm", "fused_add_norm", "residual_in_fp32", "bimamba_type",
        "if_bimamba", "mixer_type", "if_devide_out", "momentum", "norm_target",
        "num_heads", "mlp_ratio",
    ]
    for key in keys:
        if key in checkpoint_config:
            setattr(model_cfg, key, checkpoint_config[key])