import torch
import numpy as np
import cv2
import os
import os.path as osp
from PIL import Image
import pandas as pd
import time
import natsort
import scipy
import lap
import warnings
import argparse
import pickle

from tqdm import tqdm
from torchvision import transforms
from sklearn.metrics.pairwise import cosine_similarity
from scipy.spatial.distance import cdist, pdist, squareform, cosine
from sklearn.cluster import KMeans
from scipy.sparse.linalg import eigsh
from scipy.sparse import csgraph
from collections import defaultdict, Counter
from cython_bbox import bbox_overlaps as bbox_ious

from models.experimental import attempt_load
from utils.general import non_max_suppression_kpt
from utils.plots import output_to_keypoint


def make_parser():
    parser = argparse.ArgumentParser("MCMT Tracking Demo!")

    parser.add_argument('result_path', help='SCMT Tracking result path')
    parser.add_argument('homography_path', help='Homography file')
    parser.add_argument('pose_weight', help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--save_txt_path', default='./MCMT_result', help='Device used for inference')
    return parser


def load_SCMT_tracklet(gts, vid_width, vid_height):
    result = []
    for gt in gts:
        data_dict = {}
        for row in gt:
            row_values = row.strip().split(',')
            key = int(row_values[0])
            track_id=int(row.split(',')[1])
            x = int(max(0, int(float(row.split(',')[2]))))
            y = int(max(0, int(float(row.split(',')[3]))))

            if int(float(row.split(',')[4]))+x >vid_width:
                obj_w=vid_width-x-1
            else:
                obj_w=int(float(row.split(',')[4]))
            if int(float(row.split(',')[5]))+y >vid_height:
                obj_h=vid_height-y-1
            else:
                obj_h=int(float(row.split(',')[5]))

            values = [x,y,x+obj_w,y+obj_h,float(row.split(',')[6]),track_id]
            
            if key in data_dict:
                data_dict[key].append(values)
            else:
                data_dict[key] = [values]
        result.append(data_dict)
    return result


def check_dict_keys(arr, key_list, channel_id):
    """
    Checks if all elements in the numpy array arr are keys in the given key_list.
    """
    result = []
    key_set = list(key_list)
    for elem in list(arr):
        if elem not in key_set:
            result.append([int(channel_id), int(elem)])
    return result


def make_square(img):
    """
    Given an image as a cv2 object, add padding to make it square.
    """
    height, width = img.shape[:2]
    square_size = max(width, height)
    # Create a new image of the appropriate size, with a black background
    new_img = np.zeros((square_size, square_size, 3), dtype=np.uint8)
    # Determine where to paste the original image in the new image
    x_offset = (square_size - width) // 2
    y_offset = (square_size - height) // 2
    # Paste the original image in the center of the new image
    new_img[y_offset:y_offset+height, x_offset:x_offset+width] = img
    return new_img


def run_keypoint(frame_bbox, frame_img, keypoint_model, DEVICE):
    thres = 0.5
    padding_size = 10
    square_size = 256
    H,W,_ = frame_img.shape
    
    padding_height = frame_bbox[3]-frame_bbox[1]+padding_size
    center_x = (frame_bbox[0]+frame_bbox[2])/2

    patch_x1 = int(max(0, frame_bbox[0]-padding_size))
    patch_y1 = int(max(0, frame_bbox[1]-padding_size))
    patch_x2 = int(min(W - 1, frame_bbox[2]+padding_size))
    patch_y2 = int(min(H - 1, frame_bbox[3]+padding_size))

    patch = frame_img[patch_y1:patch_y2, patch_x1:patch_x2, :]
    patch_squared = make_square(patch)

    resized_patch = cv2.resize(cv2.cvtColor(patch_squared, cv2.COLOR_BGR2RGB), (square_size,square_size))
    resized_patch = torch.tensor(np.array([transforms.ToTensor()(resized_patch).numpy()]))
    resized_patch = resized_patch.to(DEVICE)  #convert image data to device
    resized_patch = resized_patch.float() #convert image to float precision (cpu)

    with torch.no_grad():  #get predictions
        output_data, _ = keypoint_model(resized_patch)

    output_data = non_max_suppression_kpt(output_data,   #Apply non max suppression
                0.25,   # Conf. Threshold.
                0.65, # IoU Threshold.
                nc=keypoint_model.yaml['nc'], # Number of classes.
                nkpt=keypoint_model.yaml['nkpt'], # Number of keypoints.
                kpt_label=True)
    output = output_to_keypoint(output_data)

    state_bit = False # leg non-exist
    avg_foot_coor = None
    max_area = 10
    for o in output:
        obj_area = o[4]*o[5]
        if (obj_area>max_area) and (o[54]>thres and o[57]>thres) and (o[18]>thres or o[21]>thres or o[24]>thres or o[27]>thres):
            max_area = obj_area
            state_bit = True # leg exist
            
            patch_height, patch_width = patch.shape[:2]
            square_size = max(patch_width, patch_height)
            if patch_height > patch_width :
                c_x = ((patch_x1+patch_x2)/2) - ((patch_y2 - patch_y1)/2) + ((o[52] + o[55])/2)/256 * (patch_y2 - patch_y1)
                bottom_y = patch_y1 + ((o[53]+o[56])/2)/256 * (patch_y2 - patch_y1)
                
                avg_foot_coor = (int(c_x), int(bottom_y))
            else:
                state_bit = False
            
    return state_bit, avg_foot_coor


def extract_global_position(avg_foot_point, H):
    # format (x,y,raw_track_id)
    c_x = avg_foot_point[0]
    bottom_y = avg_foot_point[1]
    point1 = np.array([[c_x, bottom_y, 1]])
    point2 = np.dot(H, point1.T)
    point2 /= point2[2]
    return int(point2[0]),int(point2[1])


def separate_features(data):
    position_features = []
    velocity_features = []
    for row in data:
        position = row[2][:2]
        velocity = row[2][2:]
        position_features.append(position)
        velocity_features.append(velocity)
    return np.array(position_features), np.array(velocity_features)

def normalize_features(features):
    mean = np.mean(features, axis=0)
    std = np.std(features, axis=0)
    return np.array((features - mean) / std)
    

def euclidean_distance_1d(x, y):
    return np.sqrt(np.sum((x - y) ** 2))


def euclidean_distance_2d(p1, p2):
    return int(np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2))


