# TODO:
# 参考yolov8-火灾检测，未完待续
...

import ultralytics
ultralytics.checks()

# -*- coding: utf-8 -*-
# @Author : tianL.R
# @Email : rtl1312@163.com
# @Time : 2023.11.26
import time

import numpy as np
from PIL import Image
from keras.applications.vgg16 import VGG16
from keras.applications.vgg16 import preprocess_input
from keras.preprocessing import image
from numpy import linalg


class VGG16Net:
    def __init__(self):
        self.input_shape = (224, 224, 3)
        self.weight = 'imagenet'
        self.pooling = 'max'
        self.model_vgg = VGG16(weights=self.weight,
                               input_shape=(self.input_shape[0], self.input_shape[1], self.input_shape[2],),
                               pooling=self.pooling,
                               include_top=False)
        self.model_vgg.predict(np.zeros((1, 224, 224, 3)))

    def detection(self, img_path):
        """
        提取VGG16最后一层卷积特征
        """
        # img = image.load_img(img_path, target_size=(self.input_shape[0], self.input_shape[1]))
        img = img_path.resize((self.input_shape[0], self.input_shape[1]))
        img = image.img_to_array(img)
        img = np.expand_dims(img, axis=0)
        img = preprocess_input(img)
        feat = self.model_vgg.predict(img)
        norm_feat = feat[0] / linalg.norm(feat[0])
        return norm_feat.tolist()


if __name__ == '__main__':
    img1 = '333.jpg'
    img2 = '555.jpg'
    img1 = Image.open(img1)
    img2 = Image.open(img2)

    vgg = VGG16Net()
    queryVec1 = np.array(vgg.detection(img1))
    queryVec2 = np.array(vgg.detection(img2))
    scores = np.dot(queryVec1, queryVec2)
    score2 = queryVec1.dot(queryVec2) / (np.linalg.norm(queryVec1) * np.linalg.norm(queryVec2))
    print(scores)
    print(score2)

MODEL_ROOT_PATH = "/root/huggingface"

MODEL_PATH = {
    "embed_model": {
		...
        "m3e-base": "/root/huggingface/m3e-base", # 修改m3e模型路径
		...
    },
    # TODO: add all supported llm models
    "llm_model": {
		...
        "chatglm2-6b": "/root/huggingface/chatglm2-6b", # 修改chatglm2-6b模型路径
		...
    },
}

EMBEDDING_MODEL = "m3e-base" # 可以尝试最新的嵌入式sota模型：bge-large-zh-v1.5
LLM_MODEL = "chatglm2-6b"

from keras.applications.vgg16 import VGG16
from keras.applications.vgg19 import VGG19

import torch
from collections import namedtuple
from torchvision import models
import torch.nn as nn
import torch.nn.functional as F


# VGG16神经网络定义
class VGG16(torch.nn.Module):
    """Vgg16 Net"""
    def __init__(self, requires_grad=False):
        super(VGG16, self).__init__()
        vgg_pretrained_features = models.vgg16(pretrained=True).features
        self.slice1 = torch.nn.Sequential()
        self.slice2 = torch.nn.Sequential()
        self.slice3 = torch.nn.Sequential()
        self.slice4 = torch.nn.Sequential()

        for x in range(4):
            self.slice1.add_module(str(x), vgg_pretrained_features[x])

        for x in range(4, 9):
            self.slice2.add_module(str(x), vgg_pretrained_features[x])

        for x in range(9, 16):
            self.slice3.add_module(str(x), vgg_pretrained_features[x])

        for x in range(16, 23):
            self.slice4.add_module(str(x), vgg_pretrained_features[x])

        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h = self.slice1(X)
        h_relu1_2 = h
        h = self.slice2(h)
        h_relu2_2 = h
        h = self.slice3(h)
        h_relu3_3 = h
        h = self.slice4(h)
        h_relu4_3 = h

        vgg_outputs = namedtuple("VggOutputs", ["relu1_2", "relu2_2", "relu3_3", "relu4_3"])
        output = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3)

        return output


class TransformerNet(torch.nn.Module):
    def __init__(self):
        super(TransformerNet, self).__init__()
        self.model = nn.Sequential(
            ConvBlock(3, 32, kernel_size=9, stride=1),
            ConvBlock(32, 64, kernel_size=3, stride=2),
            ConvBlock(64, 128, kernel_size=3, stride=2),
            ResidualBlock(128),
            ResidualBlock(128),
            ResidualBlock(128),
            ResidualBlock(128),
            ResidualBlock(128),
            ConvBlock(128, 64, kernel_size=3, upsample=True),
            ConvBlock(64, 32, kernel_size=3, upsample=True),
            ConvBlock(32, 3, kernel_size=9, stride=1, normalize=False, relu=False),
        )

    def forward(self, x):
        return self.model(x)


