/*
  Copyright 2013--2018 James E. McClure, Virginia Polytechnic & State University
  Copyright Equnior ASA

  This file is part of the Open Porous Media project (OPM).
  OPM is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.
  OPM is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
  You should have received a copy of the GNU General Public License
  along with OPM.  If not, see <http://www.gnu.org/licenses/>.
*/
#include "analysis/filters.h"
#include "math.h"
#include "ProfilerApp.h"

void Mean3D(const Array<double> &Input, Array<double> &Output) {
    PROFILE_START("Mean3D");
    // Perform a 3D Mean filter on Input array
    int i, j, k;

    int Nx = int(Input.size(0));
    int Ny = int(Input.size(1));
    int Nz = int(Input.size(2));

    for (k = 1; k < Nz - 1; k++) {
        for (j = 1; j < Ny - 1; j++) {
            for (i = 1; i < Nx - 1; i++) {
                double MeanValue = Input(i, j, k);
                // next neighbors
                MeanValue += Input(i + 1, j, k) + Input(i, j + 1, k) +
                             Input(i, j, k + 1) + Input(i - 1, j, k) +
                             Input(i, j - 1, k) + Input(i, j, k - 1);
                MeanValue += Input(i + 1, j + 1, k) + Input(i - 1, j + 1, k) +
                             Input(i + 1, j - 1, k) + Input(i - 1, j - 1, k);
                MeanValue += Input(i + 1, j, k + 1) + Input(i - 1, j, k + 1) +
                             Input(i + 1, j, k - 1) + Input(i - 1, j, k - 1);
                MeanValue += Input(i, j + 1, k + 1) + Input(i, j - 1, k + 1) +
                             Input(i, j + 1, k - 1) + Input(i, j - 1, k - 1);
                MeanValue +=
                    Input(i + 1, j + 1, k + 1) + Input(i - 1, j + 1, k + 1) +
                    Input(i + 1, j - 1, k + 1) + Input(i - 1, j - 1, k + 1);
                MeanValue +=
                    Input(i + 1, j + 1, k - 1) + Input(i - 1, j + 1, k - 1) +
                    Input(i + 1, j - 1, k - 1) + Input(i - 1, j - 1, k - 1);
                Output(i, j, k) = MeanValue / 27.0;
            }
        }
    }
    PROFILE_STOP("Mean3D");
}

void Med3D(const Array<float> &Input, Array<float> &Output) {
    PROFILE_START("Med3D");
    // Perform a 3D Median filter on Input array with specified width
    int i, j, k, ii, jj, kk;
    int imin, jmin, kmin, imax, jmax, kmax;

    float *List;
    List = new float[27];

    int Nx = int(Input.size(0));
    int Ny = int(Input.size(1));
    int Nz = int(Input.size(2));

    for (k = 1; k < Nz - 1; k++) {
        for (j = 1; j < Ny - 1; j++) {
            for (i = 1; i < Nx - 1; i++) {

                // Just use a 3x3x3 window (hit recursively if needed)
                imin = i - 1;
                jmin = j - 1;
                kmin = k - 1;
                imax = i + 2;
                jmax = j + 2;
                kmax = k + 2;

                // Populate the list with values in the window
                int Number = 0;
                for (kk = kmin; kk < kmax; kk++) {
                    for (jj = jmin; jj < jmax; jj++) {
                        for (ii = imin; ii < imax; ii++) {
                            List[Number++] = Input(ii, jj, kk);
                        }
                    }
                }
                // Sort the first 5 entries and return the median
                for (ii = 0; ii < 14; ii++) {
                    for (jj = ii + 1; jj < 27; jj++) {
                        if (List[jj] < List[ii]) {
                            float tmp = List[ii];
                            List[ii] = List[jj];
                            List[jj] = tmp;
                        }
                    }
                }
                // Return the median
                Output(i, j, k) = List[13];
            }
        }
    }
    PROFILE_STOP("Med3D");
}

int NLM3D(const Array<float> &Input, Array<float> &Mean,
          const Array<float> &Distance, Array<float> &Output, const int d,
          const float h) {
    PROFILE_START("NLM3D");
    // Implemenation of 3D non-local means filter
    // 		d determines the width of the search volume
    // 		h is a free parameter for non-local means (i.e. 1/sigma^2)
    // 		Distance is the signed distance function
    // 		If Distance(i,j,k) > THRESHOLD_DIST then don't compute NLM

    float THRESHOLD_DIST = float(d);
    float weight, sum;
    int i, j, k, ii, jj, kk;
    int imin, jmin, kmin, imax, jmax, kmax;
    int returnCount = 0;

    int Nx = int(Input.size(0));
    int Ny = int(Input.size(1));
    int Nz = int(Input.size(2));

    // Compute the local means
    for (k = 1; k < Nz - 1; k++) {
        for (j = 1; j < Ny - 1; j++) {
            for (i = 1; i < Nx - 1; i++) {

                imin = std::max(0, i - d);
                jmin = std::max(0, j - d);
                kmin = std::max(0, k - d);
                imax = std::min(Nx - 1, i + d);
                jmax = std::min(Ny - 1, j + d);
                kmax = std::min(Nz - 1, k + d);

                // Populate the list with values in the window
                sum = 0;
                weight = 0;
                for (kk = kmin; kk < kmax; kk++) {
                    for (jj = jmin; jj < jmax; jj++) {
                        for (ii = imin; ii < imax; ii++) {
                            sum += Input(ii, jj, kk);
                            weight++;
                        }
                    }
                }

                Mean(i, j, k) = sum / weight;
            }
        }
    }

    // Compute the non-local means
    for (k = 1; k < Nz - 1; k++) {
        for (j = 1; j < Ny - 1; j++) {
            for (i = 1; i < Nx - 1; i++) {

                if (fabs(Distance(i, j, k)) < THRESHOLD_DIST) {
                    // compute the expensive non-local means
                    sum = 0;
                    weight = 0;

                    imin = std::max(0, i - d);
                    jmin = std::max(0, j - d);
                    kmin = std::max(0, k - d);
                    imax = std::min(Nx - 1, i + d);
                    jmax = std::min(Ny - 1, j + d);
                    kmax = std::min(Nz - 1, k + d);

                    for (kk = kmin; kk < kmax; kk++) {
                        for (jj = jmin; jj < jmax; jj++) {
                            for (ii = imin; ii < imax; ii++) {
                                float tmp = Mean(i, j, k) - Mean(ii, jj, kk);
                                sum += exp(-tmp * tmp * h) * Input(ii, jj, kk);
                                weight += exp(-tmp * tmp * h);
                            }
                        }
                    }

                    returnCount++;
                    //Output(i,j,k) = Mean(i,j,k);
                    Output(i, j, k) = sum / weight;
                } else {
                    // Just return the mean
                    Output(i, j, k) = Mean(i, j, k);
                }
            }
        }
    }
    // Return the number of sites where NLM was applied
    PROFILE_STOP("NLM3D");
    return returnCount;
}