import time

__author__ = 'jared'

from json import loads, dumps
import sys

try:
    import numpy as np
    has_numpy = True
except ImportError as e:
    has_numpy = False

try:
    import matplotlib.pyplot as plt
    from matplotlib import ticker
    from matplotlib.ticker import MultipleLocator, FormatStrFormatter, LinearLocator
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import cm
    from matplotlib.ticker import LinearLocator, FormatStrFormatter
    has_matplotlib = True
except ImportError as e:
    has_matplotlib = False

def tracker_output_to_tuples(input_filename):
    def text_to_list(filename):
        result = []
        text = open(filename).read()
        for line in text.split('\n'):
            this_point = line.split(' ')[:3]
            if len(this_point) != 3:
                # print 'Ignoring line "%s" since there are not 2 spaces in that line.' % line
                continue
            frame, x, y = this_point
            if x == '-':
                continue
            result.append(this_point)
        return result

    frame_tuples = text_to_list(input_filename)
    f_list, x_list, y_list = [map(int, lst) for lst in zip(*frame_tuples)]
    y_ave = round(float(sum(y_list))/float(len(y_list)))
    print 'y_ave = %f' % y_ave
    print max(f_list),'frames'
    return zip(f_list, x_list)


def get_roi_data(input_filename):
    with open(input_filename) as f:
        txt = f.read()
        result = loads(txt)
    return result


def interp(t1, x1, t2, x2, t_mid):
    previous_location_guess = x1
    linear_guess = x1 + ((t_mid - t1) / float(t2 - t1)) * (x2 - x1)
    return linear_guess


def roi_data_to_dict(data):
    result = {}
    tt, xx = zip(*data)
    # first few frames
    for i in range(tt[0]+1):
        result[i] = xx[0]
    t1 = tt[0]
    x1 = xx[0]
    for i in range(1, len(tt)):
        t2 = tt[i]
        x2 = xx[i]

        t_itt = t1
        while t2 - t_itt > 1: # fill in where no data from tracker, guess location with interp
            t_itt += 1
            result[t_itt] = interp(t1, x1, t2, x2, t_itt)
        result[t2] = x2
        t1 = t2
        x1 = x2
    return result


'''
Fills in the pan as a function of frame given the locations to pan and the roi as function of frame (as a dict)
Returns the pan as a dict mapping frames to locations.
'''
def pan_location_points_to_dict(pan_location_points, roi_dict):
    result = {}
    p = 0.4
    #fill in first few points
    if not pan_location_points:
        return result
    first_pan_point = pan_location_points[0][0]
    t0, x0 = first_pan_point
    for i in range(t0):
        result[i] = x0
    for pan1, pan2 in zip(pan_location_points, pan_location_points[1:]):
        point1, point2 = pan1
        point3, point4 = pan2
        #fill in pan data
        interp_points = zip(*get_pan_location_as_points(point1, point2, p))
        for point in interp_points:
            t_i, x_i = point
            result[t_i] = x_i
        t2, x2 = point2
        t3, x3 = point3
        #fill in line to next pan
        for ti in range(t2, t3):
            result[ti] = x2
    #fill in last points
    point1, point2 = pan_location_points[-1]
    interp_points = zip(*get_pan_location_as_points(point1, point2, p))
    for point in interp_points:
        t_i, x_i = point
        result[t_i] = x_i
    tf, xf = point2
    while tf <= max(roi_dict.keys()):
        result[tf] = xf
        tf += 1
    return result


