"""
Scalable, Detailed and Mask-free Universal Photometric Stereo Network (CVPR2023)
# Copyright (c) 2023 Satoshi Ikehata
# All rights reserved.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

def divide_tensor_spatial(x, block_size=256, method='tile_stride'):
    assert x.dim() == 4, "Input tensor must have 4 dimensions [B, C, H, W]"
    B, C, H, W = x.shape    
    assert H == W, "Height and Width must be equal"
    assert H % block_size == 0 and W % block_size ==0, "The tensor size cannot be divided by the block size"
    mosaic_scale = H // block_size
    
    if method == 'tile_stride':
        """ decomposing x into K x K of (Hc, Wc) non-overlapped blocks (grid)"""           
        
        K = mosaic_scale * mosaic_scale
        fold_params_grid = dict(kernel_size=(mosaic_scale, mosaic_scale), stride=(mosaic_scale, mosaic_scale), padding=(0,0), dilation=(1,1))
        unfold_grid = nn.Unfold(**fold_params_grid)   
        tensor_grids = unfold_grid(x) # (B, C * K, Hm * Hm)
        tensor_grids = tensor_grids.reshape(B, C, K, block_size, block_size).permute(0, 2, 1, 3, 4) # (B, K, C, Hm, Hm)
        return tensor_grids
    
    if method == 'tile_block':    
        tensor_blocks = x.view(B, C, mosaic_scale, block_size, mosaic_scale, block_size)
        tensor_blocks = tensor_blocks.permute(0, 2, 4, 1, 3, 5) # (B, mc, mc, C, Hm, Wm)
        tensor_blocks = tensor_blocks.contiguous().view(B, mosaic_scale**2, C, block_size, block_size) ## (B, K, C, Hm, Hm)
        return tensor_blocks
    
    return -1

def merge_tensor_spatial(x, method='tile_stride'):
    
    K, N, feat_dim, Hm, Wm = x.shape
    mosaic_scale = int(math.sqrt(K))

    if method == 'tile_stride':
        x = x.reshape(K, N, feat_dim, -1)
        fold_params_grid = dict(kernel_size=(mosaic_scale, mosaic_scale), stride=(mosaic_scale, mosaic_scale), padding=(0,0), dilation=(1,1))
        fold_grid = nn.Fold(output_size=(Hm * mosaic_scale, Wm * mosaic_scale), **fold_params_grid) #  downsample based on the encoder     
        x = x.permute(1, 2, 0, 3).reshape(N, feat_dim * K, -1) 
        x = fold_grid(x)
        return x

    if method == 'tile_block':
        x = x.permute(1, 0, 2, 3, 4).reshape(N, mosaic_scale, mosaic_scale, feat_dim, Hm, Wm)
        x = x.permute(0, 3, 1, 4, 2, 5)
        x = x.reshape(N, feat_dim, mosaic_scale * Hm, mosaic_scale * Wm)
        return x

def divide_overlapping_patches(input_tensor, patch_size, margin):
    B, C, W, _ = input_tensor.shape
    stride = patch_size - margin
    padded_W = ((W - patch_size + stride - 1) // stride) * stride + patch_size
    pad = padded_W - W

    padded_tensor = F.pad(input_tensor, (0, pad, 0, pad), mode='constant', value=0)

    patches = F.unfold(padded_tensor, kernel_size=patch_size, stride=stride)
    patches = patches.view(B, C, patch_size, patch_size, -1).permute(0, 4, 1, 2, 3)

    return patches

def merge_overlappnig_patches(patches, patch_size, margin, original_size):
    B, _, C, _, _ = patches.shape
    stride = patch_size - margin
    W = original_size[2]

    patches = patches.permute(0, 2, 3, 4, 1).contiguous().view(B, C * patch_size * patch_size, -1)
    output = F.fold(patches, (W, W), kernel_size=patch_size, stride=stride)  

    weight = torch.ones(patches.size()).to(patches.device)
    weight = F.fold(weight, (W, W), kernel_size=patch_size, stride=stride)

    output = output / weight
    return output