import torchvision.models.feature_extraction
import torchvision
import os
import torch
import onnx
import onnxruntime
import numpy as np

from .config_extractor import MODEL_CONFIG

os.environ["KMP_DUPLICATE_LIB_OK"] = "True"


class FeatureExtractor:
    """Class for extracting features from images using a pre-trained model"""

    def __init__(self, base_model, onnx_path=None):
        # set the base model
        self.base_model = base_model
        # get the number of features
        self.feat_dims = MODEL_CONFIG[base_model]["feat_dims"]
        # get the feature layer name
        self.feat_layer = MODEL_CONFIG[base_model]["feat_layer"]

        # Set default ONNX path if not provided
        if onnx_path is None:
            onnx_path = f"model/{base_model}_feature_extractor.onnx"

        self.onnx_path = onnx_path
        self.onnx_session = None

        # Initialize transforms (needed for both ONNX and PyTorch)
        _, self.transforms = self.init_model(base_model)

        # Check if ONNX model exists
        if os.path.exists(onnx_path):
            print(f"Loading existing ONNX model from {onnx_path}")
            self.onnx_session = onnxruntime.InferenceSession(onnx_path)
        else:
            print(
                f"ONNX model not found at {onnx_path}. Initializing PyTorch model and converting to ONNX..."
            )
            # Initialize PyTorch model
            self.model, _ = self.init_model(base_model)
            self.model.eval()
            self.device = torch.device("cpu")
            self.model.to(self.device)

            # Create directory if it doesn't exist
            os.makedirs(os.path.dirname(onnx_path), exist_ok=True)

            # Convert to ONNX
            self.convert_to_onnx(onnx_path)

            # Load the newly created ONNX model
            self.onnx_session = onnxruntime.InferenceSession(onnx_path)
            print(f"Successfully created and loaded ONNX model from {onnx_path}")

    def init_model(self, base_model):
        """Initialize the model for feature extraction

        Args:
            base_model: str, the name of the base model

        Returns:
            model: torch.nn.Module, the feature extraction model
            transforms: torchvision.transforms.Compose, the image transformations
        """
        if base_model not in MODEL_CONFIG:
            raise ValueError(f"Invalid base model: {base_model}")

        # get the model and weights
        weights = MODEL_CONFIG[base_model]["weights"]
        model = torchvision.models.feature_extraction.create_feature_extractor(
            MODEL_CONFIG[base_model]["model"](weights=weights),
            [MODEL_CONFIG[base_model]["feat_layer"]],
        )
        # get the image transformations
        transforms = weights.transforms()
        return model, transforms

    def extract_features(self, img):
        """Extract features from an image

        Args:
            img: PIL.Image, the input image

        Returns:
            output: torch.Tensor, the extracted features
        """
        # apply transformations
        x = self.transforms(img)
        # add batch dimension
        x = x.unsqueeze(0)

        # Convert to numpy for ONNX Runtime
        x_numpy = x.numpy()
        # Run inference with ONNX Runtime
        print("Running inference with ONNX Runtime")
        output = self.onnx_session.run(
            None,
            {'input': x_numpy}
        )[0]
        # Convert back to torch tensor
        output = torch.from_numpy(output)

        return output

    def convert_to_onnx(self, save_path):
        """Convert the model to ONNX format and save it

        Args:
            save_path: str, the path to save the ONNX model

        Returns:
            None
        """
        # Create a dummy input tensor
        dummy_input = torch.randn(1, 3, 224, 224, device=self.device)

        # Export the model
        torch.onnx.export(
            self.model,
            dummy_input,
            save_path,
            export_params=True,
            opset_version=14,
            do_constant_folding=True,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size'},
                'output': {0: 'batch_size'}
            }
        )

        # Verify the exported model
        onnx_model = onnx.load(save_path)
        onnx.checker.check_model(onnx_model)
        print(f"ONNX model saved to {save_path}")