def get_cleaned_data(roi_data, ave_width=5, thresh=80):
    roi_dict = roi_data_to_dict(roi_data)
    ref_dict = roi_data_to_dict(roi_data)
    frames = sorted(roi_dict.keys())
    #smooth the data
    for i in range(0, len(frames), ave_width):
        frame_slice = frames[i:i + ave_width]
        values = [roi_dict[frame] for frame in frame_slice]
        average = sum(values) / float(ave_width)
        for j in range(ave_width):
            frame = frames[j]
            roi_dict[frame] = average

    # mean filter
    import copy
    for k in range(10):
        for i in range(len(frames)):
            ave = ref_dict[frames[i]]
            cnt = 1
            for j in range(1, 3*ave_width):
                if i-j > 0:
                    cnt+=1
                    ave+=ref_dict[frames[i-j]]
                else:
                    cnt+=1
                    ave+=ref_dict[frames[i]]
                if i+j < len(frames)-1:
                    cnt+=1
                    ave+=ref_dict[frames[i+j]]
                else:
                    cnt+=1
                    ave+=ref_dict[frames[i]]
            roi_dict[frames[i]] = float(ave)/cnt
        ref_dict = copy.copy(roi_dict)
    xx = sorted(roi_dict.keys())
    yy = [roi_dict[x] for x in xx]
    plt.plot(xx, yy, 'b')
    cleaned_data = [[frames[0], roi_dict[frames[0]]]]
    i = 0
    mxi = 0 # index of the frame whose roi value is maximal
    mni = 0
    while i < len(frames):
        prev_frame, prev_val = cleaned_data[-1]
        val = roi_dict[frames[i]]  # x value at frame
        min_val, max_val = roi_dict[frames[mni]], roi_dict[frames[mxi]]
        if val > max_val:
            mxi = i
            max_val = val
        if val < min_val:
            mni = i
            min_val = val
        if abs(max_val - min_val) > thresh:  # threshold reached, add a "cleaned point"
            if min_val < prev_val < max_val:  # if end points are not local max and min respectively
                if frames[mxi] < frames[mni]:  # then grab the max value in the range if it comes first
                    i = mxi
                else:  # otherwise grab min value in the range
                    i = mni
            point = [frames[i], roi_dict[frames[i]]]
            cleaned_data.append(point)
            mxi = i
            mni = i
        i += 1
    cleaned_data.append([frames[-1], roi_dict[frames[-1]]])
    return cleaned_data


'''
Given a list of x,y values of the form [[x1,y1],[x2,y2]], returns the
filtered list of local extrema. This list includes the first and last
values in the original list.
'''
def find_extrema(data):
    result = [data[0]]
    for i in range(1, len(data) - 1):
        x_prev, y_prev = data[i - 1]
        x_itt, y_itt = data[i]
        x_next, y_next = data[i + 1]
        is_local_max = y_itt > y_prev and y_itt > y_next
        is_local_min = y_itt < y_prev and y_itt < y_next
        if is_local_max or is_local_min:
            result.append([x_itt, y_itt])
    result.append(data[-1])
    return result


'''
Given a list of intervals like [[1,6],[3,9]], returns the list of
merged intervals, such as [[1,9]].
'''
def merge_intervals(intervals):
    intervals = sorted(intervals)
    result = [intervals[0]]
    for i in range(1, len(intervals)):
        stack_top = result[-1]
        this_interval = intervals[i]
        if this_interval[0] > stack_top[1]:
            result.append(this_interval)
        elif this_interval[1] > stack_top[1]:
            stack_top[1] = this_interval[1]
    return result


'''
Given extrema and roi_data, find intervals around each extrema on which the
roi_data are within (-y_thresh, y_thresh). Further, the intervals should not
be more than 2*x_thresh+1 wide.
The intervals are of the form [i1, i2], where i1 refers to the index of
roi_data where the interval starts
'''
def get_intervals(extrema_points, roi_data, xthresh=25*4, ythresh=80):
    data_iterator = 0
    result = []
    for extreme_point in extrema_points:
        px, py = extreme_point
        dx, dy = roi_data[data_iterator]
        if dx > px:
            data_iterator = 0
            dx, dy = roi_data[data_iterator]
        while dx < px: # move the roi_data pointer to the extreme points frame
            data_iterator += 1
            dx, dy = roi_data[data_iterator]
        left = 0
        right = 0
        # go right
        dy = py
        while abs(dy - py) < ythresh and right < xthresh and data_iterator + right < len(roi_data) - 1:
            right += 1
            dx, dy = roi_data[data_iterator + right]
        #go left
        dy = py
        while abs(dy - py) < ythresh and left < xthresh and data_iterator - left > 0:
            left += 1
            dx, dy = roi_data[data_iterator - left]
        interval = [data_iterator - left, data_iterator + right]
        result.append(interval)
    return result


