import os 
import time
import json
import queue
from multiprocessing import Process, Queue
from multiprocessing.pool import Pool
from tqdm import tqdm

import pandas as pd 
import numpy as np 
import matplotlib.pyplot as plt 

DATA_PATH = 'GQA/'
REL_PATH = 'full_relations_df.pkl'
IMG_SIZE = (500, 500)
NUM_PROCESSES = 20
NUM_SAMPLES = 100


def bbox_to_mask(x, y, w, h, img_size=IMG_SIZE, name=None, visualize=False):
    img = np.zeros(img_size)
    mask_w = np.ones(np.clip(w, 0, img_size[1]-x))
    
    for j in range(y, np.clip(y+h, 0, img_size[0])):
        img[j][x:x+w] = mask_w
    
    if visualize:   
        fig = plt.figure(figsize=(img_size[0] // 80, img_size[1] // 80))
        plt.imshow(img, cmap='gray')
        if name: 
            plt.title(name)
        plt.axis('off')
        plt.show()
    
    return img

def get_all_relations_df(data):
	
	print(f'Length of scenegraph data set: {len(data)}')
	start = time.time()
 
	df = pd.DataFrame(columns=['image_id', 'relation', 'from', 'to', 'obj_loc', 'obj_w', 'obj_h', 'obj_center', 
                               'rel_obj_loc', 'rel_obj_w', 'rel_obj_h'])       
	
	for img_id in data.keys():
		all_objects = data.get(str(img_id)).get('objects').items()
  
  		# get all object names 
		all_objects_dict = {id_num: (obj_dict.get('name'), obj_dict.get('x'), obj_dict.get('y'), obj_dict.get('w'), obj_dict.get('h')) 
                      		for (id_num, obj_dict) in all_objects}
		
		# get all relations
		for obj in all_objects:    
			id_num, obj_dict = obj
			name = obj_dict.get('name')
			x, y, width, height  = obj_dict.get('x'), obj_dict.get('y'), obj_dict.get('w'), obj_dict.get('h')
			center = [x + width / 2, y + height / 2]
			
			for relation in obj_dict.get('relations'):
				rel = relation.get('name')
				rel_obj, rel_x, rel_y, rel_w, rel_h = all_objects_dict.get(relation.get('object'))
				
				
				temp = pd.DataFrame.from_dict([{'image_id': img_id, 'relation': rel, 'from': name, 'to': rel_obj, 
												'obj_loc': [x, y], 'obj_w': width, 'obj_h': height, 'center': center,
												'rel_obj_loc': [rel_x, rel_y], 'rel_obj_w': rel_w, 'rel_obj_h': rel_h}])
				
				df = pd.concat([df,  temp], ignore_index=True)
				#print(f'{df.iloc[-1]["from"]} {df.iloc[-1].relation} {df.iloc[-1].to}')
    
	out_path = 'all_relations.pkl'
	df.to_pickle(out_path)
	print(f'Saved df to {out_path}')
 
	end = time.time()
	elapsed = end - start 
	print(f'Took {int(elapsed // 60)}:{int(elapsed % 60)} min:s for all {len(df)} relations --> {elapsed / len(df):.2f}s / relation')


def generate_query_mask(df, rel, i, img_center=np.array([250, 250]), uni_size=np.array([50, 50])):
    
    # uni_obj only needed for visualization in the end 
    uni_obj = bbox_to_mask(img_center[0] - (uni_size[0] // 2), img_center[1] - (uni_size[1] // 2), 
                       50, 50, img_size=(500, 500))
    
    temp_df = df.loc[df.relation == rel]
    print(f'[{i}] Number of "{rel}" samples: {len(temp_df)}')

    query_mask = np.zeros((500, 500), dtype=np.uint8)
    counter = 0 
    num_discard = 0 
    
    for idx in range(len(temp_df)):       
        if counter >= NUM_SAMPLES: 
            print(f'[{i}] Reached {counter} samples for relation "{rel}":')
            break 
        
        img_id = temp_df.iloc[idx].image_id
        img_size = (data.get(img_id)['height'], data.get(img_id)['width'])

        # get relative object info and generate binary mask
        obj_loc = temp_df.iloc[idx].rel_obj_loc
        width = temp_df.iloc[idx].rel_obj_w
        height = temp_df.iloc[idx].rel_obj_h

        # get mask info and generate binary mask
        mask_loc = temp_df.iloc[idx].obj_loc
        mask_w = temp_df.iloc[idx].obj_w
        mask_h = temp_df.iloc[idx].obj_h
        
        if obj_loc[0] > img_size[1] or obj_loc[1] > img_size[0] or mask_loc[0] > img_size[1] or mask_loc[1] > img_size[0]:
            #print('error in bounding box -- discard sample')
            continue 
        
        obj = bbox_to_mask(obj_loc[0], obj_loc[1], width, height, img_size=img_size)
        mask = bbox_to_mask(mask_loc[0], mask_loc[1], mask_w, mask_h, img_size=img_size)

        img = obj*2 + mask 
        img_transformed = np.zeros((1000, 1000), dtype=np.uint8)
        
        # scale image first 
        scale_x, scale_y = uni_size[0] / width, uni_size[1] / height
        scale_mat = np.array([[scale_y, 0, 0], [0, scale_x, 0], [0, 0, 1]])
        
        if scale_x > 5 or scale_y > 5: 
            num_discard += 1
            #print(f'Scale is too high! x: {scale_x}, y: {scale_y} -- discard sample')
            continue 
        
        for i, row in enumerate(img):
            for j, col in enumerate(row):
                pixel_data = img[i, j]
                input_coords = np.array([i, j, 1])
                i_out, j_out, _ = scale_mat @ input_coords

                if i_out > 0 and i_out < 1000 and j_out > 0 and j_out < 1000 and pixel_data > 0: 
                    # new indices must be within new image -- discard others 
                    img_transformed[int(i_out), int(j_out)] = pixel_data                      
        
        if not len(np.where(img_transformed >= 2)[0]) > 0: 
            # no data in transformed image -- discard sample
            continue
            
        # find new (x, y) location of object
        new_loc = sorted([[y, x] for (y, x) in zip(*np.where(img_transformed >= 2))])[0]
        new_center = [new_loc[0] + uni_size[0] // 2, new_loc[1] + uni_size[1] // 2]
        
        # move object to center 
        move_x, move_y = img_center - new_center 
        move_mat = np.array([[1, 0, move_x], [0, 1, move_y], [0, 0, 1]])

        img_moved = np.zeros((500, 500), dtype=np.uint8)
        for i, row in enumerate(img_transformed):
            for j, col in enumerate(row):
                pixel_data = img_transformed[i, j]
                input_coords = np.array([i, j, 1])
                i_out, j_out, _ = move_mat @ input_coords

                if i_out > 0 and i_out < 500 and j_out > 0 and j_out < 500 and pixel_data > 0: 
                    # new indices must be within new image -- discard others 
                    img_moved[int(i_out), int(j_out)] = pixel_data
                    
        # extract relative object mask and add to query mask 
        mask_transformed = np.where(img_moved==1, img_moved, 0) + np.where(img_moved==3, img_moved, 0)
        query_mask += mask_transformed
        counter += 1
        
    if counter > 0: 
        query_mask = query_mask / counter
        rel_name = '_'.join(rel.split(' ')) 
        np.save(f'relations/{rel_name}.npy', query_mask)
        print(f'[{i}] Saved query mask to: relations/{rel_name}.npy')
        
        if num_discard > 0: 
            print(f'[{i}] Discarded {num_discard} samples, because scaling was too high.')

        plt.figure(figsize=(3,3))
        plt.imshow(uni_obj*0.1+ query_mask, cmap='gray')
        plt.title(rel)
        plt.axis('off')
        plt.savefig(f'relations/{rel_name}.png', bbox_inches='tight', dpi=300)
        plt.clf()
        
    else: 
        print(f'[{i}] Could not generate query mask for "{rel}"')
    
def run_process(tasks, df):
    while True:
        try:
            '''
                try to get task from the queue. get_nowait() function will 
                raise queue.Empty exception if the queue is empty. 
                queue(False) function would do the same task also.
            '''
            task = tasks.get_nowait()
            i = list(df.relation.unique()).index(task)
        except queue.Empty:
            break
        else:
            ''' no exception has been raised '''
            print(f'[{i}] Starting relation #{i}: {task}')
            print()
            generate_query_mask(df, task, i)
            time.sleep(.5)
    
    return True


# task executed in a worker process
def get_relations_task(img_id):
    width, height = data.get(str(img_id))['width'], data.get(str(img_id))['height']
    all_objects = data.get(str(img_id)).get('objects').items()
  
    # get all object names 
    all_objects_dict = {id_num: (obj_dict.get('name'), obj_dict.get('x'), obj_dict.get('y'), obj_dict.get('w'), obj_dict.get('h')) 
                        for (id_num, obj_dict) in all_objects}
    
    all_relations = []
    
    # get all relations
    for obj in all_objects:    
        id_num, obj_dict = obj
        name = obj_dict.get('name')
        x, y, obj_w, obj_h  = obj_dict.get('x'), obj_dict.get('y'), obj_dict.get('w'), obj_dict.get('h')
        center = [x + width / 2, y + height / 2]

        for relation in obj_dict.get('relations'):
            rel = relation.get('name')
            rel_obj, rel_x, rel_y, rel_w, rel_h = all_objects_dict.get(relation.get('object'))
            
            all_relations.append({'image_id': img_id, 'width': width, 'height': height, 'relation': rel, 
                                'from': name, 'to': rel_obj, 'obj_loc': [x, y], 'obj_w': obj_w, 'obj_h': obj_h, 
                                'obj_center': center,'rel_obj_loc': [rel_x, rel_y], 'rel_obj_w': rel_w, 'rel_obj_h': rel_h})
    
    return all_relations



if __name__ == '__main__': 
    

    path = os.path.join(DATA_PATH, 'train_sceneGraphs.json')
    assert os.path.exists(path), f'{path} does not exist!'

    with open(os.path.join(DATA_PATH, 'train_sceneGraphs.json'), 'r') as f:
        data = json.load(f)
        
    print(f'Length of scenegraph data set: {len(data)}')

    if not os.path.exists(REL_PATH):    
        print('Generating dataframe of all relations...')
        # generate list of relations pkl -- use multiprocessing!
        # create and configure the process pool
        with Pool(processes=NUM_PROCESSES) as pool:
            
            df = pd.DataFrame(columns=['image_id', 'width', 'height', 'relation', 'from', 'to', 'obj_loc', 'obj_w', 
                                       'obj_h', 'obj_center', 'rel_obj_loc', 'rel_obj_w', 'rel_obj_h'])       
            
            # execute tasks in order
            for i, result in enumerate(tqdm(pool.map(get_relations_task, list(data.keys()), chunksize=100))):
                temp = pd.DataFrame.from_dict(result)
                df = pd.concat([df,  temp], ignore_index=True)
                if i % 10000 == 0: 
                    df.to_pickle('temp_' + REL_PATH)
                    print(f'Saved df to {"temp_" + REL_PATH}')
        
        df.to_pickle(REL_PATH)
        print(f'Saved df to {REL_PATH}')
    else: 
        df = pd.read_pickle(REL_PATH)
    
    print(f'Number of relations: {len(df.relation.unique())}')
    print(df.relation.unique())

    # generate query mask for each relation
    #for i, rel in enumerate(df.relation.unique()): 
    #	generate_query_mask(df, rel, i)

    print('Generating a query mask for each relation...')
    # generate query mask for each relation -- use multiprocessing 
    tasks = Queue()
    procs = []
    
    # only use relations with at least 1000 samples 
    rel_lst = df.relation.value_counts()[df.relation.value_counts() > 1000].index.to_list()

    for rel in rel_lst: 
        tasks.put(rel)
        
    # creating processes -- run only NUM_PROCESSES processes at the same time 
    for _ in range(NUM_PROCESSES):
        p = Process(target=run_process, args=(tasks, df,))
        procs.append(p)
        p.start()
        
    # completing all processes
    for p in procs:
        p.join()
        

