import torch
import yaml
import logging
import onnx
from google.protobuf.json_format import MessageToDict

# import self-made modules
import utils
import pruning
import quantization
import ImageDataLoader
import test_suite_model
import train_maskrcnn
# ---------------------------------------------------------------------------- #
def main():
    #logging setup
    utils.setup_logging()

    # Load the config file
    with open("config.yaml", "r") as file:
        config = yaml.safe_load(file)

    # Load the model

    model = utils.load_model_to_pytorch(pretrained = True) #pretrained uses the pytorch pretrained model

# ---------------------------------------------------------------------------- #
    #Data, fine-tune the model on a small number of images and their annotations
    train_dataloader  = ImageDataLoader.return_dataloader(config["datasets"]["train_data_path"],
                                                          config["datasets"]["train_annotation_path"])
    val_dataloader  = ImageDataLoader.return_dataloader(config["datasets"]["val_data_path"],
                                                        config["datasets"]["val_annotation_path"])
    test_dataloader = ImageDataLoader.return_dataloader(config["datasets"]["test_data_path"],
                                                        config["datasets"]["test_annotation_path"])

    #Check if the dataloader is working correctly by visualizing a sample
    utils.visualize_one_dataloader_sample(train_dataloader)

    model,train_losses_head, val_losses_head = train_maskrcnn.train_finetune_head(model, train_dataloader, val_dataloader, device=config["device"], num_epochs=10) #train only the detection head
    #model, train_losses_backbone, val_losses_backbone = train_maskrcnn.train_backbone(model, train_dataloader, val_dataloader, device=config["device"], num_epochs=5) #train the whole model

    #utils.plot_combined_losses(train_losses_head, train_losses_backbone,val_losses_head, val_losses_backbone)

# ---------------------------------------------------------------------------- #
    # Set up naming components
    name_components = []

    # Apply compression techniques
    if config["compression"]["prune"]["enabled"]:
        logging.info("Applying pruning")

        amount = config["compression"]["prune"]["amount"]
        model = pruning.apply_local_pruning(model, amount)
        name_components.append(f"pruned_{amount}")

    if config["compression"]["quantize"]["enabled"]:
        logging.info("Applying quantization")

        bit_precision = config["compression"]["quantize"]["bit_precision"]
        model = quantization.apply_quantization(model, bit_precision)
        name_components.append(f"quantized_{bit_precision}bit")

    if config["compression"]["distill"]["enabled"]:
        logging.info("Applying distillation")

        name_components.append("distilled")
        pass

# ---------------------------------------------------------------------------- #
    # Generate & save the ONNX model name based on the compression techniques
    base_model_name = "MRCNN_base"
    model_name = f"{base_model_name}_{'_'.join(name_components)}"
    model_full_file_path = f"./models/{model_name}.onnx" #use the name without the .onnx extension

    torch_input = torch.randn(1, 3, 1216, 1368)
    utils.save_model_onnx(model_full_file_path,model, torch_input)
    """
    IMPORTANT:
    input: [batch, channels, height, width]
    output: boxes, labels, scores, masks
    """
# ---------------------------------------------------------------------------- # WHERE IT GOES WRONG

    # Run inference on a single image
    image_path = "./images/train/img_ (1).bmp"
    onnx_model_path = "./models/MRCNN_base_.onnx"

    utils.print_model_info("./models/MRCNN_base_.onnx")

    results = utils.run_inference(image_path, onnx_model_path, conf_threshold=0.5)
    utils.visualize_detections(results)

# ---------------------------------------------------------------------------- #
    # Run the test suite for both cpu and gpu
    # metrics consist of: accuracy, F1, precision, recall, jaccard index, auc_roc
    #model_full_file_path = "./models/MRCNN_base_.onnx"
    avg_cpu_inference_time, cpu_metrics, all_ground_truths_cpu, all_predictions_cpu = test_suite_model.run_cpu_tests(model_full_file_path, test_dataloader)
    avg_gpu_inference_time, gpu_metrics, all_ground_truths_gpu, all_predictions_gpu = test_suite_model.run_gpu_tests(model_full_file_path, test_dataloader)
    print(avg_cpu_inference_time, cpu_metrics, all_ground_truths_cpu, all_predictions_cpu)

    # Extract metrics
    cpu_accuracy,cpu_f1_score,cpu_precision,cpu_recall,cpu_jaccard_index,cpu_auc_roc = cpu_metrics["accuracy"],cpu_metrics["precision"],cpu_metrics["recall"],cpu_metrics["jaccard_index"],cpu_metrics["auc_roc"]
    gpu_accuracy,gpu_f1_score,gpu_precision,gpu_recall,gpu_jaccard_index,gpu_auc_roc = gpu_metrics["accuracy"],gpu_metrics["precision"],gpu_metrics["recall"],gpu_metrics["jaccard_index"],gpu_metrics["auc_roc"]
# ---------------------------------------------------------------------------- #
    # save the results
    results_file_path = f"./results/{model_name}" #use the name without the .onnx extension
    #utils.write_out_results(results_file_path, config, inference_time, accuracy_score)

    logging.info("program finished")

if __name__ == "__main__":
    main()