import torch
import torch.nn as nn
import torch.optim as optim

import torchvision
from torchvision import models
from torchvision import transforms
from torchvision.utils import save_image

from collections import OrderedDict
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import cv2

import warnings
import os
from nukki import generate_result_filename, get_result_image_path, save_uploaded_file

from stylegan_model import G_mapping
from stylegan_model import G_synthesis
    
#전역변수설정
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
resolution = 1024
weight_file = 'C:/Users/remil/바탕 화면/AI프로젝트/venv/weights/karras2019stylegan-ffhq-1024x1024.pt'

# 특정 경로에서 이미지를 읽어 torch.Tensor 객체로 만드는 함수
def image_reader(image_path, resize=None):
    with open(image_path, "rb") as f: # 특정 경로에서 이미지 불러오기
        image = Image.open(f)
        image = image.convert("RGB") # RGB 색상 이미지로 사용
    # 미리 정해 놓은 해상도에 맞게 크기 변환
    if resize != None:
        image = image.resize((resize, resize))
    transform = transforms.Compose([
        transforms.ToTensor() # [0, 1] 사이의 값을 가지는 Tensor 형태로 변형
    ])
    image = transform(image)
    image = image.unsqueeze(0) # 배치(batch) 목적의 차원 추가 (N, C, H, W)
    return image


# torch.Tensor 형태의 이미지를 화면에 출력하는 함수
def imshow(tensor):
    # matplotlib는 CPU 기반이므로 CPU로 옮기기
    image = tensor.cpu().clone()
    # torch.Tensor에서 사용되는 배치 목적의 차원(dimension)을 제거
    image = image.squeeze(0)
    gray_scale = False # 흑백 이미지 여부
    if image.shape[0] == 1:
        gray_scale = True
    # PIL 객체로 변경
    image = transforms.ToPILImage()(image)
    # 이미지를 화면에 출력(matplotlib는 [0, 1] 사이의 값이라고 해도 정상적으로 출력)
    if gray_scale: # 흑백인 경우 흑백 색상으로 출력
        plt.imshow(image, cmap='gray')
    else:
        plt.imshow(image)
    plt.show()


# 화면에 출력되는 이미지 크기 조절
plt.rcParams['figure.figsize'] = [12, 8]
plt.rcParams['figure.dpi'] = 80




# VGG를 활용한 특징 추출기(Feature Extractor)
class FeatureExtractor(torch.nn.Module):
    def __init__(self, n_layers):
        super(FeatureExtractor, self).__init__()
        extractor = models.vgg16(pretrained=True).features

        # 각 레이어까지의 출력 값을 계산 (n_layers=[1, 3, 13, 20])
        index = 0
        self.layers = nn.ModuleList([])
        for i in range(len(n_layers)):
            # 해당 레이어까지의 출력 값을 내보낼 수 있도록 하기
            self.layers.append(torch.nn.Sequential())
            for j in range(index, n_layers[i] + 1):
                self.layers[i].add_module(str(j), extractor[j])
            index = n_layers[i] + 1

        # 모델을 학습할 필요는 없으므로 기울기 추적 중지
        for param in self.parameters():
            param.requires_grad = False

    # 각 레이어까지의 출력 값들을 리스트에 담아 반환
    def forward(self, x):
        result = []
        for i in range(len(self.layers)):
            x = self.layers[i](x)
            result.append(x)

        return result
    
    # 생성된 이미지(generated_image)와 대상 이미지(target_image)에 대한 손실(loss)을 계산
def loss_function(generated_image, target_image, feature_extractor):
    library()
    MSE = nn.MSELoss(reduction='mean')
    mse_loss = MSE(generated_image, target_image) # 손실(loss) 값 계산

    # VGG 네트워크의 입력은 256이므로, 크기를 256으로 바꾸는 업샘플링(upsampling)을 이용합니다.
    upsample2d = torch.nn.Upsample(scale_factor=256 / resolution, mode='bilinear')
    real_features = feature_extractor(upsample2d(target_image))
    generated_features = feature_extractor(upsample2d(generated_image))

    perceptual_loss = 0
    # 활성화 맵(activation map)의 개수만큼 반복하며
    for i in range(len(real_features)):
        perceptual_loss += MSE(real_features[i], generated_features[i]) # 손실(loss) 값 계산

    return mse_loss, perceptual_loss

