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 srlut 
from common import layers
from common import losses
from models.base import SRNetBase

class SRNet(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SRNet, self).__init__()
        self.scale = scale 
        self.stage1_S = layers.UpscaleBlock(
            hidden_dim=hidden_dim, 
            layers_count=layers_count, 
            upscale_factor=self.scale
        )
        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(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(b*c, 1, h, w)
        x = self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage1_S)
        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):
        stage_lut = lut.transfer_4_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = srlut.SRLut.init_from_numpy(stage_lut)
        return lut_model


class SRNetChebyKan(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SRNetChebyKan, self).__init__()
        self.scale = scale 
        self.stage1_S = layers.UpscaleBlockChebyKAN(
            hidden_dim=hidden_dim, 
            layers_count=layers_count, 
            upscale_factor=self.scale
        )
        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(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(b*c, 1, h, w)
        x = self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage1_S)
        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):
        stage_lut = lut.transfer_4_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = srlut.SRLut.init_from_numpy(stage_lut)
        return lut_model


class SRNetY(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SRNetY, self).__init__()
        self.scale = scale       
        self.stage1_S = layers.UpscaleBlock(
            hidden_dim=hidden_dim, 
            layers_count=layers_count, 
            upscale_factor=self.scale
        )
        self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
        self.rgb_to_ycbcr = layers.RgbToYcbcr()
        self.ycbcr_to_rgb = layers.YcbcrToRgb()

    def forward(self, x, config=None):
        b,c,h,w = x.shape
        x = self.rgb_to_ycbcr(x)
        y = x[:,0:1,:,:]
        cbcr = x[:,1:,:,:]
        cbcr_scaled = F.interpolate(cbcr, size=[h*self.scale, w*self.scale], mode='bilinear')

        x = y.view(b, 1, h, w)
        output = self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage1_S)
        output = torch.cat([output, cbcr_scaled], dim=1)
        output = self.ycbcr_to_rgb(output).clamp(0, 255)       
        return output

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        stage_lut = lut.transfer_4_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = srlut.SRLutY.init_from_numpy(stage_lut)
        return lut_model


