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 settings

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 sample_data(in_file, outfile, x=0.1):
    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))
    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()

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

    bound_min = np.min(feat_arr,axis=0)
    bound_max = np.max(feat_arr, axis=0)

    print('ranges:',bound_min,bound_max)

    mult = 2
    nbins = [mult*64,mult*8,mult*64]

    H, edges = np.histogramdd(feat_arr, bins=nbins)
    #print(H,feat_arr)
    # try thresholding
    th = np.max(H)/3 # heuristic
    H[H < th] = 0

    print('Histo:',H.shape, edges[0].size, edges[1].size, edges[2].size, edges[0][0])
    #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 density information for visualization
    out_file = 'density_out_th.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()

    # now write out a sampled data from the original locations
    # no samples come from the empty bins
    print('Sarting the sampling procedure:')

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

    filtered = vtk.vtkPolyData()
    Points = vtk.vtkPoints()

    ## insert the points

    tot_pts_after_th = np.sum(H)

    for z in range(nbins[2]):
        for y in range(nbins[1]):
            for x in range(nbins[0]):
                if H[x,y,z]>0:
                    expected_val = H[x,y,z]*target_points/(1.0*tot_pts_after_th)
                    cur_count = int(expected_val)
                    frac_part = expected_val-cur_count
                    # generate a random number to check to see if one more sample is needed or not
                    if random.random() < frac_part:
                        cur_count = cur_count+1
                    if cur_count>0:
                        for i in range(cur_count):
                            samp_x = random.uniform(edges[0][x+1],edges[0][x])
                            samp_y = random.uniform(edges[1][y + 1], edges[1][y])
                            samp_z = random.uniform(edges[2][z + 1], edges[2][z])
                            Points.InsertNextPoint(samp_x,samp_y,samp_z)


    filtered.SetPoints(Points)
    multiblock.SetBlock(0, filtered)

    w = vtk.vtkXMLMultiBlockDataWriter()
    w.SetInputData(multiblock)
    filename = "sampled_"+outfile
    w.SetFileName(filename)
    w.Write()
    #
    # print("total non-zero pieces:",pc_cnt)


def compute_entropy_from_histogram(in_file, outfile, fraction=0.1):
    # H, edges = np.histogramdd(feat_arr, bins=nbins)
    bound_min = np.load('npy_files/bound_min_' + outfile + '.npy')
    bound_max = np.load('npy_files/bound_max_' + outfile + '.npy')
    nbins = np.load('npy_files/nbins_' + outfile + '.npy')
    H = np.load('npy_files/hist_count_' + outfile + '.npy')
    edges = np.load('npy_files/hist_edges_' + outfile + '.npy')
    print(np.shape(H), np.shape(edges))
    tot_pts = np.sum(H)
    print("Total size:", tot_pts)
    target_points = int(tot_pts * fraction)

    # compute entropy from the histogram
    entropy_val = 0

    for z in range(nbins[2]):
        for y in range(nbins[1]):
            for x in range(nbins[0]):
                if H[x,y,z]>0:
                    val = H[x,y,z]/(1.0*tot_pts)
                    entropy_val = entropy_val - val*math.log(val,2)

    return entropy_val



