import sys
import os
import random
sys.path.append("/home/shanwei-luo/userdata/mmdetection")
from mmdet.apis import (async_inference_detector, inference_detector,
                        init_detector, show_result_pyplot)

class_AOI_name = {"bu_pi_pei":"1","fang_xiang_fan":"2","err.txt_c_not_f":"3", "shang_xi_bu_lia":"4"}

config_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test61/AD_dsxw_test61.py'
checkpoint_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test61/epoch_36.pth'

config_file_2 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test62/AD_dsxw_test62.py'
checkpoint_file_2 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test62/epoch_44.pth'

imgs_ok_path = '/home/shanwei-luo/userdata/datasets/dsxw_test_SMD12_2022_0301_0314/ok/'
imgs_ng_path = '/home/shanwei-luo/userdata/datasets/dsxw_test_SMD12_2022_0301_0314/ng/'

model_1 = init_detector(config_file_1, checkpoint_file_1, device='cuda:0')
model_2 = init_detector(config_file_2, checkpoint_file_2, device='cuda:1')

imgs_ok = os.listdir(imgs_ok_path)
imgs_ng = os.listdir(imgs_ng_path)

count_label_ok = len(imgs_ok)
count_label_ng = len(imgs_ng)
print(count_label_ok,count_label_ng)
imgs_labels = []
imgs_name = []
for img in imgs_ok:
    img_name = img.split("@")
    if img_name[2] in class_AOI_name.keys():
        count_label_ok -= 1
        continue
    imgs_labels.append(0)
    imgs_name.append(imgs_ok_path+img)

for img in imgs_ng:
    img_name = img.split("@")
    if img_name[2] in class_AOI_name.keys():
        count_label_ng -= 1
        continue
    imgs_labels.append(1)
    imgs_name.append(imgs_ng_path+img)

print(count_label_ok,count_label_ng, len(imgs_labels))
print("before infer")
index = 0
num = len(imgs_name)
results_1 = []
results_2 = []
step = 256
while index<num:
    index += step
    if index < num:
        results_1_tmp = inference_detector(model_1, imgs_name[index-step:index])
        results_2_tmp = inference_detector(model_2, imgs_name[index-step:index])
    else:
        results_1_tmp = inference_detector(model_1, imgs_name[index-step:num])
        results_2_tmp = inference_detector(model_2, imgs_name[index-step:num])
    results_1 += results_1_tmp
    results_2 += results_2_tmp
    print(len(results_1), len(results_2))
print("after infer")

imgs_results_1 = []
for result in results_1:
    res_predict = 0
    for i in result:
        if i.shape[0]>0:
            res_predict = 1
    imgs_results_1.append(res_predict)

imgs_results_2 = []
for result in results_2:
    res_predict = 0
    for i in result:
        if i.shape[0]>0:
            res_predict = 1
    imgs_results_2.append(res_predict)

print(len(imgs_results_1), len(imgs_results_2))

imgs_results = []
for r1, r2 in zip(imgs_results_1, imgs_results_2):
    if r1 == 1 or r2 == 1:
        imgs_results.append(1)
    else:
        imgs_results.append(0)

print(len(imgs_results_1), len(imgs_results_2), len(imgs_results))

count_ng = 0
count_ok = 0
for i in range(len(imgs_labels)):
    if imgs_labels[i]==0 and imgs_results_1[i]==0:
        count_ok += 1
    if imgs_labels[i]==1 and imgs_results_1[i]==1:
        count_ng += 1
    if imgs_labels[i]==1 and imgs_results_1[i]==0:
        print(imgs_name[i])

print(count_ok/count_label_ok, count_ng/count_label_ng)

count_ng = 0
count_ok = 0
for i in range(len(imgs_labels)):
    if imgs_labels[i]==0 and imgs_results_2[i]==0:
        count_ok += 1
    if imgs_labels[i]==1 and imgs_results_2[i]==1:
        count_ng += 1
    if imgs_labels[i]==1 and imgs_results_2[i]==0:
        print(imgs_name[i])

print(count_ok/count_label_ok, count_ng/count_label_ng)

count_ng = 0
count_ok = 0
for i in range(len(imgs_labels)):
    if imgs_labels[i]==0 and imgs_results[i]==0:
        count_ok += 1
    if imgs_labels[i]==1 and imgs_results[i]==1:
        count_ng += 1
    if imgs_labels[i]==1 and imgs_results[i]==0:
        print(imgs_name[i])
#print(count_ok, count_ng)
#print(count_ng/label_ng)
print(count_ok/count_label_ok, count_ng/count_label_ng)