#include "analysis/uCT.h"
#include "analysis/analysis.h"
#include "analysis/eikonal.h"
#include "analysis/filters.h"
#include "analysis/imfilter.h"


template<class T>
inline int sign( T x )
{
    if ( x==0 )
        return 0;
    return x>0 ? 1:-1;
}


inline float trilinear( float dx, float dy, float dz, float f1, float f2,
    float f3, float f4, float f5, float f6, float f7, float f8 )
{
    double f, dx2, dy2, dz2, h0, h1;
    dx2 = 1.0 - dx;
    dy2 = 1.0 - dy;
    dz2 = 1.0 - dz;
    h0  = ( dx * f2 + dx2 * f1 ) * dy2 + ( dx * f4 + dx2 * f3 ) * dy;
    h1  = ( dx * f6 + dx2 * f5 ) * dy2 + ( dx * f8 + dx2 * f7 ) * dy;
    f   = h0 * dz2 + h1 * dz;
    return ( f );
}


void InterpolateMesh( const Array<float> &Coarse, Array<float> &Fine )
{
	PROFILE_START("InterpolateMesh");

	// Interpolate values from a Coarse mesh to a fine one
	// This routine assumes cell-centered meshes with 1 ghost cell

	// Fine mesh
	int Nx = int(Fine.size(0))-2;
	int Ny = int(Fine.size(1))-2;
	int Nz = int(Fine.size(2))-2;

	// Coarse mesh
	int nx = int(Coarse.size(0))-2;
	int ny = int(Coarse.size(1))-2;
	int nz = int(Coarse.size(2))-2;

	// compute the stride
	int hx = Nx/nx;
	int hy = Ny/ny;
	int hz = Nz/nz;
    ASSERT(nx*hx==Nx);
    ASSERT(ny*hy==Ny);
    ASSERT(nz*hz==Nz);

	// value to map distance between meshes (since distance is in voxels)
	//  usually hx=hy=hz (or something very close)
	//  the mapping is not exact
	//  however, it's assumed the coarse solution will be refined
	//  a good guess is the goal here!
	float mapvalue = sqrt(hx*hx+hy*hy+hz*hz);

	// Interpolate to the fine mesh
	for (int k=-1; k<Nz+1; k++){
        int k0 = floor((k-0.5*hz)/hz);
        int k1 = k0+1;
        int k2 = k0+2;
        float dz = ( (k+0.5) - (k0+0.5)*hz ) / hz;
        ASSERT(k0>=-1&&k0<nz+1&&dz>=0&&dz<=1);
		for (int j=-1; j<Ny+1; j++){
            int j0 = floor((j-0.5*hy)/hy);
            int j1 = j0+1;
            int j2 = j0+2;
            float dy = ( (j+0.5) - (j0+0.5)*hy ) / hy;
            ASSERT(j0>=-1&&j0<ny+1&&dy>=0&&dy<=1);
			for (int i=-1; i<Nx+1; i++){
                int i0 = floor((i-0.5*hx)/hx);
                int i1 = i0+1;
                int i2 = i0+2;
                float dx = ( (i+0.5) - (i0+0.5)*hx ) / hx;
                ASSERT(i0>=-1&&i0<nx+1&&dx>=0&&dx<=1);
                float val = trilinear( dx, dy, dz,
                    Coarse(i1,j1,k1), Coarse(i2,j1,k1), Coarse(i1,j2,k1), Coarse(i2,j2,k1),
                    Coarse(i1,j1,k2), Coarse(i2,j1,k2), Coarse(i1,j2,k2), Coarse(i2,j2,k2) );
                Fine(i+1,j+1,k+1) = mapvalue*val;
			}
		}
	}
	PROFILE_STOP("InterpolateMesh");
}


// Smooth the data using the distance
void smooth( const Array<float>& VOL, const Array<float>& Dist, float sigma, Array<float>& MultiScaleSmooth, fillHalo<float>& fillFloat )
{
    for (size_t i=0; i<VOL.length(); i++) {
		// use exponential weight based on the distance
		float dst = Dist(i);
		float tmp = exp(-(dst*dst)/(sigma*sigma));
		float value = dst>0 ? -1:1;
		MultiScaleSmooth(i) = tmp*VOL(i) + (1-tmp)*value;
	}
	fillFloat.fill(MultiScaleSmooth);
}


// Segment the data
void segment( const Array<float>& data, Array<char>& ID, float tol )
{
    ASSERT(data.size()==ID.size());
    for (size_t i=0; i<data.length(); i++) {
        if ( data(i) > tol )
            ID(i) = 0;
        else
            ID(i) = 1;
    }
}


