from time import sleep

import numpy as np
import torch
from torch.utils.data import DataLoader
from torch import nn, optim
from tqdm import tqdm

from model.loss import YoloLoss
from util.Augmentation import SSDAugmentation
import VOCDataset
from model.yolo import Yolo

with open('data/test.txt') as f:
    val_lines = f.readlines()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model=Yolo(n_classes=20,image_size=416,anchors=None,nms_thresh=0.45).to(device)
model.darknet.load_state_dict(torch.load('data/darknet_weights.pth',map_location=device))
for param in model.parameters():
    param.requires_grad = True

dataset = VOCDataset.VOCDataset(val_lines, transform=SSDAugmentation([416,416]), train=True)
dataloader = DataLoader(dataset=dataset,
                        batch_size=4,
                        shuffle=True,
                        num_workers=0,
                        collate_fn=VOCDataset.collate_fn
                        )

anchors=[
        [116, 90], [156, 198], [373, 326],
        [30, 61], [62, 45], [59, 119],
        [10, 13], [16, 30], [33, 23],
]
anchor_masks=[[0,1,2],[3,4,5],[6,7,8]]
num_classes=20

# 优化器
darknet_params = model.darknet.parameters()
other_params = [i for i in model.parameters()
                if i not in darknet_params]
optimizer = optim.SGD(
    [
        {"params": darknet_params, 'initial_lr': 1e-3, 'lr':1e-3},
        {'params': other_params, 'initial_lr': 0.01, 'lr': 0.01}
    ],
    momentum=0.9,
    weight_decay=4e-5
)

lr_schedule = optim.lr_scheduler.StepLR(optimizer, 1, 0.1, last_epoch=0)

# 损失函数
yololoss = YoloLoss(anchors, anchor_masks, num_classes, overlap_thresh=0.5)

for epoch in range(0,2):
    loss=0
    for iteration, (images, targets) in enumerate(dataloader):
        print(epoch,iteration,'------------------------------------')
        optimizer.zero_grad()

        preds = model(images.to(device))

        loc_loss, conf_loss, cls_loss = yololoss(preds, targets)
        cur_loss = loc_loss + cls_loss + conf_loss
        #print(cur_loss)

        loss += cur_loss

        cur_loss.backward()
        optimizer.step()

        # 学习率退火
        lr_schedule.step()






# p_mask=[[1,0,1],
#         [1,1,0],
#         [0,1,0]]
# cls=[[[1,11],[2,22],[3,33]],
#      [[4,44],[5,55],[6,66]],
#      [[7,77],[8,88],[9,99]]]
# p_mask=torch.tensor(p_mask,dtype=torch.int)
# cls=torch.tensor(cls,dtype=torch.int)
# m=p_mask==1
#
# print(cls[...,1])
# print(m)
# print(cls[m])


# for epoch in range(0,100):
#     with tqdm(total=10, desc=f'Epoch {epoch + 1}/100', postfix=dict,mininterval=0.3) as pbar:
#         for iteration in range(0, 10):
#             sleep(1)
#             pbar.set_postfix(**{'loss': 10 / (iteration + 1)})
#             pbar.update(1)


# darknet_params = model.darknet.parameters()
# other_params = [i for i in model.parameters()if i not in darknet_params]
# optimizer = optim.SGD(
#     [
#         {"params": darknet_params, 'initial_lr': 0.1, 'lr': 0.2},
#         {'params': other_params, 'initial_lr': 0.3, 'lr': 0.4}
#     ],
#     momentum=0.9,
#     weight_decay=0.001
# )
#
# print(optimizer.param_groups[0]['initial_lr'])
# print(optimizer.param_groups[1]['initial_lr'])