import torch
import torchvision
from torch import nn

def create_densenet161(output_shape: int = 4,
                       seed: int = 42):

  """
  Creates an DenseNet161 as a feature extractor setting all parameters as not trainable
  and changing the shape of the classifier layer equals to 'output_shape'.
  Also creates a corresponding 'torchivision.transforms' trasnformer to preprocess data

  Args:
    output_shape (int): number of classes
    seed (int): sedd for reproducibility

  Returns:
    model: DenseNet161 with 'output_shape' output neurons
    model_transformer: DenseNet161 transformer
  """

  if seed:
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

  model = torchvision.models.densenet161(weights = torchvision.models.DenseNet161_Weights.DEFAULT)
  model_transformer = torchvision.models.DenseNet161_Weights.DEFAULT.transforms()

  for param in model.parameters():
    param.requires_grad = False

  model.classifier = nn.Linear(in_features = 2208, out_features = output_shape, bias=True)

  return model, model_transformer