# app.py
# -*- coding: utf-8 -*-
"""
雨效增强器（Rain FX）- Gradio 应用（增强版：支持手动控制雨丝粗细）
功能：
- 为上传图片叠加“下雨”效果
- 可设置：风向（°）、雨的大小（长度&粗细一体化）
- 新增：可选“手动设置雨丝粗细（像素）”，做极细/特粗风格
- 保持原图分辨率与尺寸不变
- 界面美观、布局合理，适合直接部署到 Hugging Face Spaces（CPU）

实现思路（核心算法简述）：
1) 生成稀疏随机噪声作为“雨滴种子”（密度与强度相关）。
2) 通过自定义的“运动模糊卷积核”（可旋转的细线核）在指定角度方向拉伸成雨线。
3) 对雨线图做归一化与轻微模糊，使线条更自然。
4) 与原图采用加权叠加（仅对雨线区域变亮），最终得到雨景效果。
5) 全流程不改变输入图片的宽高（分辨率）与像素比例。

注意：
- 角度约定：0° 代表“垂直向下”的雨；正值表示“向右偏”，负值表示“向左偏”。
- “雨的大小”同时影响雨线长度与粗细；你也可以调“雨量”控制密度/亮度强度。
- 本增强版将默认粗细区间提升为 2~5 px，并提供可选“手动粗细（1~9 px）”。
"""

import math
import random
from typing import Optional

import cv2  # 建议使用 opencv-python-headless
import gradio as gr
import numpy as np
from PIL import Image


# ----------------------------
# 工具函数：创建可旋转的运动模糊卷积核（雨线核）
# ----------------------------
def motion_kernel(length: int, angle_deg: float, thickness: int = 1) -> np.ndarray:
    """
    创建一个 size=length x length 的卷积核，其中包含一条可旋转的直线，用于产生“运动模糊”效果。
    - length: 雨线长度（卷积核尺寸，需为奇数以保证中心对齐）
    - angle_deg: 旋转角度，0 表示竖直向下；正值向右倾斜，负值向左倾斜
    - thickness: 线条粗细（像素）

    返回：
        归一化后的 float32 卷积核（和为 1）
    """
    # 确保 length 为 >=3 的奇数，避免旋转后中心漂移
    length = int(max(3, length))
    if length % 2 == 0:
        length += 1

    k = np.zeros((length, length), dtype=np.float32)

    # 在正中画一条竖线（先默认竖直方向），稍后整体旋转到目标角度
    center = length // 2
    cv2.line(
        k,
        (center, 0),               # 起点（上中）
        (center, length - 1),      # 终点（下中）
        color=1.0,
        thickness=int(max(1, thickness)),
        lineType=cv2.LINE_AA
    )

    # 围绕中心旋转到指定角度
    M = cv2.getRotationMatrix2D((center, center), angle_deg, 1.0)
    k_rot = cv2.warpAffine(
        k, M, (length, length),
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0
    )

    s = k_rot.sum()
    if s <= 1e-6:
        # 极端情况下（很细的线且角度接近 45° 时），避免全零
        k_rot[center, center] = 1.0
        s = 1.0
    k_rot /= s
    return k_rot.astype(np.float32)


