import argparse
import numpy as np
import cv2
import os
import matplotlib.pyplot as plt
from MovingObjectDetector.BackgroundModel import BackgroundModel
from MovingObjectDetector.DetectionRefinement import DetectionRefinement
import TrainNetwork.BaseFunctions as basefunctions
import timeit
from SimpleTracker.KalmanFilter import KalmanFilter
from copy import copy
from MovingObjectDetector.BaseFunctions import TimePropagate, TimePropagate_, draw_error_ellipse2d
import hdf5storage


model_folder = "../Models/"
model_binary, aveImg_binary, model_regression, aveImg_regression = basefunctions.ReadModels(model_folder)
model = (model_binary, aveImg_binary, model_regression, aveImg_regression)

imagefolder = "E:/WAMI-test/WASABI-AngelFire-test/"
num_of_template = 5
IDX = 26

# Read image folder
filenames = os.listdir(imagefolder)
filenames.sort()
print(str(len(filenames)) + " images in the folder...")
# Load background
images = []
for i in range(num_of_template):
    ReadImage = cv2.imread(imagefolder + filenames[i+IDX], cv2.IMREAD_GRAYSCALE)
    images.append(ReadImage)
bgt = BackgroundModel(num_of_template=num_of_template, templates=images)

# Read input image
input_image = cv2.imread(imagefolder + filenames[num_of_template+IDX], cv2.IMREAD_GRAYSCALE)
print("Filename: " + filenames[i+IDX])
Hs = bgt.doCalculateHomography(input_image)

CompensatedImages = bgt.doMotionCompensationAndValidArea(input_image, Hs, input_image.shape)
CandiateCentres, BackgroundSubtractionProperties, BackgroundSubtractionLabels = bgt.doBackgroundSubtraction(
    input_image, thres=8)
print("background subtraction finished...")
dr = DetectionRefinement(input_image, bgt.getCompensatedImages(), CandiateCentres,
                         BackgroundSubtractionProperties, BackgroundSubtractionLabels, model)
Detections1, Detections2, RefinedCentres = dr.do_refine_bs()
Detections1_for_img = [[ele["centre"][1], ele["centre"][0]] for ele in Detections1]
Detections1_for_img = np.int64(np.asarray(Detections1_for_img))
Detections2_for_img = [[ele["centre"][1], ele["centre"][0]] for ele in Detections2]
Detections2_for_img = np.int64(np.asarray(Detections2_for_img))

output_image = copy(input_image)
output_image = cv2.cvtColor(output_image, cv2.COLOR_GRAY2RGB)
for thisDetection in Detections1_for_img:
    if (thisDetection[0] > 0) and (thisDetection[0] < input_image.shape[1]) and (thisDetection[1] > 0)\
            and (thisDetection[1] < input_image.shape[0]):
        cv2.circle(output_image, (thisDetection[0], thisDetection[1]), 6, (0, 200, 0), 1)
for thisDetection in Detections2_for_img:
    if (thisDetection[0] > 0) and (thisDetection[0] < input_image.shape[1]) and (thisDetection[1] > 0)\
            and (thisDetection[1] < input_image.shape[0]):
        cv2.circle(output_image, (thisDetection[0], thisDetection[1]), 6, (0, 0, 200), 1)
print("Draw image finished...")



