import os
from torchvision.models import resnet50, resnet101,ResNet101_Weights
from torchvision.models.resnet import Bottleneck, BasicBlock
import torch.nn.init as init
import torch.nn as nn
import torch


class ResNet50Classifier(nn.Module):
    def __init__(self, num_classes: int = 14):
        super(ResNet50Classifier, self).__init__()
        self.resnet = resnet50(pretrained=False, num_classes=num_classes)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if not self._load():
            for m in self.modules():
                # print(m)
                if isinstance(m, Bottleneck) and m.bn3.weight is not None:
                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
                elif isinstance(m, BasicBlock) and m.bn2.weight is not None:
                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]
                elif 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.Linear):
                    init.xavier_normal_(m.weight)
                    if m.bias is not None:
                        init.zeros_(m.bias)

    def forward(self, x):
        x = self.resnet(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", "ResNet50Classifier.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


class ResNet101Classifier(nn.Module):
    def __init__(self, num_classes: int = 14):
        super(ResNet101Classifier, self).__init__()
        # self.resnet = resnet101(pretrained=False, num_classes=num_classes)
        
        # ======== ONLY USE THIS FOR TESTING ========
        self.resnet = resnet101(weights=ResNet101_Weights.DEFAULT)
        # Replace the final fully connected layer with one that outputs the desired number of classes
        in_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Linear(in_features, num_classes)
        # ======== ONLY USE THIS FOR TESTING ========
        

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        # if not self._load():
        #     for m in self.modules():
        #         # print(m)
        #         if isinstance(m, Bottleneck) and m.bn3.weight is not None:
        #             nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
        #         elif isinstance(m, BasicBlock) and m.bn2.weight is not None:
        #             nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]
        #         elif 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.Linear):
        #             init.xavier_normal_(m.weight)
        #             if m.bias is not None:
        #                 init.zeros_(m.bias)

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

    def _load(self, filename: str = None) -> bool:
        return False
        if filename is None:
            current_work_dir = os.path.dirname(__file__)
            filename = os.path.join(current_work_dir, "best_pth", "ResNet101Classifier.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