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

class SDYLutx1(nn.Module):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SDYLutx1, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
        self._extract_pattern_D = PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_Y = PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]], center=[0,0], window_size=3)
        self.stageS = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stageD = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stageY = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))

    @staticmethod
    def init_from_numpy(
        stageS, stageD, stageY
    ):   
        scale = int(stageS.shape[-1])
        quantization_interval = 256//(stageS.shape[0]-1)
        lut_model = SDYLutx1(quantization_interval=quantization_interval, scale=scale)
        lut_model.stageS = nn.Parameter(torch.tensor(stageS).type(torch.float32))
        lut_model.stageD = nn.Parameter(torch.tensor(stageD).type(torch.float32))
        lut_model.stageY = nn.Parameter(torch.tensor(stageY).type(torch.float32))
        return lut_model

    def forward_stage(self, x, scale, percieve_pattern, lut):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = select_index_4dlut_tetrahedral(index=x, lut=lut)
        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.view(b*c, 1, h, w).type(torch.float32)
        output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stageS) 
        output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stageD)
        output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stageY)
        output /= 3
        output = round_func(output)
        output = output.view(b, c, h*self.scale, w*self.scale)
        return output

    def __repr__(self):
        return f"{self.__class__.__name__}" + \
               f"\n  stageS size: {self.stageS.shape}" + \
               f"\n  stageD size: {self.stageD.shape}" + \
               f"\n  stageY size: {self.stageY.shape}"

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

class SDYLutx2(nn.Module):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SDYLutx2, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
        self._extract_pattern_D = PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_Y = PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]], center=[0,0], window_size=3)
        self.stage1_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage1_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage1_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage2_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stage2_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stage2_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))

    @staticmethod
    def init_from_numpy(
        stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y
    ):   
        scale = int(stage2_S.shape[-1])
        quantization_interval = 256//(stage2_S.shape[0]-1)
        lut_model = SDYLutx2(quantization_interval=quantization_interval, scale=scale)
        lut_model.stage1_S = nn.Parameter(torch.tensor(stage1_S).type(torch.float32))
        lut_model.stage1_D = nn.Parameter(torch.tensor(stage1_D).type(torch.float32))
        lut_model.stage1_Y = nn.Parameter(torch.tensor(stage1_Y).type(torch.float32))
        lut_model.stage2_S = nn.Parameter(torch.tensor(stage2_S).type(torch.float32))
        lut_model.stage2_D = nn.Parameter(torch.tensor(stage2_D).type(torch.float32))
        lut_model.stage2_Y = nn.Parameter(torch.tensor(stage2_Y).type(torch.float32))
        return lut_model

    def forward_stage(self, x, scale, percieve_pattern, lut):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = select_index_4dlut_tetrahedral(index=x, lut=lut)
        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.view(b*c, 1, h, w).type(torch.float32)
        output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
        output += self.forward_stage(x, 1, self._extract_pattern_S, self.stage1_S) 
        output += self.forward_stage(x, 1, self._extract_pattern_D, self.stage1_D)
        output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y)
        output /= 3
        output = round_func(output)
        x = output
        output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage2_S) 
        output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage2_D)
        output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage2_Y)
        output /= 3
        output = round_func(output)
        output = output.view(b, c, h*self.scale, w*self.scale)
        return output

    def __repr__(self):
        return f"{self.__class__.__name__}" + \
               f"\n  stage1_S size: {self.stage1_S.shape}" + \
               f"\n  stage1_D size: {self.stage1_D.shape}" + \
               f"\n  stage1_Y size: {self.stage1_Y.shape}" + \
               f"\n  stage2_S size: {self.stage2_S.shape}" + \
               f"\n  stage2_D size: {self.stage2_D.shape}" + \
               f"\n  stage2_Y size: {self.stage2_Y.shape}"

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


