import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from common.utils import round_func
from pathlib import Path
from common import lut
from . import rclut
from common import layers 

# class ReconstructedConvCentered(nn.Module):
#     def __init__(self, hidden_dim,  window_size=7):
#         super(ReconstructedConvCentered, self).__init__()
#         self.window_size = window_size
#         self.projection1 = torch.nn.Parameter(torch.rand((window_size**2, hidden_dim))/window_size)
#         self.projection2 = torch.nn.Parameter(torch.rand((window_size**2, hidden_dim))/window_size)

#     def pixel_wise_forward(self, x):
#         x = (x-127.5)/127.5
#         out = torch.einsum('bwk,wh,wh -> bwk', x, self.projection1, self.projection2) 
#         out = torch.tanh(out)
#         out = out*127.5 + 127.5
#         return out 

#     def forward(self, x):
#         original_shape = x.shape  
#         x = F.pad(x, pad=[self.window_size//2]*4, mode='replicate')
#         x = F.unfold(x, self.window_size)
#         x = self.pixel_wise_forward(x)     
#         x = x.mean(1)  
#         x = x.reshape(*original_shape)    
#         x = round_func(x) 
#         return x

#     def __repr__(self):
#         return f"{self.__class__.__name__} projection1: {self.projection1.shape} projection2: {self.projection2.shape}"

# class RCBlockCentered(nn.Module):
#     def __init__(self, hidden_dim = 32, window_size=3, dense_conv_layer_count=4, upscale_factor=4):
#         super(RCBlockCentered, self).__init__()   
#         self.window_size = window_size    
#         self.rc_conv = ReconstructedConvCentered(hidden_dim=hidden_dim, window_size=window_size)
#         self.dense_conv_block = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=dense_conv_layer_count, upscale_factor=upscale_factor)

#     def forward(self, x):
#         b,c,hs,ws = x.shape
#         x = self.rc_conv(x)
#         x = F.pad(x, pad=[0,1,0,1], mode='replicate') 
#         x = self.dense_conv_block(x)
#         return x

# class RCNetCentered_3x3(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetCentered_3x3, self).__init__()
#         self.hidden_dim = hidden_dim
#         self.layers_count = layers_count
#         self.scale = scale       
#         self.stage = RCBlockCentered(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=3)
            
#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         x = self.stage(x)
#         x = x.view(b, c, h*self.scale, w*self.scale)
#         return x

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         window_size = self.stage.rc_conv.window_size
#         rc_conv_luts = lut.transfer_rc_conv(self.stage.rc_conv, quantization_interval=quantization_interval).reshape(window_size,window_size,-1)
#         dense_conv_lut = lut.transfer_2x2_input_SxS_output(self.stage.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
#         lut_model = rclut.RCLutCentered_3x3.init_from_numpy(rc_conv_luts, dense_conv_lut)
#         return lut_model

# class RCNetCentered_7x7(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetCentered_7x7, self).__init__()
#         self.hidden_dim = hidden_dim
#         self.layers_count = layers_count
#         self.scale = scale       
#         window_size = 7
#         self.stage = RCBlockCentered(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=window_size)
            
#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         x = self.stage(x)
#         x = x.view(b, c, h*self.scale, w*self.scale)
#         return x

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         window_size = self.stage.rc_conv.window_size
#         rc_conv_luts = lut.transfer_rc_conv(self.stage.rc_conv, quantization_interval=quantization_interval).reshape(window_size,window_size,-1)
#         dense_conv_lut = lut.transfer_2x2_input_SxS_output(self.stage.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
#         lut_model = rclut.RCLutCentered_7x7.init_from_numpy(rc_conv_luts, dense_conv_lut)
#         return lut_model


# class ReconstructedConvRot90(nn.Module):
#     def __init__(self, hidden_dim,  window_size=7):
#         super(ReconstructedConvRot90, self).__init__()
#         self.window_size = window_size
#         self.projection1 = torch.nn.Parameter(torch.rand((window_size**2, hidden_dim))/window_size)
#         self.projection2 = torch.nn.Parameter(torch.rand((window_size**2, hidden_dim))/window_size)

#     def pixel_wise_forward(self, x):
#         x = (x-127.5)/127.5
#         out = torch.einsum('bik,ij,ij -> bik', x, self.projection1, self.projection2) 
#         out = torch.tanh(out)
#         out = out*127.5 + 127.5
#         return out 

#     def forward(self, x):
#         original_shape = x.shape
#         x = F.pad(x, pad=[0,self.window_size-1,0,self.window_size-1], mode='replicate')
#         x = F.unfold(x, self.window_size)
#         x = self.pixel_wise_forward(x)        
#         x = x.mean(1)  
#         x = x.reshape(*original_shape)  
#         x = round_func(x) # quality likely suffer from this
#         return x

