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

class HDBLut(nn.Module):
    def __init__(
        self, 
        quantization_interval,
        scale
    ):
        super(HDBLut, self).__init__()     
        assert scale == 4
        self.scale = scale
        self.quantization_interval = quantization_interval

        self.stage1_3H = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*3 + (2,2)).type(torch.float32))
        self.stage1_3D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*3 + (2,2)).type(torch.float32))
        self.stage1_3B = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*3 + (2,2)).type(torch.float32))
        self.stage1_2H = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*2 + (2,2)).type(torch.float32))
        self.stage1_2D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*2 + (2,2)).type(torch.float32))

        self.stage2_3H = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*3 + (2,2)).type(torch.float32))
        self.stage2_3D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*3 + (2,2)).type(torch.float32))
        self.stage2_3B = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*3 + (2,2)).type(torch.float32))
        self.stage2_2H = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*2 + (2,2)).type(torch.float32))
        self.stage2_2D = nn.Parameter(torch.randint(0, 255, size=(256//quantization_interval+1,)*2 + (2,2)).type(torch.float32))

        self._extract_pattern_3H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1],[0,2]], center=[0,0], window_size=3)
        self._extract_pattern_3D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1],[2,2]], center=[0,0], window_size=3)
        self._extract_pattern_3B = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,2],[2,1]], center=[0,0], window_size=3)
        self._extract_pattern_2H = layers.PercievePattern(receptive_field_idxes=[[0,0],[0,1]], center=[0,0], window_size=2)
        self._extract_pattern_2D = layers.PercievePattern(receptive_field_idxes=[[0,0],[1,1]], center=[0,0], window_size=2)

    @staticmethod
    def init_from_numpy(
        stage1_3H, stage1_3D, stage1_3B, stage1_2H, stage1_2D, 
        stage2_3H, stage2_3D, stage2_3B, stage2_2H, stage2_2D 
    ):   
        # quantization_interval = 256//(stage1_3H.shape[0]-1)
        quantization_interval = 16
        lut_model = HDBLut(quantization_interval=quantization_interval, scale=4)
        lut_model.stage1_3H = nn.Parameter(torch.tensor(stage1_3H).type(torch.float32))
        lut_model.stage1_3D = nn.Parameter(torch.tensor(stage1_3D).type(torch.float32))
        lut_model.stage1_3B = nn.Parameter(torch.tensor(stage1_3B).type(torch.float32))
        lut_model.stage1_2H = nn.Parameter(torch.tensor(stage1_2H).type(torch.float32))
        lut_model.stage1_2D = nn.Parameter(torch.tensor(stage1_2D).type(torch.float32))

        lut_model.stage2_3H = nn.Parameter(torch.tensor(stage2_3H).type(torch.float32))
        lut_model.stage2_3D = nn.Parameter(torch.tensor(stage2_3D).type(torch.float32))
        lut_model.stage2_3B = nn.Parameter(torch.tensor(stage2_3B).type(torch.float32))
        lut_model.stage2_2H = nn.Parameter(torch.tensor(stage2_2H).type(torch.float32))
        lut_model.stage2_2D = nn.Parameter(torch.tensor(stage2_2D).type(torch.float32))
        return lut_model

    def forward_stage(self, x, scale, percieve_pattern, lut):
        b,c,h,w = x.shape
        print(np.prod(x.shape))
        x = percieve_pattern(x)
        shifts = torch.tensor([lut.shape[0]**d for d in range(len(lut.shape)-2)], device=x.device).flip(0).reshape(1,1,len(lut.shape)-2)
        print(x.shape, x.min(), x.max())
        x = torch.sum(x * shifts, dim=-1)
        print(x.shape)
        lut = torch.clamp(lut, 0, 255)
        lut = lut.reshape(-1, scale, scale)
        x = x.flatten().type(torch.int64)
        x = lut[x]

        x = round_func(x)
        x = x.reshape(b, c, h, w, scale, scale)
        print(x.shape)
        # raise RuntimeError
        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)
        lsb = x % 16
        msb = x - lsb        
        output_msb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
        output_lsb = torch.zeros([b*c, 1, h*2, w*2], dtype=x.dtype, device=x.device)
        for rotations_count in range(4):
           rotated_msb = torch.floor_divide(torch.rot90(msb, k=rotations_count, dims=[2, 3]), 16)
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage1_3H), k=-rotations_count, dims=[2, 3])
           output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage1_3D), k=-rotations_count, dims=[2, 3])
           output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage1_3B), k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage1_2H), k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage1_2D), k=-rotations_count, dims=[2, 3])
        output_msb /= 4*3
        output_lsb /= 4*2
        print(output_msb.min(), output_msb.max())
        print(output_lsb.min(), output_lsb.max())

        output_msb = output_msb + output_lsb
        x = output_msb
        lsb = x % 16
        msb = x - lsb        
        output_msb = torch.zeros([b*c, 1, h*4, w*4], dtype=x.dtype, device=x.device)
        output_lsb = torch.zeros([b*c, 1, h*4, w*4], dtype=x.dtype, device=x.device)
        print("STAGE2", msb.min(), msb.max(), lsb.min(), lsb.max())
        for rotations_count in range(4):
           rotated_msb = torch.floor_divide(torch.rot90(msb, k=rotations_count, dims=[2, 3]), 16)
           rotated_lsb = torch.rot90(lsb, k=rotations_count, dims=[2, 3])
           output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3H, self.stage2_3H), k=-rotations_count, dims=[2, 3])
           output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3D, self.stage2_3D), k=-rotations_count, dims=[2, 3])
           output_msb += torch.rot90(self.forward_stage(rotated_msb, 2, self._extract_pattern_3B, self.stage2_3B), k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2H, self.stage2_2H), k=-rotations_count, dims=[2, 3])
           output_lsb += torch.rot90(self.forward_stage(rotated_lsb, 2, self._extract_pattern_2D, self.stage2_2D), k=-rotations_count, dims=[2, 3])
        output_msb /= 4*3
        output_lsb /= 4*2
        output_msb = output_msb + output_lsb
        x = output_msb
        x = x.reshape(b, c, h*self.scale, w*self.scale)
        return x        
    
    def __repr__(self):
        return f"{self.__class__.__name__}" + \
               f"\n  stage1_3H size: {self.stage1_3H.shape}" + \
               f"\n  stage1_3D size: {self.stage1_3D.shape}" + \
               f"\n  stage1_3B size: {self.stage1_3B.shape}" + \
               f"\n  stage1_2H size: {self.stage1_2H.shape}" + \
               f"\n  stage1_2D size: {self.stage1_2D.shape}" + \
               f"\n  stage2_3H size: {self.stage2_3H.shape}" + \
               f"\n  stage2_3D size: {self.stage2_3D.shape}" + \
               f"\n  stage2_3B size: {self.stage2_3B.shape}" + \
               f"\n  stage2_2H size: {self.stage2_2H.shape}" + \
               f"\n  stage2_2D size: {self.stage2_2D.shape}"