import gradio as gr
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
from art.attacks.evasion import FastGradientMethod
from art.estimators.classification import PyTorchClassifier
from PIL import Image
import numpy as np
import os
import io
from blind_watermark import WaterMark

# Pretrained ResNet50 모델 불러오기 (ImageNet 사전 훈련)
model = models.resnet50(pretrained=True)

# CIFAR-10에 맞춰 마지막 분류 레이어 수정
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)

# 모델을 GPU로 이동
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 손실 함수와 옵티마이저 설정
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# PyTorchClassifier 생성
classifier = PyTorchClassifier(
    model=model,
    loss=criterion,
    optimizer=optimizer,
    input_shape=(3, 64, 64),
    nb_classes=10,
)

# 이미지 전처리 함수
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0).to(device)

# FGSM 공격 적용 및 이미지 처리 함수
def generate_adversarial_image(image, eps_value):
    img_tensor = preprocess_image(image)

    # FGSM 공격 설정
    attack = FastGradientMethod(estimator=classifier, eps=eps_value)

    # 적대적 예제 생성
    adv_img_tensor = attack.generate(x=img_tensor.cpu().numpy())
    adv_img_tensor = torch.tensor(adv_img_tensor).to(device)

    # 적대적 이미지 변환
    adv_img_np = adv_img_tensor.squeeze(0).cpu().numpy()
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    adv_img_np = (adv_img_np * std[:, None, None]) + mean[:, None, None]
    adv_img_np = np.clip(adv_img_np, 0, 1)
    adv_img_np = adv_img_np.transpose(1, 2, 0)

    # PIL 이미지로 변환
    adv_image_pil = Image.fromarray((adv_img_np * 255).astype(np.uint8))
    
    return adv_image_pil

# 워터마크 삽입 함수
def apply_watermark(image_pil, wm_text="123", password_img=123, password_wm=456):
    bwm = WaterMark(password_img=password_img, password_wm=password_wm)

    # 이미지 바이트 데이터를 임시 파일로 저장
    temp_image_path = "temp_image.png"
    image_pil.save(temp_image_path)

    # temp_image_path 경로로 워터마크 삽입 처리
    bwm.read_img(temp_image_path)
    bwm.read_wm(wm_text, mode='str')
    
    # 워터마크 삽입
    output_path = "watermarked_image.png"
    bwm.embed(output_path)

    # 삽입된 워터마크 이미지 파일을 다시 읽어서 PIL 이미지로 변환
    result_image = Image.open(output_path)

    # 임시 파일 삭제
    os.remove(temp_image_path)
    os.remove(output_path)

    return result_image

# 전체 이미지 처리 함수
def process_image(image, eps_value):
    # 적대적 이미지 생성
    adv_image = generate_adversarial_image(image, eps_value)

    # 워터마크 적용
    watermarked_image = apply_watermark(adv_image)

    return watermarked_image

# Gradio 인터페이스 정의
gr.Interface(
    fn=process_image,
    inputs=[gr.Image(type="pil"), gr.components.Slider(0.1, 1.0, step=0.1, value=0.3, label="Epsilon")],
    outputs="image"
).launch()