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 common import layers
from pathlib import Path
from . import sdylut

class SDYNetx1(nn.Module):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SDYNetx1, self).__init__()
        self.scale = scale   
        s_pattern = [[0,0],[0,1],[1,0],[1,1]]
        d_pattern = [[0,0],[2,0],[0,2],[2,2]]
        y_pattern = [[0,0],[1,1],[1,2],[2,1]]
        self.stage1_S = layers.UpscaleBlock(receptive_field_idxes=s_pattern, center=[0,0], window_size=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
        self.stage1_D = layers.UpscaleBlock(receptive_field_idxes=d_pattern, center=[0,0], window_size=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
        self.stage1_Y = layers.UpscaleBlock(receptive_field_idxes=y_pattern, center=[0,0], window_size=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)

    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, -1])
            rb,rc,rh,rw = rotated.shape       
            output += torch.rot90(self.stage1_S(rotated), k=-rotations_count, dims=[-2, -1])   
            output += torch.rot90(self.stage1_D(rotated), k=-rotations_count, dims=[-2, -1])
            output += torch.rot90(self.stage1_Y(rotated), k=-rotations_count, dims=[-2, -1])

        output /= 4*3
        x = output
        x = round_func(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):
        stageS = lut.transfer_2x2_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        stageD = lut.transfer_2x2_input_SxS_output(self.stage1_D, quantization_interval=quantization_interval, batch_size=batch_size)
        stageY = lut.transfer_2x2_input_SxS_output(self.stage1_Y, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = sdylut.SDYLutx1.init_from_lut(stageS, stageD, stageY)
        return lut_model

class SDYNetx2(nn.Module):
    def __init__(self, hidden_dim = 64, layers_count = 4, scale = 4):
        super(SDYNetx2, self).__init__()
        self.scale = scale   
        s_pattern = [[0,0],[0,1],[1,0],[1,1]]
        d_pattern = [[0,0],[2,0],[0,2],[2,2]]
        y_pattern = [[0,0],[1,1],[1,2],[2,1]]
        self.stage1_S = layers.UpscaleBlock(receptive_field_idxes=s_pattern, center=[0,0], window_size=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
        self.stage1_D = layers.UpscaleBlock(receptive_field_idxes=d_pattern, center=[0,0], window_size=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
        self.stage1_Y = layers.UpscaleBlock(receptive_field_idxes=y_pattern, center=[0,0], window_size=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=1)
        self.stage2_S = layers.UpscaleBlock(receptive_field_idxes=s_pattern, center=[0,0], window_size=2, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
        self.stage2_D = layers.UpscaleBlock(receptive_field_idxes=d_pattern, center=[0,0], window_size=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)
        self.stage2_Y = layers.UpscaleBlock(receptive_field_idxes=y_pattern, center=[0,0], window_size=3, hidden_dim=hidden_dim, layers_count=layers_count, upscale_factor=scale)

    def forward(self, x):
        b,c,h,w = x.shape
        x = x.view(b*c, 1, h, w)
        output_1 = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
        output_1 += self.stage1_S(x)
        output_1 += self.stage1_D(x)
        output_1 += self.stage1_Y(x)
        for rotations_count in range(1,4):
            rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
            output_1 += torch.rot90(self.stage1_S(rotated), k=-rotations_count, dims=[-2, -1])
            output_1 += torch.rot90(self.stage1_D(rotated), k=-rotations_count, dims=[-2, -1])
            output_1 += torch.rot90(self.stage1_Y(rotated), k=-rotations_count, dims=[-2, -1])              
        output_1 /= 4*3
        x = round_func(output_1)
        output_2 = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output_2 += self.stage2_S(x)
        output_2 += self.stage2_D(x)
        output_2 += self.stage2_Y(x)
        for rotations_count in range(1,4):
            rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
            output_2 += torch.rot90(self.stage2_S(rotated), k=-rotations_count, dims=[-2, -1])
            output_2 += torch.rot90(self.stage2_D(rotated), k=-rotations_count, dims=[-2, -1])
            output_2 += torch.rot90(self.stage2_Y(rotated), k=-rotations_count, dims=[-2, -1])
        output_2 /= 4*3
        x = round_func(output_2)
        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):
        stage1_S = lut.transfer_2x2_input_SxS_output(self.stage1_S, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_D = lut.transfer_2x2_input_SxS_output(self.stage1_D, quantization_interval=quantization_interval, batch_size=batch_size)
        stage1_Y = lut.transfer_2x2_input_SxS_output(self.stage1_Y, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_S = lut.transfer_2x2_input_SxS_output(self.stage2_S, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_D = lut.transfer_2x2_input_SxS_output(self.stage2_D, quantization_interval=quantization_interval, batch_size=batch_size)
        stage2_Y = lut.transfer_2x2_input_SxS_output(self.stage2_Y, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = sdylut.SDYLutx2.init_from_lut(stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y)
        return lut_model