#include "localfitmethod.h"

#include "../../structures/image2d.h"
#include "../../msio/pngfile.h"

#include "../../util/rng.h"

#include "thresholdtools.h"

#include <cmath>
#include <iostream>

LocalFitMethod::LocalFitMethod() : _weights(0) {}

LocalFitMethod::~LocalFitMethod() { ClearWeights(); }

void LocalFitMethod::Initialize(const TimeFrequencyData &input) {
  ClearWeights();
  _original = input.GetSingleImage();
  _background2D =
      Image2D::CreateZeroImagePtr(_original->Width(), _original->Height());
  _background = TimeFrequencyData(input.ComplexRepresentation(),
                                  input.Polarizations()[0], _background2D);
  _mask = input.GetSingleMask();
  if (_hSquareSize * 2 > _original->Width())
    _hSquareSize = _original->Width() / 2;
  if (_vSquareSize * 2 > _original->Height())
    _vSquareSize = _original->Height() / 2;
  switch (_method) {
    case None:
    case Average:
    case Median:
    case Minimum:
      break;
    case GaussianWeightedAverage:
    case FastGaussianWeightedAverage:
      InitializeGaussianWeights();
      break;
  }
}

void LocalFitMethod::ClearWeights() {
  if (_weights) {
    for (unsigned y = 0; y < _vSquareSize * 2 + 1; ++y) delete[] _weights[y];
    delete[] _weights;
    _weights = 0;
  }
}

void LocalFitMethod::InitializeGaussianWeights() {
  _weights = new num_t *[_vSquareSize * 2 + 1];
  for (int y = 0; y < (int)(_vSquareSize * 2 + 1); ++y) {
    _weights[y] = new num_t[_hSquareSize * 2 + 1];
    for (int x = 0; x < (int)(_hSquareSize * 2 + 1); ++x) {
      _weights[y][x] =
          RNG::EvaluateGaussian2D(x - (int)_hSquareSize, y - (int)_vSquareSize,
                                  _hKernelSize, _vKernelSize);
    }
  }
}

unsigned LocalFitMethod::TaskCount() {
  if (_method == FastGaussianWeightedAverage)
    return 1;
  else
    return _original->Height();
}

void LocalFitMethod::PerformFit(unsigned taskNumber) {
  if (_mask == 0) throw std::runtime_error("Mask has not been set!");
  if (_method == FastGaussianWeightedAverage) {
    CalculateWeightedAverageFast();
  } else {
    unsigned y = taskNumber;
    for (unsigned x = 0; x < _original->Width(); ++x)
      _background2D->SetValue(x, y, CalculateBackgroundValue(x, y));
  }
}

long double LocalFitMethod::CalculateBackgroundValue(unsigned x, unsigned y) {
  ThreadLocal local;
  local.image = this;
  local.currentX = x;
  local.currentY = y;

  if (local.currentY >= _vSquareSize)
    local.startY = local.currentY - _vSquareSize;
  else
    local.startY = 0;
  local.endY = local.currentY + _vSquareSize;
  if (local.endY >= _original->Height()) local.endY = _original->Height() - 1;

  if (local.currentX >= _hSquareSize)
    local.startX = local.currentX - _hSquareSize;
  else
    local.startX = 0;
  local.endX = local.currentX + _hSquareSize;
  if (local.endX >= _original->Width()) local.endX = _original->Width() - 1;
  local.emptyWindows = 0;

  switch (_method) {
    case None:
    case FastGaussianWeightedAverage:
      return 0.0;
    case Median:
      return CalculateMedian(x, y, local);
    case Minimum:
      return CalculateMinimum(x, y, local);
    case Average:
      return CalculateAverage(x, y, local);
    case GaussianWeightedAverage:
      return CalculateWeightedAverage(x, y, local);
    default:
      throw std::runtime_error(
          "The LocalFitMethod was not initialized before a fit was executed.");
  }
}

long double LocalFitMethod::CalculateAverage(unsigned x, unsigned y,
                                             ThreadLocal &local) {
  long double sum = 0.0;
  unsigned long count = 0;
  for (unsigned yi = local.startY; yi <= local.endY; ++yi) {
    for (unsigned xi = local.startX; xi <= local.endX; ++xi) {
      if (!_mask->Value(xi, yi) && std::isfinite(_original->Value(xi, yi))) {
        sum += _original->Value(xi, yi);
        ++count;
      }
    }
  }
  if (count != 0)
    return sum / (long double)count;
  else
    return _original->Value(x, y);
}

long double LocalFitMethod::CalculateMedian(unsigned x, unsigned y,
                                            ThreadLocal &local) {
  // unsigned maxSize = (local.endY-local.startY)*(local.endX-local.startX);
  std::vector<long double> orderList;
  unsigned long count = 0;
  for (unsigned yi = local.startY; yi <= local.endY; ++yi) {
    for (unsigned xi = local.startX; xi <= local.endX; ++xi) {
      if (!_mask->Value(xi, yi) && std::isfinite(_original->Value(xi, yi))) {
        orderList.push_back(_original->Value(xi, yi));
        ++count;
      }
    }
  }
  if (count == 0) {
    local.emptyWindows++;
    return _original->Value(x, y);
  } else if (count % 2 == 1) {
    std::nth_element(orderList.begin(), orderList.begin() + count / 2,
                     orderList.end());
    long double mOdd = orderList[count / 2];
    return mOdd;
  } else {
    std::nth_element(orderList.begin(), orderList.begin() + count / 2,
                     orderList.end());
    long double mOdd = orderList[count / 2];
    std::nth_element(orderList.begin(), orderList.begin() + (count / 2 - 1),
                     orderList.end());
    long double mEven = orderList[count / 2 - 1];
    return (mOdd + mEven) * 0.5L;
  }
}

