from .visual import CocoPart
import numpy as np
from helpers import *
from default_params import *


def match_ip(ip_set, new_ips, lstm_set, num_matched, consecutive_frames=DEFAULT_CONSEC_FRAMES):
    len_ip_set = len(ip_set)
    added = [False for _ in range(len_ip_set)]

    new_len_ip_set = len_ip_set
    for new_ip in new_ips:
        if not is_valid(new_ip):
            continue
        # assert valid_candidate_hist(new_ip)
        cmin = [MIN_THRESH, -1]
        for i in range(len_ip_set):
            if not added[i] and dist(last_ip(ip_set[i])[0], new_ip) < cmin[0]:
                # here add dome condition that last_ip(ip_set[0] >-5 or someting)
                cmin[0] = dist(last_ip(ip_set[i])[0], new_ip)
                cmin[1] = i

        if cmin[1] == -1:
            ip_set.append([None for _ in range(consecutive_frames - 1)] + [new_ip])
            lstm_set.append([None, 0, 0, 0])  # Initial hidden state of lstm is None
            new_len_ip_set += 1

        else:
            added[cmin[1]] = True
            pop_and_add(ip_set[cmin[1]], new_ip, consecutive_frames)

    new_matched = num_matched

    removed_indx = []
    removed_match = []

    for i in range(len(added)):
        if not added[i]:
            pop_and_add(ip_set[i], None, consecutive_frames)
        if ip_set[i] == [None for _ in range(consecutive_frames)]:
            if i < num_matched:
                new_matched -= 1
                removed_match.append(i)

            new_len_ip_set -= 1
            removed_indx.append(i)

    for i in sorted(removed_indx, reverse=True):
        ip_set.pop(i)
        lstm_set.pop()

    return new_matched, new_len_ip_set, removed_match


def extend_vector(p1, p2, l):
    p1 += (p1-p2)*l/(2*np.linalg.norm((p1-p2), 2))
    p2 -= (p1-p2)*l/(2*np.linalg.norm((p1-p2), 2))
    return p1, p2


def perp(a):
    b = np.empty_like(a)
    b[0] = -a[1]
    b[1] = a[0]
    return b

# line segment a given by endpoints a1, a2
# line segment b given by endpoints b1, b2
# return


def seg_intersect(a1, a2, b1, b2):
    da = a2-a1
    db = b2-b1
    dp = a1-b1
    dap = perp(da)
    denom = np.dot(dap, db)
    num = np.dot(dap, dp)
    return (num / denom.astype(float))*db + b1


def get_kp(kp):
    threshold1 = 5e-3

    # dict of np arrays of coordinates
    inv_pend = {}
    # print(type(kp[CocoPart.LEar]))
    numx = (kp[CocoPart.LEar][2]*kp[CocoPart.LEar][0] + kp[CocoPart.LEye][2]*kp[CocoPart.LEye][0] +
            kp[CocoPart.REye][2]*kp[CocoPart.REye][0] + kp[CocoPart.REar][2]*kp[CocoPart.REar][0])
    numy = (kp[CocoPart.LEar][2]*kp[CocoPart.LEar][1] + kp[CocoPart.LEye][2]*kp[CocoPart.LEye][1] +
            kp[CocoPart.REye][2]*kp[CocoPart.REye][1] + kp[CocoPart.REar][2]*kp[CocoPart.REar][1])
    den = kp[CocoPart.LEar][2] + kp[CocoPart.LEye][2] + kp[CocoPart.REye][2] + kp[CocoPart.REar][2]

    if den < HEAD_THRESHOLD:
        inv_pend['H'] = None
    else:
        inv_pend['H'] = np.array([numx/den, numy/den])

    if all([kp[CocoPart.LShoulder], kp[CocoPart.RShoulder],
            kp[CocoPart.LShoulder][2] > threshold1, kp[CocoPart.RShoulder][2] > threshold1]):
        inv_pend['N'] = np.array([(kp[CocoPart.LShoulder][0]+kp[CocoPart.RShoulder][0])/2,
                                  (kp[CocoPart.LShoulder][1]+kp[CocoPart.RShoulder][1])/2])
    else:
        inv_pend['N'] = None

    if all([kp[CocoPart.LHip], kp[CocoPart.RHip],
            kp[CocoPart.LHip][2] > threshold1, kp[CocoPart.RHip][2] > threshold1]):
        inv_pend['B'] = np.array([(kp[CocoPart.LHip][0]+kp[CocoPart.RHip][0])/2,
                                  (kp[CocoPart.LHip][1]+kp[CocoPart.RHip][1])/2])
    else:
        inv_pend['B'] = None

    if kp[CocoPart.LKnee] is not None and kp[CocoPart.LKnee][2] > threshold1:
        inv_pend['KL'] = np.array([kp[CocoPart.LKnee][0], kp[CocoPart.LKnee][1]])
    else:
        inv_pend['KL'] = None

    if kp[CocoPart.RKnee] is not None and kp[CocoPart.RKnee][2] > threshold1:
        inv_pend['KR'] = np.array([kp[CocoPart.RKnee][0], kp[CocoPart.RKnee][1]])
    else:
        inv_pend['KR'] = None

    if inv_pend['B'] is not None:
        if inv_pend['N'] is not None:
            height = np.linalg.norm(inv_pend['N'] - inv_pend['B'], 2)
            LS, RS = extend_vector(np.asarray(kp[CocoPart.LShoulder][:2]),
                                   np.asarray(kp[CocoPart.RShoulder][:2]), height/4)
            LB, RB = extend_vector(np.asarray(kp[CocoPart.LHip][:2]),
                                   np.asarray(kp[CocoPart.RHip][:2]), height/3)
            ubbox = (LS, RS, RB, LB)

            if inv_pend['KL'] is not None and inv_pend['KR'] is not None:
                lbbox = (LB, RB, inv_pend['KR'], inv_pend['KL'])
            else:
                lbbox = ([0, 0], [0, 0])
                #lbbox = None
        else:
            ubbox = ([0, 0], [0, 0])
            #ubbox = None
            if inv_pend['KL'] is not None and inv_pend['KR'] is not None:
                lbbox = (np.array(kp[CocoPart.LHip][:2]), np.array(kp[CocoPart.RHip][:2]),
                         inv_pend['KR'], inv_pend['KL'])
            else:
                lbbox = ([0, 0], [0, 0])
                #lbbox = None
    else:
        ubbox = ([0, 0], [0, 0])
        lbbox = ([0, 0], [0, 0])
        #ubbox = None
        #lbbox = None
    # condition = (inv_pend["H"] is None) and (inv_pend['N'] is not None and inv_pend['B'] is not None)
    # if condition:
    #     print("half disp")

    return inv_pend, ubbox, lbbox


