import sys
import os
import random
import numpy as np
import shutil
sys.path.append("/home/shanwei-luo/userdata/mmdetection")
from mmdet.apis import (async_inference_detector, inference_detector,
                        init_detector, show_result_pyplot)
import argparse
import pandas as pd
#python get_score_csv.py --config_file /home/shanwei-luo/teamdata/anomaly_detection_active_learning/model/work_dirs/AD_dsxw_test66_06_10/AD_dsxw_test66_06_10.py --checkpoint_file /home/shanwei-luo/teamdata/anomaly_detection_active_learning/model/work_dirs/AD_dsxw_test66_06_10/latest.pth --images_path /home/shanwei-luo/teamdata/anomaly_detection_active_learning/data0422/unlabel_11_12/images/ --test_batch_size 128 --result_path test_unlabel_11_12.csv
def parse_args():
    parser = argparse.ArgumentParser(description='get best threshold')
    parser.add_argument('--config_file', help='config')
    parser.add_argument('--checkpoint_file', help='checkpoint')
    parser.add_argument('--images_path', help='images')
    parser.add_argument('--test_batch_size', help='images')
    parser.add_argument('--result_path', help='result_path')
    args = parser.parse_args()
    return args

args = parse_args()

config_file_1 = args.config_file
checkpoint_file_1 = args.checkpoint_file

img_path = args.images_path

model_1 = init_detector(config_file_1, checkpoint_file_1, device='cuda:0')

imgs = os.listdir(img_path)
imgs_name = []
for img in imgs:
    imgs_name.append(img_path+img)
print(len(imgs_name))
print("before infer")
index = 0
num = len(imgs_name)
results_1 = []
results_feat = []
step = int(args.test_batch_size)
while index<num:
    index += step
    if index < num:
        results_1_tmp, results_feat_tmp = inference_detector(model_1, imgs_name[index-step:index], feat=True)
    else:
        results_1_tmp, results_feat_tmp = inference_detector(model_1, imgs_name[index-step:num], feat=True)
    results_1 += results_1_tmp
    results_feat += results_feat_tmp
    print(len(results_1))
print("after infer")
result_score = []
for result in results_1:
    max_v = 0
    for i in result:
        for j in range(i.shape[0]):
            if max_v < i[j, 4]:
                max_v = i[j, 4]
    result_score.append(max_v)
'''np_result = np.array(result_score)
print(np_result.shape)
np.save('/home/shanwei-luo/userdata/datasets/PCBA_dataset_v15_MLOPS/result.npy', np_result)'''

result_path = args.result_path
dataframe = pd.DataFrame({'Image_Name':imgs,'score':result_score, 'feature':results_feat, "model_name":"cascade_mlops_baseline"})
dataframe.to_csv(result_path)