# -*- coding: utf-8 -*-
# @Author: <NAME>
# @Date:   2018-11-09 10:06:59
# @Last Modified by:   <NAME>
# @Last Modified time: 2018-11-20 21:26:08
import os
import random
import sys; sys.path.append('../')
import torch
import torch.nn as nn
import numpy as np
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import argparse
from data_loader import TrainDataLoader
from PIL import Image, ImageOps, ImageStat, ImageDraw
from net import SiameseRPN
from torch.nn import init

parser = argparse.ArgumentParser(description='PyTorch SiameseRPN Training')

parser.add_argument('--train_path', default='/home/ly/chz/ILSVRC2015/Data/VID/train/ILSVRC2015_VID_train_0000', metavar='DIR',help='path to dataset')

parser.add_argument('--weight_dir', default='/home/song/srpn/weight', metavar='DIR',help='path to weight')

parser.add_argument('--checkpoint_path', default='/home/ly/chz/weights-0690000.pth.tar', help='resume')

parser.add_argument('--max_epoches', default=100, type=int, metavar='N', help='number of total epochs to run')

parser.add_argument('--max_batches', default=1, type=int, metavar='N', help='number of batch in one epoch')

parser.add_argument('--init_type',  default='xavier', type=str, metavar='INIT', help='init net')

parser.add_argument('--lr', default=0.001, type=float, metavar='LR', help='initial learning rate')

parser.add_argument('--momentum', default=0.9, type=float, metavar='momentum', help='momentum')

parser.add_argument('--weight_decay', '--wd', default=5e-5, type=float, metavar='W', help='weight decay (default: 1e-4)')

def main():
    """ dataloader """
    args = parser.parse_args()
    data_loader = TrainDataLoader(args.train_path, check = False)

    """ Model on gpu """
    model = SiameseRPN()
    model = model.cuda()
    cudnn.benchmark = True

    """ loss and optimizer """
    criterion = MultiBoxLoss()

    """ load weights """
    init_weights(model)
    if args.checkpoint_path == None:
        sys.exit('please input trained model')
    else:
        assert os.path.isfile(args.checkpoint_path), '{} is not valid checkpoint_path'.format(args.checkpoint_path)
        checkpoint = torch.load(args.checkpoint_path)
        start = checkpoint['epoch']
        model.load_state_dict(checkpoint['state_dict'])

    """ test phase """
    index_list = range(data_loader.__len__())
    for example in range(args.max_batches):
        ret = data_loader.__get__(random.choice(index_list))
        template = ret['template_tensor'].cuda()
        detection= ret['detection_tensor'].cuda()
        pos_neg_diff = ret['pos_neg_diff_tensor'].cuda()
        cout, rout = model(template, detection) #[1, 10, 17, 17], [1, 20, 17, 17]

        cout = cout.reshape(-1, 2)
        rout = rout.reshape(-1, 4)
        cout = cout.cpu().detach().numpy()
        print('cout size: {}'.format(cout.shape))
        score = 1/(1 + np.exp(cout[:,1]-cout[:,0]))
        print('score: {}, size: {}'.format(score,score.shape))
        diff   = rout.cpu().detach().numpy() #1445

        num_proposals = 15
        score_64_index = np.argsort(score)[::-1][:num_proposals]

        print('score_64_index: {}, size: {}'.format(score_64_index,score_64_index.shape))
        score64 = score[score_64_index]
        print('score: {}'.format(score64))
        diffs64 = diff[score_64_index, :]
        anchors64 = ret['anchors'][score_64_index]
        proposals_x = (anchors64[:, 0] + anchors64[:, 2] * diffs64[:, 0]).reshape(-1, 1)
        proposals_y = (anchors64[:, 1] + anchors64[:, 3] * diffs64[:, 1]).reshape(-1, 1)
        proposals_w = (anchors64[:, 2] * np.exp(diffs64[:, 2])).reshape(-1, 1)
        proposals_h = (anchors64[:, 3] * np.exp(diffs64[:, 3])).reshape(-1, 1)
        proposals = np.hstack((proposals_x, proposals_y, proposals_w, proposals_h))

        d = os.path.join(ret['tmp_dir'], '6_pred_proposals')
        if not os.path.exists(d):
            os.makedirs(d)

        detection = ret['detection_cropped_resized']
        save_path = os.path.join(ret['tmp_dir'], '6_pred_proposals', '{:04d}_1_detection.jpg'.format(example))
        detection.save(save_path)

        template = ret['template_cropped_resized']
        save_path = os.path.join(ret['tmp_dir'], '6_pred_proposals', '{:04d}_0_template.jpg'.format(example))
        template.save(save_path)

        """ 可视化 """
        draw = ImageDraw.Draw(detection)
        for i in range(num_proposals):
            x, y, w, h = proposals_x[i], proposals_y[i], proposals_w[i], proposals_h[i]
            x1, y1, x2, y2 = x-w//2, y-h//2, x+w//2, y+h//2
            draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='red')
        """ save detection template proposals"""
        save_path = os.path.join(ret['tmp_dir'], '6_pred_proposals', '{:04d}_2_proposals.jpg'.format(example))
        detection.save(save_path)

        print('save at {}'.format(save_path))
            # restore
        """
            predictions = (cout, rout)
            targets = pos_neg_diff

            closs, rloss = criterion(predictions, targets)
            loss = closs + rloss
            closses.update(closs.cpu().item())
            rlosses.update(rloss.cpu().item())
            print("epoch:{:04d} example:{:06d} lr:{:.2f} closs:{:.2f}\trloss:{:.2f}".format(epoch, example, cur_lr, closses.avg, rlosses.avg))

        if epoch % 5 == 0 :
            file_path = os.path.join(args.weight_dir, 'epoch_{:04d}_weights.pth.tar'.format(epoch))
            state = {
            'epoch' :epoch+1,
            'state_dict' :model.state_dict(),
            'optimizer' : optimizer.state_dict(),
            }
            torch.save(state, file_path)
        """


def init_weights(net, init_type='normal', gain=0.02):
    def init_func(m):
        # this will apply to each layer
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('conv')!=-1 or classname.find('Linear')!=-1):
            if init_type=='normal':
                init.normal_(m.weight.data, 0.0, gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')#good for relu
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)

            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:
            init.normal_(m.weight.data, 1.0, gain)
            init.constant_(m.bias.data, 0.0)
    #print('initialize network with %s' % init_type)
    net.apply(init_func)



class MultiBoxLoss(nn.Module):
    def __init__(self):
        super(MultiBoxLoss, self).__init__()
        self.closs = torch.nn.CrossEntropyLoss()
        self.rloss = torch.nn.SmoothL1Loss()

    def forward(self, predictions, targets):
        cout, rout = predictions
        cout = cout.reshape(1, 2, -1)
        rout = rout.reshape(-1, 4)
        class_gt, diff = targets[:,0].unsqueeze(0).long(), targets[:,1:]
        closs = self.closs(cout, class_gt)#1,2,*  1,*

        pos_index = np.where(class_gt == 1)[1]
        if pos_index.shape[0] == 0:
            rloss = torch.FloatTensor([0]).cuda()
        else:
            rout_pos = rout[pos_index]
            diff_pos = diff[pos_index]

            #print(rout_pos)
            #print(diff_pos)
            rloss = self.rloss(rout_pos, diff_pos) #16
        return closs/64, rloss/16

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def adjust_learning_rate(lr, optimizer, epoch, gamma=0.1):
    """Sets the learning rate to the initial LR decayed 0.9 every 50 epochs"""
    lr = lr * (0.9 ** (epoch // 1))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr

if __name__ == '__main__':
    main()


