import torch
import os
import yaml
from PIL import Image
from torchvision import transforms

import yaml

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm

class EfficientNetModel(nn.Module):
    def __init__(self, num_classes=8, model_name='efficientnet_b7', pretrained=True):
        super(EfficientNetModel, self).__init__()
        # self.status_size = 768
        # self.type_size = 512
        # self.num_worker = 2
        # self.target_size_status = 2
        # self.target_size_type = 8
        # self.target_size_fault = 23
        self.model = timm.create_model(model_name, pretrained=pretrained)
        self.model.classifier = nn.Linear(self.model.classifier.in_features, num_classes)
        
    def forward(self, x):
        x = self.model(x)
        return x

def get_transforms(cfg):
    return transforms.Compose([
        transforms.Resize((cfg['size'], cfg['size'])),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        ),
    ])

def load_image(image_path, transform):
    image = Image.open(image_path)
    image = transform(image).unsqueeze(0)
    return image

def inference(model, state, image, device):
    model.to(device)
    model.load_state_dict(state['model'])
    model.eval()
    image = image.to(device)
    with torch.no_grad():
        y_preds = model(image)
    return y_preds.softmax(1).cpu().numpy()

if __name__ == "__main__":
    cuda = 'cuda:0'
    image_path = '/data/si/KEPCO/src/images/7815_2075S791_LP애자_설치불량.JPG'


    device = cuda
    MODEL_DIR = 'weights/resnext50d_32x4d_fold0_best.pth'

    transform = get_transforms(cfg)
    image = load_image(image_path, transform)
    
    modelC = EfficientNetModel(num_classes=8, model_name='tf_efficientnetv2_s_in21k')
    state = torch.load('tf_efficientnetv2_s_in21k_fold0_best.pth', map_location='cuda:0')
    
    # Move model and image to the specified device
    modelC = modelC.to(device)
    image = image.to(device)
    
    modelC.load_state_dict(state['model'])
    modelC.eval()
    with torch.no_grad():
        y_preds = modelC(image)  # use 'image' instead of 'img'
    
    print(f'Predictions: {y_preds}')