import numpy
import matplotlib.pyplot as plt
from get_raw_path import *




def get_path(i):
    vid = load_vid('Clips/' + str(i) + '.mp4')
    vid = vid[:, 80:270]
    cloud = get_cloud(vid, thresh=1.5)
    cand = get_cloud_candidates(cloud, vid)
    cand = smooth_candidates(cand)
    fencers = naive_fencer_positions(cand)

    for i in range(len(vid) - 1):
        vid[i] = draw_candidates_on_frame(vid[i], cand[i])

    path = np.zeros_like(cand)

    for i, pair in enumerate(fencers):
        f1, f2 = int(pair[0]), int(pair[1])
        if f1 >= 0:
            path[i, f1] = 255
        if f2 >= 0:
            path[i, f2] = 255

        

    return path[10 : -25], vid[10 : -25]

def get_series(i):
    path, vid = get_path(i)
    
    filt = filter_isolated(path)
    left, right = split_series(filt)

    left = remove_outliers(left[0], left[1], 20)
    left = remove_outliers(left[0], left[1], 10)

    right = remove_outliers(right[0], right[1], 20)
    right = remove_outliers(right[0], right[1], 10)

    left = fill_series(left[0], left[1], path)
    right = fill_series(right[0], right[1], path)

    left = smooth_series(left)
    right = smooth_series(right)

    vid = draw_series_on_vid(vid[:-1], left, right)

    return vid, left, right, path

def draw_series_on_vid(vid, left, right):
    new_vid = copy.copy(vid)
    for i in range(len(vid)):
        new_vid[i] = draw_fencers_on_frame(vid[i],
                                       (left[i], right[i]))
    return new_vid
    
def filter_isolated(path, x_rad=5, y_rad=5):
    path = copy.copy(path)
    for i in range(path.shape[0]):
        for j in range(path.shape[1]):
            if path[i, j] == 255:
                piece = path[i - y_rad : i + y_rad, j - x_rad : j + x_rad]
                if np.sum(piece) == 255:
                    piece[:] = 0
    return path



def kill_triples(path):
    path = copy.copy(path)
    for i in range(path.shape[0]):
        if np.sum(path[i]) > 2 * 255:
            path[i] = 0
    return path


def split_series(path):
    left_t = []
    right_t = []

    left_pos = []
    right_pos = []
    
    
    for i, row in enumerate(path):
        nonzero_inds = [k for k in range(len(row)) if row[k] != 0]

        if np.sum(row) == 0:
            continue
        one_fencer = False
        for k in range(len(row) - 15):
            if np.sum(row[k : k + 15]) == np.sum(row):
                one_fencer = True

        if one_fencer:
            pos = np.average(nonzero_inds)
            if left_t == [] or right_t == []:
                continue
            if np.abs(pos - left_pos[-1]) \
               < np.abs(pos - right_pos[-1]):
                left_t.append(i)
                left_pos.append(pos)
            else:
                right_t.append(i)
                right_pos.append(pos)
            continue
            
       
        losses = np.zeros(path.shape[1])
        for split in range(len(row)):
            left_inds = np.array([k for k in nonzero_inds if k <= split])
            right_inds = np.array([k for k in nonzero_inds if k > split])

            if len(left_inds) == 0:
                left_mse = -1e-10
            else:
                left_mse = np.mean((left_inds - np.mean(left_inds))**2)

            if len(right_inds) == 0:
                right_mse = -1e-10
            else:
                right_mse = np.mean(right_inds - np.mean((right_inds))**2)
           
            losses[split] = left_mse + right_mse
        split = np.argmin(losses)

        new_left_pos = np.mean([k for k in nonzero_inds if k <= split])
        new_right_pos = np.mean([k for k in nonzero_inds if k > split])

        left_t.append(i)
        right_t.append(i)
        left_pos.append(new_left_pos)
        right_pos.append(new_right_pos)
        
    return (left_t, left_pos), (right_t, right_pos)
            
def get_missing_sections(t_arr, path):
    # inclusive exclusive
    total_length = path.shape[0]
    missing_sections = []
    started = False
    
    for i in range(total_length):
        if i not in t_arr and not started:
            a = i
            started = True

        if i in t_arr and started:
            b = i
            started = False
            missing_sections.append((a, b))
    if started:
        missing_sections.append((a, total_length))

    return missing_sections

