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

class SRLut(SRNetBase):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SRLut, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)

    @staticmethod
    def init_from_numpy(
        stage_lut
    ):   
        scale = int(stage_lut.shape[-1])
        quantization_interval = 256//(stage_lut.shape[0]-1)
        lut_model = SRLut(quantization_interval=quantization_interval, scale=scale)
        lut_model.stage_lut = nn.Parameter(torch.tensor(stage_lut).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.reshape(b*c, 1, h, w).type(torch.float32)
        x = self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage_lut)
        x = x.reshape(b, c, h*self.scale, w*self.scale)
        return x        
    
    def __repr__(self):
        return f"{self.__class__.__name__}\n  lut size: {self.stage_lut.shape}"

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

class SRLutY(SRNetBase):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SRLutY, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        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()

    @staticmethod
    def init_from_numpy( 
        stage_lut
    ):   
        scale = int(stage_lut.shape[-1])
        quantization_interval = 256//(stage_lut.shape[0]-1)
        lut_model = SRLutY(quantization_interval=quantization_interval, scale=scale)
        lut_model.stage_lut = nn.Parameter(torch.tensor(stage_lut).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 = 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')
        
        output = self.forward_stage(y, self.scale, self._extract_pattern_S, self.stage_lut)
        output = torch.cat([output, cbcr_scaled], dim=1)  
        output = self.ycbcr_to_rgb(output).clamp(0, 255)        
        return output

    def __repr__(self):
        return f"{self.__class__.__name__}\n  lut size: {self.stage_lut.shape}"

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

class SRLutR90(SRNetBase):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SRLutR90, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)

    @staticmethod
    def init_from_numpy(
        stage_lut
    ):   
        scale = int(stage_lut.shape[-1])
        quantization_interval = 256//(stage_lut.shape[0]-1)
        lut_model = SRLutR90(quantization_interval=quantization_interval, scale=scale)
        lut_model.stage_lut = nn.Parameter(torch.tensor(stage_lut).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.reshape(b*c, 1, h, w)
        output = torch.zeros([b*c, 1, h*self.scale, w*self.scale], dtype=torch.float32, device=x.device)
        output += self.forward_stage(x, self.scale, self._extract_pattern_S, self.stage_lut)
        for rotations_count in range(1, 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.stage_lut), k=-rotations_count, dims=[2, 3])
        output /= 4
        output = output.reshape(b, c, h*self.scale, w*self.scale)
        return output

    def __repr__(self):
        return f"{self.__class__.__name__}\n  lut size: {self.stage_lut.shape}"

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


class SRLutR90Y(SRNetBase):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SRLutR90Y, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        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()

    @staticmethod
    def init_from_numpy( 
        stage_lut
    ):   
        scale = int(stage_lut.shape[-1])
        quantization_interval = 256//(stage_lut.shape[0]-1)
        lut_model = SRLutR90Y(quantization_interval=quantization_interval, scale=scale)
        lut_model.stage_lut = nn.Parameter(torch.tensor(stage_lut).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 = 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')
        
        output = torch.zeros([b, 1, h*self.scale, w*self.scale], dtype=torch.float32, device=x.device)
        output += self.forward_stage(y, self.scale, self._extract_pattern_S, self.stage_lut)
        for rotations_count in range(1,4):
            rotated = torch.rot90(y, k=rotations_count, dims=[2, 3])
            output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage_lut), 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 __repr__(self):
        return f"{self.__class__.__name__}\n  lut size: {self.stage_lut.shape}"

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

class SRLutR90YCbCr(SRNetBase):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(SRLutR90YCbCr, self).__init__()     
        self.scale = scale
        self.quantization_interval = quantization_interval
        self.stage_lut = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (scale,scale)).type(torch.float32))
        self._extract_pattern_S = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2)

    @staticmethod
    def init_from_numpy( 
        stage_lut
    ):   
        scale = int(stage_lut.shape[-1])
        quantization_interval = 256//(stage_lut.shape[0]-1)
        lut_model = SRLutR90YCbCr(quantization_interval=quantization_interval, scale=scale)
        lut_model.stage_lut = nn.Parameter(torch.tensor(stage_lut).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
        y = x[:,0:1,:,:]
        cbcr = x[:,1:,:,:]
        cbcr_scaled = F.interpolate(cbcr, size=[h*self.scale, w*self.scale], mode='bilinear')
        output = torch.zeros([b, 1, h*self.scale, w*self.scale], dtype=torch.float32, device=x.device)
        output += self.forward_stage(y, self.scale, self._extract_pattern_S, self.stage_lut)
        for rotations_count in range(1,4):
            rotated = torch.rot90(y, k=rotations_count, dims=[2, 3])
            output += torch.rot90(self.forward_stage(rotated, self.scale, self._extract_pattern_S, self.stage_lut), k=-rotations_count, dims=[2, 3])
        output /= 4
        output = torch.cat([output, cbcr_scaled], dim=1).clamp(0, 255)  
        return output

    def __repr__(self):
        return f"{self.__class__.__name__}\n  lut size: {self.stage_lut.shape}"

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