##  Code for computing Mutual information and Entropy using VTK and MPI
##  author: Ayan Biswas
##  contact: ayan@lanl.gov
##
import vtk
import numpy as np
import sys
import os
import math
import time
import scipy.cluster.hierarchy as hc
from matplotlib import pyplot as plt
from mpi4py import MPI

def read_vti_file(in_file):
    print("Read the vti file", in_file)
    in_t = time.time()
    # reader = vtk.vtkDataSetReader()
    reader = vtk.vtkXMLImageDataReader()
    reader.SetFileName(in_file)
    reader.Update()
    data = reader.GetOutput()
    return data

def read_vti_fileP(in_file):
    # print("Read the vti file", in_file)
    ## read in parallel
    reader = vtk.vtkXMLImageDataReader()
    reader.SetFileName(in_file)
    reader.UpdatePiece(comm.Get_rank(), comm.Get_size(), 0)
    reader.Update()
    data = reader.GetOutput()
    return data

def create_histogram(data, var, nbins):
    data.GetPointData().SetActiveScalars(var)
    histogram = vtk.vtkImageHistogram()
    histogram.SetInputData(data)
    ranges = data.GetPointData().GetArray(var).GetRange()
    #nbins = 10
    bin_spacing = (ranges[1] - ranges[0]) / nbins
    bin_origin = ranges[0] + (ranges[1] - ranges[0]) / nbins / 2.0
    #print(ranges, bin_spacing, bin_origin)
    histogram.SetBinOrigin(bin_origin)
    histogram.SetBinSpacing(bin_spacing)
    histogram.SetNumberOfBins(nbins)
    histogram.Update()
    return histogram


def create_histogramP(data,cur_var,nbins=10):
    data.GetPointData().SetActiveScalars(cur_var)
    dims = data.GetDimensions()
    extent = data.GetExtent()

    loc_ext_max = np.zeros(3, dtype='int64')
    loc_ext_max[0] = extent[1]
    loc_ext_max[1] = extent[3]
    loc_ext_max[2] = extent[5]
    glob_ext_max = np.zeros(3, dtype='int64')

    ## keep the max of max to get the global boundary

    comm.Reduce(loc_ext_max, glob_ext_max, op=MPI.MAX, root=0)

    comm.Bcast(glob_ext_max, root=0)

    ## check whether the current block is sharing a boundary edge
    xbound = 0
    ybound = 0
    zbound = 0
    if loc_ext_max[0] == glob_ext_max[0]:
        xbound = 1
    if loc_ext_max[1] == glob_ext_max[1]:
        ybound = 1
    if loc_ext_max[2] == glob_ext_max[2]:
        zbound = 1
    # print(extent,dims,xbound,ybound,zbound)

    cur_proc_extent = np.asarray(extent)
    cur_dims = np.asarray(dims)

    activeScalar = data.GetPointData().GetArray(cur_var)
    srange = activeScalar.GetRange()

    ## now send the global range to all
    glob_min = np.zeros(1)
    glob_max = np.zeros(1)
    ## keep the min of min
    comm.Reduce(np.array(srange[0]), glob_min, op=MPI.MIN, root=0)

    ## keep the max of max
    comm.Reduce(np.array(srange[1]), glob_max, op=MPI.MAX, root=0)

    comm.Bcast(glob_min, root=0)

    comm.Bcast(glob_max, root=0)

    # print("collected range:",glob_min,glob_max)

    ## reduce the boundary layer if the proc is not holding the data of the boundary
    if xbound == 0:
        cur_proc_extent[1] = cur_proc_extent[1] - 1
        cur_dims[0] = cur_dims[0] - 1
    if ybound == 0:
        cur_proc_extent[3] = cur_proc_extent[3] - 1
        cur_dims[1] = cur_dims[1] - 1
    if zbound == 0:
        cur_proc_extent[5] = cur_proc_extent[5] - 1
        cur_dims[2] = cur_dims[2] - 1

    ## extract the region such that there are no overlaps
    voi = vtk.vtkPExtractVOI()
    # voi.SetInputConnection(reader.GetOutputPort())
    voi.SetInputData(data)
    voi.SetVOI(cur_proc_extent[0], cur_proc_extent[1], cur_proc_extent[2], cur_proc_extent[3], cur_proc_extent[4],
               cur_proc_extent[5])
    voi.Update()
    data_voi = voi.GetOutput()
    data_voi.GetPointData().SetActiveScalars(cur_var)

    dataToProcess = data_voi

    # create the local histogram
    histogram = vtk.vtkImageHistogram()
    histogram.SetInputData(dataToProcess)
    # ranges = data.GetPointData().GetArray(cur_var).GetRange()
    ranges = np.zeros(2)
    ranges[0] = glob_min
    ranges[1] = glob_max
    # print("global ranges:",ranges)
    bin_spacing = (ranges[1] - ranges[0]) / nbins
    bin_origin = ranges[0] + (ranges[1] - ranges[0]) / nbins / 2.0
    # print(ranges, bin_spacing, bin_origin)
    histogram.SetBinOrigin(bin_origin)
    histogram.SetBinSpacing(bin_spacing)
    histogram.SetNumberOfBins(nbins)
    histogram.Update()

    # local entropy
    entropy = 0
    tot_N = 1.0 * histogram.GetTotal()
    for j in range(histogram.GetNumberOfBins()):
        cur_val = histogram.GetHistogram().GetTuple1(j) / tot_N
        # print(histogram.GetHistogram().GetTuple1(j))
        if cur_val > 0:
            # print( ">0",histogram.GetHistogram().GetTuple1(j), cur_val,cur_val* math.log(cur_val, 2) )
            entropy = entropy - cur_val * math.log(cur_val, 2)
    # print("Local Entropy = ", entropy)

    local_hist = np.zeros(nbins)
    global_hist = np.zeros(nbins)

    local_N = np.zeros(1)
    global_N = np.zeros(1)

    local_N[0] = tot_N

    for j in range(nbins):
        local_hist[j] = histogram.GetHistogram().GetTuple1(j)
    # print(local_hist)

    comm.Reduce(local_hist, global_hist, op=MPI.SUM, root=0)

    comm.Reduce(local_N, global_N, op=MPI.SUM, root=0)

    comm.Bcast(global_N, root=0)

    comm.Bcast(global_hist, root=0)

    return global_N, global_hist


