
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import torch.multiprocessing as mp
from .utils import round_func
import numpy as np

##################### TRANSFER ##########################

# TODO make one function instead of 4

class Domain1DValues(Dataset):
    def __init__(self, quantization_interval=1, max_value=255):
        super(Domain1DValues, self).__init__()
        values1d = torch.arange(0, max_value+1, quantization_interval, dtype=torch.uint8)
        values1d = torch.cat([values1d, torch.tensor([max_value+1])])
        self.quantization_interval = quantization_interval
        self.values = torch.cartesian_prod(*([values1d]*1)).view(-1, 1, 2)
        
    def __getitem__(self, idx):
        if isinstance(idx, slice):
            ix1s, batch = [], []
            for i in range(idx.start, idx.stop):
                ix1, values = self.__getitem__(i)                
                ix1s.append(ix1)
                batch.append(values)
            return ix1s, batch
        else:
            v = self.values[idx]
            ix = v[0]//self.quantization_interval
            return ix[0],  v

    def __len__(self):
        return len(self.values)

    def __iter__(self):
        for i in range(len(self.values)):
            yield self.__getitem__(i)

class Domain2DValues(Dataset):
    def __init__(self, quantization_interval=1, max_value=255):
        super(Domain2DValues, self).__init__()
        values1d = torch.arange(0, max_value+1, quantization_interval, dtype=torch.uint8)
        values1d = torch.cat([values1d, torch.tensor([max_value+1])])
        self.quantization_interval = quantization_interval
        self.values = torch.cartesian_prod(*([values1d]*2)).view(-1, 1, 2)
        
    def __getitem__(self, idx):
        if isinstance(idx, slice):
            ix1s, ix2s, batch = [], [], []
            for i in range(idx.start, idx.stop):
                ix1, ix2, values = self.__getitem__(i)                
                ix1s.append(ix1)
                ix2s.append(ix2)
                batch.append(values)
            return ix1s, ix2s, batch
        else:
            v = self.values[idx]
            ix = v[0]//self.quantization_interval
            return ix[0], ix[1],  v

    def __len__(self):
        return len(self.values)

    def __iter__(self):
        for i in range(len(self.values)):
            yield self.__getitem__(i)

class Domain3DValues(Dataset):
    def __init__(self, quantization_interval=1, max_value=255):
        super(Domain3DValues, self).__init__()
        values1d = torch.arange(0, max_value+1, quantization_interval, dtype=torch.uint8)
        values1d = torch.cat([values1d, torch.tensor([max_value+1])])
        self.quantization_interval = quantization_interval
        self.values = torch.cartesian_prod(*([values1d]*3)).view(-1, 1, 3)
        
    def __getitem__(self, idx):
        if isinstance(idx, slice):
            ix1s, ix2s, ix3s, batch = [], [], [], []
            for i in range(idx.start, idx.stop):
                ix1, ix2, ix3,  values = self.__getitem__(i)                
                ix1s.append(ix1)
                ix2s.append(ix2)
                ix3s.append(ix3)
                batch.append(values)
            return ix1s, ix2s, ix3s, batch
        else:
            v = self.values[idx]
            ix = v[0]//self.quantization_interval
            return ix[0], ix[1], ix[2], v

    def __len__(self):
        return len(self.values)

    def __iter__(self):
        for i in range(len(self.values)):
            yield self.__getitem__(i)

class Domain4DValues(Dataset):
    def __init__(self, quantization_interval=1, max_value=255):
        super(Domain4DValues, self).__init__()
        values1d = torch.arange(0, max_value+1, quantization_interval, dtype=torch.uint8)
        values1d = torch.cat([values1d, torch.tensor([max_value+1])])
        self.quantization_interval = quantization_interval
        self.values = torch.cartesian_prod(*([values1d]*4)).view(-1, 1, 4)
        
    def __getitem__(self, idx):
        if isinstance(idx, slice):
            ix1s, ix2s, ix3s, ix4s, batch = [], [], [], [], []
            for i in range(idx.start, idx.stop):
                ix1, ix2, ix3, ix4, values = self.__getitem__(i)                
                ix1s.append(ix1)
                ix2s.append(ix2)
                ix3s.append(ix3)
                ix4s.append(ix4)
                batch.append(values)
            return ix1s, ix2s, ix3s, ix4s, batch
        else:
            v = self.values[idx]
            ix = v[0]//self.quantization_interval
            return ix[0], ix[1], ix[2], ix[3], v

    def __len__(self):
        return len(self.values)

    def __iter__(self):
        for i in range(len(self.values)):
            yield self.__getitem__(i)


