import torch
import torch.nn as NN
from logger import Logger
LOGGER = Logger().logger()
LOGGER.info("Started Feature Maps")
device=torch.device( "cuda" if (torch.cuda.is_available()) else 'cpu')
LOGGER.info("Running the model cuda_available = "+str(torch.cuda.is_available()))
#Author: @meetdoshi
class FeatureMaps():
    def __init__(self,arch="vgg19"):
        '''
        Init function
        @params
        arch: str {vgg11,vgg13,vgg16,vgg19,vgg19bn}
        '''
        super()
        try:
            self.model = torch.hub.load('pytorch/vision:v0.10.0',arch,pretrained=True)
        except Exception as e:
            LOGGER.error("Could not load model" + str(e))
        return
    
    def get_model(self):
        return self.model

    def get_layers(self,layers=[]):
        '''
        Function to extract layers
        @params
        layers: list
        '''
        weights = []
        for layer in layers:
            try:
                weights.append(self.model.features[layer].weight)
            except:
                LOGGER.error("Could not fetch layer "+str(layer))
        return weights

    def get_fmaps_content(self,img,layer=[21]):
        '''
        Function which will pass the image through the model and get the respective fmaps
        @params
        img: numpy image f64
        layer: list
        '''
        fmaps = []
        layer_num = 0
        for layer_i in self.model.features:
            img = layer_i(img)
            if layer_num in layer:
                fmaps.append(img)
            layer_num+=1
        return fmaps

    def get_fmaps_style(self,img,layer=[0,5,10,19,28]):
        '''
        Function which will pass the image through the model and get the respective fmaps
        @params
        img: numpy image f64
        layer: list
        '''
        fmaps = []
        layer_num = 0
        for layer_i in self.model.features:
            img = layer_i(img)
            if layer_num in layer:
                fmaps.append(img)
            layer_num+=1
        return fmaps
'''
if __name__ == "__main__":
    fmap = FeatureMaps()
    model = fmap.get_model()
    print(model.features)
    weights = fmap.get_layers([4,2,6])
    print(len(weights))
    for weight in weights:
        print(type(weight),weight.shape)
'''