def get_angle(v0, v1):
    return np.math.atan2(np.linalg.det([v0, v1]), np.dot(v0, v1))


def is_valid(ip):

    assert ip is not None

    ip = ip["keypoints"]
    return (ip['B'] is not None and ip['N'] is not None and ip['H'] is not None)


def get_rot_energy(ip0, ip1):
    t = ip1["time"] - ip0["time"]
    ip0 = ip0["keypoints"]
    ip1 = ip1["keypoints"]
    m1 = 1
    m2 = 5
    m3 = 5
    energy = 0
    den = 0
    N1 = ip1['N'] - ip1['B']
    N0 = ip0['N'] - ip0['B']
    d2sq = N1.dot(N1)
    w2sq = (get_angle(N0, N1)/t)**2
    energy += m2*d2sq*w2sq

    den += m2*d2sq
    H1 = ip1['H'] - ip1['B']
    H0 = ip0['H'] - ip0['B']
    d1sq = H1.dot(H1)
    w1sq = (get_angle(H0, H1)/t)**2
    energy += m1*d1sq*w1sq
    den += m1*d1sq

    energy = energy/(2*den)
    # energy = energy/2
    return energy


def get_angle_vertical(v):
    return np.math.atan2(-v[0], -v[1])


def get_gf(ip0, ip1, ip2):
    t1 = ip1["time"] - ip0["time"]
    t2 = ip2["time"] - ip1["time"]
    ip0 = ip0["keypoints"]
    ip1 = ip1["keypoints"]
    ip2 = ip2["keypoints"]

    m1 = 1
    m2 = 15
    g = 10
    H2 = ip2['H'] - ip2['N']
    H1 = ip1['H'] - ip1['N']
    H0 = ip0['H'] - ip0['N']
    d1 = np.sqrt(H1.dot(H1))
    theta_1_plus_2_2 = get_angle_vertical(H2)
    theta_1_plus_2_1 = get_angle_vertical(H1)
    theta_1_plus_2_0 = get_angle_vertical(H0)
    # print("H: ",H0,H1,H2)
    N2 = ip2['N'] - ip2['B']
    N1 = ip1['N'] - ip1['B']
    N0 = ip0['N'] - ip0['B']
    d2 = np.sqrt(N1.dot(N1))
    # print("N: ",N0,N1,N2)
    theta_2_2 = get_angle_vertical(N2)
    theta_2_1 = get_angle_vertical(N1)
    theta_2_0 = get_angle_vertical(N0)
    #print("theta_2_2:",theta_2_2,"theta_2_1:",theta_2_1,"theta_2_0:",theta_2_0,sep=", ")
    theta_1_0 = theta_1_plus_2_0 - theta_2_0
    theta_1_1 = theta_1_plus_2_1 - theta_2_1
    theta_1_2 = theta_1_plus_2_2 - theta_2_2

    # print("theta1: ",theta_1_0,theta_1_1,theta_1_2)
    # print("theta2: ",theta_2_0,theta_2_1,theta_2_2)

    theta2 = theta_2_1
    theta1 = theta_1_1

    del_theta1_0 = (get_angle(H0, H1))/t1
    del_theta1_1 = (get_angle(H1, H2))/t2

    del_theta2_0 = (get_angle(N0, N1))/t1
    del_theta2_1 = (get_angle(N1, N2))/t2
    # print("del_theta2_1:",del_theta2_1,"del_theta2_0:",del_theta2_0,sep=",")
    del_theta1 = 0.5 * (del_theta1_1 + del_theta1_0)
    del_theta2 = 0.5 * (del_theta2_1 + del_theta2_0)

    doubledel_theta1 = (del_theta1_1 - del_theta1_0) / 0.5*(t1 + t2)
    doubledel_theta2 = (del_theta2_1 - del_theta2_0) / 0.5*(t1 + t2)
    # print("doubledel_theta2:",doubledel_theta2)

    d1 = d1/d2
    d2 = 1
    # print("del_theta",del_theta1,del_theta2)
    # print("doubledel_theta",doubledel_theta1,doubledel_theta2)

    Q_RD1 = 0
    Q_RD1 += m1 * d1 * doubledel_theta1 * doubledel_theta1
    Q_RD1 += (m1*d1*d1 + m1*d1*d2*np.cos(theta1))*doubledel_theta2
    Q_RD1 += m1*d1*d2*np.sin(theta1)*del_theta2*del_theta2
    Q_RD1 -= m1*g*d2*np.sin(theta1+theta2)

    Q_RD2 = 0
    Q_RD2 += (m1*d1*d1 + m1*d1*d2*np.cos(theta1))*doubledel_theta1
    Q_RD2 += ((m1+m2)*d2*d2 + m1*d1*d1 + 2*m1*d1*d2*np.cos(theta1))*doubledel_theta2
    Q_RD2 -= 2*m1*d1*d2*np.sin(theta1)*del_theta2*del_theta1 + m1*d1 * \
        d2*np.sin(theta1)*del_theta1*del_theta1
    Q_RD2 -= (m1 + m2)*g*d2*np.sin(theta2) + m1*g*d1*np.sin(theta1 + theta2)

    # print("Energy: ", Q_RD1 + Q_RD2)
    return Q_RD1 + Q_RD2


