import numpy as np
import matplotlib.pyplot as plt
import cv2
from pathlib import Path
import sys

import hyperdepth as hd

sys.path.append('../')
import dataset


def get_data(n, row_from, row_to, train):
  imsizes = [(256,384)]
  focal_lengths = [160]
  dset = dataset.SynDataset(n, imsizes=imsizes, focal_lengths=focal_lengths, train=train)
  ims = np.empty((n, row_to-row_from, imsizes[0][1]), dtype=np.uint8)
  disps = np.empty((n, row_to-row_from, imsizes[0][1]), dtype=np.float32)
  for idx in range(n):
    print(f'load sample {idx} train={train}')
    sample = dset[idx]
    ims[idx] = (sample['im0'][0,row_from:row_to] * 255).astype(np.uint8)
    disps[idx] = sample['disp0'][0,row_from:row_to]
  return ims, disps



params = hd.TrainParams(
  n_trees=4,
  max_tree_depth=,
  n_test_split_functions=50,
  n_test_thresholds=10,
  n_test_samples=4096,
  min_samples_to_split=16,
  min_samples_for_leaf=8)

n_disp_bins = 20
depth_switch = 0

row_from = 100
row_to = 108
n_train_samples = 1024
n_test_samples = 32

train_ims, train_disps = get_data(n_train_samples, row_from, row_to, True)
test_ims, test_disps = get_data(n_test_samples, row_from, row_to, False)

for tree_depth in [8,10,12,14,16]:
  depth_switch = tree_depth - 4

  prefix = f'td{tree_depth}_ds{depth_switch}'
  prefix = Path(f'./forests/{prefix}/')
  prefix.mkdir(parents=True, exist_ok=True)

  hd.train_forest(params, train_ims, train_disps, n_disp_bins=n_disp_bins, depth_switch=depth_switch, forest_prefix=str(prefix / 'fr'))

  es = hd.eval_forest(test_ims, test_disps, n_disp_bins=n_disp_bins, depth_switch=depth_switch, forest_prefix=str(prefix / 'fr')) 

  np.save(str(prefix / 'ta.npy'), test_disps)
  np.save(str(prefix / 'es.npy'), es)

  # plt.figure(); 
  # plt.subplot(2,1,1); plt.imshow(test_disps[0], vmin=0, vmax=4);
  # plt.subplot(2,1,2); plt.imshow(es[0], vmin=0, vmax=4);
  # plt.show()