def create_2d_histogram(data, var1, var2, nbins):
    hist2d = vtk.vtkExtractHistogram2D()
    table = vtk.vtkTable()

    arrX = data.GetPointData().GetArray(var1)
    arrY = data.GetPointData().GetArray(var2)
    table.AddColumn(arrX)
    table.AddColumn(arrY)

    hist2d.SetInputData(table)
    hist2d.SetNumberOfBins(nbins,nbins)

    hist2d.Update()
    return hist2d

def get_global_range(data, cur_var):
    data.GetPointData().SetActiveScalars(cur_var)
    dims = data.GetDimensions()
    extent = data.GetExtent()

    loc_ext_max = np.zeros(3, dtype='int64')
    loc_ext_max[0] = extent[1]
    loc_ext_max[1] = extent[3]
    loc_ext_max[2] = extent[5]
    glob_ext_max = np.zeros(3, dtype='int64')

    ## keep the max of max to get the global boundary

    comm.Reduce(loc_ext_max, glob_ext_max, op=MPI.MAX, root=0)

    comm.Bcast(glob_ext_max, root=0)

    activeScalar = data.GetPointData().GetArray(cur_var)
    srange = activeScalar.GetRange()

    ## now send the global range to all
    glob_min = np.zeros(1)
    glob_max = np.zeros(1)
    ## keep the min of min
    comm.Reduce(np.array(srange[0]), glob_min, op=MPI.MIN, root=0)

    ## keep the max of max
    comm.Reduce(np.array(srange[1]), glob_max, op=MPI.MAX, root=0)

    comm.Bcast(glob_min, root=0)

    comm.Bcast(glob_max, root=0)

    return glob_min, glob_max