long double LocalFitMethod::CalculateMinimum(unsigned x, unsigned y,
                                             ThreadLocal &local) {
  long double minimum = 1e100;
  unsigned long count = 0;
  for (unsigned yi = local.startY; yi <= local.endY; ++yi) {
    for (unsigned xi = local.startX; xi <= local.endX; ++xi) {
      if (!_mask->Value(xi, yi) && std::isfinite(_original->Value(xi, yi)) &&
          _original->Value(xi, yi) < minimum) {
        minimum = _original->Value(xi, yi);
        ++count;
      }
    }
  }
  if (count == 0)
    return _original->Value(x, y);
  else
    return minimum;
}

long double LocalFitMethod::CalculateWeightedAverage(unsigned x, unsigned y,
                                                     ThreadLocal &local) {
  long double sum = 0.0;
  long double totalWeight = 0.0;
  for (unsigned j = local.startY; j <= local.endY; ++j) {
    for (unsigned i = local.startX; i <= local.endX; ++i) {
      if (!_mask->Value(i, j) && std::isfinite(_original->Value(i, j))) {
        long double weight =
            _weights[j - y + _vSquareSize][i - x + _hSquareSize];
        sum += _original->Value(i, j) * weight;
        totalWeight += weight;
      }
    }
  }
  if (totalWeight != 0.0)
    return sum / totalWeight;
  else {
    sum = 0.0;
    totalWeight = 0.0;
    for (unsigned j = local.startY; j <= local.endY; ++j) {
      for (unsigned i = local.startX; i <= local.endX; ++i) {
        if (std::isfinite(_original->Value(i, j))) {
          long double weight =
              _weights[j - y + _vSquareSize][i - x + _hSquareSize];
          sum += _original->Value(i, j) * weight;
          totalWeight += weight;
        }
      }
    }
    if (totalWeight != 0.0)
      return sum / totalWeight;
    else {
      sum = 0.0;
      totalWeight = 0.0;
      for (unsigned j = local.startY; j <= local.endY; ++j) {
        for (unsigned i = local.startX; i <= local.endX; ++i) {
          long double weight =
              _weights[j - y + _vSquareSize][i - x + _hSquareSize];
          sum += _original->Value(i, j) * weight;
          totalWeight += weight;
        }
      }
      return sum / totalWeight;
    }
  }
}

void LocalFitMethod::CalculateWeightedAverageFast() {
  *_background2D = *_original;
  ThresholdTools::SetFlaggedValuesToZero(_background2D.get(), _mask.get());
  PerformGaussianConvolution(_background2D);
  Image2DPtr flagWeights = CreateFlagWeightsMatrix();
  PerformGaussianConvolution(flagWeights);
  ElementWiseDivide(_background2D, flagWeights);
}

void LocalFitMethod::ElementWiseDivide(Image2DPtr leftHand,
                                       Image2DCPtr rightHand) {
  for (unsigned y = 0; y < leftHand->Height(); ++y) {
    for (unsigned x = 0; x < leftHand->Width(); ++x) {
      if (rightHand->Value(x, y) == 0.0)
        leftHand->SetValue(x, y, 0.0);
      else
        leftHand->SetValue(x, y,
                           leftHand->Value(x, y) / rightHand->Value(x, y));
    }
  }
}

Image2DPtr LocalFitMethod::CreateFlagWeightsMatrix() {
  Image2DPtr image =
      Image2D::CreateUnsetImagePtr(_mask->Width(), _mask->Height());
  for (unsigned y = 0; y < image->Height(); ++y) {
    for (unsigned x = 0; x < image->Width(); ++x) {
      if (!_mask->Value(x, y) && std::isfinite(_original->Value(x, y)))
        image->SetValue(x, y, 1.0);
      else
        image->SetValue(x, y, 0.0);
    }
  }
  return image;
}

void LocalFitMethod::PerformGaussianConvolution(Image2DPtr input) {
  // Guassian convolution can be separated in two 1D convolution
  // because of properties of the 2D Gaussian function.
  Image2DPtr temp =
      Image2D::CreateZeroImagePtr(input->Width(), input->Height());
  for (int i = -_hSquareSize; i <= (int)_hSquareSize; ++i) {
    num_t gaus = _weights[_vSquareSize][i + _hSquareSize];
    for (unsigned y = 0; y < input->Height(); ++y) {
      unsigned xStart = i >= 0 ? 0 : -i;
      unsigned xEnd = i <= 0 ? input->Width() : input->Width() - i;
      for (unsigned x = xStart; x < xEnd; ++x) {
        if (std::isfinite(input->Value(x + i, y)))
          temp->AddValue(x, y, input->Value(x + i, y) * gaus);
      }
    }
  }

  input->SetAll(0.0);
  for (int j = -_vSquareSize; j <= (int)_vSquareSize; ++j) {
    num_t gaus = _weights[j + _vSquareSize][_hSquareSize];
    unsigned yStart = j >= 0 ? 0 : -j;
    unsigned yEnd = j <= 0 ? input->Height() : input->Height() - j;
    for (unsigned y = yStart; y < yEnd; ++y) {
      for (unsigned x = 0; x < input->Width(); ++x) {
        if (std::isfinite(temp->Value(x, y + j)))
          input->AddValue(x, y, temp->Value(x, y + j) * gaus);
      }
    }
  }
}
