from img_util import *
from sklearn.metrics.cluster import adjusted_rand_score
from sklearn.metrics.cluster import fowlkes_mallows_score
from sklearn.metrics import jaccard_score
import time
import random
import json


from RFMS_seg import *
from quickshift_seg import *
from ms_seg import *
from mspp_seg import *



save = False
mapping = get_mapping()

ids = get_all_id()

methods = [run_quickshift_seg,run_MS_seg,run_BMS_seg,run_RFMS_seg,run_mspp_seg]



ARI_scores = {}
FM_scores = {}
IOU_score = {}
_time = {}
for method in methods:
	ARI_scores[method.__name__] = []
	FM_scores[method.__name__] = []
	IOU_score[method.__name__] = []
	_time[method.__name__] = 0

num = 0
for _id in ids:
	print('===',num,"-",_id,'===')
	for method in methods:
		print(method.__name__)
		start_time = time.time()
		img, labels = load_img(_id)
		width, height, count, data, target = create_dataset(img,labels,False)
		new_img, pred = method(img, width, height, count, data, target)
		ARI = adjusted_rand_score(pred,target)
		FM = fowlkes_mallows_score(pred,target)
		IOU = jaccard_score(pred, target, average='macro')
		ARI_scores[method.__name__].append(ARI) 
		FM_scores[method.__name__].append(FM)
		IOU_score[method.__name__].append(IOU)
		end_time = time.time()
		elapsed_time = end_time - start_time
		_time[method.__name__]+=elapsed_time
		if save == True:
			save_img(str(_id),method.__name__,new_img,mapping)
	num+=1

print('1====')
print(ARI_scores)
print('2====')
print(FM_scores)
print('3====')
print(IOU_score)
print('4====')
print(_time)

with open('result/'+'ARI.json', 'w') as file:
    json.dump(ARI_scores, file, indent=4)

with open('result/'+'FM.json', 'w') as file:
    json.dump(FM_scores, file, indent=4)

with open('result/'+'IOU.json', 'w') as file:
    json.dump(IOU_score, file, indent=4)

with open('result/'+'Time.json', 'w') as file:
    json.dump(_time, file, indent=4)

print('==== End of Experiment ====')

