#     def __repr__(self):
#         return f"{self.__class__.__name__} projection1: {self.projection1.shape} projection2: {self.projection2.shape}"

# class RCBlockRot90(nn.Module):
#     def __init__(self, hidden_dim = 32, window_size=3, dense_conv_layer_count=4, upscale_factor=4):
#         super(RCBlockRot90, self).__init__()   
#         self.window_size = window_size    
#         self.rc_conv = ReconstructedConvRot90(hidden_dim=hidden_dim, window_size=window_size)
#         self.dense_conv_block = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=dense_conv_layer_count, upscale_factor=upscale_factor)

#     def forward(self, x):
#         b,c,hs,ws = x.shape
#         x = self.rc_conv(x)
#         x = F.pad(x, pad=[0,1,0,1], mode='replicate') 
#         x = self.dense_conv_block(x)
        
#         return x

# class RCNetRot90_3x3(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetRot90_3x3, self).__init__()
#         self.hidden_dim = hidden_dim
#         self.layers_count = layers_count
#         self.scale = scale       
#         self.stage = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=3)

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         window_size = self.stage.rc_conv.window_size
#         rc_conv_luts = lut.transfer_rc_conv(self.stage.rc_conv, quantization_interval=quantization_interval).reshape(window_size,window_size,-1)
#         dense_conv_lut = lut.transfer_2x2_input_SxS_output(self.stage.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
#         lut_model = rclut.RCLutRot90_3x3.init_from_numpy(rc_conv_luts, dense_conv_lut)
#         return lut_model

#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             rotated_prediction = self.stage(rotated)
#             unrotated_prediction = torch.rot90(rotated_prediction, k=-rotations_count, dims=[2, 3])
#             output += unrotated_prediction
#         output /= 4
#         output = output.view(b, c, h*self.scale, w*self.scale)
#         return output

# class RCNetRot90_7x7(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetRot90_7x7, self).__init__()
#         self.hidden_dim = hidden_dim
#         self.layers_count = layers_count
#         self.scale = scale       
#         self.stage = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=7)

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         window_size = self.stage.rc_conv.window_size
#         rc_conv_luts = lut.transfer_rc_conv(self.stage.rc_conv, quantization_interval=quantization_interval).reshape(window_size,window_size,-1)
#         dense_conv_lut = lut.transfer_2x2_input_SxS_output(self.stage.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
#         lut_model = rclut.RCLutRot90_7x7.init_from_numpy(rc_conv_luts, dense_conv_lut)
#         return lut_model

#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             rotated_prediction = self.stage(rotated)
#             unrotated_prediction = torch.rot90(rotated_prediction, k=-rotations_count, dims=[2, 3])
#             output += unrotated_prediction
#         output /= 4
#         output = output.view(b, c, h*self.scale, w*self.scale)
#         return output

# class RCNetx1(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetx1, self).__init__()
#         self.scale = scale  
#         self.hidden_dim = hidden_dim     
#         self.stage1_3x3 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=3)
#         self.stage1_5x5 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=5)
#         self.stage1_7x7 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=7)

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage1_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage1_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage1_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage1_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage1_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage1_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
        
#         lut_model = rclut.RCLutx1.init_from_numpy(
#             rc_conv_luts_3x3=rc_conv_luts_3x3, dense_conv_lut_3x3=dense_conv_lut_3x3,
#             rc_conv_luts_5x5=rc_conv_luts_5x5, dense_conv_lut_5x5=dense_conv_lut_5x5,
#             rc_conv_luts_7x7=rc_conv_luts_7x7, dense_conv_lut_7x7=dense_conv_lut_7x7
#         )
#         return lut_model

#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_7x7(rotated), k=-rotations_count, dims=[2, 3])
            
#         output /= 3*4
#         output = output.view(b, c, h*self.scale, w*self.scale)
#         return output


# class RCNetx2(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetx2, self).__init__()
#         self.scale = scale  
#         self.hidden_dim = hidden_dim     
#         self.stage1_3x3 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=3)
#         self.stage1_5x5 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=5)
#         self.stage1_7x7 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=7)
#         self.stage2_3x3 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=3)
#         self.stage2_5x5 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=5)
#         self.stage2_7x7 = RCBlockRot90(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=7)

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         s1_rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage1_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         s1_dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage1_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s1_rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage1_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         s1_dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage1_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s1_rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage1_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         s1_dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage1_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage2_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         s2_dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage2_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage2_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         s2_dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage2_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage2_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         s2_dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage2_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
        
