import librosa
import librosa.display
import numpy as np
import matplotlib.pyplot as plt
import scipy.io
from scipy.signal import butter, lfilter
import pandas as pd
import math
import os
from scipy.interpolate import interp1d


extension = '.npy'
svname = '35C_10GPM_1HP_T'

# Find the indices where testn is not NaN
sr = 100000
ssr = 10000
n_fft = 1024
hop_length = 512
freq_thres = 5000.0
eps = 1e-20
N_SEGMENTS = 101


def _freq_mask(sr, n_fft, fmax):
    freqs = librosa.fft_frequencies(sr=sr, n_fft=n_fft)
    mask = (np.abs(freqs) <= fmax)
    return freqs[mask], mask

def _avg_power_over_time(seg, fmask):
    S = librosa.stft(seg, n_fft=n_fft, hop_length=hop_length)
    P = (np.abs(S[fmask, :]) ** 2)              # (Fmask, T)
    return np.maximum(np.mean(P, axis=1), eps)   # (Fmask,)

def process_bg(bg_doc):
    file_name = f"26824_1_{bg_doc}"
    file_path = f"{base_name}{file_name}{extension}"
    M = np.load(file_path)
    return M[1, :].astype(np.float64)

def process_range(mat_indices, bg_y, svname, snr_tag=None):
    # snr_tag=None -> clean; '1000' -> loads *_1000.npy
    tag_txt = '' if (snr_tag is None) else f'_{snr_tag}'
    save_folder_path = f'/Path/to/folder/for/hydrophone spectrum/{svname}'
    os.makedirs(save_folder_path, exist_ok=True)

    freqs_m, fmask = _freq_mask(sr, n_fft, freq_thres)
    max_f = np.max(np.abs(freqs_m))
    y_axis = (freqs_m / max_f).reshape(-1, 1)

    # background mean power per frequency bin (0–5 kHz)
    bg_seg = bg_y[:ssr] if bg_y.size >= ssr else bg_y
    P_bg = _avg_power_over_time(bg_seg, fmask)  # (Fmask,)

    for idx in mat_indices:
        file_path = f"26824_1_{idx}{tag_txt}{extension}"
        if not os.path.exists(file_path):
            print("Missing:", file_path); continue

        M = np.load(file_path)
        y = M[1, :].astype(np.float64)
        total_chunks = y.size // ssr

        for nf in range(N_SEGMENTS):
            if nf < total_chunks and (nf*ssr + ssr) <= y.size and ssr >= n_fft:
                seg = y[nf*ssr : nf*ssr + ssr]
                if seg.size >= n_fft:
                    P_sig = _avg_power_over_time(seg, fmask)      # (Fmask,)
                    R = P_sig / (P_bg + eps)
                    Srel_db_mean = 10.0 * np.log10(np.maximum(R, eps))
                else:
                    Srel_db_mean = np.zeros_like(freqs_m)
            else:
                Srel_db_mean = np.zeros_like(freqs_m)

            # plot (no titles/labels), keep original name pattern
            plt.figure(figsize=(12, 12))
            plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
            H = np.tile(Srel_db_mean[:, None], (1, 20))
            x_axis = np.arange(20) * (hop_length / sr)
            plt.pcolormesh(x_axis, y_axis, H, shading='gouraud', cmap='gray', vmin=-30, vmax=30)
            plt.ylim(0, 1); plt.xlim(0, 0.09)
            out_png = os.path.join(save_folder_path, f'STFT_rmbg_ML_{idx}_{nf}.png')
            plt.savefig(out_png, dpi=150); plt.clf(); plt.close()


bg_2025_012 = [2273,2279]
background = [bg_2025_012]
tar_2025_012 = [[list(range(2275,2279))],[list(range(2281,2286))]]
target = [tar_2025_012]

for k in range(0,1):
    for j in range(0,1): #(0,1) for prediction [35C_10GPM (IPF condition in Figure 7)]; (1,2) for training
        bg_y = process_bg(background[k][j])
        process_range(target[k][j][0], bg_y, svname, snr_tag=None)    # loads 26824_1_<SSSidx>.npy







