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

class PercievePattern():
    """
    Coordinates scheme: [channel, height, width]. Channel can be ommited for all channels.
    Examples: 
    1. receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]]
    2. receptive_field_idxes=[[0,0,0],[0,1,0],[1,1,0],[1,1]]
    """
    def __init__(self, receptive_field_idxes=[[0,0],[0,1],[1,0],[1,1]], center=[0,0], window_size=2, channels=1, mode='replicate'):
        assert window_size >= (np.max([x for y in receptive_field_idxes for x in y])+1)
        tmp = []
        for coords in receptive_field_idxes:
            if len(coords) < 3:
                for i in range(channels):
                    tmp.append([i,] + coords)
            if len(coords) == 3:
                tmp.append(coords)
        receptive_field_idxes = np.array(sorted(tmp))
        self.mode = mode
        self.window_size = window_size
        self.center = center
        self.receptive_field_idxes = [receptive_field_idxes[i,0]*self.window_size*self.window_size + receptive_field_idxes[i,1]*self.window_size + receptive_field_idxes[i,2] for i in range(len(receptive_field_idxes))]
        assert len(np.unique(self.receptive_field_idxes) == len(self.receptive_field_idxes)), "Duplicated coordinates found. Coordinates scheme: [channel, height, width]."

    def __call__(self, x):
        b,c,h,w = x.shape
        if not self.mode is None:
            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=self.mode
            )
        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, stage):
        super(UpscaleBlock, self).__init__() 
        self.stage = stage

    def forward(self, x, percieve_pattern):
        b,c,h,w = x.shape
        upscale_factor = self.stage.upscale_factor
        out_channels = self.stage.out_channels
        x = percieve_pattern(x)
        x = self.stage(x)
        x = x.reshape(b, h, w, out_channels, upscale_factor, upscale_factor)
        x = x.permute(0, 3, 1, 4, 2, 5)
        x = x.reshape(b, out_channels, h*upscale_factor, w*upscale_factor)
        return x 

class LinearUpscaleBlockNet(nn.Module):
    def __init__(self, in_features=4, out_channels=1, hidden_dim = 32, layers_count=4, upscale_factor=1, input_max_value=255, output_max_value=255):
        super(LinearUpscaleBlockNet, self).__init__()   
        assert layers_count > 0     
        self.in_features = in_features
        self.upscale_factor = upscale_factor 
        self.out_channels = out_channels
        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=out_channels*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
        x = round_func(x)
        return x  

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        stage = lut.transfer_N_input_SxS_output(self, quantization_interval=quantization_interval, batch_size=batch_size)
        lut_model = LinearUpscaleBlockLut.init_from_numpy(stage)
        return lut_model

class LinearUpscaleBlockLut(nn.Module):
    def __init__(self, quantization_interval, upscale_factor):
        super(LinearUpscaleBlockLut, self).__init__()  
        self.out_channels = 1
        self.upscale_factor = upscale_factor
        self.quantization_interval = quantization_interval
        self.stage = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*4 + (upscale_factor,upscale_factor)).type(torch.float32))

    @staticmethod
    def init_from_numpy(
        stage
    ):   
        upscale_factor = int(stage.shape[-1])
        quantization_interval = 256//(stage.shape[0]-1)
        lut_model = LinearUpscaleBlockLut(quantization_interval=quantization_interval, upscale_factor=upscale_factor)
        lut_model.stage = nn.Parameter(torch.tensor(stage).type(torch.float32))
        return lut_model

    def forward(self, x):
        x = lut.select_index_4dlut_tetrahedral(x, self.stage)
        return x

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