#         lut_model = rclut.RCLutx2.init_from_numpy(
#             s1_rc_conv_luts_3x3=s1_rc_conv_luts_3x3, s1_dense_conv_lut_3x3=s1_dense_conv_lut_3x3,
#             s1_rc_conv_luts_5x5=s1_rc_conv_luts_5x5, s1_dense_conv_lut_5x5=s1_dense_conv_lut_5x5,
#             s1_rc_conv_luts_7x7=s1_rc_conv_luts_7x7, s1_dense_conv_lut_7x7=s1_dense_conv_lut_7x7,
#             s2_rc_conv_luts_3x3=s2_rc_conv_luts_3x3, s2_dense_conv_lut_3x3=s2_dense_conv_lut_3x3,
#             s2_rc_conv_luts_5x5=s2_rc_conv_luts_5x5, s2_dense_conv_lut_5x5=s2_dense_conv_lut_5x5,
#             s2_rc_conv_luts_7x7=s2_rc_conv_luts_7x7, s2_dense_conv_lut_7x7=s2_dense_conv_lut_7x7
#         )
#         return lut_model

#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_7x7(rotated), k=-rotations_count, dims=[2, 3])            
#         output /= 3*4
#         x = output
#         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_7x7(rotated), k=-rotations_count, dims=[2, 3])            
#         output /= 3*4
#         output = output.view(b, c, h*self.scale, w*self.scale)
#         return output

# class RCNetx2Centered(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetx2Centered, self).__init__()
#         self.scale = scale  
#         self.hidden_dim = hidden_dim     
#         self.stage1_3x3 = RCBlockCentered(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=3)
#         self.stage1_5x5 = RCBlockCentered(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=5)
#         self.stage1_7x7 = RCBlockCentered(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=7)
#         self.stage2_3x3 = RCBlockCentered(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=3)
#         self.stage2_5x5 = RCBlockCentered(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=5)
#         self.stage2_7x7 = RCBlockCentered(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=7)

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         s1_rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage1_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         s1_dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage1_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s1_rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage1_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         s1_dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage1_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s1_rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage1_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         s1_dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage1_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage2_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         s2_dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage2_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage2_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         s2_dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage2_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage2_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         s2_dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage2_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
        
#         lut_model = rclut.RCLutx2Centered.init_from_numpy(
#             s1_rc_conv_luts_3x3=s1_rc_conv_luts_3x3, s1_dense_conv_lut_3x3=s1_dense_conv_lut_3x3,
#             s1_rc_conv_luts_5x5=s1_rc_conv_luts_5x5, s1_dense_conv_lut_5x5=s1_dense_conv_lut_5x5,
#             s1_rc_conv_luts_7x7=s1_rc_conv_luts_7x7, s1_dense_conv_lut_7x7=s1_dense_conv_lut_7x7,
#             s2_rc_conv_luts_3x3=s2_rc_conv_luts_3x3, s2_dense_conv_lut_3x3=s2_dense_conv_lut_3x3,
#             s2_rc_conv_luts_5x5=s2_rc_conv_luts_5x5, s2_dense_conv_lut_5x5=s2_dense_conv_lut_5x5,
#             s2_rc_conv_luts_7x7=s2_rc_conv_luts_7x7, s2_dense_conv_lut_7x7=s2_dense_conv_lut_7x7
#         )
#         return lut_model

#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_7x7(rotated), k=-rotations_count, dims=[2, 3])            
#         output /= 3*4
#         x = output
#         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_7x7(rotated), k=-rotations_count, dims=[2, 3])            
#         output /= 3*4
#         output = output.view(b, c, h*self.scale, w*self.scale)
#         return output

# class ReconstructedConvRot90Unlutable(nn.Module):
#     def __init__(self, hidden_dim,  window_size=7):
#         super(ReconstructedConvRot90Unlutable, self).__init__()
#         self.window_size = window_size
#         self.projection1 = torch.nn.Parameter(torch.rand((window_size**2, hidden_dim))/window_size)
#         self.projection2 = torch.nn.Parameter(torch.rand((window_size**2, hidden_dim))/window_size)

#     def pixel_wise_forward(self, x):
#         x = (x-127.5)/127.5
#         out = torch.einsum('bik,ij,ij -> bik', x, self.projection1, self.projection2) 
#         out = torch.tanh(out)
#         out = out*127.5 + 127.5
#         return out 

