import vtk
import numpy as np
import random
import platform
import sys
import os
import math
import time
import scipy.cluster.hierarchy as hc
from matplotlib import pyplot as plt
# from mpi4py import MPI
from sklearn.cluster import KMeans
from scipy.ndimage.measurements import label
from scipy.ndimage.morphology import generate_binary_structure
import argparse


def getVtkImageData(origin, dimensions, extents, spacing):
    localDataset = vtk.vtkImageData()
    localDataset.SetOrigin(origin)
    localDataset.SetDimensions(dimensions)
    localDataset.SetExtent(extents)
    localDataset.SetSpacing(spacing)
    # print(origin,dimensions,spacing)
    return localDataset


def bubble_based_sampling(in_file, outfile, x=0.1, bound_min=None, bound_max=None, nbins=None):

    reader = vtk.vtkXMLMultiBlockDataReader()
    # reader = vtk.vtkXMLPolyDataReader()
    reader.SetFileName(in_file)
    # reader.UpdatePiece(comm.Get_rank(), comm.Get_size(), 0)
    reader.Update()

    # print("total points:",data.GetNumberOfPoints(),data.GetNumberOfElements(0))
    NBlocks = reader.GetOutput().GetNumberOfBlocks()
    print("total blocks:", NBlocks)
    tot_pts = 0
    if NBlocks == 1:
        print('Here 1')
        mpData = reader.GetOutput().GetBlock(0)
        numPieces = mpData.GetNumberOfPieces()
        print("pieces:", numPieces)

        multiblock = vtk.vtkMultiBlockDataSet()
        multiblock.SetNumberOfBlocks(numPieces)

        # exit()

        pc_cnt = 0
        # tot_pts = 0

        for i in range(numPieces):
            data = mpData.GetPiece(i)

            if data is not None:
                # filtered = vtk.vtkPolyData()
                # Points = vtk.vtkPoints()

                pts = data.GetPoints()

                tot_pts = tot_pts + data.GetNumberOfPoints()

                # for i in range(local_pts):
                #     loc = pts.GetPoint(i)
                # #     # feat_arr[i, :] = np.asarray(loc)

        print("Total size:", tot_pts)
        target_points = int(tot_pts * x)

        feat_arr = np.zeros((tot_pts, 3), dtype='float')
        cur_count = 0
        for i in range(numPieces):
            data = mpData.GetPiece(i)

            if data is not None:
                # filtered = vtk.vtkPolyData()
                # Points = vtk.vtkPoints()
                pts = data.GetPoints()
                local_pts = data.GetNumberOfPoints()

                for i in range(local_pts):
                    loc = pts.GetPoint(i)
                    feat_arr[cur_count + i, :] = np.asarray(loc)
                cur_count = cur_count + data.GetNumberOfPoints()

    else:
        print('Here 2')
        mpData = reader.GetOutput()
        numPieces = mpData.GetNumberOfBlocks()
        print("pieces:", numPieces)

        multiblock = vtk.vtkMultiBlockDataSet()
        multiblock.SetNumberOfBlocks(numPieces)

        # exit()

        pc_cnt = 0
        # tot_pts = 0

        for i in range(numPieces):
            data = mpData.GetBlock(i)

            if data is not None:
                # filtered = vtk.vtkPolyData()
                # Points = vtk.vtkPoints()

                pts = data.GetPoints()

                tot_pts = tot_pts + data.GetNumberOfPoints()

                # for i in range(local_pts):
                #     loc = pts.GetPoint(i)
                # #     # feat_arr[i, :] = np.asarray(loc)

        print("Total size:", tot_pts)
        target_points = int(tot_pts * x)

        feat_arr = np.zeros((tot_pts, 3), dtype='float')
        cur_count = 0
        for i in range(numPieces):
            data = mpData.GetBlock(i)

            if data is not None:
                # filtered = vtk.vtkPolyData()
                # Points = vtk.vtkPoints()
                pts = data.GetPoints()
                local_pts = data.GetNumberOfPoints()

                for i in range(local_pts):
                    loc = pts.GetPoint(i)
                    feat_arr[cur_count + i, :] = np.asarray(loc)
                cur_count = cur_count + data.GetNumberOfPoints()

    print("read all the points, Total size:", np.shape(feat_arr))

    if np.asarray(bound_min).any() == None:
        bound_min = np.min(feat_arr, axis=0)
    if np.asarray(bound_max).any() == None:
        bound_max = np.max(feat_arr, axis=0)

    print('ranges:', bound_min, bound_max)

    ## Soumya: change this to add the nbins as command line args
    # mult = 2
    # nbins = [mult * 64, mult * 8, mult * 64]

    H, edges = np.histogramdd(feat_arr, bins=nbins, range=[[bound_min[0], bound_max[0]], [bound_min[1], bound_max[1]],
                                                           [bound_min[2], bound_max[2]]])

    # print(edges)

    print(np.shape(H), np.shape(edges),'non_zero:',np.count_nonzero(H),'total:',np.sum(H))

    np.save(outfile + '/histogram.npy', H)

    ## for comparison Ayan

    H_my = np.zeros_like(H)

