import numpy as np
import cv2

class tracker():
    prev_l_center = 0
    prev_r_center = 0

    def __init__(self, Mywindow_width, Mywindow_height, Mymargin, My_ym =1, My_xm = 1, Mysmooth_factor=15):
        self.recent_centers = []
        self.window_width = Mywindow_width
        self.window_height = Mywindow_height
        self.margin = Mymargin
        self.ym_per_pix = My_ym
        self.xm_per_pix = My_xm
        self.smooth_factor = Mysmooth_factor

    def find_window_centroids(self, warped):
        window_width = self.window_width
        window_height = self.window_height
        margin = self.margin

        window_centroids = []
        window = np.ones(window_width)

        l_sum = np.sum(warped[int(3*warped.shape[0]/4):,:int(warped.shape[1]/2)], axis=0)
        l_center = np.argmax(np.convolve(window, l_sum)) - window_width / 2
        r_sum = np.sum(warped[int(3*warped.shape[0]/4):,int(warped.shape[1]/2):], axis=0)
        r_center = np.argmax(np.convolve(window, r_sum)) - window_width / 2 + int(warped.shape[1]/2)

        if tracker.prev_l_center == 0:
            tracker.prev_l_center = l_center    # if first frame assign to class variable
        elif abs(tracker.prev_l_center - l_center) > self.margin*2:
            l_center = tracker.prev_l_center    # if center is too far away from prev center, assign it to prev center

        if tracker.prev_r_center == 0:
            tracker.prev_r_center = r_center    # if first frame assign to class variable
        elif abs(tracker.prev_r_center - r_center) > self.margin*2:
            r_center = tracker.prev_r_center    # if center is too far away from prev center, assign it to prev center

        #print( " ")
        window_centroids.append((l_center, r_center))
        signal_threshold = 10000.0
        for level in range(1, (int)(warped.shape[0]/window_height)):
            image_layer = np.sum(warped[int(warped.shape[0]-(level+1)*window_height):int(warped.shape[0]-level*window_height), :], axis = 0)
            conv_signal = np.convolve(window, image_layer)
            offset = window_width/2

            l_min_index = int(max(l_center+offset-margin,0))
            l_max_index = int(min(l_center+offset+margin, warped.shape[1]))
            # if conv signal too low, do not move center
            if np.max(conv_signal[l_min_index:l_max_index]) >= signal_threshold:
                l_center = np.argmax(conv_signal[l_min_index:l_max_index])+l_min_index-offset

            #print(l_center, np.max(conv_signal[l_min_index:l_max_index]))

            r_min_index = int(max(r_center+offset-margin, 0))
            r_max_index = int(min(r_center+offset+margin, warped.shape[1]))
            # if conv signal too low, do not move center
            if np.max(conv_signal[r_min_index:r_max_index]) >= signal_threshold:
                r_center = np.argmax(conv_signal[r_min_index:r_max_index])+r_min_index-offset

            #print(r_center, np.max(conv_signal[r_min_index:r_max_index]))

            window_centroids.append((l_center, r_center))

        self.recent_centers.append(window_centroids)
        return np.average(self.recent_centers[-self.smooth_factor:], axis=0)