// Remove disconnected phases
void removeDisconnected( Array<char>& ID, const Domain& Dm )
{
    // Run blob identification to remove disconnected volumes
    BlobIDArray GlobalBlobID;
    DoubleArray SignDist(ID.size());
    DoubleArray Phase(ID.size());
    for (size_t i=0; i<ID.length(); i++) {
        SignDist(i) = (2*ID(i)-1);
        Phase(i) = 1;
    }
    ComputeGlobalBlobIDs( ID.size(0)-2, ID.size(1)-2, ID.size(2)-2,
        Dm.rank_info, Phase, SignDist, 0, 0, GlobalBlobID, Dm.Comm );
    for (size_t i=0; i<ID.length(); i++) {
        if ( GlobalBlobID(i) > 0 )
            ID(i) = 0;
        ID(i) = GlobalBlobID(i);
    }
}


// Solve a level (without any coarse level information)
void solve( const Array<float>& VOL, Array<float>& Mean, Array<char>& ID,
    Array<float>& Dist, Array<float>& MultiScaleSmooth, Array<float>& NonLocalMean, 
    fillHalo<float>& fillFloat, const Domain& Dm, int nprocx )
{
    PROFILE_SCOPED(timer,"solve");
    // Compute the median filter on the sparse array
    Med3D( VOL, Mean );
    fillFloat.fill( Mean );
    segment( Mean, ID, 0.01 );
    // Compute the distance using the segmented volume
	Eikonal3D( Dist, ID, Dm, ID.size(0)*nprocx );
	fillFloat.fill(Dist);
    smooth( VOL, Dist, 2.0, MultiScaleSmooth, fillFloat );
    // Compute non-local mean
	int depth = 5;
	float sigsq=0.1;
	int nlm_count = NLM3D( MultiScaleSmooth, Mean, Dist, NonLocalMean, depth, sigsq);
	fillFloat.fill(NonLocalMean);
}


// Refine a solution from a coarse grid to a fine grid
void refine( const Array<float>& Dist_coarse, 
    const Array<float>& VOL, Array<float>& Mean, Array<char>& ID,
    Array<float>& Dist, Array<float>& MultiScaleSmooth, Array<float>& NonLocalMean, 
    fillHalo<float>& fillFloat, const Domain& Dm, int nprocx, int level )
{
    PROFILE_SCOPED(timer,"refine");
    int ratio[3] = { int(Dist.size(0)/Dist_coarse.size(0)),
                     int(Dist.size(1)/Dist_coarse.size(1)),
                     int(Dist.size(2)/Dist_coarse.size(2)) };
    // Interpolate the distance from the coarse to fine grid
	InterpolateMesh( Dist_coarse, Dist );
    // Compute the median filter on the array and segment
    Med3D( VOL, Mean );
    fillFloat.fill( Mean );
    segment( Mean, ID, 0.01 );
    // If the ID has the wrong distance, set the distance to 0 and run a simple filter to set neighbors to 0
    for (size_t i=0; i<ID.length(); i++) {
        char id = Dist(i)>0 ? 1:0;
        if ( id != ID(i) )
            Dist(i) = 0;
    }
    fillFloat.fill( Dist );
    std::function<float(int,const float*)> filter_1D = []( int N, const float* data )
    {
        bool zero = data[0]==0 || data[2]==0;
        return zero ? data[1]*1e-12 : data[1];
    };
    std::vector<imfilter::BC> BC(3,imfilter::BC::replicate);
    std::vector<std::function<float(int,const float*)>> filter_set(3,filter_1D);
    Dist = imfilter::imfilter_separable<float>( Dist, {1,1,1}, filter_set, BC );
    fillFloat.fill( Dist );
    // Smooth the volume data
    float lambda = 2*sqrt(double(ratio[0]*ratio[0]+ratio[1]*ratio[1]+ratio[2]*ratio[2]));
    smooth( VOL, Dist, lambda, MultiScaleSmooth, fillFloat );
    // Compute non-local mean
	int depth = 3;
	float sigsq = 0.1;
	int nlm_count = NLM3D( MultiScaleSmooth, Mean, Dist, NonLocalMean, depth, sigsq);
	fillFloat.fill(NonLocalMean);
    segment( NonLocalMean, ID, 0.001 );
    for (size_t i=0; i<ID.length(); i++) {
        char id = Dist(i)>0 ? 1:0;
        if ( id!=ID(i) || fabs(Dist(i))<1 )
    	    Dist(i) = 2.0*ID(i)-1.0;
    }
    // Remove disconnected domains
    //removeDisconnected( ID, Dm );
    // Compute the distance using the segmented volume
    if ( level > 0 ) {
    	//Eikonal3D( Dist, ID, Dm, ID.size(0)*nprocx );
        //CalcDist3D( Dist, ID, Dm );
        CalcDistMultiLevel( Dist, ID, Dm );
	    fillFloat.fill(Dist);
    }
}


