#!/usr/bin/python3
import numpy as np
import matplotlib.pyplot as plt
from scipy.io import wavfile
from scipy.signal import butter, sosfiltfilt, hilbert, medfilt
from scipy.ndimage import binary_closing, binary_opening
from sklearn.cluster import KMeans

MORSE_CODE_DICT = {
    '.-': 'A', '-...': 'B', '-.-.': 'C', '-..': 'D', '.': 'E',
    '..-.': 'F', '--.': 'G', '....': 'H', '..': 'I', '.---': 'J',
    '-.-': 'K', '.-..': 'L', '--': 'M', '-.': 'N', '---': 'O',
    '.--.': 'P', '--.-': 'Q', '.-.': 'R', '...': 'S', '-': 'T',
    '..-': 'U', '...-': 'V', '.--': 'W', '-..-': 'X', '-.--': 'Y',
    '--..': 'Z', '-----': '0', '.----': '1', '..---': '2', '...--': '3',
    '....-': '4', '.....': '5', '-....': '6', '--...': '7', '---..': '8',
    '----.': '9', '.-.-.-': '.', '--..--': ',', '..--..': '?', '.----.': "'",
    '-.-.--': '!', '-..-.': '/', '-.--.': '(', '-.--.-': ')', '.-...': '&',
    '---...': ':', '-.-.-.': ';', '-...-': '=', '.-.-.': '+', '-....-': '-',
    '..--.-': '_', '.-..-.': '"', '...-..-': '$', '.--.-.': '@'
}

def decode_clean_signal(wav_path, bandwidth_hz=100):
    sample_rate, data = wavfile.read(wav_path)
    if data.ndim > 1:
        data = data[:, 0]
    
    data_float = data.astype(np.float64)
    n_samples = len(data_float)
    
    # 1. Bandpass-Filterung um die Peak-Frequenz
    fft_spec = np.fft.rfft(data_float)
    freqs = np.fft.rfftfreq(n_samples, d=1/sample_rate)
    cw_mask = (freqs >= 300) & (freqs <= 1200)
    cw_center = freqs[cw_mask][np.argmax(np.abs(fft_spec[cw_mask]))]
    
    sos = butter(N=4, Wn=[cw_center - bandwidth_hz/2, cw_center + bandwidth_hz/2], btype='bandpass', fs=sample_rate, output='sos')
    filtered = sosfiltfilt(sos, data_float)
    
    # 2. Einhüllende erzeugen & glätten
    envelope = np.abs(hilbert(filtered))
    smooth_samples = int(sample_rate * 0.025) # 25ms Glättung
    if smooth_samples % 2 == 0: smooth_samples += 1
    smooth_env = medfilt(envelope, kernel_size=smooth_samples)
    
    # 3. Absolute Schwelle setzen (30% des Maximalpegels)
    # Bei deinen Peaks (~25000) liegt die Schwelle damit perfekt bei ~7500
    threshold = np.max(smooth_env) * 0.30
    digital = (smooth_env > threshold).astype(int)
    
    # 4. Entprellen / Verkleben (Lücken < 20ms schließen, Impulse < 15ms verwerfen)
    close_samples = int(sample_rate * 0.020)
    digital = binary_closing(digital, structure=np.ones(close_samples)).astype(int)
    open_samples = int(sample_rate * 0.015)
    digital_clean = binary_opening(digital, structure=np.ones(open_samples)).astype(int)
    
    # 5. Pulse und Pausen messen
    diff = np.diff(np.concatenate(([0], digital_clean, [0])))
    rises = np.where(diff == 1)[0]
    falls = np.where(diff == -1)[0]
    
    if len(rises) == 0:
        print("Kein Signal erkannt.")
        return ""

    pulse_durs = [(f - r) / sample_rate for r, f in zip(rises, falls)]
    pause_durs = [(rises[i+1] - falls[i]) / sample_rate for i in range(len(falls)-1)]
    
    # 6. Dit/Dah Schwelle ermitteln
    pulse_arr = np.array(pulse_durs).reshape(-1, 1)
    kmeans = KMeans(n_clusters=2, random_state=0, n_init=10).fit(pulse_arr)
    centers = np.sort(kmeans.cluster_centers_.flatten())
    dit_dah_thresh = (centers[0] + centers[1]) / 2
    
    # Dit-Dauer als Basis für Pausenzeit nutzen (1 Dit = T, Zeichenpause ~ 3T, Wortpause ~ 7T)
    dit_duration = centers[0]
    char_thresh = dit_duration * 2.0
    word_thresh = dit_duration * 4.5

    # Dit-Dauer entspricht dem Mittelwert des ersten Clusters (Zentrum 0)
    dit_duration_sec = centers[0]
    
    # WPM Berechnen (PARIS-Standard)
    wpm = 1.2 / dit_duration_sec
    
    # Alternativ: Geschwindigkeit in Zeichen pro Minute (CPM)
    cpm = wpm * 5

    print(f"[+] Ermittelte Dit-Dauer: {dit_duration_sec * 1000:.1f} ms")
    print(f"[+] Morsegeschwindigkeit: {wpm:.1f} WPM ({cpm:.0f} BpM / CPM)")
    
    # 7. Dekodierung
    decoded_text = ""
    current_char = ""
    
    for i in range(len(pulse_durs)):
        is_dit = pulse_durs[i] < dit_dah_thresh
        current_char += "." if is_dit else "-"
        
        if i < len(pause_durs):
            p = pause_durs[i]
            if p >= word_thresh:
                decoded_text += MORSE_CODE_DICT.get(current_char, "?") + " "
                current_char = ""
            elif p >= char_thresh:
                decoded_text += MORSE_CODE_DICT.get(current_char, "?")
                current_char = ""
                
    if current_char:
        decoded_text += MORSE_CODE_DICT.get(current_char, "?")
        
    print("\n================ DECODIERTER TEXT ================")
    print(decoded_text.strip())
    print("==================================================\n")
    return decoded_text.strip()

# Aufruf:
decode_clean_signal("filtered_morse.wav")