##    for i in range(tot_pts):
##        loc = feat_arr[i,:]
##        # decide which bin this point goes to
##        x_id = int(nbins[0] * (loc[0] - bound_min[0]) / (bound_max[0] - bound_min[0]))
##        if x_id == nbins[0]:
##            x_id = x_id - 1
##        y_id = int(nbins[1] * (loc[1] - bound_min[1]) / (bound_max[1] - bound_min[1]))
##        if y_id == nbins[1]:
##            y_id = y_id - 1
##        z_id = int(nbins[2] * (loc[2] - bound_min[2]) / (bound_max[2] - bound_min[2]))
##        if z_id == nbins[2]:
##            z_id = z_id - 1
##        H_my[x_id, y_id, z_id] = H_my[x_id, y_id, z_id] + 1
##
##    print ("my 1 substract val: " + str(np.sum(abs(H-H_my))))

    # print(edges)
    # try thresholding

    th = np.max(H) / 3.0  # heuristic

    ###  for debug: soumya
    F_orig = np.copy(H)
    F_orig[H < th] = 1
    F_orig[H >= th] = 0
    np.save(outfile + '/orig_F.npy', F_orig)

    # th = 1
    H[H < th] = 0

    print('After TH: non_zero:', np.count_nonzero(H), 'total:', np.sum(H))

    print('Histo:', H.shape, edges[0].size, edges[1].size, edges[2].size, edges[0][0])

    ## my changes Ayan

    print('We are working with Threshold:',th)

    np.save(outfile+'/histogram_th.npy', H)

    ## go through the particles again

    ## create another histogram to keep track of if atleast one is picked or not
    H_book = np.zeros_like(H)

    ## generate new histogram type array and fill it with random index which will
    # be used later to pick randomized particels from each bin: soumya
    H_randsamp = np.zeros_like(H)
    for kk in range(np.shape(H_randsamp)[2]):
        for jj in range(np.shape(H_randsamp)[1]):
            for ii in range(np.shape(H_randsamp)[0]):
                if H[ii,jj,kk]>0:
                    H_randsamp[ii,jj,kk] = random.randint(1,H[ii,jj,kk])

    print('After new rand H: non_zero:', np.count_nonzero(H_randsamp))               

    pt_list = []

    multiblock = vtk.vtkMultiBlockDataSet()
    multiblock.SetNumberOfBlocks(numPieces)

    if 1:
        H_my = np.zeros_like(H)
        print('Here 1')
        mpData = reader.GetOutput().GetBlock(0)
        numPieces = mpData.GetNumberOfPieces()
        print("pieces:", numPieces)

        multiblock = vtk.vtkMultiBlockDataSet()
        multiblock.SetNumberOfBlocks(numPieces)

        # exit()

        pc_cnt = 0
        tot_pts = 0

        for i in range(numPieces):
            data = mpData.GetPiece(i)

            if data is not None:
                filtered = vtk.vtkPolyData()
                Points = vtk.vtkPoints()
                pts = data.GetPoints()
                local_pts = data.GetNumberOfPoints()

                for i in range(local_pts):
                    loc = np.asarray(pts.GetPoint(i))
                    # decide which bin this point goes to
                    x_id = int(nbins[0]*(loc[0]-bound_min[0])/(bound_max[0] - bound_min[0]))
                    if x_id==nbins[0]:
                        x_id = x_id-1
                    y_id = int(nbins[1] * (loc[1] - bound_min[1]) / (bound_max[1] - bound_min[1]))
                    if y_id == nbins[1]:
                        y_id = y_id - 1
                    z_id = int(nbins[2] * (loc[2] - bound_min[2]) / (bound_max[2] - bound_min[2]))
                    if z_id == nbins[2]:
                        z_id = z_id - 1

                    H_my[x_id, y_id, z_id] = H_my[x_id, y_id, z_id] + 1 ## Ayan check to see if the histogram is same

                    H_randsamp[x_id,y_id,z_id] = H_randsamp[x_id,y_id,z_id] - 1 ## soumya: reduce random value by 1

                    ## now picks one point from each non-zero bin randomly
                    if H_book[x_id,y_id,z_id]==0 and H[x_id,y_id,z_id]>0.00005 and H_randsamp[x_id,y_id,z_id]==0.0: 
                        Points.InsertNextPoint(loc)
                        H_book[x_id, y_id, z_id] = 1
                        pt_list.append(loc)
                filtered.SetPoints(Points)

                # for i in range(tot_components):
                #     # print(p.GetArrayName(i))
                #     val_arr = vtk.vtkFloatArray()
                #     val_arr.SetNumberOfComponents(1)
                #     val_arr.SetName(p.GetArrayName(i))
                #     orig_arr = p.GetArray(i)
                #     for cnt in range(len(sample_i)):
                #         val_arr.InsertNextValue(orig_arr.GetTuple1(sample_i[cnt]))
                #
                #     # print(orig_arr.GetTuple1(0))
                #     filtered.GetPointData().AddArray(val_arr)

                multiblock.SetBlock(pc_cnt, filtered)
                pc_cnt = pc_cnt + 1


        w = vtk.vtkXMLMultiBlockDataWriter()
        w.SetInputData(multiblock)
        filename = outfile
        w.SetFileName(outfile+'/bubble_based_samples.vtm')
        w.Write()

        print("total non-zero pieces:", pc_cnt)
        print('Total points gathered:',len(pt_list))

