import torch


def _get_intermediate_features(resnet, input_):
    """ Receives a ResNet50 network and a valid input and extracts intermediate features.

    Extracts features right after each conv or fc layer (plus one right before the last 
    fc) so 51 in total.
    
    Arguments:
        resnet (nn.Module): A ResNet50 from torchvision.models.
        input_ (torch.Tensor): Input to the Resnet (N x C x H x W)
    
    Returns:
        x (torch.Tensor): Output of the resnet (same as resnet(input_))
        features (list): A list with 51 feature representations. Each element is a 
            (N x num_features).
            
    Note:
        To make sense of it, check the original Resnet50 forward. I essentially just 
        copied everything here.
    """
    features = []

    x = resnet.conv1(input_)
    features.append(x.mean(dim=(-1, -2)).cpu().numpy())
    x = resnet.bn1(x)
    x = resnet.relu(x)
    x = resnet.maxpool(x)

    for layer in [resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4]:
        for bottleneck in layer:
            identity = x

            out = bottleneck.conv1(x)
            features.append(out.mean(dim=(-1, -2)).cpu().numpy())
            out = bottleneck.bn1(out)
            out = bottleneck.relu(out)

            out = bottleneck.conv2(out)
            features.append(out.mean(dim=(-1, -2)).cpu().numpy())
            out = bottleneck.bn2(out)
            out = bottleneck.relu(out)

            out = bottleneck.conv3(out)
            features.append(out.mean(dim=(-1, -2)).cpu().numpy())
            out = bottleneck.bn3(out)

            if bottleneck.downsample is not None:
                identity = bottleneck.downsample(x)

            out += identity
            out = bottleneck.relu(out)

            #return out
            x = out

    x = resnet.avgpool(x)
    x = torch.flatten(x, 1)
    features.append(x.cpu().numpy())
    x = resnet.fc(x)
    features.append(x.cpu().numpy())

    return x, features