import os, sys
import torch
import torchvision
from torch import nn

sys.path.append(os.getcwd())
from algorithms.arch.resnet import loadResnetBackbone
import utilities.runUtils as rutl


##================= CLassifier Wrapper =========================================

class ClassifierNet(nn.Module):
    def __init__(self, arch, fc_layer_sizes=[512,1000],
                    feature_dropout=0, classifier_dropout=0,
                    feature_freeze = False, feature_bnorm = False,
                    torch_pretrain=None):
        super().__init__()
        rutl.START_SEED(7)

        self.fc_layer_sizes = fc_layer_sizes

        # Feature Extractor
        self.backbone,self.feat_outsize = loadResnetBackbone(arch=arch,
                                            torch_pretrain=torch_pretrain,
                                            freeze=feature_freeze)
        fx_layers = []
        if feature_bnorm:
            fx_layers.append(nn.BatchNorm1d(self.feat_outsize, affine=False))
        fx_layers.append(nn.Dropout(p=feature_dropout))

        self.featx_proc = nn.Sequential(*fx_layers)

        # Classifier
        sizes = [self.feat_outsize] + list(self.fc_layer_sizes)
        layers = []
        for i in range(len(sizes) - 2):
            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=False))
            layers.append(nn.LayerNorm(sizes[i + 1]))
            layers.append(nn.ReLU(inplace=True))
            layers.append(nn.Dropout(p=classifier_dropout))
        layers.append(nn.Linear(sizes[-2], sizes[-1], bias=False))

        self.classifier = nn.Sequential(*layers)


    def forward(self, x):
        x = self.backbone(x)
        x = self.featx_proc(x)
        out = self.classifier(x)

        return out



if __name__ == "__main__":

    from torchinfo import summary

    model = ClassifierNet(arch='efficientnet_b0', fc_layer_sizes=[64,8],
                    feature_dropout=0, classifier_dropout=0,
                    torch_pretrain=None)
    summary(model, (1, 3, 200, 200))
    print(model)