##        print ("my substract val: " + str(np.sum(abs(H-H_my))))


        ###############################
        ## for debuggin: soumya
        H_temp, edges_temp = np.histogramdd(np.asarray(pt_list), bins=nbins, range=[[bound_min[0], bound_max[0]], [bound_min[1], bound_max[1]],
                                                           [bound_min[2], bound_max[2]]])

        np.save(outfile + '/histogram_temp.npy', H_temp)
        print np.max(H_temp), np.min(H_temp)

        ###  for debug: soumya
        th = np.max(H_temp) / 3.0  # heuristic
        
        F = np.copy(H_temp)
        F[H_temp < th] = 1
        F[H_temp >= th] = 0

        np.save(outfile + '/sampled_F.npy', F)

        print ("subtract val: " + str(np.sum(abs(F-F_orig))))

        aa = F-F_orig
        aa = aa.flatten()
        aa = aa.tolist()
        aa = set(aa)
        print aa

    return

use_common_hist = 0
## run for all files
# input_path = path
# tot_files = np.size(os.listdir(input_path))
# all_files = os.listdir(input_path)
# print(tot_files, all_files)
#
if use_common_hist:
    bound_min = np.load('npy_files/bound_min_' + name_string + '.npy')
    bound_max = np.load('npy_files/bound_max_' + name_string + '.npy')
else:
    bound_min = None
    bound_max = None


############################
## Parse input parameters

def parse_bound(inpbounds):
    bounds = inpbounds.split(' ')
    bounds = np.asarray(bounds)
    range_min = [np.float(bounds[0]), np.float(bounds[2]), np.float(bounds[4])]
    range_max = [np.float(bounds[1]), np.float(bounds[3]), np.float(bounds[5])]
    return np.asarray(range_min), np.asarray(range_max)


def parse_nbins(inputNbins):
    nbins = inputNbins.split(' ')
    nbins = np.asarray(nbins)
    return nbins.astype(np.int)


parser = argparse.ArgumentParser()
parser.add_argument('--input', action="store", required=True, help="input file or files")
parser.add_argument('--output', action="store", required=True, help="name of the output file")
parser.add_argument('--bounds', action="store", required=False, help="global spatial bound of the data")
parser.add_argument('--percentage', action="store", required=True, help="sampling fraction (0.0-1.0)")
parser.add_argument('--nbins', action="store", required=True, help="number of bins")

args = parser.parse_args()

inputstring = getattr(args, 'input')
outFile = getattr(args, 'output')
inpbounds = getattr(args, 'bounds')
sample_fraction = float(getattr(args, 'percentage'))
inputNbins = getattr(args, 'nbins')

bound_min = np.zeros(3)
bound_max = np.zeros(3)
bound_min, bound_max = parse_bound(inpbounds)

nbins = np.zeros(3)
nbins = parse_nbins(inputNbins)

## call the function now
bubble_based_sampling(inputstring, outFile, sample_fraction, bound_min, bound_max, nbins)

## the 'sample_fraction' is not being used right now for this script. The bubble-based sampling gives the lowest possible samples without disturbing the bubbles.

## Sample Command:
# vtkpython bubble_based_sampling.py --input=/media/biswas/29d42253-dbfd-4e49-bc64-61662de409e0/MFIX/fcc-data/vtm_data/timestep_150.vtm --output=study/ --percentage=0.05 --nbins='128 16 128' --bounds='7.44e-5 0.15 7.44e-5 0.0031256 7.44e-5 0.0507256'

## sample run command: vtkpython bubble_based_sampling.py --input=/home/soumya/Shared_Vbox/Data/MFIX_fcc_vtm/MFIX_fcc_vtm_150.vtm --output=/home/soumya/Shared_Vbox/pythonVtkScripts/particle_sampling/for_cinema/not_in_cinema_repo/out --bounds='7.44e-5 0.15 7.44e-5 0.0031256 7.44e-5 0.0507256' --percentage=0.05 --nbins='128 16 128'

## Sample run command : vtkpython bubble_based_sampling.py --input=/media/sdutta/disk_1/Data/MFIX_fcc_vtm/MFIX_fcc_vtm_150.vtm --output=/media/sdutta/disk_2/MFIX_sampling_output/ --percentage=0.05 --nbins='128 16 128' --bounds='7.44e-5 0.15 7.44e-5 0.0031256 7.44e-5 0.0507256'