// Remove regions that are likely noise by shrinking the volumes by dx,
// removing all values that are more than dx+delta from the surface, and then
// growing by dx+delta and intersecting with the original data
void filter_final( Array<char>& ID, Array<float>& Dist,
    fillHalo<float>& fillFloat, const Domain& Dm,
    Array<float>& Mean, Array<float>& Dist1, Array<float>& Dist2 )
{
    PROFILE_SCOPED(timer,"filter_final");
	int rank;
	MPI_Comm_rank(Dm.Comm,&rank);
    int Nx = Dm.Nx-2;
    int Ny = Dm.Ny-2;
    int Nz = Dm.Nz-2;
    // Calculate the distance
    CalcDistMultiLevel( Dist, ID, Dm );
    fillFloat.fill(Dist);
    // Compute the range to shrink the volume based on the L2 norm of the distance
    Array<float> Dist0(Nx,Ny,Nz);
    fillFloat.copy(Dist,Dist0);
    float tmp = 0;
    for (size_t i=0; i<Dist0.length(); i++)
        tmp += Dist0(i)*Dist0(i);
    tmp = sqrt( sumReduce(Dm.Comm,tmp) / sumReduce(Dm.Comm,(float)Dist0.length()) );
    const float dx1 = 0.3*tmp;
    const float dx2 = 1.05*dx1;
    if (rank==0)
        printf("   %0.1f %0.1f %0.1f\n",tmp,dx1,dx2);
    // Update the IDs/Distance removing regions that are < dx of the range
    Dist1 = Dist;
    Dist2 = Dist;
    Array<char> ID1 = ID;
    Array<char> ID2 = ID;
    for (size_t i=0; i<ID.length(); i++) {
        ID1(i) = Dist(i)<-dx1 ? 1:0;
        ID2(i) = Dist(i)> dx1 ? 1:0;
    }
    //Array<float> Dist1 = Dist;
    //Array<float> Dist2 = Dist;
    CalcDistMultiLevel( Dist1, ID1, Dm );
    CalcDistMultiLevel( Dist2, ID2, Dm );
    fillFloat.fill(Dist1);
    fillFloat.fill(Dist2);
    // Keep those regions that are within dx2 of the new volumes
    Mean = Dist;
    for (size_t i=0; i<ID.length(); i++) {
        if ( Dist1(i)+dx2>0 && ID(i)<=0 ) {
            Mean(i) = -1;
        } else if ( Dist2(i)+dx2>0 && ID(i)>0 ) {
            Mean(i) = 1;
        } else {
            Mean(i) = Dist(i)>0 ? 0.5:-0.5;
        }
    }
    // Find regions of uncertainty that are entirely contained within another region
    fillHalo<double> fillDouble(Dm.Comm,Dm.rank_info,Nx,Ny,Nz,1,1,1,0,1);
    fillHalo<BlobIDType> fillInt(Dm.Comm,Dm.rank_info,Nx,Ny,Nz,1,1,1,0,1);
    BlobIDArray GlobalBlobID;
    DoubleArray SignDist(ID.size());
    for (size_t i=0; i<ID.length(); i++)
        SignDist(i) = fabs(Mean(i))==1 ? -1:1;
    fillDouble.fill(SignDist);
    DoubleArray Phase(ID.size());
    Phase.fill(1);
    ComputeGlobalBlobIDs( Nx, Ny, Nz, Dm.rank_info, Phase, SignDist, 0, 0, GlobalBlobID, Dm.Comm );
    fillInt.fill(GlobalBlobID);
    int N_blobs = maxReduce(Dm.Comm,GlobalBlobID.max()+1);
    std::vector<float> mean(N_blobs,0);
    std::vector<int> count(N_blobs,0);
    for (int k=1; k<=Nz; k++) {
        for (int j=1; j<=Ny; j++) {
            for (int i=1; i<=Nx; i++) {
                int id = GlobalBlobID(i,j,k);
                if ( id >= 0 ) {
                    if ( GlobalBlobID(i-1,j,k)<0 ) {
                        mean[id] += Mean(i-1,j,k);
                        count[id]++;
                    }
                    if ( GlobalBlobID(i+1,j,k)<0 ) {
                        mean[id] += Mean(i+1,j,k);
                        count[id]++;
                    }
                    if ( GlobalBlobID(i,j-1,k)<0 ) {
                        mean[id] += Mean(i,j-1,k);
                        count[id]++;
                    }
                    if ( GlobalBlobID(i,j+1,k)<0 ) {
                        mean[id] += Mean(i,j+1,k);
                        count[id]++;
                    }
                    if ( GlobalBlobID(i,j,k-1)<0 ) {
                        mean[id] += Mean(i,j,k-1);
                        count[id]++;
                    }
                    if ( GlobalBlobID(i,j,k+1)<0 ) {
                        mean[id] += Mean(i,j,k+1);
                        count[id]++;
                    }
                }
            }
        }
    }
    mean = sumReduce(Dm.Comm,mean);
    count = sumReduce(Dm.Comm,count);
    for (size_t i=0; i<mean.size(); i++)
        mean[i] /= count[i];
    /*if (rank==0) {
        for (size_t i=0; i<mean.size(); i++)
            printf("%i %0.4f\n",i,mean[i]);
    }*/
    for (size_t i=0; i<Mean.length(); i++) {
        int id = GlobalBlobID(i);
        if ( id >= 0 ) {
            if ( fabs(mean[id]) > 0.95 ) {
                // Isolated domain surrounded by one domain
                GlobalBlobID(i) = -2;
                Mean(i) = sign(mean[id]);
            } else {
                // Boarder volume, set to liquid
                Mean(i) = 1;
            }
        }
    }
    // Perform the final segmentation and update the distance
    fillFloat.fill(Mean);
    segment( Mean, ID, 0.01 );
    CalcDistMultiLevel( Dist, ID, Dm );
    fillFloat.fill(Dist);
}