# ----------------------------
# 核心：为图片添加雨效
# ----------------------------
def add_rain_effect(
    img: Image.Image,
    angle_deg: float = 0.0,
    size_level: int = 50,
    rain_intensity: int = 50,
    seed: Optional[int] = 0,
    use_thickness_override: bool = False,   # 新增：是否手动指定粗细
    thickness_px: int = 3                   # 新增：手动粗细像素（1~9）
) -> Image.Image:
    """
    为输入 PIL 图片添加雨效（不改变分辨率）。

    参数：
    - img: 输入图片（PIL.Image），支持 RGB/RGBA/JPEG/PNG 等
    - angle_deg: 风向角度（°）。0=垂直下落；>0 向右偏；<0 向左偏
    - size_level: 雨的大小（1~100），同时影响雨线长度与粗细（默认映射 2~5 px）
    - rain_intensity: 雨量（0~100），影响雨滴密度与亮度
    - seed: 随机种子（相同种子可复现；0 表示随机）
    - use_thickness_override: 是否启用“手动粗细像素”
    - thickness_px: 手动粗细（像素），仅在 use_thickness_override=True 时生效

    返回：
    - PIL.Image：与输入同尺寸的结果图（保持分辨率不变）
    """
    if img is None:
        raise gr.Error("请先上传一张图片。")

    # 统一到 RGBA，便于保留透明通道（若原图无 alpha 也不会改变分辨率）
    pil_rgba = img.convert("RGBA")
    w, h = pil_rgba.size

    # numpy 格式（OpenCV 使用 BGR 或灰度）
    rgba = np.array(pil_rgba)
    bgr = cv2.cvtColor(rgba, cv2.COLOR_RGBA2BGR)
    alpha_ch = rgba[:, :, 3]  # 原图 alpha，稍后原样带回

    # 随机数发生器（可复现）
    rng = np.random.default_rng(None if (seed is None or int(seed) == 0) else int(seed))

    # ---------- 参数映射 ----------
    # 雨量 -> 噪声密度（越大越密）
    # 经验范围：0.0015 ~ 0.035（避免全屏白）
    p_noise = float(np.interp(rain_intensity, [0, 100], [0.0015, 0.035]))

    # 雨大小 -> 雨线长度与（默认）粗细
    # 长度：7 ~ 48 像素；默认粗细：*提升为* 2 ~ 5 像素（较原版更粗）
    line_len = int(np.interp(size_level, [1, 100], [7, 48]))
    mapped_thick = int(np.interp(size_level, [1, 100], [2, 5]))

    # 若启用手动粗细，则覆盖默认映射；并进行安全夹取
    if use_thickness_override:
        line_thick = int(np.clip(thickness_px, 1, 9))
    else:
        line_thick = int(np.clip(mapped_thick, 1, 9))

    # 雨量 -> 叠加强度（亮度权重），避免过曝
    blend_alpha = float(np.interp(rain_intensity, [0, 100], [0.15, 0.65]))

    # ---------- 生成稀疏噪声（雨滴种子） ----------
    # True/False 的布尔图，True 代表一个潜在雨滴种子
    noise_mask = (rng.random((h, w)) < p_noise).astype(np.uint8) * 255  # 0/255 灰度

    # ---------- 构建雨线卷积核并拉伸成条 ----------
    kernel = motion_kernel(length=line_len, angle_deg=angle_deg, thickness=line_thick)
    streaks = cv2.filter2D(noise_mask, ddepth=-1, kernel=kernel)

    # 归一化到 0~255，并轻微平滑，降低颗粒感
    streaks = cv2.normalize(streaks, None, 0, 255, cv2.NORM_MINMAX)
    streaks = cv2.GaussianBlur(streaks, (3, 3), 0)

    # 可选：轻微对比增强，让雨线更有层次
    # 使用简单的 gamma 矫正（gamma < 1 变亮）
    gamma = 0.9
    streaks_f = (streaks / 255.0) ** gamma
    streaks = np.clip(streaks_f * 255.0, 0, 255).astype(np.uint8)

    # 叠加为 3 通道白色雨线层（黑底）
    rain_bgr = cv2.cvtColor(streaks, cv2.COLOR_GRAY2BGR)

    # ---------- 与原图叠加（仅雨线区域亮化），不改变原图尺寸 ----------
    # 使用 addWeighted：rain_bgr 的黑底对背景几乎无影响，白色线条按权重亮化
    out_bgr = cv2.addWeighted(bgr, 1.0, rain_bgr, blend_alpha, 0.0)

    # 还原回 RGBA，并带回原 alpha 通道（不改变尺寸）
    out_rgba = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGBA)
    out_rgba[:, :, 3] = alpha_ch  # 保持输入透明度
    out_img = Image.fromarray(out_rgba)

    return out_img