def get_data_voi(data):
        dims = data.GetDimensions()
        extent = data.GetExtent()

        loc_ext_max = np.zeros(3, dtype='int64')
        loc_ext_max[0] = extent[1]
        loc_ext_max[1] = extent[3]
        loc_ext_max[2] = extent[5]
        glob_ext_max = np.zeros(3, dtype='int64')

        ## keep the max of max to get the global boundary

        comm.Reduce(loc_ext_max, glob_ext_max, op=MPI.MAX, root=0)

        comm.Bcast(glob_ext_max, root=0)

        ## check whether the current block is sharing a boundary edge
        xbound = 0
        ybound = 0
        zbound = 0
        if loc_ext_max[0] == glob_ext_max[0]:
            xbound = 1
        if loc_ext_max[1] == glob_ext_max[1]:
            ybound = 1
        if loc_ext_max[2] == glob_ext_max[2]:
            zbound = 1
        # print(extent,dims,xbound,ybound,zbound)

        cur_proc_extent = np.asarray(extent)
        cur_dims = np.asarray(dims)

        ## reduce the boundary layer if the proc is not holding the data of the boundary
        if xbound == 0:
            cur_proc_extent[1] = cur_proc_extent[1] - 1
            cur_dims[0] = cur_dims[0] - 1
        if ybound == 0:
            cur_proc_extent[3] = cur_proc_extent[3] - 1
            cur_dims[1] = cur_dims[1] - 1
        if zbound == 0:
            cur_proc_extent[5] = cur_proc_extent[5] - 1
            cur_dims[2] = cur_dims[2] - 1

        ## extract the region such that there are no overlaps
        voi = vtk.vtkPExtractVOI()
        # voi.SetInputConnection(reader.GetOutputPort())
        voi.SetInputData(data)
        voi.SetVOI(cur_proc_extent[0], cur_proc_extent[1], cur_proc_extent[2], cur_proc_extent[3], cur_proc_extent[4],
                   cur_proc_extent[5])
        voi.Update()
        data_voi = voi.GetOutput()
        return data_voi

def create_2d_histogramP(data, var1, var2, nbins):
    hist2d = vtk.vtkExtractHistogram2D()
    ## get only the voi that is needed
    data_voi = get_data_voi(data)
    data_to_process = data_voi
    table = vtk.vtkTable()
    arrX = data_to_process.GetPointData().GetArray(var1)
    arrY = data_to_process.GetPointData().GetArray(var2)
    table.AddColumn(arrX)
    table.AddColumn(arrY)

    hist2d.SetInputData(table)
    # hist2d.SetInputData(data)
    hist2d.SetNumberOfBins(nbins,nbins)
    hist2d.SetUseCustomHistogramExtents(True)
    ## now get the global extent
    glob_min1, glob_max1 = get_global_range(data_to_process, var1)
    glob_min2, glob_max2 = get_global_range(data_to_process, var2)

    a = np.zeros(4)
    a[0] = glob_min1
    a[1] = glob_max1
    a[2] = glob_min2
    a[3] = glob_max2
    hist2d.SetCustomHistogramExtents(a)

    hist2d.Update()

    ext_ptr = hist2d.GetHistogramExtents()

    b = vtk.vtkDoubleArray()
    b.SetNumberOfTuples(4)
    b.SetVoidArray(ext_ptr, 1, 1)

    image_data = hist2d.GetOutputHistogramImage()
    nbins_2d = hist2d.GetNumberOfBins()

    local_N = np.zeros(1)
    global_N = np.zeros(1)

    tot_N = 0
    for i in range(nbins_2d[0]):
        for j in range(nbins_2d[1]):
            tot_N = tot_N + image_data.GetScalarComponentAsFloat(i,j,0,0)
                #print(image_data.GetScalarComponentAsFloat(i,j,0,0))

    local_N[0] = tot_N

    ## now communicate the local 2d histograms across procs

    local_hist = np.zeros(nbins_2d)
    global_hist = np.zeros(nbins_2d)

    for i in range(nbins_2d[0]):
        for j in range(nbins_2d[1]):
            local_hist[i,j] = image_data.GetScalarComponentAsFloat(i, j, 0, 0)

    comm.Reduce(local_hist, global_hist, op=MPI.SUM, root=0)

    comm.Reduce(local_N, global_N, op=MPI.SUM, root=0)

    return hist2d, global_N, global_hist


