import os
import numpy as np
import torch
from torch.utils.data import Dataset
from utils.laserscan import LaserScan

EXTENSIONS_SCAN = ['.bin']
EXTENSIONS_LABEL = ['.npy']


def is_scan(filename):
  return any(filename.endswith(ext) for ext in EXTENSIONS_SCAN)


def is_label(filename):
  return any(filename.endswith(ext) for ext in EXTENSIONS_LABEL)


class Kitti(Dataset):

  def __init__(self, root,    # directory where data is
               sequences,     # sequences for this data (e.g. [1,3,4,6])
               sensor,              # sensor to parse scans from
               max_points=150000,   # max number of points present in dataset
               gt=True):            # send ground truth?
    # save deats
    self.root = os.path.join(root, "sequences")
    self.sequences = sequences
    self.sensor = sensor
    self.sensor_img_H = sensor["img_prop"]["height"]
    self.sensor_img_W = sensor["img_prop"]["width"]
    self.sensor_img_means = torch.tensor(sensor["img_means"],
                                         dtype=torch.float)
    self.sensor_img_stds = torch.tensor(sensor["img_stds"],
                                        dtype=torch.float)
    self.sensor_fov_up = sensor["fov_up"]
    self.sensor_fov_down = sensor["fov_down"]
    self.max_range = sensor["max_range"]
    self.max_points = max_points
    self.gt = gt

    print(self.root)
    # make sure directory exists
    if os.path.isdir(self.root):
      print("Sequences folder exists! Using sequences from %s" % self.root)
    else:
      raise ValueError("Sequences folder doesn't exist! Exiting...")

    # make sure sequences is a list
    assert(isinstance(self.sequences, list))

    # placeholder for filenames
    self.scan_files = []
    self.mask_files = []

    # fill in with names, checking that all sequences are complete
    for seq in self.sequences:
      # to string
      seq = '{0:02d}'.format(int(seq))

      print("parsing seq {}".format(seq))

      # get paths for each
      scan_path = os.path.join(self.root, seq, "velodyne")
      mask_path = os.path.join(self.root, seq, "masks")

      # get files
      scan_files = [os.path.join(dp, f) for dp, dn, fn in os.walk(
          os.path.expanduser(scan_path)) for f in fn if is_scan(f)]
      mask_files = [os.path.join(dp, f) for dp, dn, fn in os.walk(
          os.path.expanduser(mask_path)) for f in fn if is_label(f)]

      # each mask is generated by a pair of scans
      # so labels = scans-1
      if self.gt:
        assert(len(scan_files)-1 == len(mask_files))

      # extend list
      self.scan_files.extend(scan_files)
      self.mask_files.extend(mask_files)

    # sort for correspondance
    self.scan_files.sort()
    self.mask_files.sort()

    print("Using {} scans from sequences {}".format(len(self.scan_files),
                                                    self.sequences))

  def __getitem__(self, index):
    # get item in tensor shape
    prev_scan_file = self.scan_files[index]
    curr_scan_file = self.scan_files[index+1]
    if self.gt:
      mask_file = self.mask_files[index]
      mask = np.load(mask_file)
      mask = np.expand_dims(mask, axis=0)

    scan = LaserScan(project=True,
                   H=self.sensor_img_H,
                   W=self.sensor_img_W,
                   fov_up=self.sensor_fov_up,
                   fov_down=self.sensor_fov_down)

    # open and obtain scan
    scan.open_scan(prev_scan_file)
    proj_range_prev = torch.from_numpy(scan.proj_range).clone()
    proj_xyz_prev = torch.from_numpy(scan.proj_xyz).clone()
    proj_remission_prev = torch.from_numpy(scan.proj_remission).clone()
    proj_prev = torch.cat([proj_range_prev.unsqueeze(0).clone(),
                           proj_xyz_prev.clone().permute(2, 0, 1),
                           proj_remission_prev.unsqueeze(0).clone()])
    proj_prev = (proj_prev- self.sensor_img_means[:, None, None]) / self.sensor_img_stds[:, None, None]
    proj_mask_prev = torch.from_numpy(scan.proj_mask).unsqueeze(0).clone()
    proj_prev = proj_prev * proj_mask_prev.float()

    scan.open_scan(curr_scan_file)
    proj_range_curr = torch.from_numpy(scan.proj_range).clone()
    proj_xyz_curr = torch.from_numpy(scan.proj_xyz).clone()
    proj_remission_curr = torch.from_numpy(scan.proj_remission).clone()
    proj_curr = torch.cat([proj_range_prev.unsqueeze(0).clone(),
                           proj_xyz_prev.clone().permute(2, 0, 1),
                           proj_remission_prev.unsqueeze(0).clone()])
    proj_curr = (proj_curr- self.sensor_img_means[:, None, None]) / self.sensor_img_stds[:, None, None]

    proj_mask = torch.from_numpy(scan.proj_mask).unsqueeze(0).clone()
    proj_curr = proj_curr * proj_mask.float()
    if self.gt:
      proj_labels = torch.from_numpy(mask).float().clone()
      # proj_labels = proj_labels * proj_mask
    else:
      proj_labels = []

    # get name and sequence
    path_norm = os.path.normpath(curr_scan_file)
    path_split = path_norm.split(os.sep)
    path_seq = path_split[-3]
    path_name = path_split[-1].replace(".bin", ".label")
    # print("path_norm: ", path_norm)
    # print("path_seq", path_seq)
    # print("path_name", path_name)
    proj = [proj_prev, proj_curr]

    return proj, proj_mask, proj_labels, path_seq, path_name

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

