##  Code for computing 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
from mpi4py import MPI
from vtk.util import numpy_support

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

## select the input vti file
file_name1 = "/media/biswas/ssd2/AsteriodImpact/Res300x300x300/oceans11.lanl.gov/deepwaterimpact/yA31/300x300x300-AllScalars_resolution/"
file_name2 = ".vti"
name_string = 'pv_insitu_300x300x300_43537'
input = file_name1 + name_string + file_name2

## select the variable
var_name = 'prs'

## read in parallel
reader = vtk.vtkXMLImageDataReader()
reader.SetFileName(input)
reader.UpdatePiece(comm.Get_rank(), comm.Get_size(), 0)

reader.Update()
data = reader.GetOutput()
cur_var = var_name
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)
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()

# 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)

if rank == 0:
    #print( global_hist)
    ## compute entropy
    entropy = 0
    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)