def average_y_data(interval, points):
    i1, i2 = interval
    y_values = zip(*points[i1:i2])[1]
    return sum(y_values) / float(i2-i1)


# ~ def average_y_data2(interval,data):
#~ a, b=interval
#~ yvals = zip(*data[a:b])[1]
#~ return (max(yvals)-min(yvals))/2.0


'''
Given non-overlapping intervals and roi_data, find panning points. Note that intervals
are of the form [i1,i2] where i1 is the index in the roi_data where we should start
panning.
'''
def get_pan_locations(intervals, roi_data, clipped_frame_width):
    #we must get the average y-value of the data over each interval to know
    #where to pan to.
    result = []
    for i in range(1, len(intervals)):
        a, b = intervals[i - 1:i + 1]
        left_point = min(a[1], b[0])  # a[1]
        right_point = max(a[1], b[0])  # b[0]
        if left_point is not right_point:
            x1, y1 = roi_data[left_point]
            y1 = average_y_data(a, roi_data)
            x2, y2 = roi_data[right_point]
            y2 = average_y_data(b, roi_data)  # could also pan to max or min value in data...
            # moving_right = (y2-y1 > 0)
            # if moving_right:
            #     y2-=0.1*clipped_frame_width
            # else:
            #     y2+=0.1*clipped_frame_width
            result.append([[x1, y1], [x2, y2]])
    return result


'''
Helper method for plotting. This takes in pan points of the form [t, x]
and returns a list of t and x values for plotting. It uses the panning
function which we have developed, which accelerates for p percent of the
time and decelerates the rest of the time.
'''
def get_pan_location_as_points(pan_point1, pan_point2, p):
    assert 0 <= p <= 1
    local_pan_t = []
    local_pan_x = []
    t1, x1 = pan_point1
    t2, x2 = pan_point2
    DX = float(x2 - x1)
    DT = float(t2 - t1)
    t_mid = t1 + p * DT # (t1 + t2) * p
    tt = np.arange(t1, int(round(t_mid)), 1)
    c = DX / (p * DT ** 2)
    xx = x1 + c * (tt - t1) ** 2
    local_pan_t += list(tt)
    local_pan_x += list(xx)

    tt = np.arange(int(round(t_mid)), t2, 1)
    c0 = p * DX
    c1 = 2 * DX / DT
    c2 = (tt - t_mid)
    c3 = DX / ((1 - p) * DT ** 2)
    xx = x1 + c0 + c1 * c2 - c3 * (c2 ** 2)
    local_pan_t += list(tt)
    local_pan_x += list(xx)
    return [local_pan_t, local_pan_x]


def format_axes(major_ticks_location, minor_ticks_location, scale, ax):
    major_locator = MultipleLocator(major_ticks_location)
    major_formatter = ticker.FuncFormatter(lambda x, pos: '%d' % (x / float(scale)))
    minor_locator = MultipleLocator(minor_ticks_location)
    ax.xaxis.set_major_locator(major_locator)
    ax.xaxis.set_major_formatter(major_formatter)
    ax.xaxis.set_minor_locator(minor_locator)


