from transformers import PreTrainedModel
from .configuration_efficientnetv25 import EfficientNetV25Config
import torch, sys, os
from huggingface_hub import hf_hub_download

class EfficientNetV25ForImageClassification(PreTrainedModel):
    config_class = EfficientNetV25Config

    def __init__(self, config):
        super().__init__(config)

        repo_id = '/'.join(config.url.split('/')[3:5])
        file_name = config.url.split('/')[-1]
        path = f"./models/{file_name}"
        if not os.path.exists(path):
            hf_hub_download(repo_id=repo_id, filename=file_name, local_dir="./models")
        
        self.model = torch.load(path)
        self.input_size = config.input_size
        shape = [2] + self.input_size
        example_inputs = torch.randn(shape)
        example_inputs = (example_inputs - example_inputs.min()) / (example_inputs.max() - example_inputs.min())

        self.num_classes = config.num_classes
        if self.num_classes != 1000:
            self.model.classifier = torch.nn.Linear(in_features=1984, out_features=self.num_classes, bias=True)
        
        traced_model = torch.jit.trace(self.model, example_inputs)
        traced_model.save(file_name)

        self.model = torch.jit.load(file_name)

    def forward(self, tensor, labels=None):
        logits = self.model(tensor)
        if labels is not None:
            loss = torch.nn.cross_entropy(logits, labels)
            return {"loss": loss, "logits": logits}
        return {"logits": logits}