import gradio as gr
from PIL import Image
import numpy as np
import random
import spaces
import torch
import torch.nn.functional as F
import config
from extras.expansion import FooocusExpansion
from feifeilib.feifeimodload import feifeimodload
import re

expansion = FooocusExpansion()
pipe = feifeimodload()
with open("artist.txt", "r") as file:
    artists = file.readlines()

MAX_SEED = np.iinfo(np.int32).max

# 去除每行末尾的换行符
artists = [artist.strip() for artist in artists]


@spaces.GPU()
def feifeitexttoimg(
        prompt,
        quality_select,
        sharpened_select,
        styles_Radio,
        FooocusExpansion_select,
        seed=42,
        randomize_seed=False,
        width=1024,
        height=1024,
        num_inference_steps=4,
        guidance_scale=3.5,
        num_strength=0.35,
        progress=gr.Progress(track_tqdm=True),
):
    guidance_scale = 3.5
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)
    if not prompt:
        prompt = "the photo is a 18 yo jpop girl is looking absolutely adorable and gorgeous, with a playful and mischievous grin, her eyes twinkling with joy. art by __artist__ and __artist__"
    if "__artist__" in prompt:
        # 随机选择艺术家
        selected_artists = random.sample(artists, len(artists))

        # 使用正则表达式替换所有的 __artist__
        def replace_artists(match):
            return selected_artists.pop(0)

        prompt = re.sub(r"__artist__", replace_artists, prompt)
        # print("__artist__ " + prompt)
    if quality_select:
        prompt += ", masterpiece, best quality, very aesthetic, absurdres"
        # print("111 " + prompt)
    if styles_Radio:
        for style_name in styles_Radio:
            for style in config.style_list:
                if style["name"] == style_name:
                    prompt += style["prompt"].replace("{prompt}", "the ")
        # print("222 " + prompt)
    if FooocusExpansion_select:
        prompt = expansion(prompt, seed)
        # print("333 " + prompt)
    # print("000 " + prompt)

    # 生成图像
    image = pipe(
        prompt="",
        prompt_2=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=guidance_scale,
        output_type="pil",
    ).images[0]

    if sharpened_select:
        # 将PIL图像转换为NumPy数组
        image_np = np.array(image)

        # 将NumPy数组转换为PyTorch张量
        image_tensor = (torch.tensor(image_np).permute(
            2, 0, 1).unsqueeze(0).float().to("cuda"))

        # 定义锐化滤镜，并调整中心值
        strength = num_strength
        sharpen_kernel = (torch.tensor(
            [
                [0, -1 * strength, 0],
                [-1 * strength, 1 + 4 * strength, -1 * strength],
                [0, -1 * strength, 0],
            ],
            dtype=torch.float32,
        ).unsqueeze(0).unsqueeze(0).to("cuda"))

        # 分别对每个通道应用卷积核
        sharpened_channels = []
        for i in range(3):
            channel_tensor = image_tensor[:, i:i + 1, :, :]
            sharpened_channel = F.conv2d(channel_tensor,
                                         sharpen_kernel,
                                         padding=1)
            sharpened_channels.append(sharpened_channel)

        # 合并通道
        sharpened_image_tensor = torch.cat(sharpened_channels, dim=1)

        # 将增强后的图像转换回PIL格式
        sharpened_image_np = (sharpened_image_tensor.squeeze(0).permute(
            1, 2, 0).cpu().numpy())
        sharpened_image_np = np.clip(sharpened_image_np, 0,
                                     255).astype(np.uint8)
        image = Image.fromarray(sharpened_image_np)

    return image, seed