class Parser():
  # standard conv, BN, relu
  def __init__(self,
               root,              # directory for data
               train_sequences,   # sequences to train
               valid_sequences,   # sequences to validate.
               test_sequences,    # sequences to test (if none, don't get)
               sensor,            # sensor to use
               max_points,        # max points in each scan in entire dataset
               batch_size,        # batch size for train and val
               workers,           # threads to load data
               gt=True,           # get gt?
               shuffle_train=True):  # shuffle training set?
    super(Parser, self).__init__()

    # if I am training, get the dataset
    self.root = root
    self.train_sequences = train_sequences
    self.valid_sequences = valid_sequences
    self.test_sequences = test_sequences
    self.sensor = sensor
    self.max_points = max_points
    self.batch_size = batch_size
    self.workers = workers
    self.gt = gt
    self.shuffle_train = shuffle_train

    # Data loading code
    self.train_dataset = Kitti(root=self.root,
                               sequences=self.train_sequences,
                               sensor=self.sensor,
                               max_points=max_points,
                               gt=self.gt)

    self.trainloader = torch.utils.data.DataLoader(self.train_dataset,
                                                   batch_size=self.batch_size,
                                                   shuffle=self.shuffle_train,
                                                   num_workers=self.workers,
                                                   pin_memory=True,
                                                   drop_last=True)
    print(len(self.trainloader))
    assert len(self.trainloader) > 0
    self.trainiter = iter(self.trainloader)

    self.valid_dataset = Kitti(root=self.root,
                                       sequences=self.valid_sequences,
                                       sensor=self.sensor,
                                       max_points=max_points,
                                       gt=self.gt)

    self.validloader = torch.utils.data.DataLoader(self.valid_dataset,
                                                   batch_size=self.batch_size,
                                                   shuffle=False,
                                                   num_workers=self.workers,
                                                   pin_memory=True,
                                                   drop_last=True)
    assert len(self.validloader) > 0
    self.validiter = iter(self.validloader)

    if self.test_sequences:
      self.test_dataset = Kitti(root=self.root,
          sequences=self.test_sequences,
          sensor=self.sensor,
          max_points=max_points,
          gt=False)

      self.testloader = torch.utils.data.DataLoader(self.test_dataset,
          batch_size=self.batch_size,
          shuffle=False,
          num_workers=self.workers,
          pin_memory=True,
          drop_last=True)
      assert len(self.testloader) > 0
      self.testiter = iter(self.testloader)

  def get_train_batch(self):
    scans = self.trainiter.next()
    return scans

  def get_train_set(self):
    return self.trainloader

  def get_valid_batch(self):
    scans = self.validiter.next()
    return scans

  def get_valid_set(self):
    return self.validloader

  def get_test_batch(self):
    scans = self.testiter.next()
    return scans

  def get_test_set(self):
    return self.testloader

  def get_train_size(self):
    return len(self.trainloader)

  def get_valid_size(self):
    return len(self.validloader)

  def get_test_size(self):
    return len(self.testloader)

import os
import numpy as np
import torch
from torch.utils.data import Dataset
from utils.laserscan import LaserScan

EXTENSIONS_SCAN = ['.bin']
EXTENSIONS_LABEL = ['.npy']


def is_scan(filename):
  return any(filename.endswith(ext) for ext in EXTENSIONS_SCAN)


def is_label(filename):
  return any(filename.endswith(ext) for ext in EXTENSIONS_LABEL)