def get_height_bbox(ip):
    bbox = ip["box"]
    assert isinstance(bbox, np.ndarray)
    diff_box = bbox[1] - bbox[0]
    return diff_box[1]


def get_ratio_bbox(ip):
    bbox = ip["box"]
    assert isinstance(bbox, np.ndarray)
    diff_box = bbox[1] - bbox[0]
    if diff_box[1] == 0:
        diff_box[1] += 1e5*diff_box[0]
    assert(np.any(diff_box > 0))
    ratio = diff_box[0]/diff_box[1]
    return ratio


def get_ratio_derivative(ip0, ip1):
    ratio_der = None
    time = ip1["time"] - ip0["time"]
    diff_box = ip1["features"]["ratio_bbox"] - ip0["features"]["ratio_bbox"]
    assert time != 0
    ratio_der = diff_box/time

    return ratio_der


def match_ip2(matched_ip_set, unmatched_ip_set, new_ips, re_matrix, gf_matrix, consecutive_frames=DEFAULT_CONSEC_FRAMES):
    len_matched_ip_set = len(matched_ip_set)
    added_matched = [False for _ in range(len_matched_ip_set)]
    len_unmatched_ip_set = len(unmatched_ip_set)
    added_unmatched = [False for _ in range(len_unmatched_ip_set)]
    for new_ip in new_ips:
        if not is_valid(new_ip):
            continue
        cmin = [MIN_THRESH, -1]
        connected_set = None
        connected_added = None
        for i in range(len_matched_ip_set):
            if not added_matched[i] and dist(last_ip(matched_ip_set[i])[0], new_ip) < cmin[0]:
                # here add dome condition that last_ip(ip_set[0] >-5 or someting)
                cmin[0] = dist(last_ip(matched_ip_set[i])[0], new_ip)
                cmin[1] = i
                connected_set = matched_ip_set
                connected_added = added_matched
        for i in range(len_unmatched_ip_set):
            if not added_unmatched[i] and dist(last_ip(unmatched_ip_set[i])[0], new_ip) < cmin[0]:
                # here add dome condition that last_ip(ip_set[0] >-5 or someting)
                cmin[0] = dist(last_ip(unmatched_ip_set[i])[0], new_ip)
                cmin[1] = i
                connected_set = unmatched_ip_set
                connected_added = added_unmatched

        if cmin[1] == -1:
            unmatched_ip_set.append([None for _ in range(consecutive_frames - 1)] + [new_ip])
            # re_matrix.append([])
            # gf_matrix.append([])

        else:
            connected_added[cmin[1]] = True
            pop_and_add(connected_set[cmin[1]], new_ip, consecutive_frames)

    i = 0
    while i < len(added_matched):
        if not added_matched[i]:
            pop_and_add(matched_ip_set[i], None, consecutive_frames)
            if matched_ip_set[i] == [None for _ in range(consecutive_frames)]:
                matched_ip_set.pop(i)
                # re_matrix.pop(i)
                # gf_matrix.pop(i)
                added_matched.pop(i)
                continue
        i += 1