def sample_data_from_histogram(in_file, outfile, fraction=0.1):
    # H, edges = np.histogramdd(feat_arr, bins=nbins)
    bound_min = np.load('npy_files/bound_min_' + outfile + '.npy')
    bound_max = np.load('npy_files/bound_max_' + outfile + '.npy')
    nbins = np.load('npy_files/nbins_' + outfile + '.npy')
    H = np.load('npy_files/hist_count_' + outfile + '.npy')
    edges = np.load('npy_files/hist_edges_' + outfile + '.npy')
    print(np.shape(H), np.shape(edges))
    tot_pts = np.sum(H)
    print("Total size:", tot_pts)
    target_points = int(tot_pts * fraction)
    #print(H,feat_arr)
    # try thresholding
    th = np.max(H)/3 # heuristic
    H[H < th] = 0

    print('Histo:',H.shape, edges[0].size, edges[1].size, edges[2].size, edges[0][0])
    #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()


    # now write out a sampled data from the original locations
    # no samples come from the empty bins
    print('Sarting the sampling procedure:')

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

    filtered = vtk.vtkPolyData()
    Points = vtk.vtkPoints()

    ## insert the points

    tot_pts_after_th = np.sum(H)

    for z in range(nbins[2]):
        for y in range(nbins[1]):
            for x in range(nbins[0]):
                if H[x,y,z]>0:
                    expected_val = H[x,y,z]*target_points/(1.0*tot_pts_after_th)
                    cur_count = int(expected_val)
                    frac_part = expected_val-cur_count
                    # generate a random number to check to see if one more sample is needed or not
                    if random.random() < frac_part:
                        cur_count = cur_count+1
                    if cur_count>0:
                        for i in range(cur_count):
                            samp_x = random.uniform(edges[0][x+1],edges[0][x])
                            samp_y = random.uniform(edges[1][y + 1], edges[1][y])
                            samp_z = random.uniform(edges[2][z + 1], edges[2][z])
                            Points.InsertNextPoint(samp_x,samp_y,samp_z)


    filtered.SetPoints(Points)
    multiblock.SetBlock(0, filtered)

    w = vtk.vtkXMLMultiBlockDataWriter()
    w.SetInputData(multiblock)
    directory = "sampled_"+str(fraction)
    if not os.path.exists(directory):
        os.makedirs(directory)
    filename = directory+"/sampled_"+outfile+'.vtm'
    w.SetFileName(filename)
    w.Write()
    #
    # print("total non-zero pieces:",pc_cnt)

def create_histogram(in_file, outfile, x=0.1, bound_min=None, bound_max=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))
    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()

    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)

    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]]])
    np.save('npy_files/hist_count_'+outfile,H)
    np.save('npy_files/hist_edges_'+outfile, edges)
    np.save('npy_files/bound_min_'+outfile,bound_min)
    np.save('npy_files/bound_max_'+outfile, bound_max)
    np.save('npy_files/nbins_'+outfile,nbins)
    print(edges)