def cosine_distance(x, y):
    return cosine(x, y)
    
    
def compute_similarity(tracklet1, tracklet2):
    cam1, track_id, feature1 = tracklet1
    cam2, track_id, feature2 = tracklet2
    if cam1 == cam2:
        return 0.0, 0.0
    
    pos_feature1 = feature1[:2].reshape(1, -1)
    pos_feature2 = feature2[:2].reshape(1, -1)
    
    vel_feature1 = feature1[2:].reshape(1, -1)
    vel_feature2 = feature2[2:].reshape(1, -1)
    
    pos_sim = cosine_similarity(pos_feature1, pos_feature2)[0]
    vel_sim = cosine_similarity(vel_feature1, vel_feature2)[0]
    return pos_sim, vel_sim


def find_n_for_clustering(data):
    err_dist = 100
    tracklets = []
    distance_valid_bit = True
    for camera in data:
        camera_tracklets = {}
        for track_id, coords in camera.items():
            if len(coords) == 3:
                camera_tracklets[track_id] = coords[-1]
        tracklets.append(camera_tracklets)
    
    clusters = defaultdict(set)
    cluster_id = 1
    for i in range(len(tracklets)):
        camera_i = tracklets[i]
        for track_id_i, coord_i in camera_i.items():
            for j in range(len(tracklets)):
                camera_j = tracklets[j]
                for track_id_j, coord_j in camera_j.items():
                    if i == j and track_id_i == track_id_j:
                        continue
                    distance = euclidean_distance_2d(coord_i, coord_j)
                    if i == j and distance < err_dist*1.1:
                        distance_valid_bit = False
                    elif i != j and distance < err_dist:
                        clusters[cluster_id].add((i+1, track_id_i))
                        clusters[cluster_id].add((j+1, track_id_j))
                        cluster_id += 1
    # Merge clusters with common elements
    merged_clusters = []
    for key, value in clusters.items():
        for cluster in merged_clusters:
            if value.intersection(cluster):
                cluster.update(value)
                break
        else:
            merged_clusters.append(value)

    return len(merged_clusters), distance_valid_bit, merged_clusters


