# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import division

import glob
import os
import numpy as np
from scipy import spatial
import open3d as o3d

import nnabla as nn
import nnabla.logger as logger
import nnabla.monitor as M
import nnabla.functions as F
import nnabla.parametric_functions as PF
from nnabla.utils.data_iterator import data_iterator, data_iterator_simple
from nnabla.utils.data_source import DataSource
import utils

from args import get_args


class PointCloudDataSource(DataSource):

    def __init__(self, fpath, knn=50, test_rate=0.25, test=False, shuffle=True, rng=None):
        super(PointCloudDataSource, self).__init__(shuffle=shuffle)
        self.knn = knn
        self.test_rate = 0.25
        self.rng = np.random.RandomState(313) if rng is None else rng

        # Split info
        pcd = self._read_dataset(fpath)
        total_size = len(pcd.points)
        test_size = int(total_size * test_rate)
        indices = self.rng.permutation(total_size)
        test_indices = indices[:test_size]
        train_indices = indices[test_size:]
        indices = test_indices if test else train_indices
        self._size = test_size if test else total_size - test_size
        # Points
        points = np.asarray(pcd.points)
        self._points = self._preprocess(points)[indices]
        # Normals
        normals = np.asarray(pcd.normals)
        self._normals = normals[indices] if self.has_normals(
            normals) else normals
        # Radius
        self._radius = self._compute_radius(self._points, self.knn)
        self._variables = ('points', 'normals', 'radius')
        self.METHOD_NAME()

        logger.info("Data size = {}".format(self._size))

    def has_normals(self, normals):
        return False if normals.shape[0] == 0 else True

    def _preprocess(self, points):
        return utils.normalize(points)

    def _compute_radius(self, points, knn):
        if knn < 0:
            logger.info("Radius is not computed.")
            return
        # KDTree
        logger.info(
            "Constructing KDTree and querying {}-nearest neighbors".format(self.knn))
        tree = spatial.cKDTree(points, compact_nodes=True)
        # exclude self by adding 1
        dists, indices = tree.query(points, k=knn + 1)
        return dists[:, -1].reshape(dists.shape[0], 1)

    def _read_dataset(self, fpath):
        pcd = utils.read_pcd(fpath)
        return pcd

    def _get_data(self, position):
        points = self._points[self._indices[position]]
        normals = self._normals[self._indices[position]
                                ] if self.has_normals(self._normals) else [0.0]
        radius = self._radius[self._indices[position]]
        return points, normals, radius

    @property
    def points(self):
        return self._points

    @property
    def normals(self):
        return self._normals

    @property
    def radius(self):
        return self._radius

    def METHOD_NAME(self):
        self._indices = self.rng.permutation(self._size) \
            if self._shuffle else np.arange(self._size)
        return super(PointCloudDataSource, self).METHOD_NAME()


def point_cloud_data_source(fpath, knn=50, test_rate=0.25, test=False, shuffle=True, rng=None):
    return PointCloudDataSource(fpath, knn, test_rate, test, shuffle, rng)


def point_cloud_data_iterator(data_source, batch_size):
    return data_iterator(data_source, batch_size=batch_size,
                         with_memory_cache=False,
                         with_file_cache=False)


def create_pcd_dataset_from_mesh(fpath):
    mesh = utils.read_mesh(fpath)
    pcd = mesh.sample_points_poisson_disk(len(mesh.vertices),
                                          use_triangle_normal=True,
                                          seed=412)
    dpath = "/".join(fpath.split("/")[:-1])
    fname = fpath.split("/")[-1]
    fname = "{}_pcd.ply".format(os.path.splitext(fname)[0])
    fpath = os.path.join(dpath, fname)
    logger.info("PointCloud data ({}) is being created.".format(fpath))
    utils.write_pcd(fpath, pcd)


def main():
    args = get_args()
    if args.command == "create":
        create_pcd_dataset_from_mesh(args.mesh_data_path)


if __name__ == '__main__':
    main()