def transfer_1_input_SxS_output(block, quantization_interval=16, batch_size=2**10, max_value=255):
    bucket_count = (max_value+1)//quantization_interval    
    scale = block.upscale_factor if hasattr(block, 'upscale_factor') else 1
    lut = np.full((bucket_count+1, scale, scale), dtype=np.uint8, fill_value=max_value) # 4DLUT for simple input window 2x2
    domain_values = Domain1DValues(quantization_interval=quantization_interval, max_value=max_value)
    domain_values_loader = DataLoader(
        domain_values, 
        batch_size=batch_size if quantization_interval >= 16 else 2**16, 
        pin_memory=True, 
        num_workers=1 if quantization_interval >= 16 else mp.cpu_count(),
        shuffle=False,
    )
    counter = 0
    for idx, (ix1s, batch) in enumerate(domain_values_loader):
        inputs = batch.type(torch.float32).cuda()
        with torch.no_grad():
            outputs = block(inputs)
        lut[ix1s, :, :] = outputs.reshape(len(ix1s), scale, scale).cpu().numpy().astype(np.uint8)
        counter += inputs.shape[0]
        print(f"\r {block.__class__.__name__} {counter}/{len(domain_values)}", end="    ")
    print()
    return lut

def transfer_2_input_SxS_output(block, quantization_interval=16, batch_size=2**10, max_value=255):
    bucket_count = (max_value+1)//quantization_interval    
    scale = block.upscale_factor if hasattr(block, 'upscale_factor') else 1
    lut = np.full((bucket_count+1, bucket_count+1, scale, scale), dtype=np.uint8, fill_value=max_value) # 4DLUT for simple input window 2x2
    domain_values = Domain2DValues(quantization_interval=quantization_interval, max_value=max_value)
    domain_values_loader = DataLoader(
        domain_values, 
        batch_size=batch_size if quantization_interval >= 16 else 2**16, 
        pin_memory=True, 
        num_workers=1 if quantization_interval >= 16 else mp.cpu_count(),
        shuffle=False,
    )
    counter = 0
    for idx, (ix1s, ix2s, batch) in enumerate(domain_values_loader):
        inputs = batch.type(torch.float32).cuda()
        with torch.no_grad():
            outputs = block(inputs)
        lut[ix1s, ix2s, :, :] = outputs.reshape(len(ix1s), scale, scale).cpu().numpy().astype(np.uint8)
        counter += inputs.shape[0]
        print(f"\r {block.__class__.__name__} {counter}/{len(domain_values)}", end="    ")
    print()
    return lut

def transfer_3_input_SxS_output(block, quantization_interval=16, batch_size=2**10, max_value=255):
    bucket_count = (max_value+1)//quantization_interval    
    scale = block.upscale_factor if hasattr(block, 'upscale_factor') else 1
    lut = np.full((bucket_count+1, bucket_count+1, bucket_count+1, scale, scale), dtype=np.uint8, fill_value=max_value) # 4DLUT for simple input window 2x2
    domain_values = Domain3DValues(quantization_interval=quantization_interval, max_value=max_value)
    domain_values_loader = DataLoader(
        domain_values, 
        batch_size=batch_size if quantization_interval >= 16 else 2**16, 
        pin_memory=True, 
        num_workers=1 if quantization_interval >= 16 else mp.cpu_count(),
        shuffle=False,
    )
    counter = 0
    for idx, (ix1s, ix2s, ix3s, batch) in enumerate(domain_values_loader):
        inputs = batch.type(torch.float32).cuda()
        with torch.no_grad():
            outputs = block(inputs)
        lut[ix1s, ix2s, ix3s, :, :] = outputs.reshape(len(ix1s), scale, scale).cpu().numpy().astype(np.uint8)
        counter += inputs.shape[0]
        print(f"\r {block.__class__.__name__} {counter}/{len(domain_values)}", end="    ")
    print()
    return lut

