Building/scripts_2/task3_identify.py

from __future__ import annotations

import argparse
from pathlib import Path
import sys

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np

if __package__ is None or __package__ == "":
    sys.path.append(str(Path(__file__).resolve().parents[1]))

from scripts_2.signal_utils import (
    TOP_RESPONSE_AXIS,
    TOP_RESPONSE_SENSOR,
    bandpass_filter,
    detect_decay_peaks,
    dominant_frequency_in_band,
    ensure_parent_dir,
    estimate_sampling_rate,
    read_sensor_signal,
    select_decay_peak_window,
)


def parse_args() -> argparse.Namespace:
    project_root = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(description="Identify the first natural frequency and damping ratio from free-decay data.")
    parser.add_argument(
        "--file",
        type=str,
        default=str(project_root / "downloads" / "Non_TMD" / "free_vib" / "free_decay.csv"),
        help="Path to a free-vibration CSV inside the dedicated free_vib folder.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=str(project_root / "evaluation_outputs" / "task3_identify"),
        help="Directory for plots and optional exported diagnostics.",
    )
    parser.add_argument(
        "--bandwidth",
        type=float,
        default=0.3,
        help="Half-bandwidth (Hz) of the Butterworth band-pass filter around the identified natural frequency.",
    )
    parser.add_argument(
        "--show",
        action="store_true",
        help="Display the generated figure in addition to saving it.",
    )
    return parser.parse_args()


def validate_free_vibration_source(file_path: Path) -> None:
    if "free_vib" not in {part.lower() for part in file_path.parts}:
        raise ValueError(
            f"{file_path.name} is not inside a dedicated free_vib folder. Forced-vibration tails are forbidden for task 3."
        )


