import numpy as np
from scipy.ndimage import zoom
from PIL import Image
import rasterio
import matplotlib.pyplot as plt
import torch
from tqdm.auto import tqdm
import os
import time
import pdb
import yaml

from model import FasterRCNN_V1, FasterRCNN_V2, FasterRCNN_V3

with open('parameters_for_inferences.yml', 'r') as file :
    parameters = yaml.safe_load(file)
    
DEVICE = parameters['tunning']['device']

CLASSES = ['__background__','Crater']
CHECK_PATH = 'checkpoint/best_model.pth'


COLORS_CLASS = ['dodgerblue', 'gold', 'red', 'limegreen', "fuchsia"]
NUM_CLASSES = len(CLASSES)

DATA_DIR = parameters['path']['data_directory']
OUT_DIR = parameters['path']['output_directory'] 

MODEL_NAME = parameters['tunning']['model_name']
PATCHSIZE = parameters['input']['patch_size']
MIN_SIZE = PATCHSIZE
MAX_SIZE = PATCHSIZE
detection_threshold = parameters['input']['threshold_confidence']


if not os.path.isdir(OUT_DIR): 
        os.makedirs(OUT_DIR)
        print(f"Created directory: {OUT_DIR}")


# load the best model and trained weights

# initialize the model and move to the computation device
if MODEL_NAME == 'FasterRCNN_V1' : 
    model = FasterRCNN_V1(num_classes=NUM_CLASSES, 
                            minsize = MIN_SIZE, 
                            maxsize = MAX_SIZE)

elif MODEL_NAME == 'FasterRCNN_V2' : 
    model = FasterRCNN_V2(num_classes=NUM_CLASSES, 
                            minsize = MIN_SIZE, 
                            maxsize = MAX_SIZE)

elif MODEL_NAME == 'FasterRCNN_V3' : 
    model = FasterRCNN_V3(num_classes=NUM_CLASSES, 
                            minsize = MIN_SIZE, 
                            maxsize = MAX_SIZE)



print(f"Model: {MODEL_NAME}")
print(f"send on : {DEVICE}\n")

print(f'Data from : {DATA_DIR}')
# load a model from a pth file

print(f'Model initialisation from: {CHECK_PATH}...')
checkpoint = torch.load(CHECK_PATH, map_location=DEVICE, weights_only=True)
model.load_state_dict(checkpoint['model_state_dict'])
model.to(DEVICE).eval()
print(f'Model initialized!')


# Vérifiez si le fichier est un GeoTIFF ou un fichier numpy
test_images = []
for img in os.listdir(f"{DATA_DIR}"):
    img_path = f"{DATA_DIR}/{img}"
    if img_path.endswith(".tif"):
        with rasterio.open(img_path) as raster:
            data_array = raster.read(1)  # Lecture de la première bande
            resize_data_array = zoom(data_array, (PATCHSIZE / data_array.shape[0], PATCHSIZE / data_array.shape[1]), order=0)
            test_images.append(resize_data_array)
    
    elif img.endswith((".png", ".jpg")):
        with Image.open(img_path) as img_file:
            data_array = np.array(img_file)[:,:,0]
        resize_data_array = zoom(data_array, (PATCHSIZE / data_array.shape[0], PATCHSIZE / data_array.shape[1]), order=0)
        test_images.append(resize_data_array)


    elif img.endswith(".npy"):
        # Chargement des données numpy
        images = np.load(img_path)
        for i in images : test_images.append(i)

print(f"Test instances: {len(test_images)}")

# to count the total number of images iterated through
frame_count = 0
# to keep adding the FPS for each image
total_fps = 0

n_images = len(test_images)

for idx in tqdm(range(n_images), desc = "iter"):
    # get the image file name for saving output later on
    image_name = f"image_{idx}" 

    image = np.array(test_images[idx]).reshape(1,test_images[idx].shape[0],test_images[idx].shape[1])
    orig_image = test_images[idx]#image.copy()
    # convert to tensor
    image = torch.tensor(image, dtype=torch.float).to(DEVICE)
    # add batch dimension
    image = torch.unsqueeze(image, 0)
    start_time = time.time()
    with torch.no_grad():
        outputs = model(image.to(DEVICE))
    end_time = time.time()

    # get the current fps
    fps = 1 / (end_time - start_time)
    # add `fps` to `total_fps`
    total_fps += fps
    # increment frame count
    frame_count += 1

    # load all detection to CPU for further operations
    outputs = [{k: v.to('cpu') for k, v in t.items()} for t in outputs]
    
    # carry further only if there are detected boxes
    if len(outputs[0]['boxes']) != 0:
        pred_boxes = outputs[0]['boxes'].data.numpy()
        scores = outputs[0]['scores'].data.numpy()
        # filter out boxes according to `detection_threshold`
        pred_boxes = pred_boxes[scores >= detection_threshold].astype(np.int32)
        draw_boxes = pred_boxes.copy()
        pred_classes_idx = outputs[0]['labels'].cpu().numpy()
        pred_scores = outputs[0]['scores'].cpu().numpy()

        orig_image[orig_image == 0 ] = 255
        fig = plt.figure(figsize=(10,10))
        ax = fig.add_subplot(1,1,1)
        ax.imshow(orig_image, cmap = 'binary')
        
        # draw the bounding boxes and write the class name on top of it
        for pred_box_idx, pred_box in enumerate(draw_boxes):
            pred_annot = CLASSES[pred_classes_idx[pred_box_idx]] + ' | %2.3f'%pred_scores[pred_box_idx]
            pred_bbox = plt.Rectangle( (int(pred_box[0]), int(pred_box[1])), np.abs(int(pred_box[2]) - int(pred_box[0])), np.abs(int(pred_box[1]) - int(pred_box[3])), facecolor = 'none', edgecolor = COLORS_CLASS[pred_classes_idx[pred_box_idx]-1], linewidth = 2, zorder = 1)
            ax.add_patch(pred_bbox)
            ax.annotate(pred_annot, (int(pred_box[0]), int(pred_box[1] - 3)), color = COLORS_CLASS[pred_classes_idx[pred_box_idx]-1], fontsize = 22, zorder = 1 )


        ax.set_axis_off()
        plt.tight_layout()
        plt.savefig(f"{OUT_DIR}/{str(image_name)}.jpg")

print('TEST PREDICTIONS COMPLETE')
# calculate and print the average FPS
avg_fps = total_fps / frame_count
print(f"Average FPS: {avg_fps:.3f}")