def transfer_4_input_SxS_output(block, quantization_interval=16, batch_size=2**10, max_value=255):
    bucket_count = (max_value+1)//quantization_interval    
    scale = block.upscale_factor if hasattr(block, 'upscale_factor') else 1
    lut = np.full((bucket_count+1, bucket_count+1, bucket_count+1, bucket_count+1, scale, scale), dtype=np.uint8, fill_value=max_value) # 4DLUT for simple input window 2x2
    domain_values = Domain4DValues(quantization_interval=quantization_interval, max_value=max_value)
    domain_values_loader = DataLoader(
        domain_values, 
        batch_size=batch_size if quantization_interval >= 16 else 2**16, 
        pin_memory=True, 
        num_workers=1 if quantization_interval >= 16 else mp.cpu_count(),
        shuffle=False,
    )
    counter = 0
    for idx, (ix1s, ix2s, ix3s, ix4s, batch) in enumerate(domain_values_loader):
        inputs = batch.type(torch.float32).to(list(block.parameters())[0].device)
        with torch.no_grad():
            outputs = block(inputs)            
        lut[ix1s, ix2s, ix3s, ix4s, :, :] = outputs.reshape(len(ix1s), scale, scale).cpu().numpy().astype(np.uint8)
        counter += inputs.shape[0]
        print(f"\r {block.__class__.__name__} {counter}/{len(domain_values)}", end="    ")
    print()
    return lut

def transfer_N_input_SxS_output(block, quantization_interval=16, batch_size=2**10, max_value=255):
    match block.in_features:
        case 1: return transfer_1_input_SxS_output(block, quantization_interval, batch_size, max_value)
        case 2: return transfer_2_input_SxS_output(block, quantization_interval, batch_size, max_value)
        case 3: return transfer_3_input_SxS_output(block, quantization_interval, batch_size, max_value)
        case 4: return transfer_4_input_SxS_output(block, quantization_interval, batch_size, max_value)

# TODO revive rc
# def transfer_rc_conv(rc_conv, quantization_interval=1):
#     receptive_field_pixel_count = rc_conv.window_size**2
#     bucket_count = 256//quantization_interval
#     lut = np.full((receptive_field_pixel_count, bucket_count+1), dtype=np.uint8, fill_value=255)
#     for pixel_id in range(receptive_field_pixel_count):
#         for idx, value in enumerate(range(0, 256, quantization_interval)):
#             inputs = torch.tensor([value]).type(torch.float32).view(1,1,1).cuda()
#             with torch.no_grad():
#                 outputs = rc_conv.pixel_wise_forward(inputs)            
#             lut[:,idx] = outputs.flatten().cpu().numpy().astype(np.uint8)
#             print(f"\r {rc_conv.__class__.__name__} {pixel_id*bucket_count + idx +1}/{receptive_field_pixel_count*bucket_count}", end="    ")
#     print()
#     return lut

##################### FORWARD ##########################

# def forward_rc_conv_centered(index, lut):
#     window_size = lut.shape[0]
#     index = F.pad(index, pad=[window_size//2]*4, mode='replicate') 
#     window_indexes = lut.shape[:-1]         
#     # index = index.unsqueeze(-1)
#     x = torch.zeros_like(index)
#     for i in range(window_indexes[-2]):
#         for j in range(window_indexes[-1]):
#             shift_i, shift_j = -window_indexes[-2]//2+1 + i, -window_indexes[-1]//2+1 + j
#             shifted_index = torch.roll(index, shifts=[shift_i, shift_j], dims=[-2, -1])  
#             x += select_index_1dlut_linear(ixA=shifted_index, lut=lut[i,j])
#     x /= window_indexes[-2]*window_indexes[-1] 
#     x = x.squeeze(-1)
#     x = round_func(x)
#     x = x[:,:,window_size//2:-window_size//2+1,window_size//2:-window_size//2+1]
#     return x