def run_spectral_clustering2(position_features, k):
    # Calculate the Euclidean distance similarity matrix
    position_similarity_matrix = np.zeros((len(position_features), len(position_features)))
    for i in range(len(position_features)):
        for j in range(len(position_features)):
            if i != j:
                position_similarity_matrix[i, j] = euclidean_distance_1d(position_features[i], position_features[j])
                
    position_similarity_matrix = position_similarity_matrix / np.max(position_similarity_matrix)
    position_similarity_matrix = np.abs(position_similarity_matrix - 1)
    np.fill_diagonal(position_similarity_matrix, 0)
    
    # velocity_similarity_matrix = velocity_similarity_matrix / np.max(velocity_similarity_matrix)
    # velocity_similarity_matrix = np.nan_to_num(velocity_similarity_matrix, nan=1)
    
    pos_sim_matrix = position_similarity_matrix

    # Compute the Laplacian matrix
    laplacian = csgraph.laplacian(pos_sim_matrix, normed=True)

    # Compute the first k eigenvectors
    eigvals, eigvecs = np.linalg.eig(laplacian)
    eigvecs = eigvecs[:, np.argsort(eigvals)[:k]]

    # Normalize each row of the eigenvector matrix
    normalized_eigvecs = eigvecs / np.linalg.norm(eigvecs, axis=1)[:, np.newaxis]

    # Perform k-means clustering on the rows of the normalized matrix
    kmeans = KMeans(n_clusters=k, n_init=10).fit(normalized_eigvecs)
    
    # Cluster assignments
    clusters = kmeans.labels_

    return clusters


def run_spectral_clustering(channel_id_position, k):

    matched_id_exist = []
    
    features = []
    for idx, tracklets_dic in enumerate(channel_id_position):
        for key in tracklets_dic.keys():
            if len(tracklets_dic[key])==3:
                now_x,now_y = tracklets_dic[key][-1]
                prev_x,prev_y = tracklets_dic[key][0]
                # extract motion feature
                motion_feat = np.array([now_x, now_y, now_x-prev_x, now_y-prev_y])
                # apply the additional condition for the velocity feature
                motion_feat[2:][np.abs(motion_feat[2:]) <= 10] = 1
                # normalize
                features.append((idx, key, motion_feat))
                
                matched_id_exist.append([idx,key])
    
    position_features, velocity_features = separate_features(features)

    # Calculate the Euclidean distance similarity matrix
    position_similarity_matrix = np.zeros((len(position_features), len(position_features)))
    for i in range(len(position_features)):
        for j in range(len(position_features)):
            if i != j:
                position_similarity_matrix[i, j] = euclidean_distance_1d(position_features[i], position_features[j])
    # Calculate the cosine distance similarity matrix
    velocity_similarity_matrix = np.zeros((len(velocity_features), len(velocity_features)))
    for i in range(len(velocity_features)):
        for j in range(len(velocity_features)):
            velocity_similarity_matrix[i, j] = cosine_distance(velocity_features[i], velocity_features[j])
    
    position_similarity_matrix = position_similarity_matrix / np.max(position_similarity_matrix)
    position_similarity_matrix = np.abs(position_similarity_matrix - 1)
    np.fill_diagonal(position_similarity_matrix, 0)
    
    # velocity_similarity_matrix = velocity_similarity_matrix / np.max(velocity_similarity_matrix)
    # velocity_similarity_matrix = np.nan_to_num(velocity_similarity_matrix, nan=1)
    
    pos_sim_matrix = position_similarity_matrix

    # Compute the Laplacian matrix
    laplacian = csgraph.laplacian(pos_sim_matrix, normed=True)

    # Compute the first k eigenvectors
    eigvals, eigvecs = np.linalg.eig(laplacian)
    eigvecs = eigvecs[:, np.argsort(eigvals)[:k]]

    # Normalize each row of the eigenvector matrix
    normalized_eigvecs = eigvecs / np.linalg.norm(eigvecs, axis=1)[:, np.newaxis]

    # Perform k-means clustering on the rows of the normalized matrix
    kmeans = KMeans(n_clusters=k, n_init=10).fit(normalized_eigvecs)
    
    # Cluster assignments
    clusters = kmeans.labels_

    return clusters, matched_id_exist


