Building/scripts/dataset.py

from __future__ import annotations

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

import numpy as np
import pandas as pd

try:
    from .config import CORE_FEATURE_NAMES, DataConfig, ExperimentConfig
except ImportError:
    from config import CORE_FEATURE_NAMES, 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))


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 get_middle_segment(
    time_values: np.ndarray,
    x_values: np.ndarray,
    y_values: np.ndarray,
    config: DataConfig,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    length = len(x_values)
    search_start = int(length * config.middle_segment_start_ratio)
    search_end = int(length * config.middle_segment_end_ratio)
    search_start = max(0, min(search_start, length - 1))
    search_end = max(search_start + 1, min(search_end, length))

    if search_end - search_start < config.min_segment_length:
        center = length // 2
        half = config.min_segment_length // 2
        search_start = max(0, center - half)
        search_end = min(length, search_start + config.min_segment_length)
        search_start = max(0, search_end - config.min_segment_length)

    search_length = search_end - search_start
    window_length = max(config.min_segment_length, int(length * config.steady_window_ratio))
    window_length = min(window_length, search_length)
    if search_length <= window_length:
        return (
            time_values[search_start:search_end],
            x_values[search_start:search_end],
            y_values[search_start:search_end],
        )

    stride = max(1, int(length * config.steady_window_stride_ratio))
    candidate_y = y_values[search_start:search_end]
    candidate_rms = calculate_rms(candidate_y)
    best_score = float("inf")
    best_slice = slice(search_start, search_start + window_length)

    for window_start in range(search_start, search_end - window_length + 1, stride):
        window_end = window_start + window_length
        window_y = y_values[window_start:window_end]
        split_windows = np.array_split(window_y, config.stability_subwindow_count)
        split_rms = np.asarray([calculate_rms(chunk) for chunk in split_windows], dtype=np.float64)
        split_peaks = np.asarray([float(np.max(np.abs(chunk))) for chunk in split_windows], dtype=np.float64)

        rms_cv = float(split_rms.std() / max(split_rms.mean(), config.normalization_eps))
        peak_cv = float(split_peaks.std() / max(split_peaks.mean(), config.normalization_eps))
        window_rms = calculate_rms(window_y)

        score = rms_cv + 0.35 * peak_cv - 0.05 * (window_rms / max(candidate_rms, config.normalization_eps))
        if score < best_score:
            best_score = score
            best_slice = slice(window_start, window_end)

    return time_values[best_slice], x_values[best_slice], y_values[best_slice]


def _compute_windowed_spectrum(signal: np.ndarray, sampling_rate: float) -> tuple[np.ndarray, np.ndarray]:
    signal = np.asarray(signal, dtype=np.float64).reshape(-1)
    if signal.size < 4:
        return np.asarray([], dtype=np.float64), np.asarray([], dtype=np.float64)
    centered = signal - np.mean(signal)
    window = np.hanning(signal.size)
    scale = max(np.sum(window), 1e-12)
    fft_values = np.fft.rfft(centered * window)
    freqs = np.fft.rfftfreq(signal.size, d=1.0 / sampling_rate)
    magnitudes = (2.0 / scale) * np.abs(fft_values)
    return freqs, magnitudes


def _parabolic_peak_frequency(freqs: np.ndarray, magnitudes: np.ndarray, peak_index: int) -> float:
    if peak_index <= 0 or peak_index >= magnitudes.size - 1:
        return float(freqs[peak_index])
    alpha = magnitudes[peak_index - 1]
    beta = magnitudes[peak_index]
    gamma = magnitudes[peak_index + 1]
    denominator = alpha - 2.0 * beta + gamma
    if abs(denominator) < 1e-12:
        return float(freqs[peak_index])
    offset = 0.5 * (alpha - gamma) / denominator
    bin_width = float(freqs[1] - freqs[0])
    return float(freqs[peak_index] + offset * bin_width)


def get_dominant_frequency(signal: np.ndarray, sampling_rate: float) -> float:
    freqs, magnitudes = _compute_windowed_spectrum(signal, sampling_rate)
    if magnitudes.size == 0:
        return 0.0
    magnitudes[0] = 0.0
    band_mask = (freqs >= 0.1) & (freqs <= 5.0)
    if not np.any(band_mask):
        return 0.0
    band_magnitudes = np.where(band_mask, magnitudes, 0.0)
    peak_index = int(np.argmax(band_magnitudes))
    return _parabolic_peak_frequency(freqs, magnitudes, peak_index)


def theoretical_acceleration(frequency_hz: float, amplitude_m: float) -> float:
    return float((2.0 * np.pi * frequency_hz) ** 2 * amplitude_m)


def extract_core_features(
    signal: np.ndarray,
    sampling_rate: float,
    config: DataConfig,
    known_frequency_hz: float | None = None,
) -> np.ndarray:
    frequency_hz = known_frequency_hz if known_frequency_hz is not None else get_dominant_frequency(signal, sampling_rate)
    x_rms = calculate_rms(signal)
    return np.asarray(
        [
            frequency_hz,
            frequency_hz * frequency_hz,
            x_rms,
            theoretical_acceleration(frequency_hz, config.harmonic_amplitude_m),
        ],
        dtype=np.float32,
    )


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 load_aligned_signals(file_path: Path, config: DataConfig) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
    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:
        raise ValueError(f"Missing required sensor in {file_path.name}")

    aligned = base_df.rename(columns={config.base_axis: "base_signal"}).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():
        raise ValueError(f"Remaining NaN after interpolation in {file_path.name}")

    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)
    return time_values, x_values, y_values, interpolation_count