def remove_outliers(t_arr, pos_arr):
    t_arr = copy.copy(t_arr)
    pos_arr = copy.copy(pos_arr)
    to_kill = []
    for i in range(len(pos_arr) - 3, -1, -1):
        if np.abs(pos_arr[i] - pos_arr[i+2]) < 4 and \
           np.abs(pos_arr[i+1] - pos_arr[i]) > 3:
            to_kill.append(i + 1)
    t_arr = [t_arr[i] for i in range(len(t_arr)) if i not in to_kill]
    pos_arr = [pos_arr[i] for i in range(len(pos_arr)) if i not in to_kill]
    
    return t_arr, pos_arr

def remove_outliers(t_arr, pos_arr, radius):
    t_arr = copy.copy(t_arr)
    pos_arr = copy.copy(pos_arr)
    to_kill = []
    for i in range(len(t_arr)):
        a = max(i - radius, 0)
        b = min(i + radius + 1, len(pos_arr))
        sub_pos = pos_arr[a : b]

        std = np.std(sub_pos)
        lq = np.quantile(sub_pos, 0.25)
        uq = np.quantile(sub_pos, 0.75)
        iqr = uq - lq
        
        median = np.median(sub_pos)

        if pos_arr[i] > uq + 1.5 * iqr or pos_arr[i] < lq - 1.5 * iqr:
            to_kill.append(i)

    to_kill.reverse()

    for i in to_kill:
        t_arr.pop(i)
        pos_arr.pop(i)

    return t_arr, pos_arr

    
    
            

def get_neighbours(path, y, x):
    arr = []
    for i in range(y - 1, y + 2):
        for j in range(x - 1, x + 2):
            if path[i, j] > 0 and not (i == j == 0):
                arr.append((i, j))
    return arr

def get_segment(path, y, x):
    running = True
    points = [(y, x)]
    while running:
        y += 1
        children = []
        for j in range(x - 1, x + 2):
            if path[y, j] > 0:
                children.append((y, j))

        if len(children) == 1:
            points.append(children[0])
            x = children[0][1]
        else:
            running = False
    return points

def fill_series(t_arr, pos_arr, path):
    if len(t_arr) == 0:
        return np.zeros(len(path))
    total_length = len(path)
    series = np.zeros(total_length)

    for ind, t in enumerate(t_arr[:-1]):
        if t_arr[ind + 1] == t + 1:
            series[t] = pos_arr[ind]
        else:
            s = t_arr[ind + 1]
            for i in range(t, s):
                series[i] = pos_arr[ind] + (i - t)/(s - t) * \
                            (pos_arr[ind + 1] - pos_arr[ind])

    for i in range(t_arr[0]):
        series[i] = pos_arr[0]

    for i in range(t_arr[-1], total_length):
        series[i] = pos_arr[-1]
    return series
    

def smooth_series(series, radius=5):
    new_series = copy.copy(series)
    for i in range(len(series)):
        a = max(i - radius, 0)
        b = min(i + radius + 1, len(series))
        window = series[a : b]
        if np.sum(window) == 0:
            new_series[i] = 0
        else:
            new_series[i] = np.average(window[window > 0])
    return new_series
'''
raw_left_arr = []
raw_right_arr = []
left_arr = []
right_arr = []

for i in range(1, 16):
    path, vid = get_path(i)
    filt = filter_isolated(path)
    left, right = split_series(filt)
    raw_left_arr.append(copy.copy(left))
    raw_right_arr.append(copy.copy(right))

    left = remove_outliers(left[0], left[1], 20)
    left = remove_outliers(left[0], left[1], 10)
    right = remove_outliers(right[0], right[1], 20)
    right = remove_outliers(right[0], right[1], 10)

    left_arr.append(left)
    right_arr.append(right)


for i in range(15):
    raw_left, raw_right = raw_left_arr[i], raw_right_arr[i]

    left, right = left_arr[i], right_arr[i]
    print(i + 1)
    print('left')

    plt.scatter(raw_left[0], raw_left[1])
    plt.scatter(left[0], left[1])
    plt.show()

    print('right')

    plt.scatter(raw_right[0], raw_right[1])
    plt.scatter(right[0], right[1])
    plt.show()
'''
