# voice_compressor.py
import tkinter as tk
from tkinter import filedialog, messagebox
import numpy as np
import sounddevice as sd
import threading
import time
import array
import struct

# Настройки по умолчанию
DEFAULT_DURATION = 4.0
SAMPLE_RATE_IN = 44100
SAMPLE_RATE_OUT = 8000
BIT_DEPTH = 4  # bits per sample
MAX_DURATION = 4.0  # seconds
OUTPUT_RAW = "compressed_speech.raw"
OUTPUT_INO = "PlaybackSketch.ino"

class VoiceCompressorApp:
    def __init__(self, root):
        self.root = root
        self.root.title("🔊 Речевой компрессор для Arduino")
        self.root.geometry("700x500")

        self.audio_data = None  # сырые 8-битные данные
        self.compressed = None  # 4-битный массив

        # Интерфейс
        self.label = tk.Label(root, text="Запись -> 4-битная упаковка -> Arduino", font=("Arial", 14))
        self.label.pack(pady=10)

        # Поле ввода длительности
        self.duration_frame = tk.Frame(root)
        self.duration_frame.pack(pady=5)
        tk.Label(self.duration_frame, text="Длительность (сек):").pack(side=tk.LEFT)
        self.duration_var = tk.StringVar(value=str(DEFAULT_DURATION))
        self.duration_entry = tk.Entry(self.duration_frame, textvariable=self.duration_var, width=5)
        self.duration_entry.pack(side=tk.LEFT, padx=5)

        # Кнопки
        self.btn_frame = tk.Frame(root)
        self.btn_frame.pack(pady=10)

        self.btn_record = tk.Button(self.btn_frame, text="🎤 Start", command=self.start_recording,
                                    bg="green", fg="white", font=("Arial", 12), width=12)
        self.btn_record.grid(row=0, column=0, padx=5)

        self.btn_play = tk.Button(self.btn_frame, text="▶ Play", command=self.play_compressed,
                                  state=tk.DISABLED, font=("Arial", 12), width=12)
        self.btn_play.grid(row=0, column=1, padx=5)

        self.btn_save = tk.Button(self.btn_frame, text="💾 Save", command=self.save_files,
                                  state=tk.DISABLED, font=("Arial", 12), width=12)
        self.btn_save.grid(row=0, column=2, padx=5)

        # Лог
        self.log = tk.Text(root, height=18, width=85)
        self.log.pack(pady=10)
        self.log.insert(tk.END, "Лог:\n")
        self.log.config(state=tk.DISABLED)

    def log_msg(self, msg):
        self.log.config(state=tk.NORMAL)
        self.log.insert(tk.END, msg + "\n")
        self.log.see(tk.END)
        self.log.config(state=tk.DISABLED)

    def start_recording(self):
        try:
            duration = float(self.duration_var.get())
            if not (0.5 <= duration <= MAX_DURATION):
                raise ValueError("Длительность от 0.5 до 4 сек")
        except:
            messagebox.showerror("Ошибка", "Введите число от 0.5 до 4")
            return

        self.btn_record.config(state=tk.DISABLED)
        self.status("🎙️ Запись...", "orange")
        self.log_msg(f"Запись {duration} сек...")
        threading.Thread(target=self.record_and_process, args=(duration,)).start()

    def status(self, text, color):
        if hasattr(self, 'status_label'):
            self.status_label.config(text=text, fg=color)
        else:
            self.status_label = tk.Label(self.root, text=text, fg=color)
            self.status_label.pack(pady=5)

    def record_and_process(self, duration):
        try:
            # 1. Запись
            self.log_msg("Идёт запись...")
            recording = sd.rec(int(duration * SAMPLE_RATE_IN), samplerate=SAMPLE_RATE_IN, channels=1, dtype='float32')
            sd.wait()
            audio = recording.flatten()
            self.status("⚙️ Обработка...", "orange")
            self.log_msg("Записано. Конвертируем...")

            # 2. Ресемплинг до 8000 Гц
            from scipy.signal import resample
            num_samples = int(len(audio) * SAMPLE_RATE_OUT / SAMPLE_RATE_IN)
            audio_resampled = resample(audio, num_samples)

            # 🔥 НОВОЕ: Нормализуем до максимума
            if np.max(np.abs(audio_resampled)) > 1e-6:
                audio_resampled = audio_resampled / np.max(np.abs(audio_resampled))  # [-1, 1]
                audio_resampled *= 0.95  # чуть ниже максимума, чтобы не было клиппинга

            # 3. Нормализация и квантование в 8 бит
            audio_norm = np.clip(audio_resampled, -1.0, 1.0)
            audio_8bit = ((audio_norm + 1.0) * 127.5).astype(np.uint8)  # 0..255

            # 4. Упаковка в 4 бита (по два семпла в байте)
            packed = []
            for i in range(0, len(audio_8bit), 2):
                s1 = audio_8bit[i] >> 4  # старшие 4 бита
                s2 = audio_8bit[i + 1] >> 4 if i + 1 < len(audio_8bit) else 0
                byte = (s1 << 4) | s2
                packed.append(byte)

            # 5. Сжатие RLE
            compressed_rle = self.rle_compress_4bit(packed)  # вызываем сжатие

            # Сохраняем для дальнейшей работы
            self.audio_data = audio_8bit
            self.compressed = compressed_rle  # <-- ТЕПЕРЬ СО СЖАТИЕМ!
            self.sample_rate = SAMPLE_RATE_OUT
            self.duration = duration

            self.log_msg(f"Обработано: {len(self.compressed)} байт (4 бит/семпл)")
            self.btn_play.config(state=tk.NORMAL)
            self.btn_save.config(state=tk.NORMAL)
            self.status("✅ Готово", "green")

        except Exception as e:
            self.status("❌ Ошибка", "red")
            self.log_msg(f"Ошибка: {str(e)}")
            messagebox.showerror("Ошибка", str(e))
        finally:
            self.btn_record.config(state=tk.NORMAL)

    def play_compressed(self):
        if not self.compressed:
            return

        self.status("▶ Воспроизведение...", "orange")
        self.log_msg("Воспроизведение упакованного звука...")

        # Распаковываем
        unpacked = []
        for b in self.compressed:
            s1 = (b >> 4) & 0x0F
            s2 = b & 0x0F
            unpacked.append(s1 * 16 + 8)  # обратно в 0..255
            unpacked.append(s2 * 16 + 8)

        # Преобразуем в float32 [-1, 1]
        signal = np.array(unpacked, dtype=np.float32)
        signal = (signal - 127.5) / 127.5

        # Проигрываем
        try:
            sd.play(signal, 8000)
            sd.wait()
            self.status("⏹ Готово", "green")
        except Exception as e:
            self.status("❌ Ошибка", "red")
            self.log_msg(f"Ошибка воспроизведения: {e}")

    def save_files(self):
        if not self.compressed:
            return

        # Сохраняем .raw
        with open(OUTPUT_RAW, 'wb') as f:
            f.write(self.compressed.tobytes())
        self.log_msg(f"Сохранено: {OUTPUT_RAW}")

        # Генерируем .ino
        self.generate_arduino_sketch()

        self.log_msg(f"Скетч сохранён: {OUTPUT_INO}")
        messagebox.showinfo("Успех", f"Файлы сохранены:\n{OUTPUT_RAW}\n{OUTPUT_INO}")

    def generate_arduino_sketch(self):
        data = self.compressed
        rate = self.sample_rate
        n_bytes = len(data)

        # Разбиваем массив на строки по 16 байт
        lines = []
        for i in range(0, n_bytes, 16):
            line = ", ".join(f"0x{b:02X}" for b in data[i:i+16])
            lines.append("  " + line)

        array_str = ",\n".join(lines)

        code = f'''/*
   Воспроизведение 4-битного звука на Arduino Nano
   Частота: {rate} Гц
   Длина: {self.duration:.1f} сек
   Размер: {n_bytes} байт
*/

#include <avr/pgmspace.h>

// Packed 2 samples per byte (4 bits each)
const unsigned char sound_data[] PROGMEM = {{
  {array_str}
}};

const int speakerPin = 3;
const int data_size = {n_bytes};
const int sample_rate = {rate};  // Hz

void setup() {{
  // Настройка таймера 2 на высокую частоту ШИМ (~62.5 кГц)
  TCCR2A = _BV(COM2A1) | _BV(COM2B1) | _BV(WGM21) | _BV(WGM20);
  TCCR2B = _BV(CS20);  // No prescaler
  pinMode(speakerPin, OUTPUT);

  delay(1000);
  playSound();
}}

void loop() {{
  // nothing
}}

void playSound() {{
  unsigned long next_time = micros();
  const long delay_per_sample = 1000000L / sample_rate;  // 125 мкс при 8 кГц

  for (int i = 0; i < data_size; ) {{
    byte b = pgm_read_byte(&sound_data[i]);

    if ((b & 0xF0) == 0xF0) {{
      // RLE: команда повтора
      int count = (b & 0x0F) + 3;
      byte value = pgm_read_byte(&sound_data[i + 1]) >> 4;
      i += 2;

      for (int r = 0; r < count; r++) {{
        analogWrite(speakerPin, value << 4);
        next_time += delay_per_sample;
        while (micros() < next_time);
      }}
    }} else {{
      // Обычный байт: два семпла
      byte s1 = (b >> 4) & 0x0F;
      byte s2 = b & 0x0F;

      analogWrite(speakerPin, s1 << 4);
      next_time += delay_per_sample;
      while (micros() < next_time);

      analogWrite(speakerPin, s2 << 4);
      next_time += delay_per_sample;
      while (micros() < next_time);

      i++;
    }}
  }}
  digitalWrite(speakerPin, LOW);
}}
'''

        with open(OUTPUT_INO, 'w', encoding='utf-8') as f:
            f.write(code)

    def rle_compress_4bit(self, packed_bytes):
        """Сжимает упакованные 4-bit байты с RLE"""
        result = []
        i = 0
        while i < len(packed_bytes):
            byte = packed_bytes[i]
            # Извлекаем два семпла
            s1 = (byte >> 4) & 0x0F
            s2 = byte & 0x0F

            # Проверим, сколько раз подряд идёт s1
            count1 = 1
            pos = i + 1
            val = s1
            # Попробуем продолжить цепочку
            while pos < len(packed_bytes):
                b = packed_bytes[pos]
                next_s1 = (b >> 4) & 0x0F
                if next_s1 == val:
                    count1 += 1
                    # Перейдём ко второму семплу этого байта
                    next_s2 = b & 0x0F
                    if next_s2 == val:
                        count1 += 1
                        pos += 1
                    else:
                        break
                else:
                    break

            # Если повторов ≥ 5 (экономим минимум 1 байт)
            if count1 >= 5:
                # Команда: F0 + (count - 3), затем значение
                repeat_count = count1
                cmd = 0xF0 + min(repeat_count - 3, 15)  # макс. 18 повторов
                result.append(cmd)
                result.append(val << 4)  # кладём как старшие 4 бита
                i += (repeat_count + 1) // 2  # сколько байт заменили
            else:
                result.append(byte)
                i += 1

        return array.array('B', result)

if __name__ == "__main__":
    root = tk.Tk()
    app = VoiceCompressorApp(root)
    root.mainloop()