# def forward_rc_conv_rot90(index, lut):
#     window_size = lut.shape[0]
#     index = F.pad(index, pad=[0, window_size-1]*2, mode='replicate') 
#     window_indexes = lut.shape[:-1]         
#     # index = index.unsqueeze(-1)
#     x = torch.zeros_like(index)
#     for i in range(window_indexes[-2]):
#         for j in range(window_indexes[-1]):
#             shift_i, shift_j = i, j
#             shifted_index = torch.roll(index, shifts=[shift_i, shift_j], dims=[-2, -1])  
#             x += select_index_1dlut_linear(ixA=shifted_index, lut=lut[i,j])
#     x /= window_indexes[-2]*window_indexes[-1] 
#     x = x.squeeze(-1)
#     x = round_func(x)
#     x = x[:,:,:-(window_size-1),:-(window_size-1)]
#     return x


##################### UTILS ##########################

def select_index_1dlut_linear(index, lut):
    b, hw, c = index.shape
    L = lut.shape[0]
    Q = 256/(L-1)
    msbA = torch.floor_divide(index, Q).type(torch.int64)
    msbB = msbA + 1
    msbA = msbA.flatten() 
    msbB = msbB.flatten() 
    lsb = index % Q
    outA = lut[msbA].reshape((b, hw, c))
    outB = lut[msbB].reshape((b, hw, c))
    lsb_coef = (lsb / Q).reshape((b, hw, c))
    out = outA + lsb_coef*(outB-outA) 
    return out

def select_index_3dlut_tetrahedral(index, lut):
    b, hw, c = index.shape
    lut = torch.clamp(lut, 0, 255)
    dimA, dimB, dimC = lut.shape[:4]
    q = 256/(dimA-1)
    L = dimA
    upscale = lut.shape[-1]
    weight = lut.reshape(L**3, upscale, upscale)

    msbA = torch.floor_divide(index, q).type(torch.int64)
    msbB = msbA + 1
    lsb = index % q

    img_a1 = msbA[:,:,0].reshape(b*hw, 1)
    img_b1 = msbA[:,:,1].reshape(b*hw, 1)
    img_c1 = msbA[:,:,2].reshape(b*hw, 1)
    img_a2 = msbB[:,:,0].reshape(b*hw, 1)
    img_b2 = msbB[:,:,1].reshape(b*hw, 1)
    img_c2 = msbB[:,:,2].reshape(b*hw, 1)
    fa = lsb[:,:,0].reshape(b*hw, 1)
    fb = lsb[:,:,1].reshape(b*hw, 1)
    fc = lsb[:,:,2].reshape(b*hw, 1)

    p000 = weight[img_a1 * L * L * L + img_b1 * L * L + img_c1 * L + img_d1].reshape(b*hw, upscale*upscale)
    p001 = weight[img_a1 * L * L * L + img_b1 * L * L + img_c1 * L + img_d2].reshape(b*hw, upscale*upscale)
    p010 = weight[img_a1 * L * L * L + img_b1 * L * L + img_c2 * L + img_d1].reshape(b*hw, upscale*upscale)
    p011 = weight[img_a1 * L * L * L + img_b1 * L * L + img_c2 * L + img_d2].reshape(b*hw, upscale*upscale)
    p100 = weight[img_a1 * L * L * L + img_b2 * L * L + img_c1 * L + img_d1].reshape(b*hw, upscale*upscale)
    p101 = weight[img_a1 * L * L * L + img_b2 * L * L + img_c1 * L + img_d2].reshape(b*hw, upscale*upscale)
    p110 = weight[img_a1 * L * L * L + img_b2 * L * L + img_c2 * L + img_d1].reshape(b*hw, upscale*upscale)
    p111 = weight[img_a1 * L * L * L + img_b2 * L * L + img_c2 * L + img_d2].reshape(b*hw, upscale*upscale)

    fab = fa > fb
    fbc = fb > fc
    fac = fa > fc
    fableq = fa <= fb
    fbcleq = fb <= fc
    facleq = fa <= fc

    out = torch.zeros((b*hw, upscale*upscale), dtype=weight.dtype).to(device=weight.device)

    i1 = i = torch.all(torch.cat([fab, fbc], dim=1), dim=1);                                                        
    out[i] = (q - fa[i]) * p000[i] + (fa[i] - fb[i]) * p100[i] + (fb[i] - fc[i]) * p110[i] + fc[i] * p111[i]

    i2 = i = torch.all(torch.cat([fab, fac], dim=1), dim=1);                                          
    out[i] = (q - fa[i]) * p000[i] + (fa[i] - fc[i]) * p100[i] + (fc[i] - fb[i]) * p101[i] + fb[i] * p111[i]

    i3 = i = torch.all(torch.cat([fab, ~i1[:, None], ~i2[:, None]], dim=1), dim=1);                            
    out[i] = (q - fc[i]) * p000[i] + (fc[i] - fa[i]) * p001[i] + (fa[i] - fb[i]) * p101[i] + fb[i] * p111[i]

    i4 = i = torch.all(torch.cat([~fab, ~fbcleq], dim=1), dim=1);                   
    out[i] = (q - fc[i]) * p000[i] + (fc[i] - fb[i]) * p001[i] + (fb[i] - fa[i]) * p011[i] + fa[i] * p111[i]

    i5 = i = torch.all(torch.cat([~fab, ~facleq], dim=1), dim=1);                                                
    out[i] = (q - fb[i]) * p000[i] + (fb[i] - fc[i]) * p010[i] + (fc[i] - fa[i]) * p011[i] + fa[i] * p111[i]

    i6 = i = torch.all(torch.cat([~fab, ~fbc, ~fac], dim=1), dim=1);                                  
    out[i] = (q - fb[i]) * p000[i] + (fb[i] - fa[i]) * p010[i] + (fa[i] - fc[i]) * p110[i] + fc[i] * p111[i]

    out = out.reshape((b, hw, upscale*upscale))
    out = out / q
    return out


