import os
import cv2
import matplotlib.pyplot as plt
from mmpretrain.apis import ImageClassificationInferencer
import random
from tqdm import tqdm

print("Load model")
model_six_classes_config_path = './pretrained_weights/six_classes/config2c.py'
model_six_classes_cpkt_path = './pretrained_weights/six_classes/epoch_100.pth'

cos_config_path = './pretrained_weights/work_dirs_cos/config2c.py'
cos_cpkt_path = './pretrained_weights/work_dirs_cos/epoch_100.pth'

gs_config_path = './pretrained_weights/work_dirs_gs/config2c.py'
gs_cpkt_path = './pretrained_weights/work_dirs_gs/epoch_100.pth'

insulator_config_path = './pretrained_weights/work_dirs_insulator/config2c.py'
insulator_cpkt_path = './pretrained_weights/work_dirs_insulator/epoch_100.pth'

la_config_path = './pretrained_weights/work_dirs_la/config2c.py'
la_cpkt_path = './pretrained_weights/work_dirs_la/epoch_100.pth'

pole_config_path = './pretrained_weights/work_dirs_pole/config2c.py'
pole_cpkt_path = './pretrained_weights/work_dirs_pole/epoch_100.pth'

lp_config_path = './pretrained_weights/work_dirs_lp/config2c.py'
lp_cpkt_path = './pretrained_weights/work_dirs_lp/epoch_100.pth'

model_six_classes = ImageClassificationInferencer(model_six_classes_config_path, pretrained=model_six_classes_cpkt_path)

model_COS = ImageClassificationInferencer(cos_config_path, pretrained=cos_cpkt_path)
model_GS = ImageClassificationInferencer(gs_config_path, pretrained=gs_cpkt_path)
model_Insulator = ImageClassificationInferencer(insulator_config_path, pretrained=insulator_cpkt_path)
model_LA = ImageClassificationInferencer(la_config_path, pretrained=la_cpkt_path)
model_Pole = ImageClassificationInferencer(pole_config_path, pretrained=pole_cpkt_path)
model_LP = ImageClassificationInferencer(lp_config_path, pretrained=lp_cpkt_path)

print("Load model done")
print("Load image")
IMAGE_DIR = './TTA_TEST_SET'
image_file_names = os.listdir(IMAGE_DIR)
# extract all .JPG, .jpg .PNG .png files
image_file_names = [image_file_name for image_file_name in image_file_names if
                    image_file_name.lower().endswith(('.jpg', '.png'))]
final_predictions = []
print("Load image done")
print("Start inference")

for i, image_file_name in tqdm(enumerate(image_file_names[:2000])):
    image_path = os.path.join(IMAGE_DIR, image_file_name)

    six_classes_result = model_six_classes(image_path)[0]
    detected_class = six_classes_result['pred_class']

    if detected_class == 'COS':
        damage_result = model_COS(image_path)[0]
    elif detected_class == 'GS':
        damage_result = model_GS(image_path)[0]
    elif detected_class == 'Insulator':
        damage_result = model_Insulator(image_path)[0]
    elif detected_class == 'LA':
        damage_result = model_LA(image_path)[0]
    elif detected_class == 'Pole':
        damage_result = model_Pole(image_path)[0]
    elif detected_class == 'LP':
        damage_result = model_LP(image_path)[0]

    detection_info = {}
    detection_info['image_path'] = image_path
    detection_info['detected_class'] = detected_class
    detection_info['damage_class'] = damage_result['pred_class']
    if 'Normal' in str(damage_result['pred_class']):
        detection_info['damage_level'] = 'normal'
    elif damage_result['pred_score'] > 0.9:
        detection_info['damage_level'] = 'level_1'
    elif damage_result['pred_score'] > 0.7:
        detection_info['damage_level'] = 'level_2'
    else:
        detection_info['damage_level'] = 'level_3'

    final_predictions.append(detection_info)
    print("item: ", i)

print("Inference done")

import json
with open('pred.json', 'w') as f:
    json.dump(final_predictions, f, indent=4)