'''
Gets the pan_location_points in the form [ [[x1,y1],[x2,y2]], [[x3,y3=y2],[x4,y4]], ... ]
'''
def get_pan_location_points(p_xthresh, p_ythresh, roi_data, clipped_frame_width):
    #xx, yy = zip(*roi_data)
    #plt.plot(xx, yy, 'c.')
    cleaned_data = get_cleaned_data(roi_data, ave_width=5, thresh=p_ythresh)
    #xx, yy = zip(*cleaned_data)
    #plt.plot(xx, yy, 'g*')
    extrema = find_extrema(cleaned_data)
    #xx, yy = zip(*extrema)
    #plt.plot(xx, yy, 'ro')
    overlapping_intervals = get_intervals(extrema, roi_data, xthresh=p_xthresh, ythresh=p_ythresh)
    #for i, extremepoint in enumerate(extrema):
        #begin, end = overlapping_intervals[i]
        #upper_half = extremepoint[1]+ythresh
        #lower_half = extremepoint[1]-ythresh
        #ax = plt.gca()
        #ax.fill_between(range(roi_data[begin][0], roi_data[end][0]), lower_half, upper_half, facecolor='cyan', alpha=0.25)
    intervals = merge_intervals(overlapping_intervals)
    #for begin, end in intervals:
        #upper_half = 220
        #lower_half = 200
        #ax.fill_between(range(roi_data[begin][0], roi_data[end][0]), lower_half, upper_half, facecolor='cyan', alpha=1)
    pan_location_points = get_pan_locations(intervals, roi_data, clipped_frame_width)
    return pan_location_points


# returns roi data and points where we need to pan from/to.
def make_and_return_pan_data(xthresh, ythresh, roi_filename_or_points, clipped_frame_width):
    if isinstance(roi_filename_or_points, str): # if its a filename
        roi_points = get_roi_data(roi_filename_or_points) # get points from file
    else:
        roi_points = roi_filename_or_points # if its not a filename, it is points.
    pan_location_points = get_pan_location_points(xthresh, ythresh, roi_points, clipped_frame_width)
    return roi_points, pan_location_points


def plot_points(roi_points, pan_location_points, clipped_frame_width, fps=25):

    fig, ax = plt.subplots()

    pan_t = []
    pan_x = []  # represents ALL pans x(t) values
    #plot pan points
    a,b = pan_location_points[0]
    tt, xx = get_pan_location_as_points(a, b, p=0.5)
    pan_t += tt
    pan_x += xx
    plt.plot(tt, xx, 'r')
    for pan1, pan2 in zip(pan_location_points, pan_location_points[1:]):
        a, b = pan2
        a[1] = pan1[1][1]
        tt, xx = get_pan_location_as_points(a, b, p=0.5)
        pan_t += tt
        pan_x += xx
        plt.plot(tt, xx, 'r')
    #plot between pan points
    for pan1, pan2 in zip(pan_location_points, pan_location_points[1:]):
        point1, point2 = pan1
        point3, point4 = pan2
        t1, x1 = point2
        t2, x2 = point3
        # print 'plotting from',x1,'to',x2
        tt = np.arange(t1, t2)
        xx = np.ones(len(list(tt))) * (x1 + x2) / 2.
        pan_t += list(tt)
        pan_x += list(xx)
        plt.plot(tt, xx, 'c--')
    # finally get the entire panned video as a function of frame:
    pan_t, pan_x = zip(*sorted(zip(pan_t, pan_x)))
    lower_half = [x-0.5*clipped_frame_width for x in pan_x]
    upper_half = [x+0.5*clipped_frame_width for x in pan_x]
    ax.fill_between(pan_t, lower_half, upper_half, facecolor='cyan', alpha=0.25)

    # plot initial data
    tt, xx = zip(*roi_points)
    plt.plot(tt, xx, 'b.')

    # plot clean data (lines)
    #xx, yy = zip(*get_cleaned_data(data_to_dict(roi_points), aveWidth=5, thresh=ythresh))
    #plt.plot(xx,yy,'g*')

    #Format the axes to get units right
    format_axes(5*60*fps, 60*fps, 60*fps, ax)

    plt.xlabel('time (minutes)')
    plt.ylabel('x-coordinate (pixels)')

    plt.show()