class Kitti(Dataset):

  def __init__(self, root,    # directory where data is
               sequences,     # sequences for this data (e.g. [1,3,4,6])
               sensor,              # sensor to parse scans from
               max_points=150000,   # max number of points present in dataset
               gt=True):            # send ground truth?
    # save deats
    self.root = os.path.join(root, "sequences")
    self.sequences = sequences
    self.sensor = sensor
    self.sensor_img_H = sensor["img_prop"]["height"]
    self.sensor_img_W = sensor["img_prop"]["width"]
    self.sensor_img_means = torch.tensor(sensor["img_means"],
                                         dtype=torch.float)
    self.sensor_img_stds = torch.tensor(sensor["img_stds"],
                                        dtype=torch.float)
    self.sensor_fov_up = sensor["fov_up"]
    self.sensor_fov_down = sensor["fov_down"]
    self.max_range = sensor["max_range"]
    self.max_points = max_points
    self.gt = gt

    print(self.root)
    # make sure directory exists
    if os.path.isdir(self.root):
      print("Sequences folder exists! Using sequences from %s" % self.root)
    else:
      raise ValueError("Sequences folder doesn't exist! Exiting...")

    # make sure sequences is a list
    assert(isinstance(self.sequences, list))

    # placeholder for filenames
    self.scan_files = []
    self.mask_files = []

    # fill in with names, checking that all sequences are complete
    for seq in self.sequences:
      # to string
      seq = '{0:02d}'.format(int(seq))

      print("parsing seq {}".format(seq))

      # get paths for each
      scan_path = os.path.join(self.root, seq, "velodyne")
      mask_path = os.path.join(self.root, seq, "masks")

      # get files
      scan_files = [os.path.join(dp, f) for dp, dn, fn in os.walk(
          os.path.expanduser(scan_path)) for f in fn if is_scan(f)]
      mask_files = [os.path.join(dp, f) for dp, dn, fn in os.walk(
          os.path.expanduser(mask_path)) for f in fn if is_label(f)]

      # each mask is generated by a pair of scans
      # so labels = scans-1
      if self.gt:
        assert(len(scan_files)-1 == len(mask_files))

      # extend list
      self.scan_files.extend(scan_files)
      self.mask_files.extend(mask_files)

    # sort for correspondance
    self.scan_files.sort()
    self.mask_files.sort()

    print("Using {} scans from sequences {}".format(len(self.scan_files),
                                                    self.sequences))

  def __getitem__(self, index):
    # get item in tensor shape
    prev_scan_file = self.scan_files[index]
    curr_scan_file = self.scan_files[index+1]
    if self.gt:
      mask_file = self.mask_files[index]
      mask = np.load(mask_file)
      mask = np.expand_dims(mask, axis=0)

    scan = LaserScan(project=True,
                   H=self.sensor_img_H,
                   W=self.sensor_img_W,
                   fov_up=self.sensor_fov_up,
                   fov_down=self.sensor_fov_down)

    # open and obtain scan
    scan.open_scan(prev_scan_file)
    proj_range_prev = torch.from_numpy(scan.proj_range).clone()
    proj_xyz_prev = torch.from_numpy(scan.proj_xyz).clone()
    proj_remission_prev = torch.from_numpy(scan.proj_remission).clone()
    proj_prev = torch.cat([proj_range_prev.unsqueeze(0).clone(),
                           proj_xyz_prev.clone().permute(2, 0, 1),
                           proj_remission_prev.unsqueeze(0).clone()])
    proj_prev = (proj_prev- self.sensor_img_means[:, None, None]) / self.sensor_img_stds[:, None, None]
    proj_mask_prev = torch.from_numpy(scan.proj_mask).unsqueeze(0).clone()
    proj_prev = proj_prev * proj_mask_prev.float()

    scan.open_scan(curr_scan_file)
    proj_range_curr = torch.from_numpy(scan.proj_range).clone()
    proj_xyz_curr = torch.from_numpy(scan.proj_xyz).clone()
    proj_remission_curr = torch.from_numpy(scan.proj_remission).clone()
    proj_curr = torch.cat([proj_range_prev.unsqueeze(0).clone(),
                           proj_xyz_prev.clone().permute(2, 0, 1),
                           proj_remission_prev.unsqueeze(0).clone()])
    proj_curr = (proj_curr- self.sensor_img_means[:, None, None]) / self.sensor_img_stds[:, None, None]

    proj_mask = torch.from_numpy(scan.proj_mask).unsqueeze(0).clone()
    proj_curr = proj_curr * proj_mask.float()
    if self.gt:
      proj_labels = torch.from_numpy(mask).float().clone()
      # proj_labels = proj_labels * proj_mask
    else:
      proj_labels = []

    # get name and sequence
    path_norm = os.path.normpath(curr_scan_file)
    path_split = path_norm.split(os.sep)
    path_seq = path_split[-3]
    path_name = path_split[-1].replace(".bin", ".label")
    # print("path_norm: ", path_norm)
    # print("path_seq", path_seq)
    # print("path_name", path_name)
    proj = [proj_prev, proj_curr]

    return proj, proj_mask, proj_labels, path_seq, path_name

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

