import torch
import torch.jit
import sys
import os

def check_if_jit_model(model_path):
    try:
        # Try to load the model as a TorchScript model
        torch.jit.load(model_path)
        return True, "The model is in TorchScript format."
    except Exception as e:
        return False, f"The model is not in TorchScript format. Error: {str(e)}"

# Append the path to your system path (Note: Don't use quotes)
sys.path.append("D:/Dropbox/FieldPrism/fieldprism/yolov5/weights")


# Check if the path exists and you have read permissions
model_path = "D:/Dropbox/FieldPrism/fieldprism/yolov5/weights_nano/best.pt"


if not os.path.exists(model_path):
    print(f"Model path {model_path} does not exist. Please check the path.")
    sys.exit(1)


is_jit_model, message = check_if_jit_model(model_path)
print(message)


# Load your custom model
# Load your custom model
try:
    loaded_dict = torch.load(model_path, map_location='cuda:0')  # Adjust the device as needed

    # Assuming your model architecture is defined in a class called `YourModelClass`
    # model = YourModelClass() 
    # model.load_state_dict(loaded_dict['model'])  # If the model is saved as a state dictionary

    # If the model is saved entirely (architecture + weights)
    if isinstance(loaded_dict, dict) and 'model' in loaded_dict:
        model = loaded_dict['model']
    else:
        model = loaded_dict  # Assuming the loaded object is a model
    
    # Switch the model to evaluation mode
    model.eval()

    # Create a dummy input that matches the input dimensions of the model
    dummy_input = torch.randn(1, 3, 512, 512).half().to('cuda:0')


    # Try tracing the model
    try:
        scripted_module = torch.jit.trace(model, dummy_input)
        print("Model traced successfully.")
    except Exception as e:
        print(f"An error occurred during tracing. Error: {str(e)}")

    # Try scripting the model
    try:
        scripted_module = torch.jit.script(model)
        save_path = "D:/Dropbox/FieldPrism/fieldprism/yolov5/weights/fieldprism_v_1_0.pt"
        scripted_module.save(save_path)
        print(f"Saved TorchScript model to {save_path}")
    except Exception as e:
        print(f"An error occurred during the scripting. Error: {str(e)}")
except Exception as e:
    print(f"Error in loading the model: {e}")
    sys.exit(1)


# Script the model
try:
    scripted_module = torch.jit.script(model)
    # Save the TorchScript model (make sure you have write permissions for the directory)
    save_path = "D:/Dropbox/FieldPrism/fieldprism/yolov5/weights/fieldprism_v_1_0.pt"
    scripted_module.save(save_path)
    print(f"Saved TorchScript model to {save_path}")
except Exception as e:
    print(f"An error occurred during the scripting. Error: {str(e)}")