"""

trainiq.config

~~~~~~~~~~~~~~

Dataclass-based configuration for the entire trainiq pipeline.

"""

from __future__ import annotations


from dataclasses import dataclass, field

from typing import Any



@dataclass
class trainiqConfig:

    """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 = "trainiq_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) -> "trainiqConfig":
        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})