#     def forward(self, x):
#         original_shape = x.shape
#         x = F.pad(x, pad=[0,self.window_size-1,0,self.window_size-1], mode='replicate')
#         x = F.unfold(x, self.window_size)
#         x = self.pixel_wise_forward(x)        
#         x = x.mean(1)  
#         x = x.reshape(*original_shape)  
#         # x = round_func(x) # quality likely suffer from this
#         return x

#     def __repr__(self):
#         return f"{self.__class__.__name__} projection1: {self.projection1.shape} projection2: {self.projection2.shape}"

# class RCBlockRot90Unlutable(nn.Module):
#     def __init__(self, hidden_dim = 32, window_size=3, dense_conv_layer_count=4, upscale_factor=4):
#         super(RCBlockRot90Unlutable, self).__init__()   
#         self.window_size = window_size    
#         self.rc_conv = ReconstructedConvRot90Unlutable(hidden_dim=hidden_dim, window_size=window_size)
#         self.dense_conv_block = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=dense_conv_layer_count, upscale_factor=upscale_factor)

#     def forward(self, x):
#         b,c,hs,ws = x.shape
#         x = self.rc_conv(x)
#         x = F.pad(x, pad=[0,1,0,1], mode='replicate') 
#         x = self.dense_conv_block(x)        
#         return x

# class RCNetx2Unlutable(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetx2Unlutable, self).__init__()
#         self.scale = scale  
#         self.hidden_dim = hidden_dim     
#         self.stage1_3x3 = RCBlockRot90Unlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=3)
#         self.stage1_5x5 = RCBlockRot90Unlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=5)
#         self.stage1_7x7 = RCBlockRot90Unlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=7)
#         self.stage2_3x3 = RCBlockRot90Unlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=3)
#         self.stage2_5x5 = RCBlockRot90Unlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=5)
#         self.stage2_7x7 = RCBlockRot90Unlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=7)

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         s1_rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage1_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         s1_dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage1_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s1_rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage1_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         s1_dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage1_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s1_rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage1_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         s1_dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage1_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage2_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         s2_dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage2_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage2_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         s2_dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage2_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage2_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         s2_dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage2_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
        
#         lut_model = rclut.RCLutx2.init_from_numpy(
#             s1_rc_conv_luts_3x3=s1_rc_conv_luts_3x3, s1_dense_conv_lut_3x3=s1_dense_conv_lut_3x3,
#             s1_rc_conv_luts_5x5=s1_rc_conv_luts_5x5, s1_dense_conv_lut_5x5=s1_dense_conv_lut_5x5,
#             s1_rc_conv_luts_7x7=s1_rc_conv_luts_7x7, s1_dense_conv_lut_7x7=s1_dense_conv_lut_7x7,
#             s2_rc_conv_luts_3x3=s2_rc_conv_luts_3x3, s2_dense_conv_lut_3x3=s2_dense_conv_lut_3x3,
#             s2_rc_conv_luts_5x5=s2_rc_conv_luts_5x5, s2_dense_conv_lut_5x5=s2_dense_conv_lut_5x5,
#             s2_rc_conv_luts_7x7=s2_rc_conv_luts_7x7, s2_dense_conv_lut_7x7=s2_dense_conv_lut_7x7
#         )
#         return lut_model       

#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_7x7(rotated), k=-rotations_count, dims=[2, 3])            
#         output /= 3*4
#         x = output
#         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_7x7(rotated), k=-rotations_count, dims=[2, 3])            
#         output /= 3*4
#         output = output.view(b, c, h*self.scale, w*self.scale)
#         return output



# class ReconstructedConvCenteredUnlutable(nn.Module):
#     def __init__(self, hidden_dim,  window_size=7):
#         super(ReconstructedConvCenteredUnlutable, self).__init__()
#         self.window_size = window_size
#         self.projection1 = torch.nn.Parameter(torch.rand((window_size**2, hidden_dim))/window_size)
#         self.projection2 = torch.nn.Parameter(torch.rand((window_size**2, hidden_dim))/window_size)

#     def pixel_wise_forward(self, x):
#         x = (x-127.5)/127.5
#         out = torch.einsum('bik,ij,ij -> bik', x, self.projection1, self.projection2) 
#         out = torch.tanh(out)
#         out = out*127.5 + 127.5
#         return out 

#     def forward(self, x):
#         original_shape = x.shape
#         x = F.pad(x, pad=[self.window_size//2]*4, mode='replicate')
#         x = F.unfold(x, self.window_size)
#         x = self.pixel_wise_forward(x)        
#         x = x.mean(1)  
#         x = x.reshape(*original_shape)  
#         # x = round_func(x) # quality likely suffer from this
#         return x

#     def __repr__(self):
#         return f"{self.__class__.__name__} projection1: {self.projection1.shape} projection2: {self.projection2.shape}"

