#!/usr/bin/env python3
import os
import concurrent.futures
import datetime
import numpy as np
import matplotlib.pyplot as plt
import soundfile as sf
import scipy.signal
from scipy.fft import fft, fftfreq
import shutil

CHUNK_DIR = "chunks"
PROCESSED_CHUNK_DIR = "chunks/processed"
EVENT_DIR = "events"
SAMPLE_SECONDS = 1
TOLERANCE = 2
OVERTONE_TOLERANCE = TOLERANCE * 2
THRESHOLD_BASE = 0.
THRESHOLD_OCT = THRESHOLD_BASE / 10
CLIP_PADDING_BEFORE = 1
CLIP_PADDING_AFTER = 6
TARGET_FREQ = 211
OVERTONE_FREQ = TARGET_FREQ * 2
NFFT = 32768
SKIP_SECONDS = 10

def process_chunk(filename):
    input_path = os.path.join(CHUNK_DIR, filename)
    print(f"🔍 Verarbeite {input_path}...")

    # Frequenzanalyse und Event-Erkennung
    data, samplerate = sf.read(input_path)
    if data.ndim > 1:
        data = data[:, 0]  # nur Kanal 1

    chunk_samples = int(SAMPLE_SECONDS * samplerate)
    skip_samples = int(SKIP_SECONDS * samplerate)
    padding_before = int(CLIP_PADDING_BEFORE * samplerate)
    padding_after = int(CLIP_PADDING_AFTER * samplerate)

    chunk_start_str = os.path.splitext(filename)[0]
    chunk_start_dt = datetime.datetime.strptime(chunk_start_str, "%Y%m%d-%H%M%S")

    i = 0
    last_event = -skip_samples
    while i + chunk_samples <= len(data):
        clip = data[i:i+chunk_samples]

        freqs, times, Sxx = scipy.signal.spectrogram(clip, samplerate, nperseg=NFFT)
        idx_base = np.where((freqs >= TARGET_FREQ - TOLERANCE) & (freqs <= TARGET_FREQ + TOLERANCE))[0]
        idx_oct = np.where((freqs >= OVERTONE_FREQ - OVERTONE_TOLERANCE) & (freqs <= OVERTONE_FREQ + OVERTONE_TOLERANCE))[0]
        if len(idx_base) == 0 or len(idx_oct) == 0:
            return False
        base_energy = np.mean(Sxx[idx_base])
        oct_energy = np.mean(Sxx[idx_oct])
        total_energy = np.mean(Sxx, axis=0).max()

        fft_vals = np.abs(fft(clip))
        freqs = fftfreq(len(clip), 1/samplerate)
        peak_freq = freqs[np.argmax(fft_vals)]
        is_peak_near_target = TARGET_FREQ - TOLERANCE <= peak_freq <= TARGET_FREQ + TOLERANCE

        event_detected = is_peak_near_target and base_energy > THRESHOLD_BASE * total_energy and oct_energy > THRESHOLD_OCT * total_energy

        if i - last_event >= skip_samples and event_detected:
            clip_start = max(0, i - padding_before)
            clip_end = min(len(data), i + chunk_samples + padding_after)
            clip = data[clip_start:clip_end]

            event_offset = (i - padding_before) / samplerate
            event_time_dt = chunk_start_dt + datetime.timedelta(seconds=event_offset)
            event_time = event_time_dt.strftime("%Y%m%d-%H%M%S")

            base_name = os.path.splitext(filename)[0]
            flac_out = os.path.join(EVENT_DIR, f"{base_name}_{event_time}.flac")
            png_out = os.path.join(EVENT_DIR, f"{base_name}_{event_time}.png")
            sf.write(flac_out, clip, samplerate, format='FLAC')

            plt.figure()
            plt.specgram(clip, Fs=samplerate, NFFT=NFFT, noverlap=NFFT//2, cmap='inferno', vmin=-90, vmax=-20)
            plt.title(f"Spectrogram: {base_name}_{event_time}")
            plt.xlabel("Time (s)")
            plt.ylabel("Frequency (Hz)")
            plt.colorbar(label="dB")
            plt.savefig(png_out)
            plt.close()

            print(f"Event: {event_time} peak_freq: {int(peak_freq)} base_energy: {int(base_energy)} oct_energy: {int(oct_energy)} total_energy: {int(total_energy)}")
            last_event = i
            i += skip_samples
        else:
            i += chunk_samples

    # Datei verschieben
    output_path = os.path.join(PROCESSED_CHUNK_DIR, filename)
    #shutil.move(input_path, output_path)
    print(f"✅ Verschoben nach {output_path}")

def main():
    os.makedirs(EVENT_DIR, exist_ok=True)
    os.makedirs(PROCESSED_CHUNK_DIR, exist_ok=True)

    for file in os.listdir(CHUNK_DIR):
        if file.endswith(".flac"):
            process_chunk(file)

    # with concurrent.futures.ProcessPoolExecutor() as executor:
    #     files = [f for f in os.listdir(CHUNK_DIR) if f.endswith(".flac")]
    #     executor.map(process_chunk, files)

if __name__ == "__main__":
    main()