def build_envelope_curve(
    time_values: np.ndarray,
    peak_times: np.ndarray,
    fit_amplitudes: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    if peak_times.size < 2:
        return time_values, np.zeros_like(time_values)
    envelope = np.interp(time_values, peak_times, fit_amplitudes, left=fit_amplitudes[0], right=fit_amplitudes[-1])
    return time_values, envelope


def save_task3_figure(
    file_path: Path,
    output_dir: Path,
    time_values: np.ndarray,
    raw_signal: np.ndarray,
    filtered_signal: np.ndarray,
    peak_indices: np.ndarray,
    selected_peak_indices: np.ndarray,
    selected_peak_amplitudes: np.ndarray,
    fit_amplitudes: np.ndarray,
    natural_frequency_hz: float,
    damping_ratio: float,
) -> Path:
    output_dir = ensure_parent_dir(output_dir)
    peak_times = time_values[selected_peak_indices]
    envelope_time, envelope = build_envelope_curve(time_values, peak_times, fit_amplitudes)

    fig, axes = plt.subplots(3, 1, figsize=(12, 10), sharex=False)
    fig.suptitle(f"Task 3 | Free-Vibration Identification | {file_path.name}")

    axes[0].plot(time_values, raw_signal, color="tab:gray", linewidth=1.0, label="Raw top response")
    axes[0].set_title(f"Top Response ({TOP_RESPONSE_SENSOR} / {TOP_RESPONSE_AXIS})")
    axes[0].set_ylabel("Acceleration")
    axes[0].grid(True, alpha=0.25)
    axes[0].legend()

    axes[1].plot(time_values, filtered_signal, color="tab:blue", linewidth=1.0, label="Band-pass filtered")
    axes[1].scatter(time_values[peak_indices], filtered_signal[peak_indices], color="tab:orange", s=14, label="Detected peaks")
    axes[1].set_title(f"Filtered around f_n = {natural_frequency_hz:.4f} Hz")
    axes[1].set_ylabel("Acceleration")
    axes[1].grid(True, alpha=0.25)
    axes[1].legend()

    axes[2].plot(time_values, filtered_signal, color="tab:blue", linewidth=1.0, alpha=0.65, label="Filtered signal")
    axes[2].scatter(peak_times, selected_peak_amplitudes, color="tab:red", s=28, label="Selected envelope peaks")
    axes[2].plot(envelope_time, envelope, color="tab:green", linewidth=1.8, label="Fitted positive envelope")
    axes[2].plot(envelope_time, -envelope, color="tab:green", linewidth=1.2, linestyle="--", label="Fitted negative envelope")
    axes[2].set_title(f"Selected decay segment | damping ratio zeta = {damping_ratio:.5f}")
    axes[2].set_xlabel("Time (s)")
    axes[2].set_ylabel("Acceleration")
    axes[2].grid(True, alpha=0.25)
    axes[2].legend()

    plt.tight_layout()
    figure_path = output_dir / f"{file_path.stem}_task3_identify.png"
    plt.savefig(figure_path, dpi=180, bbox_inches="tight")
    return figure_path


def analyze_free_vibration(
    file_path: Path,
    output_dir: Path,
    bandwidth_hz: float = 0.3,
    show: bool = False,
) -> dict[str, float | int | str | Path]:
    file_path = Path(file_path).resolve()
    validate_free_vibration_source(file_path)

    time_values, raw_signal = read_sensor_signal(file_path, sensor_code=TOP_RESPONSE_SENSOR, value_column=TOP_RESPONSE_AXIS)
    sampling_rate = estimate_sampling_rate(time_values)
    natural_frequency_hz, freqs, magnitudes, peak_index = dominant_frequency_in_band(
        raw_signal,
        sampling_rate=sampling_rate,
        min_hz=0.1,
        max_hz=5.0,
        zero_padding_factor=16,
    )
    filtered_signal = bandpass_filter(
        raw_signal,
        sampling_rate=sampling_rate,
        center_hz=natural_frequency_hz,
        bandwidth_hz=bandwidth_hz,
        order=4,
    )
    peak_indices = detect_decay_peaks(filtered_signal, sampling_rate=sampling_rate, frequency_hz=natural_frequency_hz)
    peak_window = select_decay_peak_window(filtered_signal, peak_indices, sampling_rate=sampling_rate)

    selected_peak_indices = np.asarray(peak_window["peak_indices"], dtype=int)
    selected_peak_amplitudes = np.asarray(peak_window["peak_amplitudes"], dtype=np.float64)
    fit_amplitudes = np.asarray(peak_window["fit_amplitudes"], dtype=np.float64)
    log_decrements = np.log(selected_peak_amplitudes[:-1] / selected_peak_amplitudes[1:])
    mean_log_decrement = float(np.mean(log_decrements))
    damping_ratio = float(mean_log_decrement / (2.0 * np.pi))
    damping_ratio_exact = float(
        mean_log_decrement / np.sqrt((2.0 * np.pi) ** 2 + mean_log_decrement**2)
    )

    figure_path = save_task3_figure(
        file_path=file_path,
        output_dir=output_dir,
        time_values=time_values,
        raw_signal=raw_signal,
        filtered_signal=filtered_signal,
        peak_indices=peak_indices,
        selected_peak_indices=selected_peak_indices,
        selected_peak_amplitudes=selected_peak_amplitudes,
        fit_amplitudes=fit_amplitudes,
        natural_frequency_hz=natural_frequency_hz,
        damping_ratio=damping_ratio,
    )
    if show:
        plt.show()
    plt.close("all")

    return {
        "file_path": str(file_path),
        "sensor_code": TOP_RESPONSE_SENSOR,
        "axis": TOP_RESPONSE_AXIS,
        "sampling_rate_hz": float(sampling_rate),
        "natural_frequency_hz": float(natural_frequency_hz),
        "fft_peak_bin_hz": float(freqs[peak_index]),
        "fft_peak_magnitude": float(magnitudes[peak_index]),
        "detected_peak_count": int(len(peak_indices)),
        "selected_peak_count": int(len(selected_peak_indices)),
        "mean_log_decrement": mean_log_decrement,
        "damping_ratio": damping_ratio,
        "damping_ratio_exact": damping_ratio_exact,
        "envelope_r2": float(peak_window["r_squared"]),
        "figure_path": str(figure_path),
    }


def print_report(result: dict[str, float | int | str | Path]) -> None:
    print("=== Task 3: Structural Dynamic Identification ===")
    print(f"Input file: {result['file_path']}")
    print(f"Top response sensor: {result['sensor_code']} / {result['axis']}")
    print(f"Sampling rate: {result['sampling_rate_hz']:.4f} Hz")
    print(f"Natural frequency f_n: {result['natural_frequency_hz']:.6f} Hz")
    print(f"Mean log decrement delta: {result['mean_log_decrement']:.6f}")
    print(f"Damping ratio zeta (delta / 2pi): {result['damping_ratio']:.6f}")
    print(f"Damping ratio exact: {result['damping_ratio_exact']:.6f}")
    print(f"Detected peaks: {result['detected_peak_count']}")
    print(f"Selected peaks for envelope: {result['selected_peak_count']}")
    print(f"Envelope linearity R^2: {result['envelope_r2']:.6f}")
    print(f"Figure saved to: {result['figure_path']}")


def main() -> None:
    args = parse_args()
    result = analyze_free_vibration(
        file_path=Path(args.file),
        output_dir=Path(args.output_dir),
        bandwidth_hz=float(args.bandwidth),
        show=bool(args.show),
    )
    print_report(result)


if __name__ == "__main__":
    main()