import gradio as gr
import numpy as np
import cv2 # Используем OpenCV для продвинутых опций
import math

# Словарь для сопоставления имен методов интерполяции с флагами OpenCV
interpolation_map = {
    'Nearest': cv2.INTER_NEAREST,      # Ближайший сосед
    'Linear': cv2.INTER_LINEAR,       # Билинейная
    'Cubic': cv2.INTER_CUBIC,        # Бикубическая
    'Area': cv2.INTER_AREA,          # Хорошо для уменьшения (усреднение по площади)
    'Lanczos4': cv2.INTER_LANCZOS4    # Ланцош (часто дает резкие результаты)
}

def resize_image_gaussian_smooth(image, target_width, target_height, interpolation_str, blur_sigma):
    """
    Изменяет размер изображения и применяет Гауссово размытие.

    Args:
        image (np.ndarray): Входное изображение в формате NumPy RGB.
        target_width (int): Целевая ширина.
        target_height (int): Целевая высота.
        interpolation_str (str): Название метода интерполяции ('Nearest', 'Linear', etc.).
        blur_sigma (float): Стандартное отклонение для Гауссова фильтра (0 = без размытия).

    Returns:
        np.ndarray: Обработанное изображение в формате NumPy RGB, или None при ошибке.
    """
    if image is None:
        gr.Warning("Сначала загрузите изображение!")
        return None
    if not (isinstance(target_width, (int, float)) and target_width > 0 and
            isinstance(target_height, (int, float)) and target_height > 0):
        gr.Warning("Ширина и высота должны быть положительными числами.")
        return None # Возвращаем None, чтобы поле вывода очистилось или показало ошибку

    target_width = int(target_width)
    target_height = int(target_height)
    blur_sigma = float(blur_sigma)

    # --- Безопасность: Ограничим максимальный размер во избежание проблем с памятью ---
    MAX_DIM = 8192
    if target_width > MAX_DIM or target_height > MAX_DIM:
         gr.Warning(f"Слишком большой размер. Максимальная сторона: {MAX_DIM}px.")
         return None
         
    h_orig, w_orig, _ = image.shape
    print(f"Original size: {w_orig}x{h_orig}")
    print(f"Target size: {target_width}x{target_height}")
    print(f"Interpolation: {interpolation_str}")
    print(f"Gaussian Sigma: {blur_sigma}")


    # 1. Выбираем флаг интерполяции OpenCV
    interpolation_flag = interpolation_map.get(interpolation_str, cv2.INTER_LANCZOS4)

    # 2. Изменяем размер изображения
    # Gradio передает изображение как NumPy array в формате RGB. OpenCV работает с BGR.
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    try:
        resized_bgr = cv2.resize(img_bgr, (target_width, target_height), interpolation=interpolation_flag)
    except Exception as e:
        gr.Error(f"Ошибка при изменении размера: {e}")
        return None


    # 3. Применяем Гауссово размытие (если sigma > 0)
    # Это имитирует сглаживающий фильтр DSR, применяемый *после* масштабирования
    # Размер ядра (ksize) должен быть нечетным. Если ksize=(0, 0), OpenCV вычисляет его из sigma.
    # Чем больше sigma, тем сильнее размытие.
    result_bgr = resized_bgr
    if blur_sigma > 0.0:
        # Рассчитаем подходящий размер ядра (должен быть нечетным)
        # Эмпирическое правило: размер ядра примерно 6*sigma. Округляем до ближайшего нечетного.
        ksize = int(math.ceil(blur_sigma * 3)) * 2 + 1 
        print(f"Calculated Gaussian kernel size: ({ksize}, {ksize}) for sigma={blur_sigma}")
        try:
             # Используем явно заданный ksize для большего контроля, как в 13-tap
             # ksize=(13, 13) было бы прямым аналогом 13-tap, но sigma дает больше гибкости
             # result_bgr = cv2.GaussianBlur(resized_bgr, ksize=(ksize, ksize), sigmaX=blur_sigma, sigmaY=blur_sigma)
             # ИЛИ: Позволим OpenCV самому выбрать ksize из sigma (более стандартный подход)
             result_bgr = cv2.GaussianBlur(resized_bgr, ksize=(0, 0), sigmaX=blur_sigma, sigmaY=blur_sigma)
        except Exception as e:
            gr.Error(f"Ошибка при применении размытия: {e}")
            return None


    # 4. Конвертируем результат обратно в RGB для отображения в Gradio
    result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)

    return result_rgb

# --- Создание интерфейса Gradio ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Масштабирование изображений с Гауссовым сглаживанием
        Загрузите изображение, выберите целевое разрешение, метод интерполяции и степень Гауссова размытия (для сглаживания).
        Это может имитировать эффект сглаживания от технологий типа NVIDIA DSR при даунскейлинге,
        или просто смягчить артефакты апскейлинга.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Входное изображение", type="numpy", height=300)
            interp_choice = gr.Dropdown(
                choices=list(interpolation_map.keys()),
                label="Метод интерполяции для масштабирования",
                value='Lanczos4' # Ланцош часто хорош для качества
            )
            target_w = gr.Number(label="Целевая ширина (px)", value=512, minimum=1, step=1, precision=0)
            target_h = gr.Number(label="Целевая высота (px)", value=512, minimum=1, step=1, precision=0)

            smoothness_sigma = gr.Slider(
                minimum=0.0, maximum=5.0, step=0.1,
                label="Сигма Гауссова размытия (0 = нет)",
                info="Применяется ПОСЛЕ масштабирования. Небольшие значения (0.3-0.8) могут сгладить пикселизацию апскейла или имитировать DSR-сглаживание даунскейла.",
                value=0.0
            )
            submit_btn = gr.Button("Обработать изображение", variant="primary")

        with gr.Column(scale=1):
            output_image = gr.Image(label="Результат", type="numpy", height=400)

    gr.Markdown(
        "**Примечание о 13-tap Gaussian filter:** DSR использует специфический фильтр. "
        "Здесь применяется стандартный Гауссов фильтр *после* изменения размера. "
        "Параметр `Сигма` контролирует степень размытия."
        "Вы можете экспериментировать с `Сигмой` около 0.5-1.0 для имитации мягкого сглаживания."
    )

    # Связываем кнопку с функцией
    submit_btn.click(
        fn=resize_image_gaussian_smooth,
        inputs=[input_image, target_w, target_h, interp_choice, smoothness_sigma],
        outputs=output_image
    )

# Запуск интерфейса
if __name__ == "__main__":
    # Установка зависимостей: pip install gradio opencv-python numpy
    demo.launch()