class SDYLutx3(nn.Module):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SDYLutx3, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
        self._extract_pattern_D = PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_Y = PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]], center=[0,0], window_size=3)
        self.stage1_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage1_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage1_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage2_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage2_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage2_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage3_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stage3_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stage3_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))

    @staticmethod
    def init_from_numpy(
        stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y, stage3_S, stage3_D, stage3_Y
    ):   
        scale = int(stage3_S.shape[-1])
        quantization_interval = 256//(stage3_S.shape[0]-1)
        lut_model = SDYLutx3(quantization_interval=quantization_interval, scale=scale)
        lut_model.stage1_S = nn.Parameter(torch.tensor(stage1_S).type(torch.float32))
        lut_model.stage1_D = nn.Parameter(torch.tensor(stage1_D).type(torch.float32))
        lut_model.stage1_Y = nn.Parameter(torch.tensor(stage1_Y).type(torch.float32))
        lut_model.stage2_S = nn.Parameter(torch.tensor(stage2_S).type(torch.float32))
        lut_model.stage2_D = nn.Parameter(torch.tensor(stage2_D).type(torch.float32))
        lut_model.stage2_Y = nn.Parameter(torch.tensor(stage2_Y).type(torch.float32))
        lut_model.stage3_S = nn.Parameter(torch.tensor(stage3_S).type(torch.float32))
        lut_model.stage3_D = nn.Parameter(torch.tensor(stage3_D).type(torch.float32))
        lut_model.stage3_Y = nn.Parameter(torch.tensor(stage3_Y).type(torch.float32))
        return lut_model

    def forward_stage(self, x, scale, percieve_pattern, lut):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = select_index_4dlut_tetrahedral(index=x, lut=lut)
        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.view(b*c, 1, h, w).type(torch.float32)
        output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
        output += self.forward_stage(x, 1, self._extract_pattern_S, self.stage1_S) 
        output += self.forward_stage(x, 1, self._extract_pattern_D, self.stage1_D)
        output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y)
        output /= 3
        output = round_func(output)
        x = output
        output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device)
        output += self.forward_stage(x, 1, self._extract_pattern_S, self.stage2_S) 
        output += self.forward_stage(x, 1, self._extract_pattern_D, self.stage2_D)
        output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage2_Y)
        output /= 3
        output = round_func(output)
        x = output
        output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage3_S) 
        output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage3_D)
        output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage3_Y)
        output /= 3
        output = round_func(output)
        output = output.view(b, c, h*self.scale, w*self.scale)
        return output

    def __repr__(self):
        return f"{self.__class__.__name__}" + \
               f"\n  stage1_S size: {self.stage1_S.shape}" + \
               f"\n  stage1_D size: {self.stage1_D.shape}" + \
               f"\n  stage1_Y size: {self.stage1_Y.shape}" + \
               f"\n  stage2_S size: {self.stage2_S.shape}" + \
               f"\n  stage2_D size: {self.stage2_D.shape}" + \
               f"\n  stage2_Y size: {self.stage2_Y.shape}" + \
               f"\n  stage3_S size: {self.stage3_S.shape}" + \
               f"\n  stage3_D size: {self.stage3_D.shape}" + \
               f"\n  stage3_Y size: {self.stage3_Y.shape}"

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


class SDYLutR90x1(nn.Module):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SDYLutR90x1, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
        self._extract_pattern_D = PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_Y = PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]], center=[0,0], window_size=3)
        self.stageS = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stageD = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stageY = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))

    @staticmethod
    def init_from_numpy(
        stageS, stageD, stageY
    ):   
        scale = int(stageS.shape[-1])
        quantization_interval = 256//(stageS.shape[0]-1)
        lut_model = SDYLutR90x1(quantization_interval=quantization_interval, scale=scale)
        lut_model.stageS = nn.Parameter(torch.tensor(stageS).type(torch.float32))
        lut_model.stageD = nn.Parameter(torch.tensor(stageD).type(torch.float32))
        lut_model.stageY = nn.Parameter(torch.tensor(stageY).type(torch.float32))
        return lut_model

    def forward_stage(self, x, scale, percieve_pattern, lut):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = select_index_4dlut_tetrahedral(index=x, lut=lut)
        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.view(b*c, 1, h, w).type(torch.float32)
        output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stageS) 
        output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stageD)
        output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stageY)
        for rotations_count in range(1, 4):
            rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
            output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stageS), k=-rotations_count, dims=[-2, -1]) 
            output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_D, self.stageD), k=-rotations_count, dims=[-2, -1])
            output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_Y, self.stageY), k=-rotations_count, dims=[-2, -1])
        output /= 4*3
        output = round_func(output)
        output = output.view(b, c, h*self.scale, w*self.scale)
        return output

    def __repr__(self):
        return f"{self.__class__.__name__}" + \
               f"\n  stageS size: {self.stageS.shape}" + \
               f"\n  stageD size: {self.stageD.shape}" + \
               f"\n  stageY size: {self.stageY.shape}"

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