def select_index_4dlut_tetrahedral(index, lut):
    b, hw, c = index.shape
    index = index.clamp(0, 255)
    lut = torch.clamp(lut, 0, 255)
    dimA, dimB, dimC, dimD = lut.shape[:4]
    q = 256/(dimA-1)
    L = dimA
    upscale = lut.shape[-1]
    weight = lut.reshape(L**4, upscale, upscale)

    msbA = torch.floor_divide(index, q).type(torch.int64)
    msbB = msbA + 1
    lsb = index % q
    img_a1 = msbA[:,:,0].reshape(b*hw, 1)
    img_b1 = msbA[:,:,1].reshape(b*hw, 1)
    img_c1 = msbA[:,:,2].reshape(b*hw, 1)
    img_d1 = msbA[:,:,3].reshape(b*hw, 1)
    img_a2 = msbB[:,:,0].reshape(b*hw, 1)
    img_b2 = msbB[:,:,1].reshape(b*hw, 1)
    img_c2 = msbB[:,:,2].reshape(b*hw, 1)
    img_d2 = msbB[:,:,3].reshape(b*hw, 1)
    fa = lsb[:,:,0].reshape(b*hw, 1)
    fb = lsb[:,:,1].reshape(b*hw, 1)
    fc = lsb[:,:,2].reshape(b*hw, 1)
    fd = lsb[:,:,3].reshape(b*hw, 1)

    p0000 = weight[img_a1 * L * L * L + img_b1 * L * L + img_c1 * L + img_d1]
    p0001 = weight[img_a1 * L * L * L + img_b1 * L * L + img_c1 * L + img_d2]
    p0010 = weight[img_a1 * L * L * L + img_b1 * L * L + img_c2 * L + img_d1]
    p0011 = weight[img_a1 * L * L * L + img_b1 * L * L + img_c2 * L + img_d2]
    p0100 = weight[img_a1 * L * L * L + img_b2 * L * L + img_c1 * L + img_d1]
    p0101 = weight[img_a1 * L * L * L + img_b2 * L * L + img_c1 * L + img_d2]
    p0110 = weight[img_a1 * L * L * L + img_b2 * L * L + img_c2 * L + img_d1]
    p0111 = weight[img_a1 * L * L * L + img_b2 * L * L + img_c2 * L + img_d2]

    p1000 = weight[img_a2 * L * L * L + img_b1 * L * L + img_c1 * L + img_d1]
    p1001 = weight[img_a2 * L * L * L + img_b1 * L * L + img_c1 * L + img_d2]
    p1010 = weight[img_a2 * L * L * L + img_b1 * L * L + img_c2 * L + img_d1]
    p1011 = weight[img_a2 * L * L * L + img_b1 * L * L + img_c2 * L + img_d2]
    p1100 = weight[img_a2 * L * L * L + img_b2 * L * L + img_c1 * L + img_d1]
    p1101 = weight[img_a2 * L * L * L + img_b2 * L * L + img_c1 * L + img_d2]
    p1110 = weight[img_a2 * L * L * L + img_b2 * L * L + img_c2 * L + img_d1]
    p1111 = weight[img_a2 * L * L * L + img_b2 * L * L + img_c2 * L + img_d2]

    out = torch.zeros((b*hw, upscale*upscale), dtype=weight.dtype).to(device=weight.device)

    p0000 = p0000.reshape(b*hw, upscale*upscale)
    p0100 = p0100.reshape(b*hw, upscale*upscale)
    p1000 = p1000.reshape(b*hw, upscale*upscale)
    p1100 = p1100.reshape(b*hw, upscale*upscale)

    p0001 = p0001.reshape(b*hw, upscale*upscale)
    p0101 = p0101.reshape(b*hw, upscale*upscale)
    p1001 = p1001.reshape(b*hw, upscale*upscale)
    p1101 = p1101.reshape(b*hw, upscale*upscale)

    p0010 = p0010.reshape(b*hw, upscale*upscale)
    p0110 = p0110.reshape(b*hw, upscale*upscale)
    p1010 = p1010.reshape(b*hw, upscale*upscale)
    p1110 = p1110.reshape(b*hw, upscale*upscale)

    p0011 = p0011.reshape(b*hw, upscale*upscale)
    p0111 = p0111.reshape(b*hw, upscale*upscale)
    p1011 = p1011.reshape(b*hw, upscale*upscale)
    p1111 = p1111.reshape(b*hw, upscale*upscale)

    fab = fa > fb;
    fac = fa > fc;
    fad = fa > fd

    fbc = fb > fc;
    fbd = fb > fd;
    fcd = fc > fd

    i1 = i = torch.all(torch.cat([fab, fbc, fcd], dim=1), dim=1);                                                        out[i] = (q - fa[i]) * p0000[i] + (fa[i] - fb[i]) * p1000[i] + (fb[i] - fc[i]) * p1100[i] + (fc[i] - fd[i]) * p1110[i] + (fd[i]) * p1111[i]
    i2 = i = torch.all(torch.cat([~i1[:, None], fab, fbc, fbd], dim=1), dim=1);                                          out[i] = (q - fa[i]) * p0000[i] + (fa[i] - fb[i]) * p1000[i] + (fb[i] - fd[i]) * p1100[i] + (fd[i] - fc[i]) * p1101[i] + (fc[i]) * p1111[i]
    i3 = i = torch.all(torch.cat([~i1[:, None], ~i2[:, None], fab, fbc, fad], dim=1), dim=1);                            out[i] = (q - fa[i]) * p0000[i] + (fa[i] - fd[i]) * p1000[i] + (fd[i] - fb[i]) * p1001[i] + (fb[i] - fc[i]) * p1101[i] + (fc[i]) * p1111[i]
    i4 = i = torch.all(torch.cat([~i1[:, None], ~i2[:, None], ~i3[:, None], fab, fbc], dim=1), dim=1);                   out[i] = (q - fd[i]) * p0000[i] + (fd[i] - fa[i]) * p0001[i] + (fa[i] - fb[i]) * p1001[i] + (fb[i] - fc[i]) * p1101[i] + (fc[i]) * p1111[i]

    i5 = i = torch.all(torch.cat([~(fbc), fab, fac, fbd], dim=1), dim=1);                                                out[i] = (q - fa[i]) * p0000[i] + (fa[i] - fc[i]) * p1000[i] + (fc[i] - fb[i]) * p1010[i] + (fb[i] - fd[i]) * p1110[i] + (fd[i]) * p1111[i]
    i6 = i = torch.all(torch.cat([~(fbc), ~i5[:, None], fab, fac, fcd], dim=1), dim=1);                                  out[i] = (q - fa[i]) * p0000[i] + (fa[i] - fc[i]) * p1000[i] + (fc[i] - fd[i]) * p1010[i] + (fd[i] - fb[i]) * p1011[i] + (fb[i]) * p1111[i]
    i7 = i = torch.all(torch.cat([~(fbc), ~i5[:, None], ~i6[:, None], fab, fac, fad], dim=1), dim=1);                    out[i] = (q - fa[i]) * p0000[i] + (fa[i] - fd[i]) * p1000[i] + (fd[i] - fc[i]) * p1001[i] + (fc[i] - fb[i]) * p1011[i] + (fb[i]) * p1111[i]
    i8 = i = torch.all(torch.cat([~(fbc), ~i5[:, None], ~i6[:, None], ~i7[:, None], fab, fac], dim=1), dim=1);           out[i] = (q - fd[i]) * p0000[i] + (fd[i] - fa[i]) * p0001[i] + (fa[i] - fc[i]) * p1001[i] + (fc[i] - fb[i]) * p1011[i] + (fb[i]) * p1111[i]

    i9 = i = torch.all(torch.cat([~(fbc), ~(fac), fab, fbd], dim=1), dim=1);                                             out[i] = (q - fc[i]) * p0000[i] + (fc[i] - fa[i]) * p0010[i] + (fa[i] - fb[i]) * p1010[i] + (fb[i] - fd[i]) * p1110[i] + (fd[i]) * p1111[i]
    i10 = i = torch.all(torch.cat([~(fbc), ~(fac), ~i9[:, None], fab, fad], dim=1), dim=1);                              out[i] = (q - fc[i]) * p0000[i] + (fc[i] - fa[i]) * p0010[i] + (fa[i] - fd[i]) * p1010[i] + (fd[i] - fb[i]) * p1011[i] + (fb[i]) * p1111[i]
    i11 = i = torch.all(torch.cat([~(fbc), ~(fac), ~i9[:, None], ~i10[:, None], fab, fcd], dim=1), dim=1);               out[i] = (q - fc[i]) * p0000[i] + (fc[i] - fd[i]) * p0010[i] + (fd[i] - fa[i]) * p0011[i] + (fa[i] - fb[i]) * p1011[i] + (fb[i]) * p1111[i]
    i12 = i = torch.all(torch.cat([~(fbc), ~(fac), ~i9[:, None], ~i10[:, None], ~i11[:, None], fab], dim=1), dim=1);     out[i] = (q - fd[i]) * p0000[i] + (fd[i] - fc[i]) * p0001[i] + (fc[i] - fa[i]) * p0011[i] + (fa[i] - fb[i]) * p1011[i] + (fb[i]) * p1111[i]

    i13 = i = torch.all(torch.cat([~(fab), fac, fcd], dim=1), dim=1);                                                    out[i] = (q - fb[i]) * p0000[i] + (fb[i] - fa[i]) * p0100[i] + (fa[i] - fc[i]) * p1100[i] + (fc[i] - fd[i]) * p1110[i] + (fd[i]) * p1111[i]
    i14 = i = torch.all(torch.cat([~(fab), ~i13[:, None], fac, fad], dim=1), dim=1);                                     out[i] = (q - fb[i]) * p0000[i] + (fb[i] - fa[i]) * p0100[i] + (fa[i] - fd[i]) * p1100[i] + (fd[i] - fc[i]) * p1101[i] + (fc[i]) * p1111[i]
    i15 = i = torch.all(torch.cat([~(fab), ~i13[:, None], ~i14[:, None], fac, fbd], dim=1), dim=1);                      out[i] = (q - fb[i]) * p0000[i] + (fb[i] - fd[i]) * p0100[i] + (fd[i] - fa[i]) * p0101[i] + (fa[i] - fc[i]) * p1101[i] + (fc[i]) * p1111[i]
    i16 = i = torch.all(torch.cat([~(fab), ~i13[:, None], ~i14[:, None], ~i15[:, None], fac], dim=1), dim=1);            out[i] = (q - fd[i]) * p0000[i] + (fd[i] - fb[i]) * p0001[i] + (fb[i] - fa[i]) * p0101[i] + (fa[i] - fc[i]) * p1101[i] + (fc[i]) * p1111[i]

    i17 = i = torch.all(torch.cat([~(fab), ~(fac), fbc, fad], dim=1), dim=1);                                            out[i] = (q - fb[i]) * p0000[i] + (fb[i] - fc[i]) * p0100[i] + (fc[i] - fa[i]) * p0110[i] + (fa[i] - fd[i]) * p1110[i] + (fd[i]) * p1111[i]
    i18 = i = torch.all(torch.cat([~(fab), ~(fac), ~i17[:, None], fbc, fcd], dim=1), dim=1);                             out[i] = (q - fb[i]) * p0000[i] + (fb[i] - fc[i]) * p0100[i] + (fc[i] - fd[i]) * p0110[i] + (fd[i] - fa[i]) * p0111[i] + (fa[i]) * p1111[i]
    i19 = i = torch.all(torch.cat([~(fab), ~(fac), ~i17[:, None], ~i18[:, None], fbc, fbd], dim=1), dim=1);              out[i] = (q - fb[i]) * p0000[i] + (fb[i] - fd[i]) * p0100[i] + (fd[i] - fc[i]) * p0101[i] + (fc[i] - fa[i]) * p0111[i] + (fa[i]) * p1111[i]
    i20 = i = torch.all(torch.cat([~(fab), ~(fac), ~i17[:, None], ~i18[:, None], ~i19[:, None], fbc], dim=1), dim=1);    out[i] = (q - fd[i]) * p0000[i] + (fd[i] - fb[i]) * p0001[i] + (fb[i] - fc[i]) * p0101[i] + (fc[i] - fa[i]) * p0111[i] + (fa[i]) * p1111[i]

    i21 = i = torch.all(torch.cat([~(fab), ~(fac), ~(fbc), fad], dim=1), dim=1);                                         out[i] = (q - fc[i]) * p0000[i] + (fc[i] - fb[i]) * p0010[i] + (fb[i] - fa[i]) * p0110[i] + (fa[i] - fd[i]) * p1110[i] + (fd[i]) * p1111[i]
    i22 = i = torch.all(torch.cat([~(fab), ~(fac), ~(fbc), ~i21[:, None], fbd], dim=1), dim=1);                          out[i] = (q - fc[i]) * p0000[i] + (fc[i] - fb[i]) * p0010[i] + (fb[i] - fd[i]) * p0110[i] + (fd[i] - fa[i]) * p0111[i] + (fa[i]) * p1111[i]
    i23 = i = torch.all(torch.cat([~(fab), ~(fac), ~(fbc), ~i21[:, None], ~i22[:, None], fcd], dim=1), dim=1);           out[i] = (q - fc[i]) * p0000[i] + (fc[i] - fd[i]) * p0010[i] + (fd[i] - fb[i]) * p0011[i] + (fb[i] - fa[i]) * p0111[i] + (fa[i]) * p1111[i]
    i24 = i = torch.all(torch.cat([~(fab), ~(fac), ~(fbc), ~i21[:, None], ~i22[:, None], ~i23[:, None]], dim=1), dim=1); out[i] = (q - fd[i]) * p0000[i] + (fd[i] - fc[i]) * p0001[i] + (fc[i] - fb[i]) * p0011[i] + (fb[i] - fa[i]) * p0111[i] + (fa[i]) * p1111[i]

    out = out.reshape((b, hw, upscale*upscale))
    out = out / q
    return out