def compute_miP(data, var1, var2):
    nbins = 10
    hist2d, global_N, global_hist = create_2d_histogramP(data, var1, var2, nbins)

    image_data = hist2d.GetOutputHistogramImage()
    nbins_2d = hist2d.GetNumberOfBins()

    tot_N1, hist1 = create_histogramP(data, var1, nbins_2d[0])
    tot_N2, hist2 = create_histogramP(data, var2, nbins_2d[1])

    #print(image_data)
    #print(nbins_2d)
    mi = 0
    if rank==0:
        # print(nbins_2d,tot_N1,tot_N2, global_N, global_hist)
        for i in range(nbins):
            p_x = hist1[i] / tot_N1
            for j in range(nbins):
                p_xy = global_hist[i,j] / global_N
                p_y = hist2[j] / tot_N2
                if p_xy > 0:
                    # print( ">0",histogram.GetHistogram().GetTuple1(j), cur_val,cur_val* math.log(cur_val, 2) )
                    mi = mi + p_xy * math.log(p_xy/p_x/p_y, 2)
        print("Mutual Information = ", mi)
    return mi

def getGlobalRange(data,cur_var):
    activeScalar = data.GetCellData().GetArray(cur_var)
    srange = activeScalar.GetRange()
    # srange = data.GetScalarRange()  # needs Update() before!
    # print("Range", srange)

    ## now send the global range to all
    glob_min = np.zeros(1)
    glob_max = np.zeros(1)
    ## keep the min of min
    comm.Reduce(np.array(srange[0]), glob_min, op=MPI.MIN, root=0)

    ## keep the max of max
    comm.Reduce(np.array(srange[1]), glob_max, op=MPI.MAX, root=0)

    comm.Bcast(glob_min, root=0)

    comm.Bcast(glob_max, root=0)

    print("collected range:", glob_min, glob_max)
    glob_range = np.zeros(2)
    glob_range[0] = glob_min
    glob_range[1] = glob_max
    return glob_range

def compute_entropy(data, var):
    nbins = 10
    histogram = create_histogram(data,var,nbins)

    # entropy
    entropy = 0
    tot_N = 1.0 * histogram.GetTotal()
    for j in range(histogram.GetNumberOfBins()):
        cur_val = histogram.GetHistogram().GetTuple1(j) / tot_N
        #print(histogram.GetHistogram().GetTuple1(j))
        if cur_val > 0:
            entropy = entropy - cur_val * math.log(cur_val, 2)
    # print("Entropy = ", entropy)
    return entropy

def compute_entropyP(data, var):
    nbins = 10
    global_N, global_hist = create_histogramP(data, var, nbins)
    entropy = 0
    if rank == 0:
        # print( global_hist)
        ## compute entropy

        tot_N = global_N
        print("tot_n: ", tot_N)
        for j in range(nbins):
            cur_val = global_hist[j] / tot_N
            # print(histogram.GetHistogram().GetTuple1(j))
            if cur_val > 0:
                # print( ">0",histogram.GetHistogram().GetTuple1(j), cur_val,cur_val* math.log(cur_val, 2) )
                entropy = entropy - cur_val * math.log(cur_val, 2)
        print("Global Entropy = ", entropy)
    # print("Entropy = ", entropy)
    return entropy


def llf(id):
    if id < n:
        return str(id)


comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# print("rank",rank)
# print("size",size)

file_name1 = "/Users/biswas/Work/y31_data/"
file_name2 = ".vti"
name_string = 'pv_insitu_300x300x300_47631'
entropy_vals = []

var1 = 'mat'
var2 = 'v02'

input_file = file_name1 + name_string + file_name2

data_in = read_vti_fileP(input_file)

ent1 = compute_entropyP(data_in, var1)
#print("Entropy = ", ent1)

ent2 = compute_entropyP(data_in, var2)
#print("Entropy = ", ent2)

mi = compute_miP(data_in, var1, var2)

global variable_names
variable_names = ['prs', 'rho', 'v02', 'v03', 'snd', 'mat', 'grd']
#variable_names = ['prs', 'rho', 'v02']

mi_mat = []
for i in range(np.shape(variable_names)[0]):
    for j in range(i+1, np.shape(variable_names)[0]):
        var1 = variable_names[i]
        var2 = variable_names[j]
        mi = compute_miP(data_in, var1, var2)
#        if rank == 0:
#            print(mi)
        mi_mat.append(mi)

if rank == 0:
    Z = hc.linkage(np.asarray(mi_mat)[:,0])

    fig = plt.figure()
    dn = hc.dendrogram(Z,labels=variable_names)
    plt.savefig('hc.pdf')
#    plt.show()