def offline_tracking(channel_loss, channel_id_mapping, channel_final_output):
    delete_key_list = []
    for idx, loss_obj in enumerate(channel_loss):
        for key in loss_obj.keys():
            if key in channel_id_mapping[idx].keys():
                global_track_id = channel_id_mapping[idx][key]
                for row in loss_obj[key]:
                    row[1] = global_track_id
                    channel_final_output[idx].append(row)
                delete_key_list.append((idx,key))
    for key_obj in delete_key_list:
        del channel_loss[key_obj[0]][key_obj[1]]
    
    return channel_final_output, channel_loss


def most_frequent_last_called(lst):
    counter = Counter(lst)
    max_count = max(counter.values())
    most_frequent = [k for k, v in counter.items() if v == max_count]
    return most_frequent[-1]


def save_txt_result(channel_final_output, tracking_result_path, channel_list):
    for idx, channel_result in enumerate(channel_final_output):
        sort_indices = np.argsort(np.array(channel_result)[:, 2])
        sorted_channel_result = list(np.array(channel_result)[sort_indices])
        
        save_dir = os.path.join(tracking_result_path, channel_list[idx])
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        
        print("save to : ", save_dir)
        # Open a file for writing
        with open(os.path.join(save_dir,'sorted_associated.txt'), 'w') as f:
            # Loop through the rows of the array and write them to the file
            for row in sorted_channel_result:
                # Convert the row to a string and add a newline character
                row_str = ','.join(str(int(elem)) for elem in row) + '\n'
                f.write(row_str)


