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 create_histogram_and_detect_bubble(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)
    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):
            if numPieces==1:
                data = mpData
            else:
                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):
            if numPieces==1:
                data = mpData
            else:
                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)

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

    print(np.shape(H),np.shape(edges))
    #print(edges)
    # try thresholding

    th = np.max(H)/3.0 # heuristic
    #th = 1
    print('We are working with Threshold:', th)

    print np.max(H), np.min(H)

    H[H < th] = 0

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

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

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

    s = np.ones((3,3,3))
    s = generate_binary_structure(3,1)
    #print(s)

    F = np.copy(H)
    F[H<th] = 1
    F[H >= th] = 0

    labeled_array, num_features = label(F, structure=s)
    ## process the labeled array to remove the air portion
    # go through each feature and print the sizes
    size_th = 5 ## work on this threshold
    cur_features = 0
    # and check if it is touching the xbound or not
    for i in range(num_features):
        feat_id = i+1
        locs = np.asarray(np.where(labeled_array==feat_id))
        feat_size = np.shape(locs)[1]
        max_x = nbins[0] - 1
        if feat_size<size_th or max_x in locs[0,:]:
            # delete the bubble
            labeled_array[labeled_array == feat_id] = 0
        else:
            cur_features = cur_features+1
            labeled_array[labeled_array == feat_id] = cur_features
    print('total features before and after',num_features,cur_features)

    print('max:', np.max(labeled_array))

    #edges = np.asarray(edges)
    xdel = edges[0][1]-edges[0][0]
    ydel = edges[1][1]-edges[1][0]
    zdel = edges[2][1]-edges[2][0]

    # plt.figure()
    # plt.hist(np.ravel(H),bins=32)
    # plt.show()

    # write out a vti file with the bubble ids for visualization
    ## Soumya: these paths should be connected to your initial command line args
    out_file = outfile + '/bubbled_out.vti'
    origin = [bound_min[0]+xdel/2.0,bound_min[1]+ydel/2.0,bound_min[2]+zdel/2.0]
    dimensions = nbins
    exts = [0,dimensions[0]-1,0,dimensions[1]-1,0,dimensions[2]-1]
    spacing = [xdel,ydel,zdel]
    dataset = getVtkImageData(origin,dimensions,exts,spacing)

    dataset.AllocateScalars(vtk.VTK_DOUBLE, 1)

    dims = dataset.GetDimensions()

    # Fill every entry of the image data with bubble ids
    for z in range(dims[2]):
        for y in range(dims[1]):
            for x in range(dims[0]):
                dataset.SetScalarComponentFromDouble(x, y, z, 0, labeled_array[x,y,z])

    ## Soumya: these paths should be connected to your initial command line args
    np.save(outfile + '/bubbled_out.npy', labeled_array)

    writer = vtk.vtkXMLImageDataWriter()
    writer.SetFileName(out_file)
    if vtk.VTK_MAJOR_VERSION <= 5:
        writer.SetInputConnection(dataset.GetProducerPort())
    else:
        writer.SetInputData(dataset)
    writer.Write()

    # write out a vti file with the density information for visualization
    ## Soumya: these paths should be connected to your initial command line args    
    out_file = outfile + '/original_out.vti'
    origin = [bound_min[0]+xdel/2.0,bound_min[1]+ydel/2.0,bound_min[2]+zdel/2.0]
    dimensions = nbins
    exts = [0,dimensions[0]-1,0,dimensions[1]-1,0,dimensions[2]-1]
    spacing = [xdel,ydel,zdel]
    dataset = getVtkImageData(origin,dimensions,exts,spacing)

    dataset.AllocateScalars(vtk.VTK_DOUBLE, 1)

    dims = dataset.GetDimensions()

    # Fill every entry of the image data with density information
    for z in range(dims[2]):
        for y in range(dims[1]):
            for x in range(dims[0]):
                dataset.SetScalarComponentFromDouble(x, y, z, 0, H[x,y,z])

    writer = vtk.vtkXMLImageDataWriter()
    writer.SetFileName(out_file)
    if vtk.VTK_MAJOR_VERSION <= 5:
        writer.SetInputConnection(dataset.GetProducerPort())
    else:
        writer.SetInputData(dataset)
    writer.Write()


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
create_histogram_and_detect_bubble(inputstring, outFile, sample_fraction, bound_min, bound_max, nbins)

## the 'sample_fraction' is not being used right now for this script. 

## Sample Command:
# vtkpython bubble_detection_mfix_fcc.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_detection_mfix_fcc.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'