class SRNetR90(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SRNetR90, self).__init__()
        self.scale = scale       
        self.stage1_S = layers.UpscaleBlock(
            hidden_dim=hidden_dim, 
            layers_count=layers_count, 
            upscale_factor=self.scale
        )
        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(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(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.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[2, 3])
        output /= 4
        output = output.reshape(b, c, h*self.scale, w*self.scale)
        return output

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        stage_lut = lut.transfer_4_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = srlut.SRLutR90.init_from_numpy(stage_lut)
        return lut_model

class SRNetChebyKanR90(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SRNetChebyKanR90, self).__init__()
        self.scale = scale 
        self.stage1_S = layers.UpscaleBlockChebyKAN(
            hidden_dim=hidden_dim, 
            layers_count=layers_count, 
            upscale_factor=self.scale
        )
        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(self, x, config=None):
        b,c,h,w = x.shape
        x = x.reshape(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.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[2, 3])
        output /= 4
        output = output.reshape(b, c, h*self.scale, w*self.scale)
        return output

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        stage_lut = lut.transfer_4_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = srlut.SRLutR90.init_from_numpy(stage_lut)
        return lut_model

class SRNetR90Y(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SRNetR90Y, self).__init__()
        self.scale = scale       
        s_pattern=[[0,0],[0,1],[1,0],[1,1]]
        self.stage1_S = layers.UpscaleBlock(
            hidden_dim=hidden_dim, 
            layers_count=layers_count, 
            upscale_factor=self.scale
        )
        self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
        self.rgb_to_ycbcr = layers.RgbToYcbcr()
        self.ycbcr_to_rgb = layers.YcbcrToRgb()

    def forward(self, x, config=None):
        b,c,h,w = x.shape
        x = self.rgb_to_ycbcr(x)
        y = x[:,0:1,:,:]
        cbcr = x[:,1:,:,:]
        cbcr_scaled = F.interpolate(cbcr, size=[h*self.scale, w*self.scale], mode='bilinear')

        x = y.view(b, 1, h, w)
        output = torch.zeros([b, 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.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[2, 3])
        output /= 4
        output = torch.cat([output, cbcr_scaled], dim=1)
        output = self.ycbcr_to_rgb(output).clamp(0, 255)       
        return output

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        stage_lut = lut.transfer_4_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = srlut.SRLutR90Yq.init_from_numpy(stage_lut)
        return lut_model


class SRNetR90Ycbcr(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SRNetR90Ycbcr, self).__init__()
        self.scale = scale       
        s_pattern=[[0,0],[0,1],[1,0],[1,1]]
        self.stage1_S = layers.UpscaleBlock(
            hidden_dim=hidden_dim, 
            layers_count=layers_count, 
            upscale_factor=self.scale
        )
        self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
        self.rgb_to_ycbcr = layers.RgbToYcbcr()
        self.ycbcr_to_rgb = layers.YcbcrToRgb()

    def forward(self, x, config=None):
        b,c,h,w = x.shape
        y = x[:,0:1,:,:]
        cbcr = x[:,1:,:,:]
        cbcr_scaled = F.interpolate(cbcr, size=[h*self.scale, w*self.scale], mode='bilinear')

        x = y.view(b, 1, h, w)
        output = torch.zeros([b, 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.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[2, 3])
        output /= 4
        output = torch.cat([output, cbcr_scaled], dim=1).clamp(0, 255)
        return output

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        stage_lut = lut.transfer_2x2_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = srlut.SRLutR90Y.init_from_numpy(stage_lut)
        return lut_model


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

        self.msb_fn = layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale,
            input_max_value=255,
            output_max_value=255
        )
        self.lsb_fn = layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale,
            input_max_value=15,
            output_max_value=255
        )
        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(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_S, self.msb_fn)           
        output_lsb = self.forward_stage(lsb, self.scale, self._extract_pattern_S, self.lsb_fn)

        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


class SRMsbLsbShiftNet(SRNetBase):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SRMsbLsbShiftNet, self).__init__()
        self.scale = scale 
        self.hidden_dim = hidden_dim
        self.layers_count = layers_count
        self.count = 4
        self.msb_fns = nn.ModuleList([layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale,
            input_max_value=255,
            output_max_value=255
        ) for x in range(self.count)])
        self.lsb_fns = nn.ModuleList([layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale,
            input_max_value=15,
            output_max_value=255
        ) for x in range(self.count)])
        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(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 i, msb_fn, lsb_fn in zip(range(self.count), self.msb_fns, self.lsb_fns):
            output_msb_s = self.forward_stage(msb, self.scale, self._extract_pattern_S, msb_fn)
            output_lsb_s = self.forward_stage(lsb, self.scale, self._extract_pattern_S, lsb_fn)
            output_msb += torch.nn.functional.pad(output_msb_s, [i, 0, i, 0], mode='replicate')[:,:,:h*self.scale,:w*self.scale]
            output_lsb += torch.nn.functional.pad(output_lsb_s, [i, 0, i, 0], mode='replicate')[:,:,:h*self.scale,:w*self.scale]
        output_msb /= self.count
        output_lsb /= self.count

        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


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

        self.msb_fn = layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale,
            input_max_value=255,
            output_max_value=255
        )
        self.lsb_fn = layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale,
            input_max_value=15,
            output_max_value=255
        )
        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(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_msb_r = self.forward_stage(rotated_msb, self.scale, self._extract_pattern_S, self.msb_fn)           
           output_lsb_r = self.forward_stage(rotated_lsb, self.scale, self._extract_pattern_S, self.lsb_fn)
           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


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

        self.msb_fns = nn.ModuleList([layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale,
            input_max_value=255,
            output_max_value=255
        ) for x in range(4)])
        self.lsb_fns = nn.ModuleList([layers.UpscaleBlock(
            in_features=4,
            hidden_dim=hidden_dim,
            layers_count=layers_count,
            upscale_factor=self.scale,
            input_max_value=15,
            output_max_value=255
        ) for x in range(4)])
        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(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), self.msb_fns, 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 += 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.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):
        raise NotImplementedError