import random
import torch  # PyTorch 核心库，用于张量操作和神经网络构建。
import torch.nn as nn  # PyTorch 的神经网络模块，包含各种层和损失函数。
import numpy as np  # 用于数值计算。
import os  # 可以处理文件路径、更改当前工作目录、访问环境变量等
from PIL import Image #读取图片数据
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from torchvision import transforms  #用来处理和转换图像，让图像变成神经网络可以训练的格式。
import time
import matplotlib.pyplot as plt
from model_utils.model import initialize_model



# 设置随机种子 方便固定和复现模型训练结果。 保证可以复现最佳模型
def seed_everything(seed): #seed_everything用来固定每次随机
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    random.seed(seed)
    np.random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
#################################################################
seed_everything(0)
###############################################


HW=224  #深度学习领域一般经典是让图片大小是224



train_transform = transforms.Compose(
    [
        transforms.ToPILImage(),  #224,224,3       模型：3,224,224
        transforms.RandomResizedCrop(224),  #放大裁切
        transforms.RandomRotation(50),  #50度以内随机旋转
        transforms.ToTensor()  #张量运算
    ]
)
val_transform = transforms.Compose(
    [
        transforms.ToPILImage(),  #224,224,3       模型：3,224,224
        transforms.ToTensor()  #张量运算
    ]
)


#数据处理
class food_Dataset(Dataset):  #食物图像数据集
    def __init__(self,path, mode = "train"):
        self.mode = mode
        if mode == "semi":
            self.X = self.read_file(path)
        else:
            self.X, self.Y = self.read_file(path)  # 给一个路径，得到数据X,Y
            self.Y = torch.LongTensor(self.Y)  # 标签转为长整形
        if mode == "train":
            self.transform = train_transform  # 将你定义的图像变换操作保存为类的一个属性，这样在数据集的 __getitem__ 方法中可以对每张图像应用这些变换。
        else:
            self.transform = val_transform

    def read_file(self, path):  # 定义一个读取路径

        if self.mode == "semi":
            file_list = os.listdir(path)  # 列出文件夹下所有文件的名字
            xi = np.zeros((len(file_list), HW, HW, 3), dtype=np.uint8)  # dtype=np.uint8是在读数据时必须是整数

            for j, img_name in enumerate(file_list):  # enumerate既可以读到下标，又可以读到下标里的值
                img_path = os.path.join(path, img_name)
                img = Image.open(img_path)
                img = img.resize((HW, HW))
                xi[j, ...] = img
            print("读到了%d个数据" % len(xi))
            return xi
        else:
            for i in tqdm(range(11)):  # 遍历11个文件夹
                file_dir = path + "/%02d" % i  # 路径到labeled
                file_list = os.listdir(file_dir)  # 列出文件夹下所有文件的名字

                xi = np.zeros((len(file_list), HW, HW, 3), dtype=np.uint8)  # dtype=np.uint8是在读数据时必须是整数
                yi = np.zeros(len(file_list), dtype=np.uint8)

                for j, img_name in enumerate(file_list):  # enumerate既可以读到下标，又可以读到下标里的值
                    img_path = os.path.join(file_dir, img_name)
                    img = Image.open(img_path)
                    img = img.resize((HW, HW))
                    xi[j, ...] = img
                    yi[j] = i

                if i == 0:
                    X = xi
                    Y = yi
                else:
                    X = np.concatenate((X, xi), axis=0)
                    Y = np.concatenate((Y, yi), axis=0)
            print("读到了%d个数据" % len(Y))
            return X, Y

    def __getitem__(self, item):
        if self.mode == "semi":
            return self.transform(self.X[item]), self.X[item]  # 取数据，并对取出的数据图像做数据增广
        else:
            return self.transform(self.X[item]), self.Y[item]

    def __len__(self):
        return len(self.X)


#半监督方式实现无标签数据打上标签的过程
class semiDataset(Dataset):
    def __int__(self, no_label_loader, model,device,thres =0.99):  #参数：无标签数据集，模型，设备gpu，置信度
        x,y = self.get_label(no_label_loader, model,device,thres)
        if x == []:
            self.flag = False
        else:
            self.flag = True
            self.X = np.array(x)
            self.Y = torch.LongTensor(y)
            self.transform = train_transform

    def get_label(self,no_label_loader, model,device,thres):
        model = model.to(device)
        pred_prob = []  #概率列表
        labels = []  #标签列表
        x = []
        y = []
        soft = nn.Softmax()
        with torch.no_grad():  #无标签数据经过模型打标签，不需要更新模型，所以不需要梯度
            for bat_x, _ in no_label_loader:
                bat_x = bat_x.to(device)
                pred = model(bat_x)
                pred_max, pred_value = pred_soft.max(1)  #返回最大值和最大值的下标
                pred_prob.extend(pred_max.cpu().numpy().toList())
                labels.extend(pred_value.cpu().numpy().toList())

        for index,prob in enumerate(pred_prob):
            if prob > thres:
                x.append(no_label_loader.dataset[index][1])
                y.append(labels[index])
        return x, y

    def __getitem__(self, item):
        return self.transform(self.X[item]),self.Y[item]
    def __len__(self):
        return len(self.X)

