Building/src_new/dataset.py

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
import re
from typing import Any

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler

try:
    from .config import DataConfig, ExperimentConfig
except ImportError:
    from config import DataConfig, ExperimentConfig


def calculate_rms(signal: np.ndarray) -> float:
    signal = np.asarray(signal, dtype=np.float64).reshape(-1)
    return float(np.sqrt(np.mean(np.square(signal))))


def extract_frequency_hz(file_name: str) -> float | None:
    match = re.search(r"(\d+(?:\.\d+)?)Hz", file_name, flags=re.IGNORECASE)
    if match is None:
        return None
    return float(match.group(1))


@dataclass
class NormalizationStats:
    x_mean: torch.Tensor
    x_std: torch.Tensor
    y_wave_mean: torch.Tensor
    y_wave_std: torch.Tensor
    aux_mean: torch.Tensor
    aux_std: torch.Tensor
    target_mean: torch.Tensor
    target_std: torch.Tensor


@dataclass
class TensorStandardScaler:
    mean: torch.Tensor
    std: torch.Tensor

    def transform(self, array: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
        if isinstance(array, torch.Tensor):
            mean = self.mean.to(array.device, dtype=array.dtype)
            std = self.std.to(array.device, dtype=array.dtype)
            return (array - mean) / std
        np_array = np.asarray(array, dtype=np.float32)
        return (np_array - self.mean.cpu().numpy()) / self.std.cpu().numpy()

    def inverse_transform(self, array: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
        if isinstance(array, torch.Tensor):
            mean = self.mean.to(array.device, dtype=array.dtype)
            std = self.std.to(array.device, dtype=array.dtype)
            return array * std + mean
        np_array = np.asarray(array, dtype=np.float32)
        return np_array * self.std.cpu().numpy() + self.mean.cpu().numpy()


@dataclass
class RMSRecord:
    file_path: Path
    split: str
    time: torch.Tensor
    x_full: torch.Tensor
    x_model_input: torch.Tensor
    y_model_target: torch.Tensor
    x_rms: torch.Tensor
    y_rms: torch.Tensor
    target_value: torch.Tensor
    aux_features: torch.Tensor
    sample_weight: float
    interpolation_count: int = 0


@dataclass
class SplitLoadReport:
    split: str
    loaded_files: list[str] = field(default_factory=list)
    skipped_files: list[tuple[str, str]] = field(default_factory=list)
    interpolated_files: dict[str, int] = field(default_factory=dict)

    def to_lines(self) -> list[str]:
        lines = [f"[{self.split}] loaded={len(self.loaded_files)} skipped={len(self.skipped_files)}"]
        for file_name, reason in self.skipped_files:
            lines.append(f"  - skipped {file_name}: {reason}")
        for file_name, count in self.interpolated_files.items():
            if count > 0:
                lines.append(f"  - interpolated {file_name}: missing_points={count}")
        return lines


class RMSRegressionDataset(Dataset):
    def __init__(
        self,
        records: list[RMSRecord],
        normalization: NormalizationStats,
        max_sequence_length: int,
    ) -> None:
        self.records = records
        self.normalization = normalization
        self.max_sequence_length = max_sequence_length
        self.sample_weights = [record.sample_weight for record in records]

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int) -> dict[str, Any]:
        record = self.records[index]
        x = record.x_model_input
        if x.shape[0] > self.max_sequence_length:
            x = x[-self.max_sequence_length :]

        pad_length = self.max_sequence_length - x.shape[0]
        if pad_length > 0:
            x = F.pad(x.transpose(0, 1), (pad_length, 0), value=0.0).transpose(0, 1)

        x_norm = (x - self.normalization.x_mean) / self.normalization.x_std
        y = record.y_model_target
        if y.shape[0] > self.max_sequence_length:
            y = y[-self.max_sequence_length :]
        aux_norm = (record.aux_features - self.normalization.aux_mean) / self.normalization.aux_std
        target_norm = (record.target_value - self.normalization.target_mean) / self.normalization.target_std
        if pad_length > 0:
            y = F.pad(y.transpose(0, 1), (pad_length, 0), value=0.0).transpose(0, 1)
        y_wave_norm = (y - self.normalization.y_wave_mean) / self.normalization.y_wave_std

        valid_length = min(record.x_model_input.shape[0], self.max_sequence_length)
        mask = torch.zeros(self.max_sequence_length, dtype=torch.float32)
        mask[-valid_length:] = 1.0

        return {
            "x": x_norm,
            "aux": aux_norm,
            "y": target_norm,
            "y_wave": y_wave_norm,
            "y_raw": record.y_rms,
            "x_rms_raw": record.x_rms,
            "frequency_hz": torch.tensor(record.aux_features[0].item(), dtype=torch.float32),
            "mask": mask,
            "file_name": record.file_path.name,
            "file_path": str(record.file_path),
            "valid_length": torch.tensor(valid_length, dtype=torch.long),
            "sample_weight": torch.tensor(record.sample_weight, dtype=torch.float32),
        }


def list_split_files(config: DataConfig, split: str) -> list[Path]:
    split_dir = getattr(config, f"{split}_dir")
    return sorted(split_dir.glob(config.csv_pattern))


def _value_frame(df: pd.DataFrame, config: DataConfig, sensor_code: str, value_column: str) -> pd.DataFrame:
    sensor_df = df.loc[df[config.code_column] == sensor_code, [config.time_column, value_column]].copy()
    sensor_df = sensor_df.sort_values(config.time_column)
    sensor_df = sensor_df.drop_duplicates(subset=config.time_column, keep="first")
    sensor_df[config.time_column] = sensor_df[config.time_column].astype("float64")
    sensor_df[value_column] = sensor_df[value_column].astype("float32")
    return sensor_df


def select_model_segment(signal: np.ndarray, config: DataConfig) -> np.ndarray:
    if not config.use_steady_state_only:
        return signal
    start_index = int(len(signal) * config.steady_state_start_ratio)
    max_start = max(0, len(signal) - config.steady_state_min_samples)
    start_index = min(start_index, max_start)
    return signal[start_index:]


def estimate_dominant_frequency(signal: np.ndarray, sampling_rate: float) -> float:
    signal = np.asarray(signal, dtype=np.float64).reshape(-1)
    if signal.size < 4:
        return 0.0
    fft_values = np.fft.rfft(signal)
    freqs = np.fft.rfftfreq(signal.size, d=1.0 / sampling_rate)
    magnitudes = np.abs(fft_values)
    magnitudes[0] = 0.0
    band_mask = (freqs >= 0.1) & (freqs <= 5.0)
    if not np.any(band_mask):
        return 0.0
    masked_magnitudes = np.where(band_mask, magnitudes, 0.0)
    return float(freqs[int(np.argmax(masked_magnitudes))])


def estimate_sampling_rate(time_values: np.ndarray) -> float:
    if time_values.size < 2:
        return 100.0
    dt = np.diff(time_values)
    dt = dt[np.isfinite(dt)]
    dt = dt[dt > 0.0]
    if dt.size == 0:
        return 100.0
    return float(1.0 / np.median(dt))


def compute_sample_weight(y_rms: float, x_rms: float, freq_hz: float, config: DataConfig) -> float:
    gain = y_rms / max(x_rms, config.normalization_eps)
    low_freq_factor = (config.low_frequency_reference_hz / max(freq_hz, config.normalization_eps)) ** config.low_frequency_emphasis_power
    weight = (gain ** config.train_weight_power) * low_freq_factor
    return float(np.clip(weight, config.train_weight_min, config.train_weight_max))


def load_split_records(config: DataConfig, split: str) -> tuple[list[RMSRecord], SplitLoadReport]:
    records: list[RMSRecord] = []
    report = SplitLoadReport(split=split)

    for file_path in list_split_files(config, split):
        df = pd.read_csv(file_path)
        base_df = _value_frame(df, config, config.base_sensor_code, config.base_axis)
        response_df = _value_frame(df, config, config.response_sensor_code, config.response_axis)

        if base_df.empty or response_df.empty:
            report.skipped_files.append((file_path.name, "missing required sensor"))
            continue

        aligned = base_df.rename(columns={config.base_axis: "base_signal"})
        aligned = aligned.merge(
            response_df.rename(columns={config.response_axis: "response_signal"}),
            on=config.time_column,
            how="left",
        )
        interpolation_count = int(aligned["response_signal"].isna().sum())
        aligned["response_signal"] = aligned["response_signal"].interpolate(
            method=config.interpolation_method,
            limit_direction="both",
        ).ffill().bfill()

        if aligned["response_signal"].isna().any():
            report.skipped_files.append((file_path.name, "remaining NaN after interpolation"))
            continue

        time_values = aligned[config.time_column].to_numpy(dtype=np.float64)
        x_values = aligned["base_signal"].to_numpy(dtype=np.float32)
        y_values = aligned["response_signal"].to_numpy(dtype=np.float32)
        if len(x_values) < config.min_sequence_length:
            report.skipped_files.append((file_path.name, f"sequence too short: {len(x_values)}"))
            continue

        x_model = select_model_segment(x_values, config)
        y_model = select_model_segment(y_values, config)
        time_model = select_model_segment(time_values, config)
        sampling_rate = estimate_sampling_rate(time_model)
        file_frequency = extract_frequency_hz(file_path.name)
        x_rms = calculate_rms(x_model)
        y_rms = calculate_rms(y_model)
        target_value = float(np.log(max(y_rms / max(x_rms, config.normalization_eps), config.normalization_eps)))
        dominant_freq = estimate_dominant_frequency(x_model, sampling_rate)
        feature_frequency = file_frequency if file_frequency is not None else dominant_freq
        sample_weight = compute_sample_weight(y_rms=y_rms, x_rms=x_rms, freq_hz=feature_frequency, config=config) if split == "train" else 1.0

        aux_features = torch.tensor(
            [feature_frequency, x_rms, float(len(x_model)) / float(config.max_sequence_length)],
            dtype=torch.float32,
        )
        records.append(
            RMSRecord(
                file_path=file_path,
                split=split,
                time=torch.tensor(time_model, dtype=torch.float64),
                x_full=torch.tensor(x_values[:, None], dtype=torch.float32),
                x_model_input=torch.tensor(x_model[:, None], dtype=torch.float32),
                y_model_target=torch.tensor(y_model[:, None], dtype=torch.float32),
                x_rms=torch.tensor([x_rms], dtype=torch.float32),
                y_rms=torch.tensor([y_rms], dtype=torch.float32),
                target_value=torch.tensor([target_value], dtype=torch.float32),
                aux_features=aux_features,
                sample_weight=sample_weight,
                interpolation_count=interpolation_count,
            )
        )
        report.loaded_files.append(file_path.name)
        if interpolation_count > 0:
            report.interpolated_files[file_path.name] = interpolation_count

    if not records:
        raise RuntimeError(f"No usable records found for split='{split}'.")
    return records, report


def fit_normalization(records: list[RMSRecord], config: DataConfig) -> NormalizationStats:
    x_all = torch.cat([record.x_model_input for record in records], dim=0)
    y_all = torch.cat([record.y_model_target for record in records], dim=0)
    aux_all = torch.stack([record.aux_features for record in records], dim=0)
    target_all = torch.cat([record.target_value for record in records], dim=0)

    def safe_std(tensor: torch.Tensor, dim: int) -> torch.Tensor:
        std = tensor.std(dim=dim, unbiased=False)
        return torch.clamp(std, min=config.normalization_eps)

    return NormalizationStats(
        x_mean=x_all.mean(dim=0),
        x_std=safe_std(x_all, dim=0),
        y_wave_mean=y_all.mean(dim=0),
        y_wave_std=safe_std(y_all, dim=0),
        aux_mean=aux_all.mean(dim=0),
        aux_std=safe_std(aux_all, dim=0),
        target_mean=target_all.mean(dim=0, keepdim=True),
        target_std=safe_std(target_all, dim=0).view(1),
    )


def build_datasets(
    config: ExperimentConfig | DataConfig,
) -> tuple[dict[str, RMSRegressionDataset], NormalizationStats, dict[str, SplitLoadReport]]:
    data_config = config.data if isinstance(config, ExperimentConfig) else config
    train_records, train_report = load_split_records(data_config, "train")
    normalization = fit_normalization(train_records, data_config)
    val_records, val_report = load_split_records(data_config, "val")
    test_records, test_report = load_split_records(data_config, "test")

    datasets = {
        "train": RMSRegressionDataset(train_records, normalization, data_config.max_sequence_length),
        "val": RMSRegressionDataset(val_records, normalization, data_config.max_sequence_length),
        "test": RMSRegressionDataset(test_records, normalization, data_config.max_sequence_length),
    }
    return datasets, normalization, {"train": train_report, "val": val_report, "test": test_report}


def build_dataloaders(
    config: ExperimentConfig | DataConfig,
) -> tuple[dict[str, DataLoader], dict[str, RMSRegressionDataset], dict[str, SplitLoadReport]]:
    data_config = config.data if isinstance(config, ExperimentConfig) else config
    datasets, _, reports = build_datasets(config)

    train_sampler = None
    train_shuffle = True
    if data_config.use_weighted_train_sampler:
        weights = torch.tensor(datasets["train"].sample_weights, dtype=torch.double)
        train_sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
        train_shuffle = False

    loaders = {
        "train": DataLoader(
            datasets["train"],
            batch_size=data_config.batch_size,
            shuffle=train_shuffle,
            sampler=train_sampler,
            num_workers=data_config.num_workers,
            pin_memory=data_config.pin_memory,
        ),
        "val": DataLoader(
            datasets["val"],
            batch_size=data_config.batch_size,
            shuffle=False,
            num_workers=data_config.num_workers,
            pin_memory=data_config.pin_memory,
        ),
        "test": DataLoader(
            datasets["test"],
            batch_size=data_config.batch_size,
            shuffle=False,
            num_workers=data_config.num_workers,
            pin_memory=data_config.pin_memory,
        ),
    }
    return loaders, datasets, reports


def get_dataloaders(
    config: ExperimentConfig | DataConfig,
) -> tuple[
    dict[str, DataLoader],
    dict[str, RMSRegressionDataset],
    dict[str, SplitLoadReport],
    TensorStandardScaler,
    TensorStandardScaler,
    TensorStandardScaler,
]:
    loaders, datasets, reports = build_dataloaders(config)
    normalization = datasets["train"].normalization
    x_scaler = TensorStandardScaler(normalization.x_mean, normalization.x_std)
    aux_scaler = TensorStandardScaler(normalization.aux_mean, normalization.aux_std)
    y_scaler = TensorStandardScaler(normalization.target_mean, normalization.target_std)
    return loaders, datasets, reports, x_scaler, aux_scaler, y_scaler


def report_to_text(reports: dict[str, SplitLoadReport]) -> str:
    lines: list[str] = []
    for split in ("train", "val", "test"):
        lines.extend(reports[split].to_lines())
    return "\n".join(lines)