class SDYLutR90x2(nn.Module):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SDYLutR90x2, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self._extract_pattern_S = PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)
        self._extract_pattern_D = PercievePattern(receptive_field_idxes=[[0,0],[2,0],[0,2],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_Y = PercievePattern(receptive_field_idxes=[[0,0],[1,1],[1,2],[2,1]], center=[0,0], window_size=3)
        self.stage1_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage1_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage1_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (1,1)).type(torch.float32))
        self.stage2_S = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stage2_D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self.stage2_Y = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))

    @staticmethod
    def init_from_numpy(
        stage1_S, stage1_D, stage1_Y, stage2_S, stage2_D, stage2_Y
    ):   
        scale = int(stage2_S.shape[-1])
        quantization_interval = 256//(stage2_S.shape[0]-1)
        lut_model = SDYLutR90x2(quantization_interval=quantization_interval, scale=scale)
        lut_model.stage1_S = nn.Parameter(torch.tensor(stage1_S).type(torch.float32))
        lut_model.stage1_D = nn.Parameter(torch.tensor(stage1_D).type(torch.float32))
        lut_model.stage1_Y = nn.Parameter(torch.tensor(stage1_Y).type(torch.float32))
        lut_model.stage2_S = nn.Parameter(torch.tensor(stage2_S).type(torch.float32))
        lut_model.stage2_D = nn.Parameter(torch.tensor(stage2_D).type(torch.float32))
        lut_model.stage2_Y = nn.Parameter(torch.tensor(stage2_Y).type(torch.float32))
        return lut_model

    def forward_stage(self, x, scale, percieve_pattern, lut):
        b,c,h,w = x.shape
        x = percieve_pattern(x)   
        x = select_index_4dlut_tetrahedral(index=x, lut=lut)
        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.view(b*c, 1, h, w).type(torch.float32)
        output = torch.zeros([b*c, 1, h, w], dtype=x.dtype, device=x.device) 
        output += self.forward_stage(x, 1, self._extract_pattern_S, self.stage1_S) 
        output += self.forward_stage(x, 1, self._extract_pattern_D, self.stage1_D)
        output += self.forward_stage(x, 1, self._extract_pattern_Y, self.stage1_Y)
        for rotations_count in range(1, 4):
            rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
            output += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_S, self.stage1_S), k=-rotations_count, dims=[-2, -1])
            output += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_D, self.stage1_D), k=-rotations_count, dims=[-2, -1])
            output += torch.rot90(self.forward_stage(rotated, 1, self._extract_pattern_Y, self.stage1_Y), k=-rotations_count, dims=[-2, -1])
        output /= 4*3
        x = round_func(output)

        output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=x.dtype, device=x.device)
        output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage2_S)
        output += self.forward_stage(x, self.scale, self._extract_pattern_D, self.stage2_D)
        output += self.forward_stage(x, self.scale, self._extract_pattern_Y, self.stage2_Y)
        for rotations_count in range(1, 4):
            rotated = torch.rot90(x, k=rotations_count, dims=[-2, -1])
            output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage2_S), k=-rotations_count, dims=[-2, -1])
            output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_D, self.stage2_D), k=-rotations_count, dims=[-2, -1])
            output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_Y, self.stage2_Y), k=-rotations_count, dims=[-2, -1])
        output /= 4*3 
        output = round_func(output) 
        output = output.view(b, c, h*self.scale, w*self.scale)
        return output

    def __repr__(self):
        return f"{self.__class__.__name__}" + \
               f"\n  stage1_S size: {self.stage1_S.shape}" + \
               f"\n  stage1_D size: {self.stage1_D.shape}" + \
               f"\n  stage1_Y size: {self.stage1_Y.shape}" + \
               f"\n  stage2_S size: {self.stage2_S.shape}" + \
               f"\n  stage2_D size: {self.stage2_D.shape}" + \
               f"\n  stage2_Y size: {self.stage2_Y.shape}"

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