def calculate_scores(roi_points, pan_location_points, clipped_frame_size):
    scores = {
        'out_of_frame_too_long': 0,
        'pan_too_much': 0,
        'pan_length_penalty': 0,
        'pan_count': 0
    }

    if not pan_location_points:
        for k in scores:
            scores[k] = -100
        return scores

    def pan_length_penalty(pan_length):
        '''
        0s to 30sec really awful - 100
        50ms to 10sec awful - 20
        100ms to 4sec bad - 9
        200ms to 2.5sec not so good  - 4
        300ms to 2sec is ok      - 1
        1 sec to 1.5 sec is best - 0
        '''

        t = float(pan_length)/fps
        if 1 < t < 1.5:
            score = 0
        elif 0.3 < t < 2:
            score = 1
        elif 0.2 < t < 3:
            score = 4
        elif 0.1 < t < 5:
            score = 9
        elif 0.05 < t < 10:
            score = 20
        else:
            score = 100
        return score

    roi_dict = roi_data_to_dict(roi_points)
    pan_dict = pan_location_points_to_dict(pan_location_points, roi_dict)
    out_of_frame_time = 0
    is_panning_previous_frame = False
    is_lecturer_in_previous_frame = False
    total_out_of_frame_time = 0
    for t_i in sorted(pan_dict.keys()):
        rx_i = roi_dict[t_i]
        px_i = pan_dict[t_i]
        dx = px_i - rx_i
        # Panning
        px_prev = pan_dict[max((t_i-1), 0)]
        is_panning = (px_i != px_prev)
        is_started_panning = is_panning and not is_panning_previous_frame
        is_stopped_panning = not is_panning and is_panning_previous_frame
        if is_started_panning:
            length_this_pan = 0
        if is_panning:
            length_this_pan += 1
            scores['pan_too_much'] += 1
        if is_stopped_panning:
            scores['pan_length_penalty'] += pan_length_penalty(length_this_pan)
            scores['pan_count'] += 1
            length_this_pan = 0
        is_panning_previous_frame = (px_i != px_prev)
        # Lecturer in frame
        is_lecturer_in_frame = (abs(dx) <= 0.5 * clipped_frame_size)
        is_lecturer_entered_frame = not is_lecturer_in_previous_frame and is_lecturer_in_frame
        is_lecturer_left_frame = is_lecturer_in_previous_frame and not is_lecturer_in_frame
        if is_lecturer_entered_frame:
            scores['out_of_frame_too_long'] += out_of_frame_time**2
            total_out_of_frame_time += out_of_frame_time
            out_of_frame_time = 0
        if not is_lecturer_in_frame:
            out_of_frame_time += 1
        if is_lecturer_left_frame:
            out_of_frame_time = 1
        is_lecturer_in_previous_frame = is_lecturer_in_frame

    scores['out_of_frame_too_long'] = scores['out_of_frame_too_long'] / max(1, float(total_out_of_frame_time))
    scores['pan_too_much'] = (100.*scores['pan_too_much']) / len(pan_dict.keys())
    scores['pan_length_penalty'] /= max(1, scores['pan_count'])

    return scores


def write_yaml(lines, filename):
    result = '%YAML:1.0\n'
    result += 'panPoints:\n'
    result += '\n'.join(lines)
    with open(filename, 'w+') as yaml_file:
        yaml_file.write(result)


def write_data_to_yaml(roi_points, pan_location_points, roi_filename, pan_filename, fps=25):
    opencv_roi_points = ['  - { frame: %d, location: %d }' % point for point in roi_points]
    write_yaml(opencv_roi_points, roi_filename)

    opencv_pan_points = []
    for a, b in pan_location_points:
        frame_start = a[0]
        frame_end = b[0]
        frames = frame_end - frame_start
        seconds = frames / float(fps)
        location = int(round(b[1]))
        opencv_pan_points.append(
            '  - { frame_start: %d, seconds: %f, location: %d }' %
            (frame_start, seconds, location)
        )
    write_yaml(opencv_pan_points, pan_filename)


