import torch
import torch.nn as nn

def calculate_conv_output_size(input_size, conv_layers):
    # Start with the input size
    size = input_size
    
    for layer in conv_layers:
        if isinstance(layer, nn.Conv2d):
            # Calculate output size for each dimension
            size = (
                (size[0] - layer.kernel_size[0] + 2*layer.padding[0]) // layer.stride[0] + 1,
                (size[1] - layer.kernel_size[1] + 2*layer.padding[1]) // layer.stride[1] + 1
            )
        elif isinstance(layer, nn.MaxPool2d):
            # If you have max pooling layers, include this calculation
            size = (
                (size[0] - layer.kernel_size + 2*layer.padding) // layer.stride + 1,
                (size[1] - layer.kernel_size + 2*layer.padding) // layer.stride + 1
            )
    
    return size

class ImageClassifier(nn.Module):
    def __init__(self, input_size=(28, 28)):
        super(ImageClassifier, self).__init__()
        
        self.conv_layers = nn.Sequential(
            nn.Conv2d(1, 32, (3,3)),
            nn.ReLU(),
            nn.Conv2d(32, 64, (3,3)),
            nn.ReLU(),
            nn.Conv2d(64, 64, (3,3)),
            nn.ReLU()
        )
        
        # Calculate the output size of the conv layers
        conv_output_size = calculate_conv_output_size(input_size, self.conv_layers)
        
        # Calculate the flattened size
        flattened_size = 64 * conv_output_size[0] * conv_output_size[1]
        
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(flattened_size, 10)
        )
        
        print(f"Convolutional output size: {conv_output_size}")
        print(f"Flattened size: {flattened_size}")

    def forward(self, x):
        x = self.conv_layers(x)
        x = self.fc_layers(x)
        return x

# Create an instance of the model
model = ImageClassifier()

# You can also verify the output size using a dummy input
dummy_input = torch.randn(1, 1, 28, 28)
output = model(dummy_input)
print(f"Output shape: {output.shape}")