def get_semi_loader(no_label_loader, model, device, thres):
    semiset = semiDataset(no_label_loader, model, device, thres)
    if semiset.flag == False:
        return None
    else:
        semi_loader = DataLoader(semiset, batch_size=16, shuffle=False)
        return semi_loader



class myModel(nn.Module):
    def __init__(self, num_class):  #num_class表示分类的个数
        super(myModel, self).__init__()
        # 3 *224 *224  -> 512*7*7 -> 拉直 ->全连接分类
        #  卷积，归一化，激活，池化

        self.conv1 = nn.Conv2d(3, 64, 3, 1, 1)  # 64*224*224
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)  # 64*112*112

        self.layer1 = nn.Sequential(
            nn.Conv2d(64, 128, 3, 1, 1),  # 128*112*112
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2)  # 128*56*56
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(128, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(2)  # 256*28*28
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.BatchNorm2d(512),
            nn.ReLU(),
            nn.MaxPool2d(2)  # 512*14*14
        )

        self.pool2 = nn.MaxPool2d(2)  # 512*7*7
        self.fc1 = nn.Linear(25088, 1000)  # 25088->1000
        self.relu2 = nn.ReLU()
        self.fc2 = nn.Linear(1000, num_class)  # 1000->11

    def forward(self, x):  #让你的数据通过你的模型
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.pool1(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.pool2(x)
        x = x.view(x.size()[0], -1)
        x = self.fc1(x)
        x = self.relu2(x)
        x = self.fc2(x)
        return x


# 这个是从回归那个test.py导入过来的,再加入acc，统计准确率
def train_val(model, train_loader, val_loader, no_label_loader,device, epochs, optimizer, loss, thres,save_path):
    model = model.to(device)
    semi_loader = None

    plt_train_loss = []
    plt_val_loss = []
    plt_train_acc = []
    plt_val_acc = []

    max_acc = 0.0


    for epoch in range(epochs):
        train_loss = 0.0
        val_loss = 0.0
        train_acc = 0.0
        val_acc = 0.0
        semi_loss = 0.0
        semi_acc = 0.0
        start_time = time.time()

        model.train()
        for batch_x, batch_y in train_loader:
            x, target = batch_x.to(device), batch_y.to(device)
            pred = model(x)
            # train_bat_loss = loss(pred, target,model)
            train_bat_loss = loss(pred, target)
            train_bat_loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            train_loss += train_bat_loss.cpu().item()

            train_acc += np.sum(np.argmax(pred.detach().cpu().numpy(), axis=1) == target.cpu().numpy())

        plt_train_loss.append(train_loss/train_loader.__len__())
        plt_train_acc.append(train_acc/train_loader.dataset.__len__())  #记录准确率


        if semi_loader!= None:
            for batch_x, batch_y in semi_loader:
                x, target = batch_x.to(device), batch_y.to(device)
                pred = model(x)
                # train_bat_loss = loss(pred, target,model)
                semi_bat_loss = loss(pred, target)
                semi_bat_loss.backward()
                optimizer.step()
                optimizer.zero_grad()
                semi_loss += train_bat_loss.cpu().item()
                semi_acc += np.sum(np.argmax(pred.detach().cpu().numpy(), axis=1) == target.cpu().numpy())
            print("半监督数据集的训练准确度为",semi_acc/train_loader.dataset.__len__())


        model.eval()
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                x, target = batch_x.to(device), batch_y.to(device)
                pred = model(x)
                # val_bat_loss = loss(pred, target,model)
                val_bat_loss = loss(pred, target)
                val_loss += val_bat_loss.cpu().item()
                val_acc += np.sum(np.argmax(pred.detach().cpu().numpy(), axis=1) == target.cpu().numpy())
        plt_val_loss.append(val_loss / val_loader.dataset.__len__())
        plt_val_acc.append(val_acc / val_loader.dataset.__len__())

        if epoch%3 == 0 and plt_val_acc[-1] > 0.6:
            semi_loader = get_semi_loader(no_label_loader, model, device, thres)

        if val_acc > max_acc:
            torch.save(model, save_path)
            max_acc = val_acc


        print('[%03d/%03d]  %2.2f secs Trainloss: %.6f Valloss: %.6f Trainacc: %.6f Valacc: %.6f' % \
              (epoch, epochs, time.time()-start_time,plt_train_loss[-1],plt_val_loss[-1], plt_train_acc[-1],plt_val_acc[-1]))
               # 打印训练结果。 注意python语法， %2.2f 表示小数位为2的浮点数， 后面可以对应。
    plt.plot(plt_train_loss)
    plt.plot(plt_val_loss)
    plt.title("loss")
    plt.legend(["train","val"])
    plt.show()

    plt.plot(plt_train_acc)
    plt.plot(plt_val_acc)
    plt.title("acc")
    plt.legend(["train","val"])
    plt.show()

# path = r"E:\shenduxuexi\food_classification\food-11\training\labeled"  #前面加r是去除转义字符

# train_path = r"E:\shenduxuexi\food_classification\food-11\training\labeled"
# val_path = r"E:\shenduxuexi\food_classification\food-11\validation"
train_path = r"E:\shenduxuexi\food_classification\food-11_sample\training\labeled"
val_path = r"E:\shenduxuexi\food_classification\food-11_sample\validation"
no_label_path = r"E:\shenduxuexi\food_classification\food-11\training\unlabeled\00"



train_set = food_Dataset(train_path, "train")
val_set = food_Dataset(val_path, "val")
no_label_set = food_Dataset(no_label_path, "semi")

train_loader = DataLoader(train_set, batch_size=16, shuffle=True)  #train_loader把它变成一批的数据，利用shuffle打乱，不让一类的图片在一块
val_loader = DataLoader(val_set, batch_size=16, shuffle=True)
no_label_loader = DataLoader(no_label_set, batch_size=16, shuffle=False)  #不用打乱，怕无标签数据和它打上的标签不对应


# model =myModel(11)
# from torchvision.models import resnet18
# model = resnet18(pretrained=True)  #pretrained=True表示不光要用大佬们的架构，也要用大佬们的参数
# in_fetures = model.fc.in_features  #提取特征多长
# model.fc = nn.Linear(in_fetures, 11)

model, _ = initialize_model("resnet18", 11, use_pretrained=True)


lr = 0.001  #学习率
loss = nn.CrossEntropyLoss()  #交叉熵损失函数，它用来衡量模型预测结果和真实标签之间的差距，告诉模型“你猜得准不准”。
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
device = "cuda" if torch.cuda.is_available() else "cpu"
save_path = "model_save/best_model.pth"
epochs = 15
thres = 0.1


train_val(model, train_loader, val_loader, no_label_loader,device, epochs, optimizer, loss, thres,save_path)

监督学习：有x，有y

✅ X：输入（Input） / 特征（Features）

你给模型看的东西，也就是“问题”或“条件”。

✅ Y：输出（Output） / 标签（Labels）

你希望模型预测的结果，也就是“答案”或“目标”。

无监督学习：只有x，没有y

半监督学习：

伪标签法。数据集中包含有标签的数据（有x，有y）和无标签的数据（只有x，没有y）。无标签的数据通过模型获得预测值，如果这个预测值大于一定的置信度，那这个预测值就相当于我们打上的标签,这回就把这个数据归入了有标签数据。

1.之前学习的回归任务主要是用：nn.linear(x,y)

2.之前学习的分类任务：那个图片分类任务主要是把3×224×224的图片经过卷积，然后拉平，再全连接，最后一步一步转换为linear(x,3)

3.深度学习， 就是把我们的输入转换成一个高维的向量。一定程度上， 可以理解为编码， 或者压缩。然后， 我们才可以用这个特征去分

类，回归，或者其他任务。

4.无监督学习举例：把不同颜色的球，放在不同的桶里。虽然我们不知道每个桶的颜色，但是我们实现了把同一颜色的球放在了同一个桶里。

自监督学习：从自己的x中提取一部分x，当成自己的y，自己监督自己。

5.

对比学习：

对抗生成网络：举例：小偷做假钞，生成器生成图像，骗过判别器；警察利用判别器，找出生成器做的"假数据"；所以生成器和判别器越来越好

gan：

生成式自监督方法：

文字里的自监督：把文字遮住一部分，输出到模型中，完整版文字当成它的标签

6.模型的训练需要梯度，梯度需要loss，loss是谁与谁之间的差距

学习方式	类比
监督学习	学生做题，老师给答案批改
无监督学习	学生看一堆题目，自己总结出有哪些题型（但不做题）
自监督学习	学生自己把题目的某个部分遮住，然后尝试填空，再翻开书对答案 → 自己出题 + 自己批改