import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from .utils import  round_func

class PercievePattern():
    def __init__(self, receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2):
        assert window_size >= (np.max(receptive_field_idxes)+1)
        receptive_field_idxes = np.array(receptive_field_idxes)
        self.window_size = window_size
        self.center = center
        self.receptive_field_idxes = [receptive_field_idxes[i,0]*self.window_size + receptive_field_idxes[i,1] for i in range(len(receptive_field_idxes))]

    def __call__(self, x):
        b,c,h,w = x.shape
        x = F.pad(
            x, 
            pad=[self.center[0], self.window_size-self.center[0]-1,
                 self.center[1], self.window_size-self.center[1]-1], 
            mode='replicate'
        )
        x = F.unfold(input=x, kernel_size=self.window_size)
        x = torch.stack([x[:,self.receptive_field_idxes[i],:] for i in range(len(self.receptive_field_idxes))], 2)
        return x

class UpscaleBlock(nn.Module):
    def __init__(self, in_features=4, hidden_dim = 32, layers_count=4, upscale_factor=1, input_max_value=255, output_max_value=255):
        super(UpscaleBlock, self).__init__()   
        assert layers_count > 0     
        self.upscale_factor = upscale_factor 
        self.hidden_dim = hidden_dim
        self.embed = nn.Linear(in_features=in_features, out_features=hidden_dim, bias=True)
        
        self.linear_projections = []
        for i in range(layers_count):
            self.linear_projections.append(nn.Linear(in_features=(i+1)*hidden_dim, out_features=hidden_dim, bias=True))
        self.linear_projections = nn.ModuleList(self.linear_projections)          

        self.project_channels = nn.Linear(in_features=(layers_count+1)*hidden_dim, out_features=upscale_factor * upscale_factor, bias=True)

        self.in_bias = self.in_scale = input_max_value/2
        self.out_bias = self.out_scale = output_max_value/2
    
    def forward(self, x):
        x = (x-self.in_bias)/self.in_scale
        x = torch.nn.functional.gelu(self.embed(x))
        for linear_projection in self.linear_projections:
            x = torch.cat([x, torch.nn.functional.gelu(linear_projection(x))], dim=2)
        x = self.project_channels(x)
        x = torch.tanh(x)         
        x = x*self.out_scale + self.out_bias
        return x  

class RgbToYcbcr(nn.Module):
    r"""Convert an image from RGB to YCbCr.

    The image data is assumed to be in the range of (0, 1).

    Returns:
        YCbCr version of the image.

    Shape:
        - image: :math:`(*, 3, H, W)`
        - output: :math:`(*, 3, H, W)`

    Examples:
        >>> input = torch.rand(2, 3, 4, 5)
        >>> ycbcr = RgbToYcbcr()
        >>> output = ycbcr(input)  # 2x3x4x5

    https://github.com/kornia/kornia/blob/2c084f8dc108b3f0f3c8983ac3f25bf88638d01a/kornia/color/ycbcr.py#L105
    """

    def forward(self, image):
        r = image[..., 0, :, :]
        g = image[..., 1, :, :]
        b = image[..., 2, :, :]

        delta = 0.5
        y  = 0.299 * r + 0.587 * g + 0.114 * b 
        cb = (b - y) * 0.564 + delta
        cr = (r - y) * 0.713 + delta
        return torch.stack([y, cb, cr], -3)

class YcbcrToRgb(nn.Module):
    r"""Convert an image from YCbCr to Rgb.

    The image data is assumed to be in the range of (0, 1).

    Returns:
        RGB version of the image.

    Shape:
        - image: :math:`(*, 3, H, W)`
        - output: :math:`(*, 3, H, W)`

    Examples:
        >>> input = torch.rand(2, 3, 4, 5)
        >>> rgb = YcbcrToRgb()
        >>> output = rgb(input)  # 2x3x4x5
    
    https://github.com/kornia/kornia/blob/2c084f8dc108b3f0f3c8983ac3f25bf88638d01a/kornia/color/ycbcr.py#L127
    """

    def forward(self, image):
        y = image[..., 0, :, :]
        cb = image[..., 1, :, :]
        cr = image[..., 2, :, :]

        delta = 0.5
        cb_shifted = cb - delta
        cr_shifted = cr - delta

        r = y + 1.403 * cr_shifted
        g = y - 0.714 * cr_shifted - 0.344 * cb_shifted
        b = y + 1.773 * cb_shifted
        return torch.stack([r, g, b], -3)