# -*- coding: utf-8 -*-
import itertools
import torch
import torch.nn as nn
import torch.nn.functional as F
from .unet3d import UNet3DModel
import trimesh
from tqdm import tqdm
from skimage import measure
from ...modules.utils import convert_module_to_f16, convert_module_to_f32


def adaptive_conv(inputs,weights):
    padding = (1, 1, 1, 1, 1, 1)
    padded_input = F.pad(inputs, padding, mode="constant", value=0)
    output = torch.zeros_like(inputs)
    size=inputs.shape[-1]
    for i in range(3):
        for j in range(3):
            for k in range(3):
                output=output+padded_input[:,:,i:i+size,j:j+size,k:k+size]*weights[:,i*9+j*3+k:i*9+j*3+k+1]
    return output

def adaptive_block(inputs,conv,weights_=None):
    if weights_ != None:
        weights = conv(weights_)
    else:
        weights = conv(inputs)
    weights = F.normalize(weights, dim=1, p=1)
    for i in range(3):
        inputs = adaptive_conv(inputs, weights)
    return inputs

class GeoDecoder(nn.Module):

    def __init__(self, 
                 n_features: int,
                 hidden_dim: int = 32, 
                 num_layers: int = 4, 
                 use_sdf: bool = False,
                 activation: nn.Module = nn.ReLU):
        super().__init__()
        self.use_sdf=use_sdf
        self.net = nn.Sequential(
            nn.Linear(n_features, hidden_dim),
            activation(),
            *itertools.chain(*[[
                nn.Linear(hidden_dim, hidden_dim),
                activation(),
            ] for _ in range(num_layers - 2)]),
            nn.Linear(hidden_dim, 8),
        )

        # init all bias to zero
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        x = self.net(x)
        return x


