import os
from torchvision.models.inception import inception_v3
import torch.nn.init as init
import torch.nn as nn
import torch


class InceptionV3Classifier(nn.Module):
    def __init__(self, num_classes: int = 14):
        super(InceptionV3Classifier, self).__init__()
        self.inception = inception_v3(
            pretrained=False, num_classes=num_classes, aux_logits=False
        )

        # Initialize weights if not loading from a file
        if not self._load():
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                    if m.bias is not None:
                        init.zeros_(m.bias)
                elif isinstance(m, nn.BatchNorm2d):
                    init.ones_(m.weight)
                    init.zeros_(m.bias)
                elif isinstance(m, nn.Linear):
                    init.xavier_normal_(m.weight)
                    if m.bias is not None:
                        init.zeros_(m.bias)

    def forward(self, x):
        x = self.inception(x)
        return x

    def _load(self, filename: str = None) -> bool:
        if filename is None:
            current_work_dir = os.path.dirname(__file__)
            filename = os.path.join(
                current_work_dir, "best_pth", "InceptionV3Classifier.pth"
            )
        if not os.path.exists(filename):
            print("Model file does not exist.")
            return False
        self.load_state_dict(torch.load(filename))
        print("Model loaded successfully.")
        return True