from __future__ import annotations

from pathlib import Path
import re

import numpy as np
import pandas as pd
from scipy.signal import butter, filtfilt, find_peaks


TOP_RESPONSE_SENSOR = "WSMS00007"
TOP_RESPONSE_AXIS = "value3"
BASE_REFERENCE_SENSOR = "WSMS00012"
BASE_REFERENCE_AXIS = "value1"
TIME_COLUMN = "time"
CODE_COLUMN = "code"


def ensure_parent_dir(path: Path) -> Path:
    path = Path(path).resolve()
    path.mkdir(parents=True, exist_ok=True)
    return path


def parse_frequency_from_filename(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:
    time_values = np.asarray(time_values, dtype=np.float64).reshape(-1)
    if time_values.size < 2:
        raise ValueError("At least two time samples are required to estimate sampling rate.")
    dt = np.diff(time_values)
    dt = dt[np.isfinite(dt)]
    dt = dt[dt > 0.0]
    if dt.size == 0:
        raise ValueError("Failed to infer a positive sampling interval from the CSV time column.")
    return float(1.0 / np.median(dt))


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 read_sensor_signal(
    file_path: Path,
    sensor_code: str = TOP_RESPONSE_SENSOR,
    value_column: str = TOP_RESPONSE_AXIS,
) -> tuple[np.ndarray, np.ndarray]:
    df = pd.read_csv(file_path)
    sensor_df = df.loc[df[CODE_COLUMN] == sensor_code, [TIME_COLUMN, value_column]].copy()
    if sensor_df.empty:
        raise ValueError(f"Sensor {sensor_code}/{value_column} was not found in {file_path.name}.")

    sensor_df = sensor_df.sort_values(TIME_COLUMN)
    sensor_df = sensor_df.drop_duplicates(subset=TIME_COLUMN, keep="first")
    sensor_df[TIME_COLUMN] = sensor_df[TIME_COLUMN].astype("float64")
    sensor_df[value_column] = sensor_df[value_column].astype("float64")

    time_values = sensor_df[TIME_COLUMN].to_numpy(dtype=np.float64)
    signal = sensor_df[value_column].to_numpy(dtype=np.float64)
    if time_values.size < 8:
        raise ValueError(f"Too few usable samples in {file_path.name}.")
    return time_values, signal


def get_middle_segment(
    time_values: np.ndarray,
    signal: np.ndarray,
    start_ratio: float = 0.20,
    end_ratio: float = 0.80,
    steady_window_ratio: float = 0.25,
    stride_ratio: float = 0.05,
    min_segment_length: int = 512,
    subwindow_count: int = 4,
) -> tuple[np.ndarray, np.ndarray]:
    signal = np.asarray(signal, dtype=np.float64).reshape(-1)
    time_values = np.asarray(time_values, dtype=np.float64).reshape(-1)
    if signal.size != time_values.size:
        raise ValueError("Time and signal arrays must have the same length.")

    length = signal.size
    search_start = int(length * start_ratio)
    search_end = int(length * 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 < min_segment_length:
        center = length // 2
        half = min_segment_length // 2
        search_start = max(0, center - half)
        search_end = min(length, search_start + min_segment_length)
        search_start = max(0, search_end - min_segment_length)

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

    stride = max(1, int(length * stride_ratio))
    candidate_signal = signal[search_start:search_end]
    candidate_rms = calculate_rms(candidate_signal)
    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_signal = signal[window_start:window_end]
        split_windows = np.array_split(window_signal, 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(), 1e-8))
        peak_cv = float(split_peaks.std() / max(split_peaks.mean(), 1e-8))
        window_rms = calculate_rms(window_signal)
        score = rms_cv + 0.35 * peak_cv - 0.05 * (window_rms / max(candidate_rms, 1e-8))
        if score < best_score:
            best_score = score
            best_slice = slice(window_start, window_end)

    return time_values[best_slice], signal[best_slice]


def compute_windowed_spectrum(
    signal: np.ndarray,
    sampling_rate: float,
    zero_padding_factor: int = 8,
) -> tuple[np.ndarray, np.ndarray]:
    signal = np.asarray(signal, dtype=np.float64).reshape(-1)
    if signal.size < 4:
        raise ValueError("Signal is too short for FFT analysis.")

    centered = signal - np.mean(signal)
    window = np.hanning(centered.size)
    fft_size = int(2 ** np.ceil(np.log2(max(centered.size * zero_padding_factor, centered.size))))
    fft_values = np.fft.rfft(centered * window, n=fft_size)
    freqs = np.fft.rfftfreq(fft_size, d=1.0 / sampling_rate)
    magnitudes = 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 = float(magnitudes[peak_index - 1])
    beta = float(magnitudes[peak_index])
    gamma = float(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 dominant_frequency_in_band(
    signal: np.ndarray,
    sampling_rate: float,
    min_hz: float = 0.1,
    max_hz: float = 5.0,
    zero_padding_factor: int = 8,
) -> tuple[float, np.ndarray, np.ndarray, int]:
    freqs, magnitudes = compute_windowed_spectrum(
        signal=signal,
        sampling_rate=sampling_rate,
        zero_padding_factor=zero_padding_factor,
    )
    band_mask = (freqs >= min_hz) & (freqs <= max_hz)
    if not np.any(band_mask):
        raise ValueError(f"No FFT bins fall inside {min_hz:.2f} Hz to {max_hz:.2f} Hz.")

    band_indices = np.flatnonzero(band_mask)
    peak_index = int(band_indices[np.argmax(magnitudes[band_mask])])
    peak_frequency = parabolic_peak_frequency(freqs, magnitudes, peak_index)
    return peak_frequency, freqs, magnitudes, peak_index


def bandpass_filter(
    signal: np.ndarray,
    sampling_rate: float,
    center_hz: float,
    bandwidth_hz: float = 0.3,
    order: int = 4,
) -> np.ndarray:
    nyquist = 0.5 * sampling_rate
    low_hz = max(0.05, center_hz - bandwidth_hz)
    high_hz = min(center_hz + bandwidth_hz, nyquist * 0.98)
    if not low_hz < high_hz:
        raise ValueError("Invalid band-pass range. Check the estimated natural frequency and sampling rate.")

    b, a = butter(order, [low_hz / nyquist, high_hz / nyquist], btype="bandpass")
    return filtfilt(b, a, np.asarray(signal, dtype=np.float64))


def detect_decay_peaks(filtered_signal: np.ndarray, sampling_rate: float, frequency_hz: float) -> np.ndarray:
    filtered_signal = np.asarray(filtered_signal, dtype=np.float64).reshape(-1)
    samples_per_cycle = max(int(round(sampling_rate / max(frequency_hz, 1e-6))), 3)
    min_peak_distance = max(1, int(0.65 * samples_per_cycle))
    peaks, _ = find_peaks(filtered_signal, distance=min_peak_distance)
    if peaks.size < 4:
        peaks, _ = find_peaks(np.abs(filtered_signal), distance=max(1, min_peak_distance // 2))
    return peaks.astype(int)


def select_decay_peak_window(
    filtered_signal: np.ndarray,
    peak_indices: np.ndarray,
    sampling_rate: float,
    min_peaks: int = 5,
    max_peaks: int = 12,
) -> dict[str, np.ndarray | float | int]:
    filtered_signal = np.asarray(filtered_signal, dtype=np.float64).reshape(-1)
    peak_indices = np.asarray(peak_indices, dtype=int).reshape(-1)
    if peak_indices.size < min_peaks:
        raise ValueError("Too few peaks were detected for log-decrement damping estimation.")

    amplitudes = np.abs(filtered_signal[peak_indices])
    tail_start = int(filtered_signal.size * 0.85)
    noise_slice = filtered_signal[tail_start:] if tail_start < filtered_signal.size else filtered_signal
    noise_floor = max(float(np.median(np.abs(noise_slice))), 1e-8)
    useful_mask = amplitudes > max(3.0 * noise_floor, 0.05 * float(amplitudes.max()))
    useful_indices = peak_indices[useful_mask]
    useful_amplitudes = amplitudes[useful_mask]
    if useful_indices.size < min_peaks:
        useful_indices = peak_indices
        useful_amplitudes = amplitudes

    best_window: dict[str, np.ndarray | float | int] | None = None
    best_score = -float("inf")
    max_window_size = min(max_peaks, useful_indices.size)

    for window_size in range(min_peaks, max_window_size + 1):
        for start in range(1, useful_indices.size - window_size + 1):
            stop = start + window_size
            idx_window = useful_indices[start:stop]
            amp_window = useful_amplitudes[start:stop]
            if np.any(amp_window <= 0.0):
                continue

            time_window = idx_window.astype(np.float64) / sampling_rate
            log_amp = np.log(amp_window)
            slope, intercept = np.polyfit(time_window, log_amp, 1)
            if slope >= 0.0:
                continue

            fit_values = slope * time_window + intercept
            ss_res = float(np.sum((log_amp - fit_values) ** 2))
            ss_tot = float(np.sum((log_amp - np.mean(log_amp)) ** 2))
            r_squared = 1.0 - ss_res / max(ss_tot, 1e-12)
            amplitude_ratio = float(amp_window[0] / max(amp_window[-1], 1e-12))
            score = r_squared + 0.02 * np.log(max(amplitude_ratio, 1.0))
            if score > best_score:
                best_score = score
                best_window = {
                    "peak_indices": idx_window,
                    "peak_amplitudes": amp_window,
                    "time_window": time_window,
                    "fit_log_amplitudes": fit_values,
                    "fit_amplitudes": np.exp(fit_values),
                    "r_squared": r_squared,
                    "noise_floor": noise_floor,
                }

    if best_window is None:
        raise ValueError("Failed to isolate a clean decay envelope segment for damping estimation.")
    return best_window