# decoder.py
import time
import numpy as np
import librosa
import soundfile as sf
from PIL import Image
from transforms import ms_to_lr
from phase_generator import (griffin_lim, griffin_lim_fast, 
                             griffin_lim_multi_scale, griffin_lim_stereo_parallel)


def _pad_spectrum(data: np.ndarray, low_pad: int, high_pad: int) -> np.ndarray:
    if low_pad > 0:
        data = np.vstack([np.zeros((low_pad, data.shape[1]), dtype=data.dtype), data])
    if high_pad > 0:
        data = np.vstack([data, np.zeros((high_pad, data.shape[1]), dtype=data.dtype)])
    return data


def image_to_audio(image_path: str, output_wav_path: str, metadata: dict) -> np.ndarray:
    n_freqs, n_frames = metadata['shape']
    n_freqs_original = metadata['shape_original'][0]
    low_cut_bin = metadata.get('low_cut_bin', 0)
    high_cut_bins_removed = metadata.get('high_cut_bins_removed', 0)
    flipped = metadata.get('flipped', False)
    
    sr = metadata['sr']
    n_fft = metadata['n_fft']
    hop_length = metadata['hop_length']
    original_length = metadata['original_length']
    mag_min = metadata.get('mag_min', -120)
    ref_mid = metadata.get('ref_mid', 1.0)
    ref_side = metadata.get('ref_side', 1.0)

    img = Image.open(image_path).convert('RGB')
    
    target_size = (n_frames, n_freqs)
    if img.size != target_size:
        img = img.resize(target_size, Image.LANCZOS)
    
    arr = np.array(img, dtype=np.float32)
    if flipped:
        arr = np.flipud(arr)

    mag_mid_norm = arr[:, :, 1] / 255.0
    mag_side_norm = arr[:, :, 2] / 255.0
    mag_mid_db = mag_mid_norm * (-mag_min) + mag_min
    mag_side_db = mag_side_norm * (-mag_min) + mag_min
    mag_mid = librosa.db_to_amplitude(mag_mid_db, ref=ref_mid)
    mag_side = librosa.db_to_amplitude(mag_side_db, ref=ref_side)

    mag_mid = _pad_spectrum(mag_mid, low_cut_bin, high_cut_bins_removed)
    mag_side = _pad_spectrum(mag_side, low_cut_bin, high_cut_bins_removed)

    t_start = time.time()
    
    # Всегда генерируем фазу
    iterations = metadata.get('phase_generate_iterations', 100)
    random_seed = metadata.get('phase_generate_random_seed', 5)
    if random_seed == -1:
        random_seed = None
    
    gl_mode = metadata.get('griffin_lim_mode', 'standard')
    parallel = metadata.get('griffin_lim_parallel', True)
    
    early_stop_threshold = metadata.get('early_stop_threshold', 0.0001)
    early_stop_patience = metadata.get('early_stop_patience', 15)
    
    if parallel and gl_mode in ('standard', 'fast', 'multi_scale'):
        print(f"\nГенерация фазы через Griffin-Lim ({gl_mode}, {iterations} итераций, параллельно)")
        
        if gl_mode == 'multi_scale':
            scale_factor = metadata.get('multi_scale_factor', 4)
            coarse_iter = metadata.get('multi_scale_coarse_iterations', -1)
            fine_iter = metadata.get('multi_scale_fine_iterations', -1)
            
            if coarse_iter == -1:
                coarse_iter = None
            if fine_iter == -1:
                fine_iter = None
        else:
            scale_factor = 4
            coarse_iter = None
            fine_iter = None
        
        phase_mid, phase_side = griffin_lim_stereo_parallel(
            mag_mid, mag_side, n_fft, hop_length,
            iterations=iterations,
            random_seed=random_seed,
            mode=gl_mode,
            scale_factor=scale_factor,
            coarse_iterations=coarse_iter,
            fine_iterations=fine_iter,
            early_stop_threshold=early_stop_threshold,
            early_stop_patience=early_stop_patience,
            num_workers=2,
            verbose=True
        )
    else:
        print(f"\nГенерация фазы через Griffin-Lim ({gl_mode}, {iterations} итераций)...")
        
        if gl_mode == 'fast':
            phase_mid = griffin_lim_fast(mag_mid, n_fft, hop_length, 
                                         iterations=iterations, 
                                         random_seed=random_seed,
                                         early_stop_threshold=early_stop_threshold,
                                         early_stop_patience=early_stop_patience)
            phase_side = griffin_lim_fast(mag_side, n_fft, hop_length, 
                                          iterations=iterations, 
                                          random_seed=random_seed + 1 if random_seed else None,
                                          early_stop_threshold=early_stop_threshold,
                                          early_stop_patience=early_stop_patience)
        elif gl_mode == 'multi_scale':
            scale_factor = metadata.get('multi_scale_factor', 4)
            coarse_iter = metadata.get('multi_scale_coarse_iterations', -1)
            fine_iter = metadata.get('multi_scale_fine_iterations', -1)
            
            if coarse_iter == -1:
                coarse_iter = None
            if fine_iter == -1:
                fine_iter = None
            
            phase_mid = griffin_lim_multi_scale(mag_mid, n_fft, hop_length, 
                                                iterations=iterations, 
                                                random_seed=random_seed,
                                                scale_factor=scale_factor,
                                                coarse_iterations=coarse_iter,
                                                fine_iterations=fine_iter,
                                                early_stop_threshold=early_stop_threshold,
                                                early_stop_patience=early_stop_patience)
            phase_side = griffin_lim_multi_scale(mag_side, n_fft, hop_length, 
                                                 iterations=iterations, 
                                                 random_seed=random_seed + 1 if random_seed else None,
                                                 scale_factor=scale_factor,
                                                 coarse_iterations=coarse_iter,
                                                 fine_iterations=fine_iter,
                                                 early_stop_threshold=early_stop_threshold,
                                                 early_stop_patience=early_stop_patience)
        else:
            phase_mid = griffin_lim(mag_mid, n_fft, hop_length, 
                                    iterations=iterations, 
                                    random_seed=random_seed,
                                    early_stop_threshold=early_stop_threshold,
                                    early_stop_patience=early_stop_patience)
            phase_side = griffin_lim(mag_side, n_fft, hop_length, 
                                     iterations=iterations, 
                                     random_seed=random_seed + 1 if random_seed else None,
                                     early_stop_threshold=early_stop_threshold,
                                     early_stop_patience=early_stop_patience)

    # Проверки на NaN и некорректные значения
    for name, ph in [("Mid", phase_mid), ("Side", phase_side)]:
        if np.any(np.isnan(ph)):
            np.nan_to_num(ph, copy=False, nan=0.0)
        if np.abs(ph).max() > np.pi * 1.1:
            np.clip(ph, -np.pi, np.pi, out=ph)

    print("Восстановление комплексных спектров...")
    D_mid = mag_mid.astype(np.complex64) * np.exp(1j * phase_mid.astype(np.float32))
    D_side = mag_side.astype(np.complex64) * np.exp(1j * phase_side.astype(np.float32))

    print("Обратное STFT...")
    y_mid = librosa.istft(D_mid, hop_length=hop_length, length=original_length, window='hann')
    y_side = librosa.istft(D_side, hop_length=hop_length, length=original_length, window='hann')

    y_left, y_right = ms_to_lr(y_mid, y_side)
    
    y_recovered = np.stack([y_left, y_right], axis=1)
    max_val = np.max(np.abs(y_recovered))
    if max_val > 0:
        y_recovered = y_recovered / max_val * 0.95

    duration = y_recovered.shape[0] / sr
    elapsed = time.time() - t_start
    print(f"Сохранение: {output_wav_path} ({duration:.2f} сек, {elapsed:.1f} сек)")
    
    y_16bit = (y_recovered * 32767).clip(-32768, 32767).astype(np.int16)
    sf.write(output_wav_path, y_16bit, sr, subtype='PCM_16')
    
    return y_recovered