import torch
import math
import torch.optim as optim
import gc

#torch.cuda.empty_cache()
#gc.collect()

class ResidualAttentionTranspose(nn.Module):
    def __init__(self, dim):
        super(ResidualAttentionTranspose, self).__init__()
        self.linear1 = ResidualLinear(dim, dim)
        self.attenton = nn.MultiheadAttention(dim, 8, batch_first=True)
                

    def forward(self, x):
        shape = x.shape
        x = x.flatten(2).transpose(1, 2)
        x = self.linear1(x)
        attn_output, attn_weights = self.attenton(x, x, x)

        return (x + attn_output).view(shape)
    
class ResidualAttentionFlatten(nn.Module):
    def __init__(self, dim):
        super(ResidualAttentionFlatten, self).__init__()
        self.linear1 = ResidualLinear(dim, dim)
        encoder_layer = nn.TransformerEncoderLayer(d_model=dim, nhead=4)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=3)
        #self.attenton = nn.MultiheadAttention(dim, 8, batch_first=True)
                

    def forward(self, x):
        shape = x.shape
        x = x.flatten(2)
        x = self.linear1(x)
        #attn_output, attn_weights = self.attenton(x, x, x)
        attn_output = self.transformer_encoder(x)


        return (x + attn_output).view(shape)

class PrintShape(nn.Module):
    def __init__(self):
        super(PrintShape, self).__init__()

    def forward(self, x):
    	#print(x.shape)
    	return x

class PositionEmbedding(nn.Module):
    def __init__(self, dim):
        super(PositionEmbedding, self).__init__()
        self.register_buffer("position_ids", torch.arange(16384).expand((1, -1)), persistent=False)
        self.position_embedding = nn.Embedding(16384, dim)
        
    def forward(self, x):
        shape = x.shape
        
        x = x.flatten(2).transpose(1, 2)
        embeddings = self.position_embedding(self.position_ids)
        x =  x + embeddings

        return x.transpose(2, 1).view(shape)

class MultiScaleConvModule(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(MultiScaleConvModule, self).__init__()
        
        # Define 3 parallel convolutional layers with different kernel sizes
        self.conv4 = nn.Conv2d(in_channels, out_channels, kernel_size=4, padding=1, stride=2)
        self.conv8 = nn.Conv2d(in_channels, out_channels, kernel_size=8, padding=3, stride=2)
        self.conv16 = nn.Conv2d(in_channels, out_channels, kernel_size=16, padding=7, stride=2)
    
        #self.fe1 = FeatureEnhancer(8, 256*256)
        #self.fe2 = FeatureEnhancer(8, 256*256)
        #self.fe3 = FeatureEnhancer(8, 256*256)

    def forward(self, x):
        # Apply each convolutional layer
        x1 = (self.conv4(x))
        #print(x1.shape)
        x2 = (self.conv8(x))
        #print(x2.shape)
        x3 = (self.conv16(x))
        #print(x3.shape)

        # Combine the outputs from each convolutional branch
        # You can choose different strategies to combine them.
        # Here, we simply sum them up.
        out = torch.cat([x1, x2, x3], dim=1)
        #print(out.shape)

        # Optionally, apply non-linearity
        out = F.relu(out)
        
        return out

class FeatureEnhancer(nn.Module):
    def __init__(self, channel_count, input_dim):
        super(FeatureEnhancer, self).__init__()
        hidden_dim = input_dim // 16

        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=channel_count, out_channels=channel_count, kernel_size=16, stride=4, padding=6),
            nn.ReLU(),
            #nn.Conv2d(in_channels=channel_count, out_channels=channel_count, kernel_size=3, stride=2, padding=1),
			)
        self.attn = ResidualAttentionFlatten(hidden_dim)
        #self.ff = nn.Linear(hidden_dim, input_dim)
        self.upconv = nn.ConvTranspose2d(channel_count, channel_count, kernel_size=4, stride=4)
        print(hidden_dim)

    def forward(self, x):
        shape = x.shape
        #print(shape)
        y = self.conv(x)
        y = self.attn(y)
        y = self.upconv(y)#(y.flatten(2))
        #print(y.shape)

        return x + y.view(shape)

