import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from common.utils import round_func
from common import lut
from pathlib import Path
from . import hdblut
from common import layers
from itertools import cycle
from models.base import SRNetBase

class HDBNet(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4, rotations = 4):
        super(HDBNet, self).__init__()
        assert scale == 4
        self.scale = scale 
        self.rotations = rotations
        self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage1_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)
        self.stage1_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)

        self.stage2_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage2_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage2_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage2_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)
        self.stage2_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)

        self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
        self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
        self._extract_pattern_2H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1]], center=[0,0], window_size=2)
        self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2)        

    def forward_stage(self, x, scale, percieve_pattern, stage):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = stage(x)
        x = round_func(x)
        x = x.reshape(b, c, h, w, scale, scale)
        x = x.permute(0,1,2,4,3,5)
        x = x.reshape(b, c, h*scale, w*scale)
        return x

    def forward(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(b*c, 1, h, w)
        lsb = x % 16
        msb = x - lsb     
        output = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
        for rotations_count in range(self.rotations):
           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msb = self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage1_3H) + \
                        self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage1_3D) + \
                        self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage1_3B)
           output_msb /= 3
           output_lsb = self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage1_2H) + \
                        self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage1_2D)
           output_lsb /= 2
           if not config is None and config.current_iter % config.display_step == 0:
                config.writer.add_histogram('s1_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
                config.writer.add_histogram('s1_output_msb', output_msb.detach().cpu().numpy(), config.current_iter) 
           output += torch.rot90(output_msb + output_lsb, k=-rotations_count, dims=[2, 3]).clamp(0, 255)
        output /= self.rotations
        x = output
        lsb = x % 16
        msb = x - lsb
        output = torch.zeros([b*c, 1, h*4, w*4], dtype=x.dtype, device=x.device)
        for rotations_count in range(self.rotations):
           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msb = self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage2_3H) + \
                        self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage2_3D) + \
                        self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage2_3B)
           output_msb /= 3
           output_lsb = self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage2_2H) + \
                        self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage2_2D)
           output_lsb /= 2
           if not config is None and config.current_iter % config.display_step == 0:
                config.writer.add_histogram('s2_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
                config.writer.add_histogram('s2_output_msb', output_msb.detach().cpu().numpy(), config.current_iter) 
           output += torch.rot90(output_msb + output_lsb, k=-rotations_count, dims=[2, 3]).clamp(0, 255)
        output /= self.rotations
        x = output
        x = x.reshape(b, c, h*self.scale, w*self.scale)
        return x

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        stage1_3H = lut.transfer_3_input_SxS_output(self.stage1_3H, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_3D = lut.transfer_3_input_SxS_output(self.stage1_3D, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_3B = lut.transfer_3_input_SxS_output(self.stage1_3B, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_2H = lut.transfer_2_input_SxS_output(self.stage1_2H, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_2D = lut.transfer_2_input_SxS_output(self.stage1_2D, quantization_interval=quantization_interval, batch_size=batch_size)

        stage2_3H = lut.transfer_3_input_SxS_output(self.stage2_3H, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_3D = lut.transfer_3_input_SxS_output(self.stage2_3D, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_3B = lut.transfer_3_input_SxS_output(self.stage2_3B, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_2H = lut.transfer_2_input_SxS_output(self.stage2_2H, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_2D = lut.transfer_2_input_SxS_output(self.stage2_2D, quantization_interval=quantization_interval, batch_size=batch_size)

        lut_model = hdblut.HDBLut.init_from_numpy(
            stage1_3H, stage1_3D, stage1_3B, stage1_2H, stage1_2D, 
            stage2_3H, stage2_3D, stage2_3B, stage2_2H, stage2_2D
        )
        return lut_model

    def get_loss_fn(self):
        def loss_fn(pred, target):
            return F.mse_loss(pred/255, target/255)
        return loss_fn

class HDBNetv2(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 2, scale = 4, rotations = 4):
        super(HDBNetv2, self).__init__()
        assert scale == 4
        self.scale = scale 
        self.rotations = rotations
        self.layers_count = layers_count
        self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage1_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)
        self.stage1_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)

        self.stage2_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage2_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage2_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=255, output_max_value=255)
        self.stage2_2H = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)
        self.stage2_2D = layers.UpscaleBlock(in_features=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=2, input_max_value=15, output_max_value=255)

        self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
        self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
        self._extract_pattern_2H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1]], center=[0,0], window_size=2)
        self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2)        

    def forward(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(b*c, 1, h, w)
        lsb = x % 16
        msb = x - lsb     
        output_msb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
        output_lsb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
        for rotations_count in range(self.rotations):
           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msbt = self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage1_3H) + \
                         self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage1_3D) + \
                         self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage1_3B)
           output_lsbt = self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage1_2H) + \
                         self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage1_2D)
           if not config is None and config.current_iter % config.display_step == 0:
                config.writer.add_histogram('s1_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
                config.writer.add_histogram('s1_output_msb', output_msb.detach().cpu().numpy(), config.current_iter) 
           output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3])
        output_msb /= self.rotations*3
        output_lsb /= self.rotations*2
        output = output_msb + output_lsb
        x = output.clamp(0, 255)
        lsb = x % 16
        msb = x - lsb
        output_msb = torch.zeros([b*c, 1, h*2*2, w*2*2], dtype=x.dtype, device=x.device)
        output_lsb = torch.zeros([b*c, 1, h*2*2, w*2*2], dtype=x.dtype, device=x.device)
        for rotations_count in range(self.rotations):
           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msbt = self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage2_3H) + \
                         self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage2_3D) + \
                         self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage2_3B)
           output_lsbt = self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage2_2H) + \
                         self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage2_2D)
           if not config is None and config.current_iter % config.display_step == 0:
                config.writer.add_histogram('s2_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
                config.writer.add_histogram('s2_output_msb', output_msb.detach().cpu().numpy(), config.current_iter) 
           output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3])
        output_msb /= self.rotations*3
        output_lsb /= self.rotations*2
        output = output_msb + output_lsb
        x = output.clamp(0, 255)
        x = x.reshape(b, c, h*self.scale, w*self.scale)
        return x

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        stage1_3H = lut.transfer_3_input_SxS_output(self.stage1_3H, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_3D = lut.transfer_3_input_SxS_output(self.stage1_3D, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_3B = lut.transfer_3_input_SxS_output(self.stage1_3B, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_2H = lut.transfer_2_input_SxS_output(self.stage1_2H, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_2D = lut.transfer_2_input_SxS_output(self.stage1_2D, quantization_interval=quantization_interval, batch_size=batch_size)

        stage2_3H = lut.transfer_3_input_SxS_output(self.stage2_3H, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_3D = lut.transfer_3_input_SxS_output(self.stage2_3D, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_3B = lut.transfer_3_input_SxS_output(self.stage2_3B, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_2H = lut.transfer_2_input_SxS_output(self.stage2_2H, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_2D = lut.transfer_2_input_SxS_output(self.stage2_2D, quantization_interval=quantization_interval, batch_size=batch_size)

        lut_model = hdblut.HDBLut.init_from_numpy(
            stage1_3H, stage1_3D, stage1_3B, stage1_2H, stage1_2D, 
            stage2_3H, stage2_3D, stage2_3B, stage2_2H, stage2_2D
        )
        return lut_model

    def get_loss_fn(self):
        def loss_fn(pred, target):
            return F.mse_loss(pred/255, target/255)
        return loss_fn

class HDBLNet(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(HDBLNet, self).__init__()
        self.scale = scale 
        self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3L = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)

        self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
        self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
        self._extract_pattern_3L = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,1]], center=[0,0], window_size=2)

    def forward_stage(self, x, scale, percieve_pattern, stage):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = stage(x)
        x = round_func(x)
        x = x.reshape(b, c, h, w, scale, scale)
        x = x.permute(0,1,2,4,3,5)
        x = x.reshape(b, c, h*scale, w*scale)
        return x
            
    def forward(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(b*c, 1, h, w)
        lsb = x % 16
        msb = x - lsb     
        output_msb = self.forward_stage(msb, self.scale, self._extract_pattern_3H, self.stage1_3H) + \
                     self.forward_stage(msb, self.scale, self._extract_pattern_3D, self.stage1_3D) + \
                     self.forward_stage(msb, self.scale, self._extract_pattern_3B, self.stage1_3B)
        output_lsb = self.forward_stage(lsb, self.scale, self._extract_pattern_3L, self.stage1_3L)
        output_msb /= 3
        if not config is None and config.current_iter % config.display_step == 0:
            config.writer.add_histogram('s1_output_msb', output_msb.detach().cpu().numpy(), config.current_iter)
            config.writer.add_histogram('s1_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)            
        output = output_msb + output_lsb
        x = output.clamp(0, 255)
        x = x.reshape(b, c, h*self.scale, w*self.scale)
        return x

    def get_loss_fn(self):
        def loss_fn(pred, target):
            return F.mse_loss(pred/255, target/255)
        return loss_fn
        

class HDBLNetR90(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(HDBLNet, self).__init__()
        self.scale = scale 
        self.stage1_3H = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3D = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3B = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3L = layers.UpscaleBlock(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)

        self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
        self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
        self._extract_pattern_3L = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,1]], center=[0,0], window_size=2)

    def forward_stage(self, x, scale, percieve_pattern, stage):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = stage(x)
        x = round_func(x)
        x = x.reshape(b, c, h, w, scale, scale)
        x = x.permute(0,1,2,4,3,5)
        x = x.reshape(b, c, h*scale, w*scale)
        return x
            
    def forward(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(b*c, 1, h, w)
        lsb = x % 16
        msb = x - lsb     
        output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        for rotations_count in range(4):
           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msbt = self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3H, self.stage1_3H) + \
                         self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3D, self.stage1_3D) + \
                         self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3B, self.stage1_3B)
           output_lsbt = self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_3L, self.stage1_3L)
           if not config is None and config.current_iter % config.display_step == 0:
                config.writer.add_histogram('s1_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
                config.writer.add_histogram('s1_output_msb', output_msb.detach().cpu().numpy(), config.current_iter)
           output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3])
        output_msb /= 4*3
        output_lsb /= 4
        output = output_msb + output_lsb
        x = output.clamp(0, 255)
        x = x.reshape(b, c, h*self.scale, w*self.scale)
        return x

    def get_loss_fn(self):
        def loss_fn(pred, target):
            return F.mse_loss(pred/255, target/255)
        return loss_fn


class HDBLNetR90KAN(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(HDBLNetR90KAN, self).__init__()
        self.scale = scale 
        self.stage1_3H = layers.UpscaleBlockChebyKAN(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3D = layers.UpscaleBlockChebyKAN(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3B = layers.UpscaleBlockChebyKAN(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)
        self.stage1_3L = layers.UpscaleBlockChebyKAN(in_features=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=self.scale)

        self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
        self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
        self._extract_pattern_3L = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,1]], center=[0,0], window_size=2)

    def forward_stage(self, x, scale, percieve_pattern, stage):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = stage(x)
        x = round_func(x)
        x = x.reshape(b, c, h, w, scale, scale)
        x = x.permute(0,1,2,4,3,5)
        x = x.reshape(b, c, h*scale, w*scale)
        return x
            
    def forward(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(b*c, 1, h, w)
        lsb = x % 16
        msb = x - lsb     
        output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        for rotations_count in range(4):
           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msbt = self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3H, self.stage1_3H) + \
                         self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3D, self.stage1_3D) + \
                         self.forward_stage(rotated_msb, self.scale, self._extract_pattern_3B, self.stage1_3B)
           output_lsbt = self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_3L, self.stage1_3L)
           if not config is None and config.current_iter % config.display_step == 0:
                config.writer.add_histogram('s1_output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
                config.writer.add_histogram('s1_output_msb', output_msb.detach().cpu().numpy(), config.current_iter)
           output_msb += torch.rot90(output_msbt, k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(output_lsbt, k=-rotations_count, dims=[2, 3])
        output_msb /= 4*3
        output_lsb /= 4
        output = output_msb + output_lsb
        x = output.clamp(0, 255)
        x = x.reshape(b, c, h*self.scale, w*self.scale)
        return x

    def get_loss_fn(self):
        def loss_fn(pred, target):
            return F.mse_loss(pred/255, target/255)
        return loss_fn


class HDBHNet(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(HDBHNet, self).__init__()
        self.scale = scale 
        self.hidden_dim = hidden_dim
        self.layers_count = layers_count

        self.msb_fns = SRNetBaseList([layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale
        ) for x in range(1)])
        self.lsb_fns = SRNetBaseList([layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale
        ) for x in range(1)])
        self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)

    def forward_stage(self, x, scale, percieve_pattern, stage):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = stage(x)
        x = round_func(x)
        x = x.reshape(b, c, h, w, scale, scale)
        x = x.permute(0,1,2,4,3,5)
        x = x.reshape(b, c, h*scale, w*scale)
        return x
            
    def forward(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(b*c, 1, h, w)

        lsb = x % 16
        msb = x - lsb

        output_msb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output_lsb = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        for rotations_count, msb_fn, lsb_fn in zip(range(4), cycle(self.msb_fns), cycle(self.lsb_fns)):
           rotated_msb = torch.rot90(msb, k=rotations_count, dims=[2, 3])
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msb_r = self.forward_stage(rotated_msb, self.scale, self._extract_pattern_S, msb_fn)           
           output_lsb_r = self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_S, lsb_fn)
           output_msb_r = round_func((output_msb_r / 255)*16) * 15 
           output_lsb_r = (output_lsb_r / 255) * 15  
           output_msb += torch.rot90(output_msb_r, k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(output_lsb_r, k=-rotations_count, dims=[2, 3])
        output_msb /= 4
        output_lsb /= 4
        if not config is None and config.current_iter % config.display_step == 0:
            config.writer.add_histogram('output_lsb', output_lsb.detach().cpu().numpy(), config.current_iter)
            config.writer.add_histogram('output_msb', output_msb.detach().cpu().numpy(), config.current_iter)
        x = output_msb + output_lsb
        x = x.reshape(b, c, h*self.scale, w*self.scale)
        return x

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        raise NotImplementedError
    
    def get_loss_fn(self):
        fourier_loss_fn = FocalFrequencyLoss()
        high_frequency_loss_fn = FourierLoss()
        def loss_fn(pred, target):
            a = fourier_loss_fn(pred/255, target/255) * 1e8
            # b = F.mse_loss(pred/255, target/255) #* 1e3
            # c = high_frequency_loss_fn(pred/255, target/255) * 1e6
            return a #+ b #+ c
        return loss_fn