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
from scipy.signal import find_peaks

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,
    dominant_frequency_in_band,
    ensure_parent_dir,
    estimate_sampling_rate,
    get_middle_segment,
    parse_frequency_from_filename,
    read_sensor_signal,
)


def parse_args() -> argparse.Namespace:
    project_root = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(description="Predict the harmonic input frequency from top-response steady-state data only.")
    parser.add_argument(
        "--file",
        type=str,
        default=str(project_root / "downloads" / "Non_TMD" / "val" / "harmonic_5mm_0.75Hz.csv"),
        help="Path to a forced-vibration CSV such as data from Non_TMD val/test.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=str(project_root / "evaluation_outputs" / "task4_predict_freq"),
        help="Directory for plots and optional exported diagnostics.",
    )
    parser.add_argument(
        "--show",
        action="store_true",
        help="Display the generated figure in addition to saving it.",
    )
    return parser.parse_args()


def save_task4_figure(
    file_path: Path,
    output_dir: Path,
    time_segment: np.ndarray,
    signal_segment: np.ndarray,
    freqs: np.ndarray,
    magnitudes: np.ndarray,
    predicted_frequency_hz: float,
    true_frequency_hz: float | None,
) -> Path:
    output_dir = ensure_parent_dir(output_dir)
    fig, axes = plt.subplots(2, 1, figsize=(12, 8))
    fig.suptitle(f"Task 4 | Forced-Vibration Frequency Prediction | {file_path.name}")

    axes[0].plot(time_segment, signal_segment, color="tab:blue", linewidth=1.0)
    axes[0].set_title(f"Steady middle segment from {TOP_RESPONSE_SENSOR} / {TOP_RESPONSE_AXIS}")
    axes[0].set_xlabel("Time (s)")
    axes[0].set_ylabel("Acceleration")
    axes[0].grid(True, alpha=0.25)

    band_mask = (freqs >= 0.1) & (freqs <= 5.0)
    axes[1].plot(freqs[band_mask], magnitudes[band_mask], color="tab:purple", linewidth=1.1, label="Windowed FFT")
    axes[1].axvline(
        predicted_frequency_hz,
        color="tab:red",
        linestyle="--",
        linewidth=1.6,
        label=f"Predicted = {predicted_frequency_hz:.5f} Hz",
    )
    if true_frequency_hz is not None:
        axes[1].axvline(
            true_frequency_hz,
            color="tab:green",
            linestyle=":",
            linewidth=1.6,
            label=f"Filename truth = {true_frequency_hz:.5f} Hz",
        )
    axes[1].set_title("High-resolution spectrum from top steady-state response")
    axes[1].set_xlabel("Frequency (Hz)")
    axes[1].set_ylabel("Magnitude")
    axes[1].grid(True, alpha=0.25)
    axes[1].legend()

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


def pick_fundamental_peak_index(
    freqs: np.ndarray,
    magnitudes: np.ndarray,
    min_hz: float = 0.1,
    max_hz: float = 5.0,
    relative_threshold: float = 0.20,
    relative_prominence: float = 0.08,
) -> int:
    band_mask = (freqs >= min_hz) & (freqs <= max_hz)
    band_indices = np.where(band_mask)[0]
    if band_indices.size == 0:
        raise ValueError("No FFT bins fall inside the requested frequency band.")

    band_magnitudes = magnitudes[band_mask]
    max_magnitude = float(np.max(band_magnitudes))
    height_threshold = relative_threshold * max_magnitude
    prominence_threshold = relative_prominence * max_magnitude
    local_peaks, _ = find_peaks(
        band_magnitudes,
        height=height_threshold,
        prominence=prominence_threshold,
    )

    if local_peaks.size == 0:
        # Fallback for very smooth spectra: choose the strongest bin in band.
        return int(band_indices[int(np.argmax(band_magnitudes))])

    # Fundamental-first rule: choose the left-most valid peak.
    return int(band_indices[int(local_peaks[0])])


def analyze_forced_vibration(
    file_path: Path,
    output_dir: Path,
    show: bool = False,
) -> dict[str, float | int | str]:
    file_path = Path(file_path).resolve()
    time_values, raw_signal = read_sensor_signal(file_path, sensor_code=TOP_RESPONSE_SENSOR, value_column=TOP_RESPONSE_AXIS)
    time_segment, signal_segment = get_middle_segment(time_values, raw_signal)
    sampling_rate = estimate_sampling_rate(time_segment)
    _, freqs, magnitudes, _ = dominant_frequency_in_band(
        signal_segment,
        sampling_rate=sampling_rate,
        min_hz=0.1,
        max_hz=5.0,
        zero_padding_factor=32,
    )
    peak_index = pick_fundamental_peak_index(
        freqs=freqs,
        magnitudes=magnitudes,
        min_hz=0.1,
        max_hz=5.0,
        relative_threshold=0.20,
        relative_prominence=0.08,
    )
    predicted_frequency_hz = float(freqs[peak_index])
    true_frequency_hz = parse_frequency_from_filename(file_path.name)
    absolute_error_hz = None if true_frequency_hz is None else abs(predicted_frequency_hz - true_frequency_hz)

    figure_path = save_task4_figure(
        file_path=file_path,
        output_dir=output_dir,
        time_segment=time_segment,
        signal_segment=signal_segment,
        freqs=freqs,
        magnitudes=magnitudes,
        predicted_frequency_hz=predicted_frequency_hz,
        true_frequency_hz=true_frequency_hz,
    )
    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),
        "segment_sample_count": int(signal_segment.size),
        "segment_start_time": float(time_segment[0]),
        "segment_end_time": float(time_segment[-1]),
        "predicted_frequency_hz": float(predicted_frequency_hz),
        "fft_peak_bin_hz": float(freqs[peak_index]),
        "fft_peak_magnitude": float(magnitudes[peak_index]),
        "true_frequency_hz": None if true_frequency_hz is None else float(true_frequency_hz),
        "absolute_error_hz": None if absolute_error_hz is None else float(absolute_error_hz),
        "figure_path": str(figure_path),
        "note": "Prediction uses only the top response sensor WSMS00007/value3. Base sensor WSMS00012/value1 is not used to infer the answer.",
    }


def print_report(result: dict[str, float | int | str]) -> None:
    print("=== Task 4: Excitation Frequency Reverse Prediction ===")
    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"Steady segment samples: {result['segment_sample_count']}")
    print(f"Steady segment time range: {result['segment_start_time']:.4f} s to {result['segment_end_time']:.4f} s")
    print(f"Predicted excitation frequency: {result['predicted_frequency_hz']:.6f} Hz")
    if result["true_frequency_hz"] is not None:
        print(f"Filename frequency truth: {result['true_frequency_hz']:.6f} Hz")
        print(f"Absolute error: {result['absolute_error_hz']:.6f} Hz")
    else:
        print("Filename frequency truth: unavailable")
    print(f"Figure saved to: {result['figure_path']}")
    print(result["note"])


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


if __name__ == "__main__":
    main()