import sounddevice as sd
import numpy as np
import pyworld as pw
import blessed
import sys
import queue
import time

# --- 定数 ---
SAMPLE_RATE = 16000
CHANNELS = 1
BLOCK_SIZE = 4096
NOISE_GATE_THRESHOLD = 0.01
OUTPUT_GAIN = 0.2
FRAME_RATE = 20.0

# --- ターミナルオブジェクト ---
term = blessed.Terminal()

# --- グローバル変数 ---
pitch_ratio = 1.0
formant_ratio = 1.0
data_queue = queue.Queue()

def get_color_from_amplitude(amplitude):
    """振幅に応じた色を返す"""
    normalized = min(1.0, amplitude * 4)
    if normalized < 0.3: return term.green
    elif normalized < 0.6: return term.yellow
    else: return term.red

def draw_waveform(data):
    """波形の描画"""
    global pitch_ratio, formant_ratio
    
    header = (f"PYWORLD VOX | Pitch: {pitch_ratio:.2f} | "
              f"Formant: {formant_ratio:.2f}")
    print(term.move_y(0) + term.center(term.bold(header)))

    num_cols = term.width
    if data is None or data.size == 0: return
    
    step = len(data) // num_cols
    if step == 0: step = 1
    sliced_data = data[::step]
    center_y = term.height // 2

    for x, value in enumerate(sliced_data):
        if x >= num_cols: break
        
        amplitude = int(value * (center_y - 2))
        color = get_color_from_amplitude(abs(value))

        for y in range(abs(amplitude)):
            pos_y = center_y - np.sign(amplitude) * y
            if 1 < pos_y < term.height:
                with term.location(x, int(pos_y)):
                    print(color('█'))

def callback(indata, outdata, frames, time, status):
    """pyworldを使った音声処理"""
    global pitch_ratio, formant_ratio, data_queue

    if status: print(status, file=sys.stderr)

    try:
        x = indata.flatten().astype(np.float64)

        if np.abs(x).max() < NOISE_GATE_THRESHOLD:
            outdata.fill(0)
            if data_queue.qsize() < 2: data_queue.put(np.zeros_like(x))
            return

        # WORLDによる分析
        f0, t = pw.harvest(x, SAMPLE_RATE)
        sp = pw.cheaptrick(x, f0, t, SAMPLE_RATE)
        ap = pw.d4c(x, f0, t, SAMPLE_RATE)

        # パラメータの操作
        modified_f0 = f0 * pitch_ratio
        
        # --- 新しいフォルマントシフト処理 ---
        if formant_ratio != 1.0:
            fftlen = sp.shape[1]
            # 各フレームに補間を適用
            def shift_spectrum(spectrum):
                new_axis = np.linspace(0, fftlen, int(fftlen / formant_ratio))
                orig_axis = np.arange(fftlen)
                shifted_spectrum = np.interp(new_axis, orig_axis, spectrum)
                if len(shifted_spectrum) < fftlen:
                    shifted_spectrum = np.pad(shifted_spectrum, (0, fftlen - len(shifted_spectrum)), 'constant')
                return shifted_spectrum[:fftlen]
            modified_sp = np.apply_along_axis(shift_spectrum, 1, sp)
        else:
            modified_sp = sp

        # WORLDによる再合成
        y = pw.synthesize(modified_f0, modified_sp, ap, SAMPLE_RATE)
        
        y = y[:len(x)]
        reshaped_data = y.reshape(outdata.shape).astype(np.float32)

        # 位相の反転とゲインの適用
        outdata[:] = -reshaped_data * OUTPUT_GAIN
        
        if data_queue.qsize() < 2:
            data_queue.put(x)

    except Exception as e:
        print(f"Error in callback: {e}", file=sys.stderr)
        outdata.fill(0)
        if data_queue.qsize() < 2:
            data_queue.put(np.zeros(BLOCK_SIZE, dtype=np.float32))

def main():
    """メインループ"""
    global pitch_ratio, formant_ratio
    try:
        with sd.Stream(samplerate=SAMPLE_RATE, blocksize=BLOCK_SIZE, channels=CHANNELS, callback=callback):
            with term.fullscreen(), term.cbreak(), term.hidden_cursor():
                running = True
                last_data = np.zeros(BLOCK_SIZE, dtype=np.float32)
                
                while running:
                    try:
                        while not data_queue.empty():
                           last_data = data_queue.get_nowait()
                    except queue.Empty:
                        pass
                    
                    print(term.clear, end='')
                    draw_waveform(last_data)
                    
                    with term.location(0, term.height - 1):
                        controls = ("Pitch [h/l] | Formant [k/j] | [Enter] Exit")
                        print(term.center(term.bold_white_on_blue(controls)))
                    
                    sys.stdout.flush()

                    key = term.inkey(timeout=1/FRAME_RATE)
                    if key:
                        if key == 'h': pitch_ratio = min(3.0, pitch_ratio + 0.05)
                        elif key == 'l': pitch_ratio = max(0.1, pitch_ratio - 0.05)
                        elif key == 'k': formant_ratio = min(1.5, formant_ratio + 0.05)
                        elif key == 'j': formant_ratio = max(0.5, formant_ratio - 0.05)
                        elif key.name == "KEY_ENTER" or key == '\n':
                            running = False
    except Exception as e:
        print(term.normal)
        print(f"\nAn error occurred: {e}")
    finally:
        print("\nVoice changer stopped.")

if __name__ == "__main__":
    main()