def main(args):
    tracking_result_path = osp.join(args.result_path)
    with open('homography_list.pkl', 'rb') as file:
        loaded_list = pickle.load(file)
    H_list = loaded_list
    
    channel_list = [file for file in os.listdir(tracking_result_path) if not file.startswith(".") and os.path.isdir(os.path.join(tracking_result_path, file))]
    channel_list=natsort.natsorted(channel_list)

    caps = []
    gts = []
    for channel in channel_list:
        video_path = os.path.join(tracking_result_path, channel, 'video.mp4')
        label_path = os.path.join(tracking_result_path, channel, 'label.txt')
        cap = cv2.VideoCapture(video_path)
        caps.append(cap)
        with open(label_path) as label_file:
            gt = label_file.readlines()
        gts.append(gt)
        
    vid_width = caps[0].get(cv2.CAP_PROP_FRAME_WIDTH)  # float
    vid_height = caps[0].get(cv2.CAP_PROP_FRAME_HEIGHT)  # float
    vid_fps = caps[0].get(cv2.CAP_PROP_FPS)
    vid_length = int(caps[0].get(cv2.CAP_PROP_FRAME_COUNT))
    
    ##### [[chnnel_id], [], []... ] -> channel_id = {frame_count : [[x,y,x2,y2,score,track_id],[],...]}
    channel_track_result = load_SCMT_tracklet(gts, vid_width, vid_height)
    
    channel_id_mapping = [{} for _ in range(len(caps))]
    channel_id_position = [{} for _ in range(len(caps))]
    channel_loss = [{} for _ in range(len(caps))]
    channel_final_output = [[] for _ in range(len(caps))]
    # define start global ID #
    new_id = 100
    ##########################
    
    keypoint_model = attempt_load(args.pose_weight, map_location=args.device)  #Load model
    keypoint_model.eval()
    
    t = tqdm(total=vid_length)
    ## Tracklet Association
    for frame_count in range(vid_length):
        new_id_exist = []
        for idx,cap in enumerate(caps):
            _, frame_img = cap.read()
            
            defined_track_id = []
            if frame_count in channel_track_result[idx].keys():
                track_bbox = channel_track_result[idx][frame_count]
                for frame_bbox in track_bbox:
                    state_bit, avg_foot_point = run_keypoint(frame_bbox, frame_img, keypoint_model, args.device)
                    if state_bit:
                        track_id = frame_bbox[-1]
                        defined_track_id.append(track_id)
                        projected_x, projected_y = extract_global_position(avg_foot_point, H_list[idx])  
                        if not track_id in channel_id_position[idx].keys(): # 모든 보여진 ID에 대하여
                            channel_id_position[idx][track_id] = [(projected_x,projected_y)]
                        elif len(channel_id_position[idx][track_id])==3:
                            del channel_id_position[idx][track_id][0]
                            channel_id_position[idx][track_id].append((projected_x,projected_y))
                        else:
                            channel_id_position[idx][track_id].append((projected_x,projected_y))
            
                delete_position_list = []
                for key in channel_id_position[idx].keys():
                    if key not in defined_track_id:
                        delete_position_list.append((idx, key))
                for obj in delete_position_list:
                    del channel_id_position[obj[0]][obj[1]]

                channel_track_id = np.array(track_bbox)[:,-1]
                key_list = list(channel_id_mapping[idx].keys())
                new_id_exist+= check_dict_keys(channel_track_id, key_list, idx)
            
        if len(new_id_exist) > 0:
            # 1. keypoint -> foot point exist in 3 framses
            keypoint_valid_bit = False
            for new_id_obj in new_id_exist:
                if new_id_obj[1] in channel_id_position[new_id_obj[0]].keys() and len(channel_id_position[new_id_obj[0]][new_id_obj[1]])==3:
                    keypoint_valid_bit = True
            
            # 2. 한 채널 내서 두 객체가 가까이 있을때 X (loss에 추가)
            N, distance_valid_bit, merged_clusters = find_n_for_clustering(channel_id_position)
    
            if keypoint_valid_bit and distance_valid_bit:
                # 3. spectral clustering
                # cluster_result, matched_id_exist = run_spectral_clustering(channel_id_position, N)
                
                matched_id_exist = []
                for idx, tracklets_dic in enumerate(channel_id_position):
                    for key in tracklets_dic.keys():
                        if len(tracklets_dic[key])==3:
                            matched_id_exist.append([idx,key])
                
                new_matched_id_exist = [tup for tup in new_id_exist if tup in matched_id_exist]
                
                ## update channel_id_mapping using global_id 
                for id_obj in new_matched_id_exist:
                    find_id_obj = (id_obj[0]+1, id_obj[1])
                    other_tuples = []
                    for group in merged_clusters:
                        if find_id_obj in group:
                            for tup in group:
                                if tup != find_id_obj:
                                    other_tuples.append((tup[0]-1,tup[1]))
                            break
                    
                    target_id_list = []
                    for tup in other_tuples:
                        if tup[1] in channel_id_mapping[tup[0]].keys():
                            target_id_list.append(channel_id_mapping[tup[0]][tup[1]])
                
                    if len(target_id_list)>0:
                        target_id = most_frequent_last_called(target_id_list)
                        
                        if any(x != target_id_list[0] for x in target_id_list):

                            # target_id_list 만 모아놓고 spectral clustering
                            dup_tuples = []
                            target_id_list_final = []
                            for idx,mapping_dic in enumerate(channel_id_mapping):
                                for key,value in mapping_dic.items():
                                    if value in target_id_list and key in channel_id_position[idx].keys():
                                        dup_tuples.append((idx, key))

                                        target_id_list_final.append(value)
                            
                            dup_position = []
                            # channel_id_position
                            for tup in dup_tuples:
                                dup_position.append(channel_id_position[tup[0]][tup[1]][-1])
                            # new motion
                            dup_position.append(channel_id_position[id_obj[0]][id_obj[1]][-1])
                            
                            cluster_k = len(list(set(target_id_list_final)))
                            cluster_result = run_spectral_clustering2(np.array(dup_position), cluster_k)
                            
                            cluster_label_list = cluster_result[:-1].copy()
                            
                            connection_count = {}

                            # Count the connections
                            for label, gt_id in zip(cluster_label_list, target_id_list_final):
                                if (label, gt_id) not in connection_count:
                                    connection_count[(label, gt_id)] = 1
                                else:
                                    connection_count[(label, gt_id)] += 1

                            # Find the most frequent gt_global_id for each label
                            label_mapping = {}
                            for key, value in connection_count.items():
                                label, gt_id = key
                                if label not in label_mapping or value > label_mapping[label][1]:
                                    label_mapping[label] = (gt_id, value)
                            
                            for x,tup in enumerate(dup_tuples):
                                channel_id_mapping[tup[0]][tup[1]] = label_mapping[cluster_result[x]][0] 
                    else:
                        target_id = None

                        
                    if target_id is not None:
                        channel_id_mapping[id_obj[0]][id_obj[1]] = target_id
                        channel_id_arr = np.array(channel_track_result[id_obj[0]][frame_count])[:,-1]
                        for c_id in channel_id_arr:
                            c_id = int(c_id)
                            if (c_id in channel_id_mapping[id_obj[0]].keys()) and (c_id != id_obj[1]) and (target_id == channel_id_mapping[id_obj[0]][c_id]):
                                print("중복 발생")
                                print(target_id)
                        
                    else:
                        channel_id_mapping[id_obj[0]][id_obj[1]] = new_id
                        new_id+=1  
                        
                    print(channel_id_mapping)
                
        # 새로운 id 있어도 클러스터링 불가 -> loss 저장 
        # 새로운 id 없으면 -> final에 저장
        # 이전 loss에서 매칭 성공 -> final에 저장
        for idx,cap in enumerate(caps):
            if frame_count in channel_track_result[idx].keys():
                track_bbox = channel_track_result[idx][frame_count]
                for frame_bbox in track_bbox:
                    global_x, global_y = -1, -1
                    track_id = frame_bbox[-1]
                    if track_id in channel_id_position[idx].keys():
                        global_x, global_y = channel_id_position[idx][track_id][-1]

                    global_track_id = -1
                    if track_id in channel_id_mapping[idx].keys():
                        global_track_id = channel_id_mapping[idx][track_id]

                    x,y,w,h = frame_bbox[0], frame_bbox[1], frame_bbox[2]-frame_bbox[0], frame_bbox[3]-frame_bbox[1]

                    if global_track_id != -1:
                        # Format : channel_final_output -> ([channel],[],...)
                        # channel -> (〈channel_idx〈global_track_id>〈frame_id>〈x> <y>〈w>〈h>〈xworld〉〈yworld〉)
                        channel_final_output[idx].append([idx, global_track_id , frame_count , x, y, w, h, global_x, global_y])
                    else:
                        if track_id in channel_loss[idx].keys():
                            channel_loss[idx][track_id].append([idx, global_track_id , frame_count , x, y, w, h, global_x, global_y])
                        else:
                            channel_loss[idx][track_id] = [[idx, global_track_id , frame_count , x, y, w, h, global_x, global_y]]
                
        channel_final_output, channel_loss= offline_tracking(channel_loss, channel_id_mapping, channel_final_output) 
        
        t.update(1)
        
        if frame_count>550:
            break
            
    save_txt_result(channel_final_output, args.save_txt_path, channel_list)
            
if __name__ == "__main__":
    args = make_parser().parse_args()
    
    main(args)