import time
import torch
from torch import nn, optim
import sys
sys.path.append("..")
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy

import d2lzh_pytorch as d2l
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(torch.__version__)
print(device)

#定义vgg块
def vgg_block(num_convs, in_channels, out_channels):
    blk = []
    for i in range(num_convs):
        if i == 0:
            blk.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))
        else:
            blk.append(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))
        blk.append(nn.ReLU())
    blk.append(nn.MaxPool2d(kernel_size=2, stride=2))
    return nn.Sequential(*blk)

#定义vgg网络
conv_arch = ((1, 1, 64), (1, 64, 128), (2, 128, 256), (2, 256, 512), (2, 512, 512))
fc_features = 512*7*7 #根据卷积层的输出算出来的
fc_hidden_units = 4096
class_count = 6
def vgg(conv_arch, fc_features, fc_hidden_units=4096):
    net = nn.Sequential()
    #卷积层部分
    for i, (num_convs, in_channels, out_channels) in enumerate(conv_arch):
        net.add_module("vgg_block_" + str(i+1), vgg_block(num_convs, in_channels, out_channels))

    #全连接层部分
    net.add_module("fc", nn.Sequential(d2l.FlattenLayer(),
                                       nn.Linear(fc_features, fc_hidden_units),
                                       nn.ReLU(),
                                       nn.Dropout(0.5),
                                       nn.Linear(fc_hidden_units, fc_hidden_units),
                                       nn.ReLU(),
                                       nn.Dropout(0.5),
                                       nn.Linear(fc_hidden_units, class_count)
                                       ))
    return net


#构造单个输入，对vgg进行测试
net = vgg(conv_arch, fc_features, fc_hidden_units)
X = torch.rand(1, 1, 224, 224)

for name, blk in net.named_children():
    X = blk(X)
    print(name, 'output shape: ', X.shape)

ratio = 8
small_conv_arch = [(1, 1, 64//ratio), (1, 64//ratio, 128//ratio), (2, 128//ratio, 256//ratio), (2, 256//ratio, 512//ratio), (2, 512//ratio, 512//ratio)]
net = vgg(small_conv_arch, fc_features//ratio, fc_hidden_units//ratio)
print(net)

batch_size = 32

"""
train_iter, test_iter = d2l.load_data_fashion_mnist(
    batch_size, resize=224, root='./Datasets/FashionMNIST'
)
"""
data_dir = 'Datasets/datasets_detcting'
#加载数据这部分需要重新编写代码，使其可以加载我们自己的图片
#Compose函数将图片的几种变换组合起来
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.Grayscale(),  # 将图片变成灰度图片
        transforms.ToTensor(),#将图片转换为Tensor,归一化至[0,1]
        #transforms.Normalize([0.485], [0.229])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256), #缩放图片，保持长宽比不变，最短边的长为256像素,
        transforms.CenterCrop(224),#从中间切出 224*224的图片
        transforms.Grayscale(),#将图片变成灰度图片
        transforms.ToTensor(),#将图片转换为Tensor,归一化至[0,1]

        #transforms.Normalize([0.485], [0.229])
    ]),
}

image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}

image, label = next(iter(image_datasets['train']))

print('image_shape:')
print(image.shape)

#print(type(image_datasets))
#print(image_datasets)


class_names = image_datasets['train'].classes

print(class_names)

dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=0)
              for x in ['train', 'val']}


dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
print(dataset_sizes)

# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))
#print(inputs)
#print(classes)

#以上是参考main函数

""""""
lr, num_epochs = 0.001, 100
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
d2l.train_ch5(net, dataloaders['train'], dataloaders['val'], batch_size, optimizer, device, num_epochs)