# 생성된 이미지(generated_image)와 대상 이미지(target_image)에 대한 손실(loss)을 계산(마스크 고려)
def loss_function_with_mask(generated_image, target_image, feature_extractor, mask):
    MSE = nn.MSELoss(reduction='mean')
    # 마스크가 씌워진 부분에 대한 MSE 손실(loss) 계산
    mse_loss = MSE(generated_image * mask.expand(1, 3, resolution, resolution), target_image * mask.expand(1, 3, resolution, resolution))

    # VGG 네트워크의 입력은 256이므로, 크기를 256으로 바꾸는 업샘플링(upsampling)을 이용합니다.
    upsample2d = torch.nn.Upsample(scale_factor=256 / resolution, mode='bilinear')
    real_features = feature_extractor(upsample2d(target_image))
    generated_features = feature_extractor(upsample2d(generated_image))

    perceptual_loss = 0
    # 활성화 맵(activation map)의 개수만큼 반복하며
    for i in range(len(real_features)):
        mask_shape = mask.shape
        feature_shape = real_features[i].shape # 현재 특징 맵(feature map) 크기 구하기 (shape[2]는 높이를 의미)
        mask = torch.nn.Upsample(scale_factor=feature_shape[2] / mask_shape[2], mode='bilinear')(mask) # 특징 맵의 크기에 맞추기
        # 마스크가 씌워진 부분에 대한 손실(loss) 계산
        perceptual_loss += MSE(real_features[i] * mask.expand(feature_shape), generated_features[i] * mask.expand(feature_shape)) # 손실(loss) 값 계산

    return mse_loss, perceptual_loss


def process_face(img1_path, img2_path, result_path):
    warnings.filterwarnings(action='ignore') # 경고(warning) 메시지 보이지 않도록 하기

    

    # StyleGAN을 구성하는 두 개의 네트워크 불러오기
    g_all = nn.Sequential(OrderedDict([
        ('g_mapping', G_mapping()),
        ('g_synthesis', G_synthesis(resolution=resolution))
    ]))
    g_all.load_state_dict(torch.load(weight_file, map_location=device))
    g_all.eval()
    g_all.to(device)

    g_mapping, g_synthesis = g_all[0], g_all[1]

    # 이미지 A 불러오기
    image_1 = image_reader(img1_path, resize=resolution)
    image_1 = image_1.to(device)

    # 이미지 B 불러오기
    image_2 = image_reader(img2_path, resize=resolution)
    image_2 = image_2.to(device)

    # 교차(crossover) 이미지 생성을 위한 마스크 이미지 불러오기
    mask_src = 'C:/Users/remil/바탕 화면/AI프로젝트/venv/images/blur_mask.png'
    mask = image_reader(mask_src, resize=resolution)
    mask = mask.to(device)
    mask = mask[:, 0, :, :].unsqueeze(0)
    mask_inverse = mask.clone()
    mask_inverse = 1 - mask

    latent = torch.zeros((1, 18, 512), requires_grad=True, device=device) # 업데이트할 latent vector 변수
    optimizer_latent = optim.Adam({latent}, lr=0.02, betas=(0.9, 0.999), eps=1e-8)

    # VGG perceptual loss를 위한 레이어 명시
    feature_extractor = FeatureExtractor(n_layers=[1, 3, 13, 20]).to(device)

    name = 'crossover'

    # 최적화(optimization) 시작
    iteration = 1500
    for i in range(iteration):
        optimizer_latent.zero_grad()

        # latent vector를 이용해 이미지 생성
        generated_image = g_synthesis(latent)
        generated_image = (generated_image + 1.0) / 2.0
        generated_image = generated_image.clamp(0, 1)

        # 손실(loss) 값 계산
        loss = 0
        mse_loss_1, perceptual_loss_1 = loss_function_with_mask(generated_image, image_1, feature_extractor, mask_inverse)
        mse_loss_2, perceptual_loss_2 = loss_function_with_mask(generated_image, image_2, feature_extractor, mask)
        loss += mse_loss_1 + perceptual_loss_1 + mse_loss_2 + perceptual_loss_2
        loss.backward() # 기울기(gradient) 계산

        optimizer_latent.step() # latent vector 업데이트

        # 주기적으로 손실(loss) 값 및 이미지 출력
        if i == 0 or (i + 1) % 100 == 0:
            print(f'[iter {i + 1}/{iteration}] loss = {loss.detach().cpu().numpy()}, saved_path = {name}_{i + 1}.png')
            save_image(generated_image, f'{name}_{i + 1}.png')
            np.save(f'{name}_latent.npy', latent.detach().cpu().numpy())

    # 결과 이미지 저장
    save_image(generated_image, result_path)
    return True