class ResidualBlock(torch.nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.block = nn.Sequential(
            ConvBlock(channels, channels, kernel_size=3, stride=1, normalize=True, relu=True),
            ConvBlock(channels, channels, kernel_size=3, stride=1, normalize=True, relu=False),
        )

    def forward(self, x):
        return self.block(x) + x


class ConvBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, upsample=False, normalize=True, relu=True):
        super(ConvBlock, self).__init__()
        self.upsample = upsample
        self.block = nn.Sequential(
            nn.ReflectionPad2d(kernel_size // 2),
            nn.Conv2d(in_channels, out_channels, kernel_size, (stride,))
        )
        self.norm = nn.InstanceNorm2d(out_channels, affine=True) if normalize else None
        self.relu = relu

    def forward(self, x):
        if self.upsample:
            x = F.interpolate(x, scale_factor=2)
        x = self.block(x)
        if self.norm is not None:
            x = self.norm(x)
        if self.relu:
            x = F.relu(x)
        return x


"""
测试模型
"""
if __name__ == '__main__':
    input1 = torch.rand([224, 3, 224, 224])
    model_x = VGG16()
    print(model_x)

def gram_matrix(y):
    (b, c, h, w) = y.size()
    features = y.view(b, c, w * h)
    features_t = features.transpose(1, 2)
    gram = features.bmm(features_t) / (c * h * w)
    return gram

parser = argparse.ArgumentParser(description="Parser 4 Training")
parser.add_argument("--style", type=str, default="images/styles/the_starry_night.jpg", help="Path 2 style image")
parser.add_argument("--dataset", type=str, help="path 2 training dataset")
parser.add_argument("--epochs", type=int, default=1, help="Number of training epochs")
parser.add_argument("--batch_size", type=int, default=4, help="Batch size 4 training")
parser.add_argument("--image_size", type=int, default=256, help="Size of training images")
parser.add_argument("--style_size", type=int, help="Size of style image")
parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate")
parser.add_argument("--lambda_img", type=float, default=1e5, help="Weight 4 image loss")
parser.add_argument("--lambda_style", type=float, default=1e10, help="Weight 4 style loss")
parser.add_argument("--model_path", type=str, help="Optional path 2 checkpoint model")
parser.add_argument("--model_checkpoint", type=int, default=1000, help="Batches 4 saving model")
parser.add_argument("--result_checkpoint", type=int, default=1000, help="Batches 4 saving image result")

def train_transform(image_size):
    transform = transforms.Compose(
        [
            transforms.Resize(int(image_size * 1.15)),
            transforms.RandomCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
        ]
    )
    return transform

def style_transform(image_size=None):
    resize = [transforms.Resize(image_size)] if image_size else []
    transform = transforms.Compose(resize + [transforms.ToTensor(), transforms.Normalize(mean, std)])
    return transform

def denormalize(tensors):
    for c in range(3):
        tensors[:, c].mul_(std[c]).add_(mean[c])
    return tensors

train_args = TrainArgs()
args = train_args.initialize().parse_args()

args.dataset = './dataset'
args.style = './images/styles/the_starry_night.jpg'
args.epochs = 2400 # epochs*(数据集/batch_size)是1000的公倍数
args.batch_size = 4
args.image_size = 256

style_name = args.style.split("/")[-1].split(".")[0]
os.makedirs(f"images/train/{style_name}_training", exist_ok=True)
os.makedirs(f"checkpoints", exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_dataset = datasets.ImageFolder(args.dataset, train_transform(args.image_size))
dataloader = DataLoader(train_dataset, batch_size=args.batch_size)
transformer = TransformerNet().to(device)
vgg = VGG16(requires_grad=False).to(device)
if args.model_path:
    transformer.load_state_dict(torch.load(args.model_path))
optimizer = Adam(transformer.parameters(), args.lr)
l2_loss = torch.nn.MSELoss().to(device)
style = style_transform(args.style_size)(Image.open(args.style))
style = style.repeat(args.batch_size, 1, 1, 1).to(device)
features_style = vgg(style)
gram_style = [gram_matrix(y) for y in features_style]
image_samples = []
for path in random.sample(glob.glob(f"{args.dataset}/*/*"), len(train_dataset)):
    image_samples += [style_transform(args.image_size)(Image.open(path).resize((224, 224)))]
image_samples = torch.stack(image_samples)

def save_result(sample):
    transformer.eval()
    with torch.no_grad():
        output = transformer(image_samples.to(device))
    image_rgb = denormalize(torch.cat((image_samples.cpu(), output.cpu()), 2))
    save_image(image_rgb, f"images/train/{style_name}_training/{sample}.jpg", nrow=4)
    transformer.train()


def save_model(sample):
    torch.save(transformer.state_dict(), f"checkpoints/{style_name}_{sample}.pth")


for epoch in range(args.epochs):
    for line in range(len(dataloader)):
        batch_i = line
        batches_done = epoch * len(dataloader) + batch_i + 1
        images = list(dataloader)[line][0]
        optimizer.zero_grad()

        images_original = images.to(device)
        images_transformed = transformer(images_original)

        features_original = vgg(images_original)
        features_transformed = vgg(images_transformed)

        img_loss = args.lambda_img * l2_loss(features_transformed.relu2_2, features_original.relu2_2)

        style_loss = 0
        for ft_y, gm_s in zip(features_transformed, gram_style):
            gm_y = gram_matrix(ft_y)
            style_loss += l2_loss(gm_y, gm_s[: images.size(0), :, :])
        style_loss *= args.lambda_style

        total_loss = img_loss + style_loss
        total_loss.backward()
        optimizer.step()
        if batches_done % args.result_checkpoint == 0:
            save_result(batches_done)
        if args.model_checkpoint > 0 and batches_done % args.model_checkpoint == 0:
            save_model(batches_done)

predict_args = PredictArgs()
args = predict_args.initialize().parse_args()
args.image_path = './images/input/001.jpg'
args.model_path = './checkpoints/the_starry_night_12000.pth'

os.makedirs("images/output", exist_ok=True)
   device = torch.device('cpu')#("cuda" if torch.cuda.is_available() else "cpu")
   transform = style_transform()
   transformer = TransformerNet().to(device)
   transformer.load_state_dict(torch.load(mod_path))
   transformer.eval()
   image_tensor = Variable(transform(Image.open(img_path))).to(device)
   image_tensor = image_tensor.unsqueeze(0)

   with torch.no_grad():
       output_image = denormalize(transformer(image_tensor)).cpu()
       
   name = img_path.split("/")[-1]
   save_image(output_image, f"images/output/output_{name}")