# This is inspired by Kolmogorov-Arnold Networks but using Chebyshev polynomials instead of splines coefficients
class ChebyKANLayer(nn.Module):
    """
    https://github.com/SynodicMonth/ChebyKAN/blob/main/ChebyKANLayer.py
    """
    def __init__(self, in_features, out_features, degree=8, input_max_value=255, output_max_value=255):
        super(ChebyKANLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.degree = degree

        self.cheby_coeffs = nn.Parameter(torch.empty(in_features, out_features, degree + 1))
        nn.init.normal_(self.cheby_coeffs, mean=0.0, std=1 / (in_features * (degree + 1)))
        self.register_buffer("arange", torch.arange(0, degree + 1, 1))

        self.in_bias = self.in_scale = input_max_value/2
        self.out_bias = self.out_scale = output_max_value/2

    def forward_all_to_all(self, x):
        # Since Chebyshev polynomial is defined in [-1, 1]
        # We need to normalize x to [-1, 1] using tanh
        b, hw, c = x.shape
        assert c == self.in_features, f"Input features count={c} is not equal to specified for this layer count={self.in_features}."
        x = (x-self.in_bias)/self.in_scale
        # x = torch.tanh(x)
        # View and repeat input degree + 1 times
        x = x.view((b, hw, self.in_features, 1)).expand(
            -1, -1, -1, self.degree + 1
        )  # shape = (batch_size, inputdim, self.degree + 1)
        # Apply acos
        x = x.acos()
        # Multiply by arange [0 .. degree]
        x *= self.arange
        # Apply cos
        x = x.cos()
        # Compute the Chebyshev interpolation
        y = torch.einsum(
            "btid,iod->btio", x, self.cheby_coeffs
        )  
        # shape = (batch_size, hw, outdim)
        # print("btio", y.shape)
        y = torch.tanh(y)
        # y = (y-1) / self.degree
    
        y = y*self.out_scale + self.out_bias
        y = y.view(b, hw, self.in_features, self.out_features)
        # y = y.clamp(0,255)
        # print("btio", y.shape)
        return y

    def forward(self, x):
        out = self.forward_all_to_all(x)
        # print("net", out.min(), out.max())
        out = (out - 127.5)/127.5
        out = out.sum(dim=-2) 
        # print("net", out.min(), out.max())
        # out = torch.tanh(out)
        out = out / self.in_features
        out = out*127.5 + 127.5
        return out

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        bucket_count = 256//quantization_interval  
        in_features = self.in_features
        out_features = self.out_features  
        domain_values = torch.cat([torch.arange(0, 256, quantization_interval, dtype=torch.uint8), torch.tensor([255])])
        inputs = domain_values.type(torch.float32).to(self.cheby_coeffs.device)
        lut = np.full((in_features, out_features, bucket_count+1), dtype=np.uint8, fill_value=255)
        # lut_min = np.full((in_features, out_features), dtype=np.uint8, fill_value=0)
        # lut_length = np.full((in_features, out_features), dtype=np.uint8, fill_value=255)
        # print(in_features, out_features)
        qmodel = ChebyKANLut(in_features=self.in_features, out_features=self.out_features, quantization_interval=quantization_interval)
        with torch.no_grad():
            for d in range(len(domain_values)):
                out = self.forward_all_to_all(inputs[d].view(1,1,1).expand(1,1,in_features)).squeeze(0).squeeze(1).cpu().numpy()
                lut[:,:,d] = out.astype(np.uint8)
            # lut_min[:,:] = lut.min(axis=-1).astype(np.uint8)
            # lut_length[:,:] = lut.max(axis=-1).astype(np.uint8)-lut.min(axis=-1).astype(np.uint8)

        qmodel.lut = nn.Parameter(torch.tensor(lut).type(torch.float32))
        # qmodel.lut_min = nn.Parameter(torch.tensor(lut_min).type(torch.float32))
        # qmodel.lut_length = nn.Parameter(torch.tensor(lut_length).type(torch.float32))
        return qmodel
    
    def __repr__(self):
        return (f"{self.__class__.__name__}\n" 
                f"cheby coefs size: {self.cheby_coeffs.shape}\n"
                f"in/out bias, scale: {self.in_bias}, {self.in_scale} / {self.out_bias}, {self.out_scale}" )


class ChebyKANLut(nn.Module):
    def __init__(self, in_features, out_features, quantization_interval=16):
        super(ChebyKANLut, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.lut = nn.Parameter(
            torch.randint(0, 255, size=(in_features, out_features, 256//quantization_interval+1)).type(torch.float32)
        )
        # self.lut_min = nn.Parameter(torch.ones((in_features, out_features)).type(torch.float32)*0)
        # self.lut_max = nn.Parameter(torch.ones((in_features, out_features)).type(torch.float32)*255)

    def forward_all_to_all(self, x):
        b,w,c = x.shape
        out = torch.zeros([b, w, self.in_features, self.out_features], dtype=x.dtype, device=x.device)
        for i in range(self.in_features):
            for j in range(self.out_features):
                out[:,:, i, j:(j+1)] = lut.select_index_1dlut_linear(x[:,:,i:(i+1)], self.lut[i,j])
        return out 

    def forward(self, x):
        out = self.forward_all_to_all(x)
        out = (out - 127.5)/127.5
        out = out.sum(dim=-2)
        # out = torch.tanh(out)
        out = out / self.in_features
        out = out*127.5 + 127.5
        return out

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



class ChebyKANUpscaleBlockNet(nn.Module):
    def __init__(self, in_features=4, out_channels=1, hidden_dim=64, layers_count=2, upscale_factor=1, input_max_value=255, output_max_value=255, degree=8):
        super(ChebyKANUpscaleBlockNet, self).__init__()   
        assert layers_count > 0     
        self.upscale_factor = upscale_factor 
        self.out_channels = out_channels
        
        if layers_count == 1:
            self.kan_layers = [
                ChebyKANLayer(in_features=in_features, out_features=out_channels*self.upscale_factor*self.upscale_factor, degree=degree)
            ]
        elif layers_count > 1:
            self.kan_layers = [
                ChebyKANLayer(in_features=in_features, out_features=hidden_dim, degree=degree),
            ]
            for i in range(layers_count-2):
                self.kan_layers.append(ChebyKANLayer(in_features=hidden_dim, out_features=hidden_dim, degree=degree))
            self.kan_layers.append(ChebyKANLayer(in_features=hidden_dim, out_features=out_channels*self.upscale_factor*self.upscale_factor, degree=degree))

        self.kan_layers = nn.ModuleList(self.kan_layers)
        
    def forward(self, x):
        for kan in self.kan_layers:
            x = kan(x)
            # print(x.min(), x.max())
        return x  

    def get_lut_model(self, quantization_interval=16, batch_size=2**10):
        lut_kan_layers = [v.get_lut_model(quantization_interval=quantization_interval) for v in self.kan_layers]
        qmodel = ChebyKANUpscaleBlockLut.init_from_numpy(lut_kan_layers=lut_kan_layers)
        qmodel.upscale_factor = self.upscale_factor
        return qmodel


class ChebyKANUpscaleBlockLut(nn.Module):
    def __init__(self, layer_count, in_features_list, out_features_list, quantization_interval, upscale_factor):
        super(ChebyKANUpscaleBlockLut, self).__init__()
        self.quantization_interval = quantization_interval
        self.upscale_factor = upscale_factor
        self.out_channels = 1
        self.lut_kan_layers = nn.ModuleList([
            ChebyKANLut(in_features=in_features, out_features=out_features, quantization_interval=quantization_interval) 
            for i, in_features, out_features in zip(range(layer_count), in_features_list, out_features_list)])

    @staticmethod
    def init_from_numpy(
        lut_kan_layers
    ):
        in_features_list = [l.lut.shape[0] for l in lut_kan_layers]
        out_features_list = [l.lut.shape[1] for l in lut_kan_layers]
        quantization_interval = 256//lut_kan_layers[0].lut.shape[2]+1 
        qmodel = ChebyKANUpscaleBlockLut(layer_count=len(lut_kan_layers), in_features_list=in_features_list, out_features_list=out_features_list, quantization_interval=quantization_interval, upscale_factor=out_features_list[-1])
        qmodel.lut_kan_layers = nn.ModuleList(lut_kan_layers)
        return qmodel

    def forward(self, x):
        for kan in self.lut_kan_layers:
            x = kan(x)
            # print(x.min(), x.max())
        return x

    def __repr__(self):
        return "\n".join([f"{self.__class__.__name__}"] + [f"lut size: {l.lut.shape}" for l in self.lut_kan_layers])