import cv2
import os
import sys
sys.path.append("/home/syc/caffe-new/python")
import caffe
import numpy as np
import random
import time
from math import *
from DataBase import DataBase

video_dir = "../../dataset/"
video_name = "woman"
video_transpose = False
video_resize = (960, 540)
bbox = 207,117,29,103
# bbox = 100,20,60,60
(x,y,w,h) = bbox

batch_size = 4

from_seq = True

scale = 1

model_dir = 'model/'
data_dir = 'data/'
proto_file = model_dir + 'solver.prototxt'
model_file =  model_dir + "ZF_faster_rcnn_final.caffemodel"
mean_file = model_dir + 'ilsvrc_2012_mean.npy'

target_size = 224

class DeepTracker:
    def __init__(self, proto_file, model_file, mean_file):
        caffe.set_device(0)
        caffe.set_mode_gpu()
        self.inited = False
        self.prob = None
        self.mean = np.array([102.9801, 115.9465, 122.7717])

    def transpose(self, frame, bbox):
        (x,y,w,h) = bbox
        pad = 200
        _frame = np.zeros((frame.shape[0] + 2*pad, frame.shape[0] + 2*pad, 3))
        _frame[pad:pad+frame.shape[0], pad:pad+frame.shape[1], :] = frame
        data = _frame[pad+y:pad+y+h, pad+x:pad+x+w, :]
        # data = cv2.resize(data, (scale*w,scale*h))
        data = data - self.mean
        data = data.transpose((2,0,1)) 
        return data


    def transform(self, bbox, delta=(0,0,-0,-0)):
        (x,y,w,h) = bbox
        cx = x + 0.5 * w
        cy = y + 0.5 * h
        w = w / exp(delta[2])
        h = h / exp(delta[3])
        cx -= delta[0] * w
        cy -= delta[1] * h
        x = cx - 0.5 * w
        y = cy - 0.5 * h
        return (int(round(x)),int(round(y)),int(round(w)),int(round(h)))

    def transform_inv(self, bbox, delta):
        (x,y,w,h) = bbox
        cx = x + 0.5 * w
        cy = y + 0.5 * h
        cx += delta[0] * w
        cy += delta[1] * h
        w = w * exp(delta[2])
        h = h * exp(delta[3])
        x = cx - 0.5 * w
        y = cy - 0.5 * h
        return (int(round(x)),int(round(y)),int(round(w)),int(round(h)))


    def IoU(self, box1, box2):
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[0]+box1[2], box2[0]+box2[2])
        y2 = min(box1[1]+box1[3], box2[1]+box2[3])
        if x1>=x2 or y1>=y2:
            return 0
        else:
            a1 = box1[2] * box2[3]
            a2 = box1[2] * box2[3]
            return (x2-x1)*(y2-y1)*1.0 / (a1 + a2)

    def saveImage(self, frame, bbox, idx):
        x,y,w,h = bbox
        cv2.imwrite('data/%d.jpg'%idx, frame[y:y+h,x:x+w,:])

    def genData(self, frame, bbox, dw , dh):
        num_pos = 0
        for i in range(batch_size):
            rand_range = 0.2
            dx = rand_range * (2*random.random()-1)
            dy = rand_range * (2*random.random()-1)
            delta = np.array([dx,dy,dw,dh])
            box_sample = self.transform(bbox, delta)
            # self.saveImage(frame,box_sample,i)
            box_large = self.transform(bbox)
            # data_sample = self.transpose(frame, box_sample)
            IoU = self.IoU(box_sample, box_large)