@dataclass
class SampleRecord:
    file_path: Path
    frequency_hz: float
    features: np.ndarray
    x_rms: float
    y_rms: float
    target_tr: float
    time_middle: np.ndarray
    x_middle: np.ndarray
    y_middle: np.ndarray
    sampling_rate: float
    interpolation_count: int = 0


@dataclass
class LoadReport:
    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"[all_data] 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


def list_harmonic_files(config: DataConfig) -> list[Path]:
    files = sorted(config.data_root.rglob(config.harmonic_pattern))
    return [file_path for file_path in files if extract_frequency_hz(file_path.name) is not None]


def build_record_from_file(file_path: Path, config: DataConfig) -> SampleRecord:
    time_values, x_values, y_values, interpolation_count = load_aligned_signals(file_path, config)
    time_middle, x_middle, y_middle = get_middle_segment(time_values, x_values, y_values, config)
    sampling_rate = estimate_sampling_rate(time_middle)
    frequency_hz = extract_frequency_hz(file_path.name)
    if frequency_hz is None:
        frequency_hz = get_dominant_frequency(x_middle, sampling_rate)

    x_rms = calculate_rms(x_middle)
    y_rms = calculate_rms(y_middle)
    target_tr = float(y_rms / max(x_rms, config.normalization_eps))
    features = extract_core_features(x_middle, sampling_rate, config, known_frequency_hz=frequency_hz)

    return SampleRecord(
        file_path=file_path,
        frequency_hz=float(frequency_hz),
        features=features,
        x_rms=float(x_rms),
        y_rms=float(y_rms),
        target_tr=target_tr,
        time_middle=np.asarray(time_middle, dtype=np.float64),
        x_middle=np.asarray(x_middle, dtype=np.float32).reshape(-1),
        y_middle=np.asarray(y_middle, dtype=np.float32).reshape(-1),
        sampling_rate=float(sampling_rate),
        interpolation_count=interpolation_count,
    )


def load_all_records(config: ExperimentConfig | DataConfig) -> tuple[list[SampleRecord], LoadReport]:
    data_config = config.data if isinstance(config, ExperimentConfig) else config
    records: list[SampleRecord] = []
    report = LoadReport()

    for file_path in list_harmonic_files(data_config):
        try:
            record = build_record_from_file(file_path, data_config)
        except ValueError as error:
            report.skipped_files.append((file_path.name, str(error)))
            continue
        records.append(record)
        report.loaded_files.append(file_path.name)
        if record.interpolation_count > 0:
            report.interpolated_files[file_path.name] = record.interpolation_count

    if not records:
        raise RuntimeError("No usable harmonic records found in data_root.")
    return records, report


def records_to_frame(records: list[SampleRecord]) -> pd.DataFrame:
    rows = []
    for record in records:
        row = {
            "file_name": record.file_path.name,
            "frequency_hz": record.frequency_hz,
            "x_rms": record.x_rms,
            "y_rms": record.y_rms,
            "target_tr": record.target_tr,
            "sampling_rate": record.sampling_rate,
        }
        for feature_name, feature_value in zip(CORE_FEATURE_NAMES, record.features.tolist()):
            row[feature_name] = float(feature_value)
        rows.append(row)
    return pd.DataFrame(rows)


def report_to_text(report: LoadReport) -> str:
    return "\n".join(report.to_lines())