config.py

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Mapping

import yaml


def _coerce_bool(value: Any, field_name: str) -> bool:
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        normalized = value.strip().lower()
        if normalized in {"true", "yes", "1", "on"}:
            return True
        if normalized in {"false", "no", "0", "off"}:
            return False
    raise ValueError(f"{field_name} must be a boolean")


def _positive_int(value: Any, field_name: str) -> int:
    parsed = int(value)
    if parsed <= 0:
        raise ValueError(f"{field_name} must be greater than zero")
    return parsed


@dataclass
class DataConfig:
    train_dir: str
    train_ann: str
    val_dir: str
    val_ann: str
    num_workers: int = 8
    batch_size: int = 4

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> DataConfig:
        num_workers = int(data.get("num_workers", 8))
        if num_workers < 0:
            raise ValueError("data.num_workers must be zero or greater")
        return cls(
            train_dir=str(data["train_dir"]),
            train_ann=str(data["train_ann"]),
            val_dir=str(data["val_dir"]),
            val_ann=str(data["val_ann"]),
            num_workers=num_workers,
            batch_size=_positive_int(data.get("batch_size", 4), "data.batch_size"),
        )


@dataclass
class ModelConfig:
    config_path: str
    weights_path: str
    lora_weights: str | None = None

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> ModelConfig:
        lora_weights = data.get("lora_weights", data.get("lora_weigths"))
        return cls(
            config_path=str(data["config_path"]),
            weights_path=str(data["weights_path"]),
            lora_weights=None if lora_weights is None else str(lora_weights),
        )

    @property
    def lora_weigths(self) -> str | None:
        """Backward-compatible alias for the misspelled pre-0.2 field."""
        return self.lora_weights


@dataclass
class EvaluationConfig:
    iou_thresholds: list[float] = field(default_factory=lambda: [0.5, 0.75])
    score_threshold: float = 0.25
    max_detections: int = 100
    metrics_output_dir: str = "evaluation_results"
    generate_visualizations: bool = True
    num_visualization_samples: int = 10

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> EvaluationConfig:
        iou_thresholds = data.get("iou_thresholds", [0.5, 0.75])
        if isinstance(iou_thresholds, str):
            iou_thresholds = [item.strip() for item in iou_thresholds.split(",")]
        parsed_thresholds = [float(value) for value in iou_thresholds]
        if not parsed_thresholds or any(not 0 < value <= 1 for value in parsed_thresholds):
            raise ValueError("training.evaluation.iou_thresholds must be within (0, 1]")

        score_threshold = float(data.get("score_threshold", 0.25))
        if not 0 <= score_threshold <= 1:
            raise ValueError("training.evaluation.score_threshold must be within [0, 1]")

        num_visualization_samples = int(data.get("num_visualization_samples", 10))
        if num_visualization_samples < 0:
            raise ValueError(
                "training.evaluation.num_visualization_samples must be zero or greater"
            )

        return cls(
            iou_thresholds=parsed_thresholds,
            score_threshold=score_threshold,
            max_detections=_positive_int(
                data.get("max_detections", 100),
                "training.evaluation.max_detections",
            ),
            metrics_output_dir=str(data.get("metrics_output_dir", "evaluation_results")),
            generate_visualizations=_coerce_bool(
                data.get("generate_visualizations", True),
                "training.evaluation.generate_visualizations",
            ),
            num_visualization_samples=num_visualization_samples,
        )


@dataclass
class TrainingConfig:
    num_epochs: int = 1000
    learning_rate: float = 1e-3
    save_dir: str = "weights"
    save_frequency: int = 100
    warmup_epochs: int = 5
    use_lora: bool = False
    use_ema: bool = False
    visualization_frequency: int = 5
    evaluate_during_training: bool = True
    evaluation_frequency: int = 10
    evaluation: EvaluationConfig | None = None

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> TrainingConfig:
        evaluation_data = data.get("evaluation")
        evaluation_config = (
            None if evaluation_data is None else EvaluationConfig.from_dict(evaluation_data)
        )
        learning_rate = float(data.get("learning_rate", 1e-3))
        if learning_rate <= 0:
            raise ValueError("training.learning_rate must be greater than zero")

        warmup_epochs = int(data.get("warmup_epochs", 5))
        if warmup_epochs < 0:
            raise ValueError("training.warmup_epochs must be zero or greater")

        return cls(
            num_epochs=_positive_int(data.get("num_epochs", 1000), "training.num_epochs"),
            learning_rate=learning_rate,
            save_dir=str(data.get("save_dir", "weights")),
            save_frequency=_positive_int(
                data.get("save_frequency", 100), "training.save_frequency"
            ),
            warmup_epochs=warmup_epochs,
            use_lora=_coerce_bool(data.get("use_lora", False), "training.use_lora"),
            use_ema=_coerce_bool(data.get("use_ema", False), "training.use_ema"),
            visualization_frequency=_positive_int(
                data.get("visualization_frequency", 5),
                "training.visualization_frequency",
            ),
            evaluate_during_training=_coerce_bool(
                data.get("evaluate_during_training", True),
                "training.evaluate_during_training",
            ),
            evaluation_frequency=_positive_int(
                data.get("evaluation_frequency", 10),
                "training.evaluation_frequency",
            ),
            evaluation=evaluation_config,
        )


class ConfigurationManager:
    @staticmethod
    def load_config(config_path: str) -> tuple[DataConfig, ModelConfig, TrainingConfig]:
        """Load and validate an AViD YAML configuration."""
        with Path(config_path).open("r", encoding="utf-8") as file_obj:
            config = yaml.safe_load(file_obj)

        if not isinstance(config, Mapping):
            raise ValueError("Invalid configuration format: root must be a mapping")

        try:
            for section_name in ("data", "model", "training"):
                if not isinstance(config[section_name], Mapping):
                    raise TypeError(f"{section_name} must be a mapping")
            data_config = DataConfig.from_dict(config["data"])
            model_config = ModelConfig.from_dict(config["model"])
            training_config = TrainingConfig.from_dict(config["training"])
        except (KeyError, ValueError, TypeError) as exc:
            raise ValueError(f"Invalid configuration format: {exc}") from exc

        return data_config, model_config, training_config

    @staticmethod
    def save_config(
        config_path: str | Path,
        data_config: DataConfig,
        model_config: ModelConfig,
        training_config: TrainingConfig,
    ) -> None:
        """Write a configuration that can be loaded again with ``load_config``."""
        destination = Path(config_path)
        destination.parent.mkdir(parents=True, exist_ok=True)
        payload = {
            "data": asdict(data_config),
            "model": asdict(model_config),
            "training": asdict(training_config),
        }
        with destination.open("w", encoding="utf-8") as file_obj:
            yaml.safe_dump(payload, file_obj, sort_keys=False)