import os
import numpy as np

from torch.utils.data import Dataset
from torch_geometric.datasets import ShapeNet

from itertools import repeat, product
from collections import defaultdict

import os
from tqdm import tqdm
import os.path as osp
import glob

import torch
from torch_geometric.data import (Data, InMemoryDataset, download_url,
                                  extract_zip)
from torch.utils.data import Dataset
from torch_geometric.read import read_txt_array
from torch_geometric.datasets import ShapeNet

from torch_geometric.read import parse_txt_array


class CustomShapeNet(InMemoryDataset):

    categories = {key: val for val, key in enumerate(['Box', 'Cone', 'Cylinder', 'Sphere'])}

    def __init__(self, root, train=True, transform=None, pre_filter=None, pre_transform=None, **kwargs):
        super(CustomShapeNet, self).__init__(root, transform, pre_transform, pre_filter)
        path = self.processed_paths[0] if train else self.processed_paths[1]
        self.data, self.slices = torch.load(path)
        print("Initialized")

    @property
    def raw_file_names(self):
        # Maybe add more data like validation sets
        return ['train', 'test']

    @property
    def processed_file_names(self):
        return [f'{x}.pt' for x in self.raw_file_names]

    def download(self):
        dir_count = len([name for name in os.listdir(self.raw_dir) if os.path.isdir(os.path.join(self.raw_dir, name))])
        print(f'{dir_count} folders have been found....')
        if dir_count:
            return dir_count
        raise IOError("No raw pointclouds have been found.")

    @property
    def num_classes(self):
        return len(self.categories)

    def _load_dataset(self):
        data, slices = None, None
        while True:
            try:
                filepath = os.path.join(self.root, self.processed_dir, f'{"train" if self.train else "test"}.pt')
                data, slices = torch.load(filepath)
                print('Dataset Loaded')
                break
            except FileNotFoundError:
                self.process()
                continue
        return data, slices

    def process(self, delimiter=' '):
        # idx = self.categories[self.category]
        # paths = [osp.join(path, idx) for path in self.raw_paths]

        datasets = defaultdict(list)
        for idx, setting in enumerate(self.raw_file_names):
            for pointcloud in tqdm(os.scandir(os.path.join(self.raw_dir, setting))):
                if not os.path.isdir(pointcloud):
                    continue
                for element in glob.glob(os.path.join(pointcloud.path, '*.dat')):
                    if os.path.split(element)[-1] not in ['pc.dat']:
                        # Assign training data to the data container
                        # Following the original logic;
                        # y should be the label;
                        # pos should be the six dimensional vector describing: !its pos not points!!
                        # x,y,z,x_rot,y_rot,z_rot
                        y_raw = os.path.splitext(element)[0].split('_')[-2]
                        with open(element,'r') as f:
                            headers = f.__next__()
                            # Check if there are no useable nodes in this file, header says 0.
                            if not int(headers.rstrip().split(delimiter)[0]):
                                continue
                            # Get the y - Label

                            # Iterate over all rows
                            src = [[float(x) if x not in ['-nan(ind)', 'nan(ind)'] else 0
                                    for x in line.rstrip().split(delimiter)[None:None]] for line in f if line != '']
                        points = torch.tensor(src, dtype=None).squeeze()
                        if not len(points.shape) > 1:
                            continue
                        # pos = points[:, :3]
                        # norm = points[:, 3:]
                        y_all = [self.categories[y_raw]] * points.shape[0]
                        y = torch.as_tensor(y_all, dtype=torch.int)
                        # points = torch.as_tensor(points, dtype=torch.float)
                        # norm = torch.as_tensor(norm, dtype=torch.float)
                        data = Data(y=y, pos=points[:, :3])
                        # , points=points, norm=points[:3], )
                        # ToDo: ANy filter to apply? Then do it here.
                        if self.pre_filter is not None and not self.pre_filter(data):
                            data = self.pre_filter(data)
                            raise NotImplementedError
                        # ToDo: ANy transformation to apply? Then do it here.
                        if self.pre_transform is not None:
                            data = self.pre_transform(data)
                            raise NotImplementedError
                        datasets[setting].append(data)

            os.makedirs(self.processed_dir, exist_ok=True)
            torch.save(self.collate(datasets[setting]), self.processed_paths[idx])

    def __repr__(self):
        return f'{self.__class__.__name__}({len(self)})'


class ShapeNetPartSegDataset(Dataset):
    """
    Resample raw point cloud to fixed number of points.
    Map raw label from range [1, N] to [0, N-1].
    """
    def __init__(self, root_dir, train=True, transform=None, npoints=1024):
        super(ShapeNetPartSegDataset, self).__init__()
        self.npoints = npoints
        self.dataset = CustomShapeNet(root=root_dir, train=train, transform=transform)

    def __getitem__(self, index):
        data = self.dataset[index]
        points, labels = data.pos, data.y

        # Resample to fixed number of points
        try:
            choice = np.random.choice(points.shape[0], self.npoints, replace=True)
        except ValueError:
            choice = []

        points, labels = points[choice, :], labels[choice]

        labels -= 1 if self.num_classes() in labels else 0   # Map label from [1, C] to [0, C-1]

        sample = {
            'points': points,  # torch.Tensor (n, 3)
            'labels': labels  # torch.Tensor (n,)
        }

        return sample

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

    def num_classes(self):
        return self.dataset.num_classes