class Parser():
  # standard conv, BN, relu
  def __init__(self,
               root,              # directory for data
               train_sequences,   # sequences to train
               valid_sequences,   # sequences to validate.
               test_sequences,    # sequences to test (if none, don't get)
               sensor,            # sensor to use
               max_points,        # max points in each scan in entire dataset
               batch_size,        # batch size for train and val
               workers,           # threads to load data
               gt=True,           # get gt?
               shuffle_train=True):  # shuffle training set?
    super(Parser, self).__init__()

    # if I am training, get the dataset
    self.root = root
    self.train_sequences = train_sequences
    self.valid_sequences = valid_sequences
    self.test_sequences = test_sequences
    self.sensor = sensor
    self.max_points = max_points
    self.batch_size = batch_size
    self.workers = workers
    self.gt = gt
    self.shuffle_train = shuffle_train

    # Data loading code
    self.train_dataset = Kitti(root=self.root,
                               sequences=self.train_sequences,
                               sensor=self.sensor,
                               max_points=max_points,
                               gt=self.gt)

    self.trainloader = torch.utils.data.DataLoader(self.train_dataset,
                                                   batch_size=self.batch_size,
                                                   shuffle=self.shuffle_train,
                                                   num_workers=self.workers,
                                                   pin_memory=True,
                                                   drop_last=True)
    print(len(self.trainloader))
    assert len(self.trainloader) > 0
    self.trainiter = iter(self.trainloader)

    self.valid_dataset = Kitti(root=self.root,
                                       sequences=self.valid_sequences,
                                       sensor=self.sensor,
                                       max_points=max_points,
                                       gt=self.gt)

    self.validloader = torch.utils.data.DataLoader(self.valid_dataset,
                                                   batch_size=self.batch_size,
                                                   shuffle=False,
                                                   num_workers=self.workers,
                                                   pin_memory=True,
                                                   drop_last=True)
    assert len(self.validloader) > 0
    self.validiter = iter(self.validloader)

    if self.test_sequences:
      self.test_dataset = Kitti(root=self.root,
          sequences=self.test_sequences,
          sensor=self.sensor,
          max_points=max_points,
          gt=False)

      self.testloader = torch.utils.data.DataLoader(self.test_dataset,
          batch_size=self.batch_size,
          shuffle=False,
          num_workers=self.workers,
          pin_memory=True,
          drop_last=True)
      assert len(self.testloader) > 0
      self.testiter = iter(self.testloader)

  def get_train_batch(self):
    scans = self.trainiter.next()
    return scans

  def get_train_set(self):
    return self.trainloader

  def get_valid_batch(self):
    scans = self.validiter.next()
    return scans

  def get_valid_set(self):
    return self.validloader

  def get_test_batch(self):
    scans = self.testiter.next()
    return scans

  def get_test_set(self):
    return self.testloader

  def get_train_size(self):
    return len(self.trainloader)

  def get_valid_size(self):
    return len(self.validloader)

  def get_test_size(self):
    return len(self.testloader)

# directory where data is # sequences for this data (e.g. [1,3,4,6]) # sensor to parse scans from # max number of points present in dataset # send ground truth? # save deats # make sure directory exists # make sure sequences is a list # placeholder for filenames # fill in with names, checking that all sequences are complete # to string # get paths for each # get files # each mask is generated by a pair of scans # so labels = scans-1 # extend list # sort for correspondance # get item in tensor shape # open and obtain scan # proj_labels = proj_labels * proj_mask # get name and sequence # print("path_norm: ", path_norm) # print("path_seq", path_seq) # print("path_name", path_name) # standard conv, BN, relu # directory for data # sequences to train # sequences to validate. # sequences to test (if none, don't get) # sensor to use # max points in each scan in entire dataset # batch size for train and val # threads to load data # get gt? # shuffle training set? # if I am training, get the dataset # Data loading code