# class RCBlockCenteredUnlutable(nn.Module):
#     def __init__(self, hidden_dim = 32, window_size=3, dense_conv_layer_count=4, upscale_factor=4):
#         super(RCBlockRot90Unlutable, self).__init__()   
#         self.window_size = window_size    
#         self.rc_conv = ReconstructedConvCenteredUnlutable(hidden_dim=hidden_dim, window_size=window_size)
#         self.dense_conv_block = layers.UpscaleBlock(hidden_dim=hidden_dim, layers_count=dense_conv_layer_count, upscale_factor=upscale_factor)

#     def forward(self, x):
#         b,c,hs,ws = x.shape
#         x = self.rc_conv(x)
#         x = F.pad(x, pad=[0,1,0,1], mode='replicate') 
#         x = self.dense_conv_block(x)        
#         return x

# class RCNetx2CenteredUnlutable(nn.Module):
#     def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
#         super(RCNetx2CenteredUnlutable, self).__init__()
#         self.scale = scale  
#         self.hidden_dim = hidden_dim     
#         self.stage1_3x3 = RCBlockCenteredUnlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=3)
#         self.stage1_5x5 = RCBlockCenteredUnlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=5)
#         self.stage1_7x7 = RCBlockCenteredUnlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=1, window_size=7)
#         self.stage2_3x3 = RCBlockCenteredUnlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=3)
#         self.stage2_5x5 = RCBlockCenteredUnlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=5)
#         self.stage2_7x7 = RCBlockCenteredUnlutable(hidden_dim=hidden_dim, dense_conv_layer_count=layers_count, upscale_factor=scale, window_size=7)

#     def get_lut_model(self, quantization_interval=16, batch_size=2**10):
#         s1_rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage1_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         s1_dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage1_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s1_rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage1_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         s1_dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage1_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s1_rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage1_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         s1_dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage1_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_3x3 = lut.transfer_rc_conv(self.stage2_3x3.rc_conv, quantization_interval=quantization_interval).reshape(3,3,-1)
#         s2_dense_conv_lut_3x3 = lut.transfer_2x2_input_SxS_output(self.stage2_3x3.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_5x5 = lut.transfer_rc_conv(self.stage2_5x5.rc_conv, quantization_interval=quantization_interval).reshape(5,5,-1)
#         s2_dense_conv_lut_5x5 = lut.transfer_2x2_input_SxS_output(self.stage2_5x5.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)

#         s2_rc_conv_luts_7x7 = lut.transfer_rc_conv(self.stage2_7x7.rc_conv, quantization_interval=quantization_interval).reshape(7,7,-1)
#         s2_dense_conv_lut_7x7 = lut.transfer_2x2_input_SxS_output(self.stage2_7x7.dense_conv_block, quantization_interval=quantization_interval, batch_size=batch_size)
        
#         lut_model = rclut.RCLutx2Centered.init_from_numpy(
#             s1_rc_conv_luts_3x3=s1_rc_conv_luts_3x3, s1_dense_conv_lut_3x3=s1_dense_conv_lut_3x3,
#             s1_rc_conv_luts_5x5=s1_rc_conv_luts_5x5, s1_dense_conv_lut_5x5=s1_dense_conv_lut_5x5,
#             s1_rc_conv_luts_7x7=s1_rc_conv_luts_7x7, s1_dense_conv_lut_7x7=s1_dense_conv_lut_7x7,
#             s2_rc_conv_luts_3x3=s2_rc_conv_luts_3x3, s2_dense_conv_lut_3x3=s2_dense_conv_lut_3x3,
#             s2_rc_conv_luts_5x5=s2_rc_conv_luts_5x5, s2_dense_conv_lut_5x5=s2_dense_conv_lut_5x5,
#             s2_rc_conv_luts_7x7=s2_rc_conv_luts_7x7, s2_dense_conv_lut_7x7=s2_dense_conv_lut_7x7
#         )
#         return lut_model      

#     def forward(self, x):
#         b,c,h,w = x.shape
#         x = x.view(b*c, 1, h, w)
#         output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage1_7x7(rotated), k=-rotations_count, dims=[2, 3])            
#         output /= 3*4
#         x = output
#         output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
#         for rotations_count in range(4):
#             rotated = torch.rot90(x, k=rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_3x3(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_5x5(rotated), k=-rotations_count, dims=[2, 3])
#             output += torch.rot90(self.stage2_7x7(rotated), k=-rotations_count, dims=[2, 3])            
#         output /= 3*4
#         output = output.view(b, c, h*self.scale, w*self.scale)
#         return output

