import logging

import tensorflow as tf
from tensorflow import keras

# class ResNet(tf.keras.Model):
#     def __init__(self, num_classes=1000):
#         super(ResNet, self).__init__()
#         self.block_1 = ResNetBlock()
#         self.block_2 = ResNetBlock()
#         self.global_pool = layers.GlobalAveragePooling2D()
#         self.classifier = Dense(num_classes)

#     def call(self, inputs):
#         x = self.block_1(inputs)
#         x = self.block_2(x)
#         x = self.global_pool(x)
#         return self.classifier(x)


class PFENet(tf.keras.Model):
    def __init__(self, layers=50, shot=1) -> None:

        assert layers in [50, 101, 152]
        assert classes > 1

        self.zoom_factor = 0.2  # zoom_factor
        self.criterion = None  # TODO set this to CategoricalCrossentropy
        self.shot = shot
        self.ppm_scales = ppm_scales = [60, 30, 15, 8]  # TODO what is this
        # TODO batchnorm somewhere?

        logging.info("Using ResNet %d" % layers)
        resnet = tf.keras.applications.ResNet50V2(
            include_top=True, #?
            weights="imagenet",
            input_tensor=None,
            input_shape=None,
            pooling=None,
            classes=1000,
            classifier_activation="softmax",
        )

        
        self.layer0 = tf.keras.models.Sequential([
            resnet.conv1, resnet.bn1, resnet.relu1, 
            resnet.conv2, resnet.bn2, resnet
    
                                                 ]