// Filter the original data
void filter_src( const Domain& Dm, Array<float>& src )
{
	PROFILE_START("Filter source data");
    int Nx = Dm.Nx-2;
    int Ny = Dm.Ny-2;
    int Nz = Dm.Nz-2;
	fillHalo<float> fillFloat(Dm.Comm,Dm.rank_info,Nx,Ny,Nz,1,1,1,0,1);
    // Perform a hot-spot filter on the data
    std::vector<imfilter::BC> BC = { imfilter::BC::replicate, imfilter::BC::replicate, imfilter::BC::replicate };
    std::function<float(const Array<float>&)> filter_3D = []( const Array<float>& data )
    {
        float min1 = std::min(data(0,1,1),data(2,1,1));
        float min2 = std::min(data(1,0,1),data(1,2,1));
        float min3 = std::min(data(1,1,0),data(1,1,2));
        float max1 = std::max(data(0,1,1),data(2,1,1));
        float max2 = std::max(data(1,0,1),data(1,2,1));
        float max3 = std::max(data(1,1,0),data(1,1,2));
        float min = std::min(min1,std::min(min2,min3));
        float max = std::max(max1,std::max(max2,max3));
        return std::max(std::min(data(1,1,1),max),min);
    };
    std::function<float(const Array<float>&)> filter_1D = []( const Array<float>& data )
    {
        float min = std::min(data(0),data(2));
        float max = std::max(data(0),data(2));
        return std::max(std::min(data(1),max),min);
    };
    //LOCVOL[0] = imfilter::imfilter<float>( LOCVOL[0], {1,1,1}, filter_3D, BC );
    std::vector<std::function<float(const Array<float>&)>> filter_set(3,filter_1D);
    src = imfilter::imfilter_separable<float>( src, {1,1,1}, filter_set, BC );
    fillFloat.fill( src );
    // Perform a gaussian filter on the data
    int Nh[3] = { 2, 2, 2 };
    float sigma[3] = { 1.0, 1.0, 1.0 };
    std::vector<Array<float>> H(3);
    H[0] = imfilter::create_filter<float>( { Nh[0] }, "gaussian", &sigma[0] );
    H[1] = imfilter::create_filter<float>( { Nh[1] }, "gaussian", &sigma[1] );
    H[2] = imfilter::create_filter<float>( { Nh[2] }, "gaussian", &sigma[2] );
    src = imfilter::imfilter_separable( src, H, BC );
    fillFloat.fill( src );
	PROFILE_STOP("Filter source data");
}