from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path


CORE_FEATURE_NAMES: tuple[str, ...] = (
    "dominant_frequency_hz",
    "frequency_squared",
    "inverse_frequency_hz",
    "log_frequency_hz",
    "input_rms",
    "input_peak_abs",
    "input_peak_to_peak",
    "crest_factor",
    "middle_length_ratio",
    "dominant_amplitude",
    "dominant_energy_ratio",
    "harmonic_fit_amplitude",
    "harmonic_fit_residual_ratio",
    "spectral_peak_prominence",
    "half_power_bandwidth_hz",
    "spectral_centroid_hz",
    "signal_mean",
)


@dataclass
class DataConfig:
    project_root: Path = field(default_factory=lambda: Path(__file__).resolve().parents[1])
    scenario: str = "Non_TMD"
    train_split_name: str = "train"
    val_split_name: str = "val"
    test_split_name: str = "test"
    csv_pattern: str = "*.csv"

    code_column: str = "code"
    time_column: str = "time"
    base_sensor_code: str = "WSMS00012"
    base_axis: str = "value1"
    response_sensor_code: str = "WSMS00007"
    response_axis: str = "value3"

    middle_segment_start_ratio: float = 0.20
    middle_segment_end_ratio: float = 0.80
    min_segment_length: int = 512
    steady_window_ratio: float = 0.25
    steady_window_stride_ratio: float = 0.05
    stability_subwindow_count: int = 4
    interpolation_method: str = "linear"
    normalization_eps: float = 1e-6

    downloads_dir: Path = field(init=False)
    scenario_dir: Path = field(init=False)
    train_dir: Path = field(init=False)
    val_dir: Path = field(init=False)
    test_dir: Path = field(init=False)

    def __post_init__(self) -> None:
        self.project_root = Path(self.project_root).resolve()
        self.downloads_dir = self.project_root / "downloads"
        self.scenario_dir = self.downloads_dir / self.scenario
        self.train_dir = self.scenario_dir / self.train_split_name
        self.val_dir = self.scenario_dir / self.val_split_name
        self.test_dir = self.scenario_dir / self.test_split_name


@dataclass
class ModelConfig:
    input_dim: int = len(CORE_FEATURE_NAMES)
    hidden_dims: tuple[int, ...] = (96, 64, 32)
    dropout: float = 0.08


@dataclass
class TrainConfig:
    epochs: int = 400
    batch_size: int = 16
    learning_rate: float = 1e-3
    weight_decay: float = 1e-4
    seed: int = 42
    device: str = "cuda"
    grad_clip_norm: float = 1.0
    lr_scheduler_patience: int = 20
    lr_scheduler_factor: float = 0.5
    min_learning_rate: float = 1e-6
    early_stop_patience: int = 50
    checkpoint_dir: str = "checkpoints_mlp"
    history_name: str = "training_history.csv"
    best_model_name: str = "best_feature_mlp.pt"


@dataclass
class LossConfig:
    relative_rms_weight: float = 1.0
    log_rms_huber_weight: float = 0.75
    mae_weight: float = 0.15


@dataclass
class ExperimentConfig:
    data: DataConfig = field(default_factory=DataConfig)
    model: ModelConfig = field(default_factory=ModelConfig)
    train: TrainConfig = field(default_factory=TrainConfig)
    loss: LossConfig = field(default_factory=LossConfig)


def make_experiment_config() -> ExperimentConfig:
    return ExperimentConfig()