code: from pathlib import Path
  import copy
  import time
  import torch.optim as optim
  import numpy as np
  import torch
  from torch.autograd import Variable
  from model import *
  from data_utils import *
  import torch.nn as nn
  from loguru import logger
  
  feature_dim = 8
  block_size = 16
  pad=2
  n_conv=3
  thresh=0.5
  debug = False
  
  def test_bottom_io():
      tsdf = [torch.from_numpy(np.random.rand(1, 1, block_size+2*pad+2*n_conv,
          block_size+2*pad+2*n_conv,
          block_size+2*pad+2*n_conv)).float().to(device)]
      prev = {(0, 0, 0): torch.from_numpy(np.random.rand(1, feature_dim,
          block_size//2+2*pad, block_size//2+2*pad, block_size//2+2*pad)
          ).float().to(device)}
      mod = BottomLevel(feature_dim, block_size=block_size)
      if device == 'cuda':
          mod.cuda()
      out = mod(tsdf, prev)
      assert type(out) == list
      assert len(out) == 1
      out = out[0]
      assert len(out) == 1
      for X in out.keys():
          assert out[X].shape == (1, 2, block_size, block_size, block_size), out[X].shape
  
  
  def test_convtrans():
      conv1 = nn.ConvTranspose3d(10, 10, kernel_size=4, stride=2, output_padding=0, padding=0, bias=False)
      dat = torch.ones(1, 10, block_size, block_size, block_size)
      y = conv1(dat)
      assert y.shape[-1] == block_size*2+2 , (y.shape, dat.shape)
  
      pad = nn.ReplicationPad3d(1)
      conv1 = nn.ConvTranspose3d(1, 1, kernel_size=3, stride=2,
                                 output_padding=1, padding=1, bias=False)
      dat = Variable(torch.ones(1, 1, 4, 4, 4))
      y = conv1(dat)
      assert y.shape[-1] == 8, y.shape
  
  
  def test_data():
      data = TsdfGenerator(64)
      vis = visdom.Visdom()
      gt, tsdf_in = data.__getitem__(0)
      assert np.abs(tsdf_in).max() < 33
  
  
  def test_ellipsoid():
      arr = ellipsoid(10, 10, 10, levelset=True)*10  # the output is ~normalized.  multiple by 10
      assert arr.shape == (23, 23, 23), arr.shape
      dist = np.sqrt(11**2*3)-10
      assert np.abs(arr[0, 0, 0]) > dist, (arr[0, 0, 0], dist)
      print(arr[0, 0, 0], dist)
  
      a, b, c = 10, 15, 25
      arr = ellipsoid(a, b, c, levelset=True)
      # if we move 1 voxel in space the sdf should also not change by more than 1
      # compare to 1.01 for numeric reasons
      assert np.all(np.abs(np.diff(arr, axis=0)) <= 1.01), np.abs(np.diff(arr, axis=0)).max()
      assert np.all(np.abs(np.diff(arr, axis=1)) <= 1.01)
      assert np.all(np.abs(np.diff(arr, axis=2)) <= 1.01)
  
  
  def test_criteria_trivial():
      data = TsdfGenerator(block_size, sigma=0.)
      gt, tsdf_in = data.__getitem_split__()
      gt = gt[None, :]  # add dim for batch
      assert np.abs(tsdf_in).max() < 33
      gt_label = np.zeros_like(gt)
      gt_label[gt >= 0] = 1
      gt_label = torch.from_numpy(gt_label.astype(int))
      criteria = OctreeCrossEntropyLoss(gt_label, block_size)
      assert len(criteria.gt_octree) == 1
      mock_out = np.concatenate((tsdf_in[None,:]<0, tsdf_in[None,:]>=0),
              axis=1).astype(float)
      mock_out=1000*(mock_out-0.5)
      mock_out = [{(0,0,0):torch.from_numpy(mock_out).float()}]
      loss = criteria(mock_out)
      assert loss.dim()==0
      assert loss < 0.01, loss
  
  def test_gt():
      pass
      #get gt, 
      #get gt_octree
      #retnder gt
      #render gt_octree
  
  def test_criteria(levels=2):
      res=2**(levels-1)*block_size
      data = TsdfGenerator(res, sigma=0.9)
      gt, tsdf_in = data.__getitem_split__()
      gt = gt[None, :]  # add dim for batch
      assert np.abs(tsdf_in).max() < res
      #labels should be symetric
      def count_label(gt, label, level=1):
          gt_label = np.zeros_like(gt)
          gt_label[gt >= 0] = 1
          gt_label = torch.from_numpy(gt_label.astype(int))
          criteria = OctreeCrossEntropyLoss(gt_label, block_size)
          gt=criteria.gt_octree[level]
          return np.count_nonzero(np.array(list(gt.values()))==label)
  
      n_outside =  count_label(gt, OUTSIDE)
      n_inside = count_label(gt, INSIDE)
      n_mixed = count_label(gt, MIXED)
      assert n_outside+n_inside+n_mixed==(2**(levels-2))**3
      rev_inside = count_label(-gt, OUTSIDE)
      assert n_inside==rev_inside, (n_inside, rev_inside)
  
  
      gt_label = np.zeros_like(gt)
      gt_label[gt >= 0] = 1
      gt_label = torch.from_numpy(gt_label.astype(int))
      criteria = OctreeCrossEntropyLoss(gt_label, block_size)
      assert len(criteria.gt_octree) == levels
      assert len(criteria.gt_octree[0]) == (2**(levels-1))**3, len(criteria.gt_octree[0])
      assert len(criteria.gt_octree[-1]) == 1, len(criteria.gt_octree[-1])
      for l, level in enumerate(criteria.gt_octree):
          for k, v in level.items():
              assert v.dim() > 0, (l, k, v)
  
  
  def test_basic_debug():
      T = torch.zeros(1,1,36,36,36)
      outplane = 16
      mod = nn.Conv3d(1, outplane, kernel_size=3, stride=1,
                  padding=0, bias=False)
      T = mod(T)
      mod = nn.BatchNorm3d(outplane)
      T = mod(T)
      mod = nn.ReLU(inplace=True)
      T = mod(T)
      mod = nn.Conv3d(outplane, outplane, kernel_size=3, stride=1, 
              padding=0, bias=False)
      T = mod(T)
      mod = nn.BatchNorm3d(outplane)
      T = mod(T)
      assert T.shape == (1,16,32,32,32)
  
  
  def test_simple_net_single_data():
      data = TsdfGenerator(block_size, sigma=0.9)
      vis = visdom.Visdom()
      gt, tsdf_in = data.__getitem__(0)
      gt = gt[None, :]  # add dim for batch
      assert np.abs(tsdf_in).max() < block_size
      gt_label = np.zeros_like(gt)
      gt_label[gt >= 0] = 1
      gt_label = torch.from_numpy(gt_label.astype(int)).to(device)
      rep_pad = nn.ReplicationPad3d(pad+n_conv)
      tsdf = [rep_pad(torch.from_numpy(copy.copy(tsdf_in)[None, :]).float().to(device))]
      #prev = {(0, 0, 0): torch.rand(1, feature_dim, block_size//2, block_size//2,
      #    block_size//2).float().to(device)}
      prev = {(0, 0, 0): torch.from_numpy(np.random.rand(1, feature_dim,
          block_size//2+2*pad, block_size//2+2*pad, block_size//2+2*pad)
          ).float().to(device)}
      #assert tsdf[0].shape == (1, 1, block_size, block_size, block_size)
      assert gt_label.shape == (1, block_size, block_size, block_size)
      criteria = OctreeCrossEntropyLoss(gt_label, block_size)
      mod = BottomLevel(feature_dim, block_size)
      if device=='cuda':
          mod.cuda()
          criteria.cuda()
      optimizer = optim.Adam(mod.parameters(), lr=0.001)  # , momentum=0.9)
      for it in range(1, 100):
          out = mod(tsdf, prev)
          assert len(out) == 1
          assert out[0][(0,0,0)].shape[1] == 2, out.shape
          loss = criteria(out)
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          if (it+1) % 10 == 0:
              sdf_ = octree_to_sdf(out, block_size)
              print('level ', np.count_nonzero(sdf_ == 1))
              err = plotVoxelVisdom(gt[0], sdf_, tsdf_in[0], vis)
              assert np.abs(tsdf_in).max() < 33
              print(err)
  
          print(it, loss)
      assert err < 2
  
  
  def test_bottom_layer( block_size = 32):
      dataset = TsdfGenerator(block_size, n_elips=1, sigma=0.9, epoch_size=1000)
      train_loader = torch.utils.data.DataLoader(dataset, batch_size=1,
                                                 num_workers=4)
  
      vis = visdom.Visdom()
      mod = BottomLevel(feature_dim, block_size)
      if device=='cuda':
          mod.cuda()
      optimizer = optim.SGD(mod.parameters(), lr=0.0001, momentum=0.9)
      m = nn.ReplicationPad3d(mod.pad+mod.n_conv)
      prev = {(0, 0, 0): torch.rand(1, feature_dim,
              block_size//2+2*pad, block_size//2+2*pad, block_size//2+2*pad
              ).float().to(device)}
      gt_label = None
      for it, (gt, tsdf_in) in enumerate(train_loader):
          assert np.abs(tsdf_in).max() < 33
          assert gt.max() > 1 and gt.min() < -1
          gt_label = torch.ones_like(gt)*INSIDE
          gt_label[gt >= 0] = OUTSIDE
          gt_label = gt_label.long().to(device)
          tsdf = [m(tsdf_in).float().to(device)]
          for T in prev.values():
              assert torch.all(torch.isfinite(T))
          for T in tsdf:
              assert torch.all(torch.isfinite(T))
          out = mod(tsdf, prev)
          assert out[0][(0,0,0)].max()>out[0][(0,0,0)].min()
          for oct in out:
              if not np.all([torch.all(torch.isfinite(o)) for o in oct.values()]):
                  import ipdb; ipdb.set_trace()
          criteria = OctreeCrossEntropyLoss(gt_label, block_size)
          if device=='cuda':
              criteria.cuda()
          loss = criteria(out)
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          print(it, loss)
          if it>1 and it%100 == 0:
              sdf_ = octree_to_sdf(out, block_size)
              err = plotVoxelVisdom(gt[0].numpy(), sdf_, tsdf_in[0][0].numpy(), vis)
              print(it, err)
      assert err < 2, err
  
  
  def test_2tier_net_single_data():
      res = block_size*2
      dataset = TsdfGenerator(res, n_elips=3, sigma=0.9, epoch_size=100)
  
      vis = visdom.Visdom()
      mod = TopLevel(feature_dim, BottomLevel(feature_dim, block_size), block_size=block_size)
      if device == 'cuda':
          mod.cuda()
  
      optimizer = optim.Adam(mod.parameters(), lr=0.01)#, momentum=0.9)
      gt, tsdf_in = dataset.__getitem__(0)
      assert np.abs(tsdf_in).max() < 33
      assert gt.max() > 1 and gt.min() < -1
      gt = torch.from_numpy(gt[None, :])
      gt_label = torch.zeros_like(gt)
      gt_label[gt >= 0] = 1
      gt_label = gt_label.long().to(device)
      criteria = OctreeCrossEntropyLoss(gt_label, block_size)
      if device == 'cuda':
          criteria.cuda()
      tsdf = torch.from_numpy(copy.copy(tsdf_in)[None, :]).float().to(device)
      for it in range(1000):
          out = mod(tsdf)
          assert len(out) == 2
          for l in out[1:]:
              for v in l.values():
                  # only level 0 can have a full bloc
                  assert v.shape[-1] < block_size, (v.shape)
          loss = criteria(out)
          assert len(out) == 2
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          print(it, loss)
          if (it+1) % 10 == 0:
              #mod.eval()
              sdf_ = octree_to_sdf(out, block_size)
              err = plotVoxelVisdom(gt[0].numpy(), sdf_, tsdf_in[0], vis)
              #mod.train()
              print(it, err)
      assert err < 2,err
  
  
  def test_4tier_data(block_size=block_size):
      res=block_size*(2**3)
      dataset = TsdfGenerator(res, n_elips=3, sigma=0.9, epoch_size=1000)
      gt, tsdf = dataset.__getitem__(0)
  
      mod = BottomLevel(feature_dim, block_size)
      for i in range(2): #add 2 mid layers
          print('adding mid layer')
          mod = MidLevel(feature_dim, feature_dim, mod, block_size,
                  thresh=thresh, budget=4)
      mod = TopLevel(feature_dim, mod, block_size=block_size)
      out = mod(torch.from_numpy(tsdf[None,:]).float())
  
  
  
  def test_2tier_net(res=64, block_size=block_size):
      dataset = TsdfGenerator(res, n_elips=1, sigma=0.9, epoch_size=10000, debug=False)
      train_loader = torch.utils.data.DataLoader(dataset, batch_size=1,
                                                 num_workers=2)
  
      vis = visdom.Visdom()
      Force = False
      if not Force and Path('model_2tier.pth').exists():
          mod = torch.load('model_2tier.pth')
      else:
          layers = []
          layers.append(BottomLevel(feature_dim, block_size))
          while block_size*2**len(layers) <= res/2:
              print('adding mid layer', len(layers))
              layers.append(MidLevel(feature_dim, feature_dim, layers[-1],
                                     block_size, thresh=0.5, budget=4))
          mod = TopLevel(feature_dim, layers[-1], block_size=block_size)
      if device == 'cuda':
          mod.cuda()
      optimizer = optim.SGD(mod.parameters(), lr=0.0001, momentum=0.95)
      for it, (gt, tsdf_in) in enumerate(train_loader):
          assert np.abs(tsdf_in).max() < res
          assert gt.max() > 1 and gt.min() < -1
          gt_label = torch.zeros_like(gt, device=device)
          gt_label[gt >= 0] = 1
          gt_label = gt_label.long().to(device)
          criteria = OctreeCrossEntropyLoss(gt_label, block_size)
          if device == 'cuda':
              criteria.cuda()
          #tsdf = tsdf_in.float().cuda()
          t_start = time.time()
          tsdf = tsdf_in.float().to(device)
          pred = mod(tsdf)
          forward_t = time.time()-t_start
          t = time.time()
          loss = criteria(pred)
          loss_t = time.time()-t
          t = time.time()
          optimizer.zero_grad()
          loss.backward()
          back_t = time.time()-t
          t = time.time()
          optimizer.step()
          step_t = time.time()-t
          t = time.time()
          print(it, loss.data)
          print('valuated ', [len(o) for o in pred])
          print('GT voxels ', np.count_nonzero([o.numel()>3 for o in criteria.gt_octree]))
          print('timing:{total:.3f}. forward {forward_t:.3f}, loss {loss_t:.3f}, back {back_t:.3f}, step {step_t:.3f}'.format(
              total=t-t_start, forward_t=forward_t, loss_t=loss_t, back_t=back_t, step_t=step_t))
          if (it+1) % 100 == 0:
              mod.eval()
              out = mod(tsdf)
              loss = criteria(out)
              for i in range(len(out)):
                  resample = (2**i)
                  print('Eval: level %d, %d/%d evaluated' % (i, len(out[i]),
                                                             (res/block_size/resample)**3))
              sdf_ = octree_to_sdf(out, block_size)
              err = plotVoxelVisdom(gt[0].numpy(), sdf_, tsdf_in[0][0].numpy(), vis)
              if loss.data<1:
                  import ipdb; ipdb.set_trace()
              mod.train()
              print(it, err)
              torch.save(mod, 'model_2tier.pth')
              if err < 2 :
                  break
      #assert err < 2
  
  def create_model(block_size, feature_dim, res):
      layers = []
      layers.append(BottomLevel(feature_dim, block_size))
      while block_size*2**len(layers) <= res/2:
          print('adding mid layer', len(layers))
          layers.append(MidLevel(feature_dim, feature_dim, layers[-1],
                                 block_size, thresh=0.1))
      mod = TopLevel(feature_dim, layers[-1], block_size=block_size)
      return mod
  
  
  def test_simple_split(res=64, block_size=block_size):
      dataset = TsdfGenerator(res, n_elips=3, sigma=0.9, epoch_size=1000, debug=True)
      vis = visdom.Visdom()
  
      mod = torch.load('model.pth')
      if device == 'cuda':
          mod.cuda()
      mod.eval()
      gt, tsdf_in = dataset.__getitem_split__()
      gt = torch.from_numpy(gt[None, :])
      tsdf_in = torch.from_numpy(tsdf_in[None, :])
  
      gt_label = torch.zeros_like(gt, device=device)
      gt_label[gt >= 0] = 1
      gt_label = gt_label.long().to(device)
      criteria = OctreeCrossEntropyLoss(gt_label, block_size)
      if device == 'cuda':
          criteria.cuda()
      tsdf = tsdf_in.float().to(device)
      pred = mod(tsdf)
      loss = criteria(pred)
      print(loss.data)
      print('evaluated ', [len(o) for o in pred])
  
      for X in pred[0]:
          X_ = tuple(np.array(X)//2)
          print (X, pred[1][X_])
          assert pred[1][X_][0,2]>0.5
      sdf_ = octree_to_sdf(pred, block_size)
      err = plotVoxelVisdom(gt[0].numpy(), sdf_, tsdf_in[0][0].numpy(), vis)
      import ipdb; ipdb.set_trace()
      for X,v in criteria.gt_octree[0].items():
          if v.numel()>1:
              assert X[2]==1 #that's how we built the space
  
  
  def test_split_subtree(padding=0):
      feat = torch.rand(1, feature_dim, block_size+2*padding,
              block_size+2*padding,
              block_size+2*padding
              ).float()
      split = split_tree(feat,padding=padding)
      assert len(split) == 8, len(split)
      assert torch.all(split[(0, 0, 0)][0, :, padding, padding, padding] ==
              feat[0, :, padding, padding, padding])
      assert torch.all(split[(1, 0, 0)][0, :, padding, padding, padding] ==
              feat[0, :, block_size//2+padding, padding, padding])
      split[(1, 0, 0)][0, 0, padding, padding, padding] = 12.13
      #this is no longer true, I don't know how to do this inplace
      #assert feat[0, 0, block_size//2, 0, 0] == 12.13
  
  def test_split_subtree_with_padding():
      padding=2
      feat = torch.rand(1, feature_dim, block_size, block_size,
              block_size).float()
      split = split_tree(feat, padding=2)
      assert len(split) == 8, len(split)
      octant = split[(0,0,0)]
      assert torch.all(octant[0, :padding, 0, 0, 0] == 0)
      assert torch.all(octant[0, -padding:, 0, 0, 0] == 0)
      assert octant.shape[-3:]==feat.shape[-3:]//2+padding*2
      assert torch.all(octant[0, padding:-padding, 0, 0, 0] == feat[0, :, 0, 0, 0])
      assert torch.all(octant[0, padding:-padding, 0, 0, 0] == feat[0, :, 0, 0, 0])
      assert torch.all(split[(1, 0, 0)][0, :, padding, padding, padding] ==
              feat[0, :, block_size//2, 0, 0])
      split[(1, 0, 0)][0, 0, 0, 0, 0] = 12.13
      assert feat[0, 0, block_size//2+padding, 0, 0] == 12.13
  
  if __name__ == '__main__':
      import sys
      logger.remove()
      logger.add(sys.stderr , format="{time} {level} {message}",  level="INFO")
  
      #test_4tier_data()
      #test_criteria_trivial()
      #test_criteria()
      #test_criteria(4)
      #test_data()
      #test_ellipsoid()
      #test_convtrans()
      #test_split_subtree()
      #test_split_subtree(padding=2)
      #test_basic_debug()
      #test_bottom_io()
      #test_simple_net_single_data()
      #test_bottom_layer()
      # TODO why does this not converge? interesting
      #test_2tier_net_single_data()
      #test_2tier_net(res=32, block_size=block_size)
      test_2tier_net(res=64, block_size=block_size)
      test_simple_split(res=64, block_size=block_size)
      import ipdb; ipdb.set_trace()
      test_2tier_net(res=128, block_size=block_size)
  
apis: torch.utils.data.DataLoader numpy.diff torch.rand torch.nn.BatchNorm3d torch.save numpy.abs torch.nn.ConvTranspose3d torch.all torch.from_numpy numpy.random.rand torch.isfinite torch.ones_like torch.nn.ReplicationPad3d torch.ones torch.load torch.zeros numpy.count_nonzero numpy.array numpy.zeros_like torch.zeros_like numpy.sqrt numpy.concatenate torch.nn.ReLU torch.nn.Conv3d