class CrossAttention(nn.Module):
    def __init__(self, q_dim, kv_dim, attn_dim, num_heads):
        super(CrossAttention, self).__init__()
        self.query_proj = nn.Linear(q_dim, attn_dim)
        self.key_proj = nn.Linear(kv_dim, attn_dim)
        self.value_proj = nn.Linear(kv_dim, attn_dim)
        self.multihead_attn = nn.MultiheadAttention(embed_dim=attn_dim, num_heads=num_heads)

    def forward(self, query, key, value):
        query = self.query_proj(query)
        key = self.key_proj(key)
        value = self.value_proj(value)
        attn_output, _ = self.multihead_attn(query, key, value)
        return attn_output

class DifferenceAE(nn.Module):
    def __init__(self, input_channels=3, attn_channels=4, num_heads=4, attn_dim=128):
        super(DifferenceAE, self).__init__()
        self.encoder = nn.Sequential(
            MultiScaleConvModule(in_channels=3, out_channels=16),
            PrintShape(),
            nn.Conv2d( in_channels=16*3, out_channels=24, kernel_size=3, stride=2, padding=1),
            FeatureEnhancer(24, 128*128),
            PrintShape(),
            nn.ReLU(),    
            nn.Conv2d( in_channels=24, out_channels=12, kernel_size=3, stride=2, padding=1),
            FeatureEnhancer(12, 64*64),
            PrintShape(),
            nn.ReLU(),    
            nn.Conv2d( in_channels=12, out_channels=8, kernel_size=3, stride=2, padding=1),
            FeatureEnhancer(8, 32*32),
            PrintShape(),
            nn.ReLU(),    
            nn.Conv2d( in_channels=8, out_channels=4, kernel_size=3, stride=1, padding=1),
            FeatureEnhancer(4, 32*32),
            PrintShape(),
            #ResidualAttentionFlatten(64*64),
            #PrintShape(),
            )

        #self.cross_attention = self.cross_attention = CrossAttention(q_dim=attn_dim, kv_dim=attn_channels*32*32, attn_dim=attn_dim, num_heads=num_heads)

        self.decoder = decoder = nn.Sequential(
            nn.ConvTranspose2d(in_channels=4, out_channels=16, kernel_size=4, stride=2, padding=1),
            PrintShape(),
            FeatureEnhancer(16, 64*64),
            nn.ReLU(),
            nn.ConvTranspose2d(in_channels=16, out_channels=12, kernel_size=4, stride=2, padding=1),
            PrintShape(),
            FeatureEnhancer(12, 128*128),
            nn.ReLU(),
            nn.ConvTranspose2d(in_channels=12, out_channels=8, kernel_size=4, stride=2, padding=1),
            PrintShape(),
            #FeatureEnhancer(8, 128*128),
            nn.ReLU(),
            nn.ConvTranspose2d(in_channels=8, out_channels=6, kernel_size=4, stride=2, padding=1 ),
            nn.ReLU(),
            #FeatureEnhancer(8, 128*128),
            PrintShape(),
            nn.ConvTranspose2d(in_channels=6, out_channels=4, kernel_size=3, stride=1, padding=1 ),
            PrintShape(),
            #FeatureEnhancer(8, 256*256),
            nn.ReLU(),
            nn.ConvTranspose2d(in_channels=4, out_channels=3, kernel_size=3, stride=1, padding=1 ),

            ).to('cuda')

    def encode(self , x):
        return self.encoder(x)
    
    def decode(self , x):
        return self.decoder(x)

    def forward(self, x):
        x = self.encode(x)
        return self.decode(x)


model = DifferenceAE().to('cuda')

#linear = ResidualAttention(768)

tensor = torch.rand(1, 3, 512, 512).to('cuda')
encoded = model.encode(tensor)
print(f'encoded: {encoded.shape}')
tensor = model.decode(encoded)
print(f'decoded: {tensor.shape}')

criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)
torch.cuda.empty_cache()


LOSS BECAME NAN