"""
automl_lib.config
~~~~~~~~~~~~~~~~~
Dataclass-based configuration for the entire AutoML pipeline.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


@dataclass
class AutoMLConfig:
    """Central configuration object.

    Every parameter has a sensible default.  Users can override any field
    either programmatically or via the CLI.
    """

    # ── Data ──────────────────────────────────────────────────────────
    data_path: str = ""
    data_type: str | None = None          # auto-detected if None
    task: str | None = None               # auto-detected if None
    target_column: str | None = None      # for tabular data
    val_split: float = 0.2
    test_split: float = 0.0
    max_samples: int | None = None        # cap dataset for quick experiments

    # ── Model ─────────────────────────────────────────────────────────
    model_family: str | None = None       # e.g. "cnn", "transformer", "xgboost"
    model_name: str | None = None         # specific model e.g. "resnet18"
    num_classes: int | None = None        # auto-detected if None
    layers: list[int] = field(default_factory=lambda: [256, 128])
    activations: str = "relu"
    dropout: float = 0.3
    pretrained: bool = True

    # ── Advanced Features ─────────────────────────────────────────────
    cv_folds: int = 1                     # >1 enables Cross-Validation
    class_weights: str | None = None      # "auto" or None
    lr_finder: bool = False               # Auto-find optimal LR before training
    ensemble: bool = False                # Enable model ensembling
    ensemble_top_n: int = 3               # Top-N models to ensemble

    # ── Training ──────────────────────────────────────────────────────
    epochs: int = 50
    batch_size: int = 32
    learning_rate: float = 1e-3
    optimizer: str = "adam"               # adam | sgd | adamw
    weight_decay: float = 1e-4
    scheduler: str | None = "cosine"      # cosine | step | None
    early_stopping_patience: int = 7
    gradient_clip: float = 1.0

    # ── Hyperparameter Tuning ─────────────────────────────────────────
    tune: bool = False
    tune_trials: int = 30
    tune_timeout: int | None = None       # seconds

    # ── Hardware ──────────────────────────────────────────────────────
    device: str | None = None             # auto-detected if None
    num_workers: int = 4
    pin_memory: bool = True
    mixed_precision: bool = True

    # ── Checkpointing & Logging ───────────────────────────────────────
    output_dir: str = "automl_output"
    checkpoint_dir: str | None = None     # defaults to output_dir/checkpoints
    save_best_only: bool = True
    log_every_n_steps: int = 10
    use_tensorboard: bool = False

    # ── Export ─────────────────────────────────────────────────────────
    export_format: str = "torchscript"    # torchscript | onnx | both
    export_path: str | None = None

    # ── Misc ──────────────────────────────────────────────────────────
    seed: int = 42
    verbose: bool = True
    extra: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        if self.checkpoint_dir is None:
            import os
            self.checkpoint_dir = os.path.join(self.output_dir, "checkpoints")

    # ── Convenience ───────────────────────────────────────────────────
    def to_dict(self) -> dict:
        from dataclasses import asdict
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "AutoMLConfig":
        import dataclasses
        valid = {f.name for f in dataclasses.fields(cls)}
        return cls(**{k: v for k, v in d.items() if k in valid})