def get_scores(roi_points, pan_location_points, clipped_frame_width, N):
    import random
    import datetime

    print 'making the pan'
    score_out_of_frame_too_long = np.zeros((N, N))
    score_pan_too_much = np.zeros((N, N))
    score_pan_length_penalty = np.zeros((N, N))
    score_pan_count = np.zeros((N, N))
    t0 = time.time()
    for xthresh in range(10, 10 * (N + 1), 10):
        for ythresh in range(10, 10 * (N + 1), 10):
            i = int(xthresh / 10. - 1)
            j = int(ythresh / 10. - 1)
            t1 = time.time()
            if random.random() < 0.1 and t1 - t0 > 5:
                t1 = time.time()
                k = (N * (xthresh - 10) / 10. + (ythresh - 10) / 10.) / (N * N)
                time_left = (1 - k) / k * (t1 - t0)
                print '%s remains as %s elapsed and %.1f%% complete.' % \
                      (str(datetime.timedelta(seconds=int(time_left))),
                       str(datetime.timedelta(seconds=int(t1 - t0))),
                       100 * k)
            roi_points, pan_location_points = make_and_return_pan_data(xthresh, ythresh, roi_points, clipped_frame_width)
            score_map = calculate_scores(roi_points, pan_location_points, clipped_frame_width)
            score_out_of_frame_too_long[i, j] = score_map['out_of_frame_too_long']
            score_pan_too_much[i, j] = score_map['pan_too_much']
            score_pan_length_penalty[i, j] = score_map['pan_length_penalty']
            score_pan_count[i, j] = score_map['pan_count']

    return score_out_of_frame_too_long, score_pan_too_much, score_pan_length_penalty, score_pan_count


def plot_scores(score_out_of_frame_too_long, score_pan_too_much, score_pan_length_penalty, score_pan_count, N):
    def plot(ax_string, X, Y, Z, zlabel):
        ax = fig.add_subplot(ax_string, projection='3d')
        surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,
                               linewidth=0, antialiased=False)
        fig.colorbar(surf, shrink=0.5, aspect=5)
        ax.zaxis.set_major_locator(LinearLocator(10))
        ax.zaxis.set_major_formatter(FormatStrFormatter('%d'))
        ax.set_xlabel('xthresh')
        ax.set_ylabel('ythresh')
        ax.set_zlabel(zlabel)

    fig = plt.figure(figsize=plt.figaspect(0.5))
    X = np.arange(10, 10 * (N + 1), 10)
    Y = np.arange(10, 10 * (N + 1), 10)
    X, Y = np.meshgrid(X, Y)

    plot(221, X, Y, score_out_of_frame_too_long, 'leave frame too long penalty')
    plot(222, X, Y, score_pan_too_much, 'panning total time')
    plot(223, X, Y, score_pan_length_penalty, 'pan length penalty')
    plot(224, X, Y, score_pan_count, 'number of pans')

    plt.show()


def write_score(filename, score):
    np.save(open(filename,'w+'), score)


def main(xthresh, ythresh, fps, clipped_frame_width):
    input_filename = '../../tracking_output.txt' # default value
    pan_yaml_filename = 'panDataPointsOut.yaml'
    roi_yaml_filename = 'initialROIDataOpenCV.yaml'
    if len(sys.argv) > 1:
        input_filename = sys.argv[1]

    json_data = tracker_output_to_tuples(input_filename)  # get data as json
    
    roi_points, pan_location_points = make_and_return_pan_data(xthresh, ythresh, json_data, clipped_frame_width)

    # print 'plotting the pan'
    # plot_points(roi_points, pan_location_points, clipped_frame_width)
    #
    write_data_to_yaml(roi_points, pan_location_points, roi_yaml_filename, pan_yaml_filename , fps=fps)

    # print 'calculating scores'
    #N=40
    #score_out_of_frame_too_long, score_pan_too_much, score_pan_length_penalty, score_pan_count =\
    #    get_scores(roi_points, pan_location_points, clipped_frame_width, N)
    
    #plot_scores(score_out_of_frame_too_long, score_pan_too_much, score_pan_length_penalty, score_pan_count, N)

    # write_score(score_out_of_frame_too_long, 'score_out_of_frame_too_long_30.npy')
    # write_score(score_pan_too_much, 'score_pan_too_much_30.npy')

    # plot_scores(score_out_of_frame_too_long, score_pan_too_much, 5)

fps = 25
clipped_frame_width = 800
xthresh = fps*10
ythresh = 280
main(xthresh, ythresh, fps, clipped_frame_width)