def detect_bubbles(in_file, outfile, x=0.1):
    #H, edges = np.histogramdd(feat_arr, bins=nbins)
    bound_min = np.load('npy_files/bound_min_'+outfile+'.npy')
    bound_max = np.load('npy_files/bound_max_'+outfile+'.npy')
    nbins = np.load('npy_files/nbins_'+outfile+'.npy')
    H = np.load('npy_files/hist_count_'+outfile+'.npy')
    edges = np.load('npy_files/hist_edges_'+outfile+'.npy')
    print(np.shape(H),np.shape(edges))
    #print(edges)
    # try thresholding

    th = np.max(H)/3 # heuristic
    H[H < th] = 0

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

    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 = 30
    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
    out_file = 'bubbled_data/bubbled_'+outfile+'.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])

    np.save('npy_files/bubbled_' + outfile, 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
    out_file = 'bubbled_data/original_'+outfile+'.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()

def detect_bubbles_refined(in_file, outfile, outfile2):
    ###### TO-DO ############
    #H, edges = np.histogramdd(feat_arr, bins=nbins)
    bound_min = np.load('npy_files/bound_min_'+outfile+'.npy')
    bound_max = np.load('npy_files/bound_max_'+outfile+'.npy')
    nbins = np.load('npy_files/nbins_'+outfile+'.npy')
    H = np.load('npy_files/hist_count_'+outfile+'.npy')
    edges = np.load('npy_files/hist_edges_'+outfile+'.npy')
    labeled_array = np.load('npy_files/bubbled_' + outfile+'.npy')
    labeled_array2 = np.load('npy_files/bubbled_' + outfile2 + '.npy')
    print(np.shape(H),np.shape(edges),np.shape(labeled_array))

    num_features = np.max(labeled_array)
    print('Number of bubbles:',num_features,np.max(labeled_array2))

    for i in range(1,num_features):
        # find locations of this bubbles for current timestep
        locs = np.where(labeled_array==i)
        # what are values at these locations
        next_vals = labeled_array2[locs]
        unique, counts = np.unique(next_vals, return_counts=True)
        uniq_val = unique[np.argmax(counts)]
        if uniq_val==0:
            uniq_val = unique[np.where(counts==np.partition(counts.flatten(), -2)[-2])][0]

        if uniq_val == 0:
            print('We have a problem!!')
            exit()

        labeled_array2[labeled_array2==uniq_val]=i


        print(uniq_val)

    # write out a vti file with the density information for visualization
    xdel = edges[0][1]-edges[0][0]
    ydel = edges[1][1]-edges[1][0]
    zdel = edges[2][1]-edges[2][0]

    out_file = 'bubbled_data/bubbled_refined_' + outfile + '.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, labeled_array2[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()


def track_bubbles_back(in_file, outfile, outfile2, corrected_outfile=None):
    # global feat_count
    # H, edges = np.histogramdd(feat_arr, bins=nbins)
    bound_min = np.load('npy_files/bound_min_' + outfile + '.npy')
    bound_max = np.load('npy_files/bound_max_' + outfile + '.npy')
    nbins = np.load('npy_files/nbins_' + outfile + '.npy')
    H = np.load('npy_files/hist_count_' + outfile + '.npy')
    edges = np.load('npy_files/hist_edges_' + outfile + '.npy')
    labeled_array = np.load('npy_files/bubbled_' + outfile + '.npy')
    if corrected_outfile is None:
        labeled_array = np.load('npy_files/bubbled_' + outfile + '.npy')
    else:
        labeled_array = corrected_outfile
    labeled_array2 = np.load('npy_files/bubbled_' + outfile2 + '.npy')
    # print(np.shape(H),np.shape(edges),np.shape(labeled_array))

    num_features1 = np.max(labeled_array)
    num_features2 = np.max(labeled_array2)
    print('Number of bubbles:', num_features1, num_features2)
    if settings.feat_count == 0:
        settings.feat_count = settings.feat_count + num_features1

    if num_features2 == num_features1:
        ## we are in the easy scenario
        print('should be normal scenario')

        for i in range(1, num_features1 + 1):
            # find locations of this bubbles for current timestep
            locs = np.where(labeled_array == i)
            # what are values at these locations
            next_vals = labeled_array2[locs]
            unique, counts = np.unique(next_vals, return_counts=True)
            uniq_val = unique[np.argmax(counts)]
            if uniq_val == 0:
                uniq_val = unique[np.where(counts == np.partition(counts.flatten(), -2)[-2])][0]

            if uniq_val == 0:
                print('We have a problem!!')
                exit()

            labeled_array2[labeled_array2 == uniq_val] = i

            print(uniq_val)

    elif num_features2 > num_features1:
        print("handle split")
        # feat_count = feat_count + num_features2-num_features1

        covered_list1 = []
        covered_list2 = []

        for k in range(1, num_features2 + 1):
            covered_list1.append(k)

        ## we handle the easy scenario first
        for i in range(1, num_features1 + 1):
            # find locations of this bubbles for current timestep
            locs = np.where(labeled_array == i)
            # what are values at these locations
            next_vals = labeled_array2[locs]
            unique, counts = np.unique(next_vals, return_counts=True)
            uniq_val = unique[np.argmax(counts)]

            if uniq_val == 0:
                uniq_val = unique[np.where(counts == np.partition(counts.flatten(), -2)[-2])][0]

            if uniq_val == 0:
                print('We have a problem!!')
                exit()

            covered_list2.append(uniq_val)

            labeled_array2[labeled_array2 == uniq_val] = i

            print(uniq_val)

        # now see which features are new
        new_feats = list(set(covered_list1) - set(covered_list2))
        print('new features:', new_feats)
        for j in new_feats:
            labeled_array2[labeled_array2 == j] = settings.feat_count + 1
            settings.feat_count = settings.feat_count + 1


    else:
        print('handle merge')

        ## we start with the easy scenario
        for i in range(1, num_features1 + 1):
            # find locations of this bubbles for current timestep
            locs = np.where(labeled_array == i)
            # what are values at these locations
            next_vals = labeled_array2[locs]
            unique, counts = np.unique(next_vals, return_counts=True)
            uniq_val = unique[np.argmax(counts)]

            if uniq_val == 0:
                if np.size(uniq_val) == 1:
                    continue
                else:
                    uniq_val = unique[np.where(counts == np.partition(counts.flatten(), -2)[-2])][0]

            if uniq_val == 0:
                print('We have a problem!!')
                exit()

            labeled_array2[labeled_array2 == uniq_val] = i

            print(uniq_val)

    # write out a vti file with the density information for visualization
    xdel = edges[0][1] - edges[0][0]
    ydel = edges[1][1] - edges[1][0]
    zdel = edges[2][1] - edges[2][0]

    out_file = 'bubbled_data/tracked_' + outfile2 + '.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, labeled_array2[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()

    return labeled_array2


def track_bubbles_back2(in_file, outfile, outfile2,corrected_outfile=None):
    #global feat_count
    #H, edges = np.histogramdd(feat_arr, bins=nbins)
    bound_min = np.load('npy_files/bound_min_'+outfile+'.npy')
    bound_max = np.load('npy_files/bound_max_'+outfile+'.npy')
    nbins = np.load('npy_files/nbins_'+outfile+'.npy')
    H = np.load('npy_files/hist_count_'+outfile+'.npy')
    edges = np.load('npy_files/hist_edges_'+outfile+'.npy')
    labeled_array = np.load('npy_files/bubbled_' + outfile+'.npy')
    if corrected_outfile is None:
        labeled_array = np.load('npy_files/bubbled_' + outfile + '.npy')
    else:
        labeled_array = corrected_outfile
    labeled_array2 = np.load('npy_files/bubbled_' + outfile2 + '.npy')
    #print(np.shape(H),np.shape(edges),np.shape(labeled_array))

    num_features1 = np.max(labeled_array)
    num_features2 = np.max(labeled_array2)
    print('Number of bubbles:', num_features1, num_features2)
    if settings.feat_count == 0:
        settings.feat_count = settings.feat_count + num_features1

    if num_features2==num_features1:
        ## we are in the easy scenario
        print('should be normal scenario')

        for i in range(1,num_features2+1):
            # find locations of this bubbles for current timestep
            locs = np.where(labeled_array2==i)
            # what are values at these locations
            prev_vals = labeled_array[locs]  ## actually prev_vals
            unique, counts = np.unique(prev_vals, return_counts=True)
            uniq_val = unique[np.argmax(counts)]
            if uniq_val==0:
                if np.size(unique)>=2:
                    uniq_val = unique[np.where(counts==np.partition(counts.flatten(), -2)[-2])][0]
                else:
                    print('a new feature is born')
                    labeled_array2[labeled_array2 ==i] = settings.feat_count + 1
                    settings.feat_count = settings.feat_count + 1

            if uniq_val == 0:
                print('We have a problem!!')
                exit()

            labeled_array2[labeled_array2==i]=uniq_val


            print(uniq_val)

    elif num_features2>num_features1:
        print("handle split")
        #feat_count = feat_count + num_features2-num_features1

        covered_list1 = []
        covered_list2 = []

        for k in range(1,num_features2+1):
            covered_list1.append(k)

        # find the feat ids for the last time step
        unique1 = list(np.unique(labeled_array))
        unique1.remove(0)
        unique1 = np.asarray(unique1)
        print(unique1)

        ## we handle the easy scenario first
        for i in unique1:
            # find locations of this bubbles for current timestep
            locs = np.where(labeled_array == i)
            # what are values at these locations
            next_vals = labeled_array2[locs]
            unique, counts = np.unique(next_vals, return_counts=True)
            uniq_val = unique[np.argmax(counts)]

            if uniq_val == 0 and np.size(unique)>=2:
                uniq_val = unique[np.where(counts == np.partition(counts.flatten(), -2)[-2])][0]

            if uniq_val == 0:
                print('destruction of a bubble!!')
            else:
                covered_list2.append(uniq_val)

                labeled_array2[labeled_array2 == uniq_val] = i

                print(uniq_val)

        # now see which features are new
        new_feats = list(set(covered_list1) - set(covered_list2))
        print('new features:',new_feats)
        for j in new_feats:
            labeled_array2[labeled_array2 == j] = settings.feat_count+1
            settings.feat_count = settings.feat_count + 1


    else:
        print('handle merge: as like inverse of split')

        ## we start with the current time step and look behind
        for i in range(1, num_features2 + 1):
            # find locations of this bubbles for current timestep
            locs = np.where(labeled_array2 == i)
            # what are values at these locations
            prev_vals = labeled_array[locs]  ## actually prev_vals
            unique, counts = np.unique(prev_vals, return_counts=True)
            uniq_val = unique[np.argmax(counts)]

            if uniq_val == 0:
                uniq_val = unique[np.where(counts == np.partition(counts.flatten(), -2)[-2])][0]


            if uniq_val == 0:
                print('We have a problem!!')
                exit()

            labeled_array2[labeled_array2 == i] = uniq_val

            print(uniq_val)




    # write out a vti file with the density information for visualization
    xdel = edges[0][1]-edges[0][0]
    ydel = edges[1][1]-edges[1][0]
    zdel = edges[2][1]-edges[2][0]

    out_file = 'bubbled_data/tracked_' + outfile2 + '.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, labeled_array2[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()

    return labeled_array2


def track_bubbles(in_file, outfile, outfile2,corrected_outfile=None):
    #global feat_count
    #H, edges = np.histogramdd(feat_arr, bins=nbins)
    bound_min = np.load('npy_files/bound_min_'+outfile+'.npy')
    bound_max = np.load('npy_files/bound_max_'+outfile+'.npy')
    nbins = np.load('npy_files/nbins_'+outfile+'.npy')
    H = np.load('npy_files/hist_count_'+outfile+'.npy')
    edges = np.load('npy_files/hist_edges_'+outfile+'.npy')
    labeled_array = np.load('npy_files/bubbled_' + outfile+'.npy')
    if corrected_outfile is None:
        labeled_array = np.load('npy_files/bubbled_' + outfile + '.npy')
    else:
        labeled_array = corrected_outfile
    labeled_array2 = np.load('npy_files/bubbled_' + outfile2 + '.npy')
    #print(np.shape(H),np.shape(edges),np.shape(labeled_array))

    num_features1 = np.max(labeled_array)
    num_features2 = np.max(labeled_array2)
    print('Number of bubbles:', num_features1, num_features2)
    if settings.feat_count == 0:
        settings.feat_count = settings.feat_count + num_features1

    unique1 = list(np.unique(labeled_array2))
    unique1.remove(0)
    unique1 = np.asarray(unique1)
    print('new feat ids:',unique1)

    for i in unique1:#range(1, num_features2 + 2):
        # find locations of this bubbles for current timestep
        locs = np.where(labeled_array2 == i)
        # what are values at these locations
        prev_vals = labeled_array[locs]  ## actually prev_vals
        unique, counts = np.unique(prev_vals, return_counts=True)
        uniq_val = unique[np.argmax(counts)]
        if uniq_val == 0:
            if np.size(unique) >= 2:
                uniq_val = unique[np.where(counts == np.partition(counts.flatten(), -2)[-2])][0]
                labeled_array2[labeled_array2 == i] = uniq_val
            else:
                print('a new feature is born')
                labeled_array2[labeled_array2 == i] = settings.feat_count + 1
                settings.feat_count = settings.feat_count + 1
        else:
            labeled_array2[labeled_array2 == i] = uniq_val

        # if uniq_val == 0:
        #     print('We have a problem!!')
        #     exit()

        #print(uniq_val,end=' ')

    ## find the feat ids after tracking
    # find the feat ids for the last time step
    unique1 = list(np.unique(labeled_array2))
    unique1.remove(0)
    unique1 = np.asarray(unique1)
    print('current feat ids:',unique1)

    np.save('npy_files/tracked_' + outfile2, labeled_array2)

    # write out a vti file with the density information for visualization
    xdel = edges[0][1]-edges[0][0]
    ydel = edges[1][1]-edges[1][0]
    zdel = edges[2][1]-edges[2][0]

    out_file = 'bubbled_data/tracked_' + outfile2 + '.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, labeled_array2[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()

    return labeled_array2


## first try to get the whole range
machine = platform.system()
# read this pvtp file
if machine == 'Linux':
    path = '/media/biswas/29d42253-dbfd-4e49-bc64-61662de409e0/MFIX/fcc-data/vtm_data/'
elif machine == 'Darwin':
    path = '/Users/biswas/Work/MFIX/fcc/'
else:
    print("Unknown machine. Add to the list")
    exit(0)

file_name2 = ".vtm"
name_string = 'timestep_185' #'test_MFIX_50000'

#in_file = '/Users/biswas/Work/MFIX/fcc/timestep_408.vtm'
in_file = path+name_string+file_name2

sample_fraction = 0.05

#sample_data(in_file,name_string+file_name2,sample_fraction)
#create_histogram(in_file,name_string,sample_fraction)
#detect_bubbles(in_file,name_string,sample_fraction)

# ## 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)
# #
# for i in sorted(os.listdir(input_path)):
#     # read files
#     filename, file_extension = os.path.splitext(i)
#     # print(filename)
#     if file_extension != ".vtm":
#         continue
#     print("opening", i)
#     # write the actual code here
#
#     #sample_fraction = 0.1
#     in_file = input_path + i
#
#     #sample_data(in_file, filename, sample_fraction)
#     #create_histogram(in_file, filename, sample_fraction)
#     detect_bubbles(in_file, filename, sample_fraction)