class Voxel_RefinerXL(nn.Module):
    def __init__(self,
                in_channels: int = 1,
                out_channels: int = 1,
                layers_per_block: int = 2,
                layers_mid_block: int = 2,
                patch_size: int = 192,
                res: int = 512,
                use_checkpoint: bool=False,
                use_fp16: bool = False):

        super().__init__()

        self.unet3d1 = UNet3DModel(in_channels=16, out_channels=8, use_conv_out=False,
                                   layers_per_block=layers_per_block, layers_mid_block=layers_mid_block, 
                                   block_out_channels=(8, 32, 128,512), norm_num_groups=4, use_checkpoint=use_checkpoint)
        self.conv_in = nn.Conv3d(in_channels, 8, kernel_size=3, padding=1)
        self.latent_mlp = GeoDecoder(32)
        self.adaptive_conv1 = nn.Sequential(nn.Conv3d(8, 8, kernel_size=3, padding=1),
                                            nn.ReLU(),
                                            nn.Conv3d(8, 27, kernel_size=3, padding=1, bias=False))
        self.adaptive_conv2 = nn.Sequential(nn.Conv3d(8, 8, kernel_size=3, padding=1),
                                            nn.ReLU(),
                                            nn.Conv3d(8, 27, kernel_size=3, padding=1, bias=False))
        self.adaptive_conv3 = nn.Sequential(nn.Conv3d(8, 8, kernel_size=3, padding=1),
                                            nn.ReLU(),
                                            nn.Conv3d(8, 27, kernel_size=3, padding=1, bias=False))
        self.mid_conv = nn.Conv3d(8, 8, kernel_size=3, padding=1)
        self.conv_out = nn.Conv3d(8, out_channels, kernel_size=3, padding=1)
        self.patch_size = patch_size
        self.res = res

        self.use_fp16 = use_fp16
        self.dtype = torch.float16 if use_fp16 else torch.float32
        if use_fp16:
            self.convert_to_fp16()

    def convert_to_fp16(self) -> None:
        """
        Convert the torso of the model to float16.
        """
        # self.blocks.apply(convert_module_to_f16)
        self.apply(convert_module_to_f16)

    def run(self,
            reconst_x,
            feat, 
            mc_threshold=0,
        ):
        batch_size = int(reconst_x.coords[..., 0].max()) + 1
        sparse_sdf, sparse_index = reconst_x.feats, reconst_x.coords
        sparse_feat = feat.feats
        device = sparse_sdf.device
        dtype = sparse_sdf.dtype
        res = self.res

        sdfs = []
        for i in range(batch_size):
            idx = sparse_index[..., 0] == i
            sparse_sdf_i, sparse_index_i = sparse_sdf[idx].squeeze(-1),  sparse_index[idx][..., 1:]
            sdf = torch.ones((res, res, res)).to(device).to(dtype)
            sdf[sparse_index_i[..., 0], sparse_index_i[..., 1], sparse_index_i[..., 2]] = sparse_sdf_i
            sdfs.append(sdf.unsqueeze(0))

        sdfs = torch.stack(sdfs, dim=0)
        feats = torch.zeros((batch_size, sparse_feat.shape[-1], res, res, res), 
                            device=device, dtype=dtype)
        feats[sparse_index[...,0],:,sparse_index[...,1],sparse_index[...,2],sparse_index[...,3]] = sparse_feat
        
        N = sdfs.shape[0]
        outputs = torch.ones([N,1,res,res,res], dtype=dtype, device=device)
        stride = 160
        patch_size = self.patch_size
        step = 3
        sdfs = sdfs.to(dtype)
        feats = feats.to(dtype)
        patchs=[]
        for i in range(step):
            for j in range(step):
                for k in tqdm(range(step)):
                    sdf = sdfs[:, :, stride * i: stride * i + patch_size,
                               stride * j: stride * j + patch_size,
                               stride * k: stride * k + patch_size]
                    crop_feats = feats[:, :, stride * i: stride * i + patch_size, 
                                       stride * j: stride * j + patch_size, 
                                       stride * k: stride * k + patch_size]
                    inputs = self.conv_in(sdf)
                    crop_feats = self.latent_mlp(crop_feats.permute(0,2,3,4,1)).permute(0,4,1,2,3)
                    inputs = torch.cat([inputs, crop_feats],dim=1)
                    mid_feat = self.unet3d1(inputs)  
                    mid_feat = adaptive_block(mid_feat, self.adaptive_conv1)
                    mid_feat = self.mid_conv(mid_feat)
                    mid_feat = adaptive_block(mid_feat, self.adaptive_conv2)
                    final_feat = self.conv_out(mid_feat)
                    final_feat = adaptive_block(final_feat, self.adaptive_conv3, weights_=mid_feat)
                    output = F.tanh(final_feat)
                    patchs.append(output)
        weights = torch.linspace(0, 1, steps=32, device=device, dtype=dtype)
        lines=[]
        for i in range(9):
            out1 = patchs[i * 3]
            out2 = patchs[i * 3 + 1]
            out3 = patchs[i * 3 + 2]
            line = torch.ones([N, 1, 192, 192,res], dtype=dtype, device=device) * 2
            line[:, :, :, :, :160] = out1[:, :, :, :, :160]
            line[:, :, :, :, 192:320] = out2[:, :, :, :, 32:160]
            line[:, :, :, :, 352:] = out3[:, :, :, :, 32:]
            
            line[:,:,:,:,160:192] = out1[:,:,:,:,160:] * (1-weights.reshape(1,1,1,1,-1)) + out2[:,:,:,:,:32] * weights.reshape(1,1,1,1,-1)
            line[:,:,:,:,320:352] = out2[:,:,:,:,160:] * (1-weights.reshape(1,1,1,1,-1)) + out3[:,:,:,:,:32] * weights.reshape(1,1,1,1,-1)
            lines.append(line)
        layers=[]
        for i in range(3):
            line1 = lines[i*3]
            line2 = lines[i*3+1]
            line3 = lines[i*3+2]
            layer = torch.ones([N,1,192,res,res], device=device, dtype=dtype) * 2
            layer[:,:,:,:160] = line1[:,:,:,:160]
            layer[:,:,:,192:320] = line2[:,:,:,32:160]
            layer[:,:,:,352:] = line3[:,:,:,32:]
            layer[:,:,:,160:192] = line1[:,:,:,160:]*(1-weights.reshape(1,1,1,-1,1))+line2[:,:,:,:32]*weights.reshape(1,1,1,-1,1)
            layer[:,:,:,320:352] = line2[:,:,:,160:]*(1-weights.reshape(1,1,1,-1,1))+line3[:,:,:,:32]*weights.reshape(1,1,1,-1,1)
            layers.append(layer)
        outputs[:,:,:160] = layers[0][:,:,:160]
        outputs[:,:,192:320] = layers[1][:,:,32:160]
        outputs[:,:,352:] = layers[2][:,:,32:]
        outputs[:,:,160:192] = layers[0][:,:,160:]*(1-weights.reshape(1,1,-1,1,1))+layers[1][:,:,:32]*weights.reshape(1,1,-1,1,1)
        outputs[:,:,320:352] = layers[1][:,:,160:]*(1-weights.reshape(1,1,-1,1,1))+layers[2][:,:,:32]*weights.reshape(1,1,-1,1,1)
        # outputs = -outputs

        meshes = []
        for i in range(outputs.shape[0]):
            vertices, faces, _, _ = measure.marching_cubes(outputs[i, 0].cpu().numpy(), level=mc_threshold, method='lewiner')
            vertices = vertices / res * 2 - 1
            meshes.append(trimesh.Trimesh(vertices, faces))
        
        return meshes