# ----------------------------
# Gradio UI
# ----------------------------
APP_TITLE = "🌧️ 雨效增强器（Rain FX）"
APP_DESC = """
**给图片一键加“下雨”效果**，支持调节**风向**与**雨的大小**，不改变原图分辨率。  
- 角度约定：`0° = 垂直向下`；正数 = 向右偏；负数 = 向左偏  
- “雨的大小”同时影响雨线长度与**默认粗细（现已提升为更粗）**  
- “雨量”影响密度与亮度强度  
- **新增**：可勾选“手动设置雨丝粗细（像素）”，范围 1~9 px
"""

CUSTOM_CSS = """
/* 简洁卡片风格 */
.gradio-container {max-width: 1080px !important; margin: auto;}
#rainfx-head h1 {margin-bottom: .25rem}
#rainfx-head p {opacity: .9}
"""

theme = gr.themes.Soft(
    primary_hue="blue",
    radius_size=gr.themes.sizes.radius_lg,
)

with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="Rain FX - Gradio") as demo:
    with gr.Column():
        gr.Markdown(f"<div id='rainfx-head'><h1>{APP_TITLE}</h1><p>{APP_DESC}</p></div>")

        with gr.Row():
            with gr.Column(scale=1):
                in_img = gr.Image(
                    label="上传图片（保持原尺寸）",
                    type="pil",
                    image_mode="RGBA",
                    height=360,
                )
                angle = gr.Slider(
                    label="风向角度（°）",
                    minimum=-80,
                    maximum=80,
                    value=0,
                    step=1,
                    info="0=垂直下落；正值向右偏；负值向左偏",
                )
                size_level = gr.Slider(
                    label="雨的大小（影响长度&默认粗细）",
                    minimum=1,
                    maximum=100,
                    value=50,
                    step=1,
                    info="长度 7~48 px；默认粗细已提升为 2~5 px",
                )
                rain_intensity = gr.Slider(
                    label="雨量（密度&亮度）",
                    minimum=0,
                    maximum=100,
                    value=50,
                    step=1,
                    info="越大越密、越明显。温和建议：40~70",
                )
                seed = gr.Number(
                    label="随机种子（0=每次随机）",
                    value=0,
                    precision=0,
                )

                # —— 新增：手动粗细控制 ——
                use_thickness_override = gr.Checkbox(
                    label="手动设置雨丝粗细（像素）",
                    value=False,
                )
                thickness_px = gr.Slider(
                    label="雨丝粗细（像素）",
                    minimum=1,
                    maximum=9,
                    value=3,
                    step=1,
                    info="仅在上方勾选时生效；数值越大线条越粗",
                )

                with gr.Row():
                    run_btn = gr.Button("✨ 生成效果", variant="primary")
                    clear_btn = gr.ClearButton([in_img], value="清空图片")

            with gr.Column(scale=1):
                out_img = gr.Image(
                    label="效果预览（可下载）",
                    type="pil",
                    image_mode="RGBA",
                    show_download_button=True,
                    height=360,
                )

    # 交互绑定（点击生成）
    run_btn.click(
        fn=add_rain_effect,
        inputs=[in_img, angle, size_level, rain_intensity, seed, use_thickness_override, thickness_px],
        outputs=[out_img],
        api_name="rainfx",
    )

    # 即改即看（任一参数变动自动刷新）
    for c in (in_img, angle, size_level, rain_intensity, seed, use_thickness_override, thickness_px):
        c.change(
            fn=add_rain_effect,
            inputs=[in_img, angle, size_level, rain_intensity, seed, use_thickness_override, thickness_px],
            outputs=[out_img],
        )

if __name__ == "__main__":
    # 本地调试时启用；部署到 Spaces 可不用改
    demo.launch(server_name="0.0.0.0", server_port=7860)