#include "../../structures/antennainfo.h"

#include "../../imaging/uvimager.h"

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

#include "../algorithms/fringestoppingfitter.h"
#include "../algorithms/sinusfitter.h"

FringeStoppingFitter::FringeStoppingFitter()
    : _originalData(0),
      _fringesToConsider(1.0),
      _minWindowSize(32),
      _maxWindowSize(128),
      _returnFittedValue(false),
      _returnMeanValue(false),
      _fringeFit(true),
      _newPhaseCentreDec(M_PInl * 0.5),
      _newPhaseCentreRA(0.0) {}

FringeStoppingFitter::~FringeStoppingFitter() {}

void FringeStoppingFitter::PerformFit(unsigned taskNumber) {
  if (_fringeFit)
    PerformDynamicFrequencyFitOnOneChannel(taskNumber, _maxWindowSize);
  else {
    size_t x = taskNumber;

    if (_fitChannelsIndividually) {
      for (size_t y = 0; y < _originalData->ImageHeight(); ++y) {
        num_t r, i;
        CalculateFitValue(*_originalData->GetRealPart(),
                          *_originalData->GetImaginaryPart(), x, y, 1, r, i);

        _realBackground->SetValue(x, y, r);
        _imaginaryBackground->SetValue(x, y, i);
      }
    } else {
      num_t r, i;
      CalculateFitValue(*_originalData->GetRealPart(),
                        *_originalData->GetImaginaryPart(), x, 0,
                        _originalData->ImageHeight(), r, i);

      for (size_t y = 0; y < _originalData->ImageHeight(); ++y) {
        _realBackground->SetValue(x, y, r);
        _imaginaryBackground->SetValue(x, y, i);
      }
    }
  }
}

void FringeStoppingFitter::PerformStaticFrequencyFitOnOneChannel(unsigned y) {
  if (_fitChannelsIndividually) {
    for (size_t x = 0; x < _originalData->ImageWidth(); ++x) {
      num_t r, i;
      CalculateFitValue(*_originalData->GetRealPart(),
                        *_originalData->GetImaginaryPart(), x, y, 1, r, i);

      _realBackground->SetValue(x, y, r);
      _imaginaryBackground->SetValue(x, y, i);
    }
  } else {
    for (size_t x = 0; x < _originalData->ImageWidth(); ++x) {
      num_t r, i;
      CalculateFitValue(*_originalData->GetRealPart(),
                        *_originalData->GetImaginaryPart(), x, 0,
                        _originalData->ImageHeight(), r, i);

      _realBackground->SetValue(x, y, r);
      _imaginaryBackground->SetValue(x, y, i);
    }
  }
}

void FringeStoppingFitter::PerformFringeStop() {
  Image2DCPtr real = _originalData->GetRealPart(),
              imaginary = _originalData->GetImaginaryPart();

  Baseline baseline(*_antenna1Info, *_antenna2Info);

  for (size_t x = 0; x < real->Width(); ++x) {
    for (size_t y = 0; y < real->Height(); ++y) {
      num_t r = real->Value(x, y);
      num_t i = imaginary->Value(x, y);
      num_t fringeCount = UVImager::GetFringeCount(0, x, y, _metaData);

      num_t cosfreq = cosn(fringeCount * 2.0L * M_PIn),
            sinfreq = sinn(fringeCount * 2.0L * M_PIn);

      num_t newR = r * cosfreq - i * sinfreq;
      i = r * sinfreq + i * cosfreq;

      _realBackground->SetValue(x, y, newR);
      _imaginaryBackground->SetValue(x, y, i);
    }
  }
}

num_t FringeStoppingFitter::CalculateMaskedAverage(const Image2D &image,
                                                   size_t x, size_t yFrom,
                                                   size_t yLength) {
  num_t average = 0.0L;
  size_t count = 0;
  for (size_t i = yFrom; i < yFrom + yLength; ++i) {
    if (!_originalMask->Value(x, i) && std::isfinite(image.Value(x, i))) {
      average += image.Value(x, i);
      ++count;
    }
  }
  return average / (long double)count;
}

num_t FringeStoppingFitter::CalculateUnmaskedAverage(const Image2D &image,
                                                     size_t x, size_t yFrom,
                                                     size_t yLength) {
  num_t average = 0.0L;
  size_t count = 0;
  for (size_t i = yFrom; i < yFrom + yLength; ++i) {
    if (std::isfinite(image.Value(x, i))) {
      average += image.Value(x, i);
      ++count;
    }
  }
  return average / (num_t)count;
}

void FringeStoppingFitter::CalculateFitValue(const Image2D &real,
                                             const Image2D &imaginary, size_t x,
                                             size_t yFrom, size_t yLength,
                                             num_t &rValue, num_t &iValue) {
  size_t windowWidth;
  size_t yMid = yFrom + yLength / 2;
  num_t estimatedFrequency = GetFringeFrequency(x, yMid);
  windowWidth = (size_t)ceil(_fringesToConsider / estimatedFrequency);
  if (windowWidth > _maxWindowSize) windowWidth = _maxWindowSize;
  if (windowWidth < _minWindowSize) windowWidth = _minWindowSize;

  size_t xLeft, xRight;
  xRight = x + windowWidth / 2;
  if (x >= (windowWidth / 2))
    xLeft = x - (windowWidth / 2);
  else {
    xLeft = 0;
    xRight += (windowWidth / 2) - x;
  }
  if (xRight > real.Width()) {
    if (xLeft > xRight - real.Width())
      xLeft -= xRight - real.Width();
    else
      xLeft = 0;
    xRight = real.Width();
  }

  num_t fringeFrequency = GetFringeFrequency(x, yMid);
  num_t *dataT = new num_t[xRight - xLeft];
  num_t *dataR = new num_t[xRight - xLeft];
  num_t *dataI = new num_t[xRight - xLeft];

  size_t index = 0;
  for (size_t i = xLeft; i < xRight; ++i) {
    dataR[index] = CalculateMaskedAverage(real, i, yFrom, yLength);
    dataI[index] = CalculateMaskedAverage(imaginary, i, yFrom, yLength);
    if (std::isfinite(dataR[index]) && std::isfinite(dataI[index])) {
      dataT[index] = i;
      ++index;
    }
  }
  if (index == 0) {
    for (size_t i = xLeft; i < xRight; ++i) {
      dataR[index] = CalculateUnmaskedAverage(real, i, yFrom, yLength);
      dataI[index] = CalculateUnmaskedAverage(imaginary, i, yFrom, yLength);
      if (std::isfinite(dataR[index]) && std::isfinite(dataI[index])) {
        dataT[index] = i;
        ++index;
      }
    }
  }

  SinusFitter fitter;
  num_t phaseR, phaseI, amplitude;

  fitter.FindPhaseAndAmplitudeComplex(phaseR, amplitude, dataR, dataI, dataT,
                                      index, fringeFrequency);
  phaseI = fmodn(phaseR + M_PIn * 0.5L, 2.0L * M_PIn);
  num_t rfiValueR =
      SinusFitter::Value(phaseR, amplitude, x, fringeFrequency, 0.0);
  num_t rfiValueI =
      SinusFitter::Value(phaseI, amplitude, x, fringeFrequency, 0.0);

  if (_returnFittedValue) {
    num_t meanR = fitter.FindMean(phaseR, amplitude, dataR, dataT, index,
                                  fringeFrequency);
    num_t meanI = fitter.FindMean(phaseI, amplitude, dataI, dataT, index,
                                  fringeFrequency);
    rValue = rfiValueR + meanR;
    iValue = rfiValueI + meanI;
  } else if (_returnMeanValue) {
    rValue = fitter.FindMean(phaseR, amplitude, dataR, dataT, index,
                             fringeFrequency);
    iValue = fitter.FindMean(phaseI, amplitude, dataI, dataT, index,
                             fringeFrequency);
  } else {
    num_t observedValueR = CalculateMaskedAverage(real, x, yFrom, yLength);
    num_t observedValueI = CalculateMaskedAverage(imaginary, x, yFrom, yLength);
    rValue = observedValueR - rfiValueR;
    iValue = observedValueI - rfiValueI;
  }

  delete[] dataR;
  delete[] dataI;
  delete[] dataT;
}

num_t FringeStoppingFitter::GetFringeFrequency(size_t x, size_t y) {
  double deltaTime;
  if (_observationTimes->size() > 1)
    deltaTime = (*_observationTimes)[1] - (*_observationTimes)[0];
  else
    deltaTime = 1.0;
  num_t observeFreq = _bandInfo->channels[y].frequencyHz;
  Baseline baseline(*_antenna1Info, *_antenna2Info);
  num_t delayRA = _fieldInfo->delayDirectionRA;
  num_t delayDec = _fieldInfo->delayDirectionDec;
  return deltaTime * UVImager::GetFringeStopFrequency(x, baseline, delayRA,
                                                      delayDec, observeFreq,
                                                      _metaData);
}

void FringeStoppingFitter::GetRFIValue(num_t &r, num_t &i, int x, int y,
                                       num_t rfiPhase, num_t rfiStrength) {
  const numl_t earthRotation = UVImager::TimeToEarthLattitude(x, _metaData);
  const Baseline baseline = _metaData->Baseline();
  const numl_t newWPos = UVImager::GetWPosition(
      _newPhaseCentreDec, _newPhaseCentreRA,
      _metaData->Band().channels[y].frequencyHz, earthRotation,
      baseline.DeltaX(), baseline.DeltaY());
  const num_t rotations =
      UVImager::GetFringeCount(0, x, y, _metaData) - newWPos;

  r = cosn(rotations * 2.0 * M_PIn + rfiPhase) * rfiStrength;
  i = -sinn(rotations * 2.0 * M_PIn + rfiPhase) * rfiStrength;
}

void FringeStoppingFitter::GetMeanValue(num_t &rMean, num_t &iMean, num_t phase,
                                        num_t amplitude, const SampleRow &real,
                                        const SampleRow &imaginary,
                                        unsigned xStart, unsigned xEnd,
                                        unsigned y) {
  rMean = 0.0;
  iMean = 0.0;
  for (unsigned t = xStart; t < xEnd; ++t) {
    num_t r, i;
    GetRFIValue(r, i, t, y, phase, amplitude);
    rMean += real.Value(t) - r;
    iMean += imaginary.Value(t) - i;
  }
  rMean /= (num_t)(xEnd - xStart);
  iMean /= (num_t)(xEnd - xStart);
}

void FringeStoppingFitter::MinimizeRFIFitError(num_t &phase, num_t &amplitude,
                                               const SampleRow &real,
                                               const SampleRow &imaginary,
                                               unsigned xStart, unsigned xEnd,
                                               unsigned y) const throw() {
  // calculate 1/N * \sum_x v(t) e^{2 i \pi \tau_g(t)}, where \tau_g(t) is the
  // number of phase rotations because of the geometric delay as function of
  // time t.

  num_t sumR = 0.0, sumI = 0.0;
  size_t n = 0;

  const num_t dx =
      _metaData->Antenna2().position.x - _metaData->Antenna1().position.x;
  const num_t dy =
      _metaData->Antenna2().position.y - _metaData->Antenna1().position.y;

  for (unsigned t = xStart; t < xEnd; ++t) {
    const num_t vR = real.Value(t);
    const num_t vI = imaginary.Value(t);

    if (std::isfinite(vR) && std::isfinite(vI)) {
      const num_t tRotation =
          UVImager::TimeToEarthLattitude(_metaData->ObservationTimes()[t]);
      const num_t tauge = UVImager::GetFringeCount(0, t, y, _metaData);
      const num_t taugeNew = UVImager::GetWPosition(
          _newPhaseCentreDec, _newPhaseCentreRA,
          _metaData->Band().channels[y].frequencyHz, tRotation, dx, dy);
      const num_t phaseShift = tauge - taugeNew;

      sumR += vR * cosn(-2.0 * M_PIn * phaseShift);
      sumR += vI * sinn(-2.0 * M_PIn * phaseShift);

      sumI += vR * sinn(-2.0 * M_PIn * phaseShift);
      sumI -= vI * cosn(-2.0 * M_PIn * phaseShift);
      ++n;
    }
  }

  sumR /= (num_t)n;
  sumI /= (num_t)n;

  phase = SinusFitter::Phase(sumR, sumI);
  amplitude = sqrtn(sumR * sumR + sumI * sumI);
}

void FringeStoppingFitter::PerformDynamicFrequencyFitOnOneChannel(unsigned y) {
  SampleRow real =
                SampleRow::MakeFromRow(_originalData->GetRealPart().get(), y),
            imaginary = SampleRow::MakeFromRow(
                _originalData->GetImaginaryPart().get(), y);
  PerformDynamicFrequencyFitOnOneRow(real, imaginary, y);
}

void FringeStoppingFitter::PerformDynamicFrequencyFitOnOneRow(
    const SampleRow &real, const SampleRow &imaginary, unsigned y) {
  num_t phase, strength;
  MinimizeRFIFitError(phase, strength, real, imaginary, 0,
                      _originalData->ImageWidth(), y);
  Logger::Debug << "Amplitude found: " << strength << " phase found: " << phase
                << '\n';
  for (size_t x = 0; x < _originalData->ImageWidth(); ++x) {
    num_t rfiR, rfiI;
    GetRFIValue(rfiR, rfiI, x, y, phase, strength);
    if (_returnFittedValue) {
      num_t rMean, iMean;
      GetMeanValue(rMean, iMean, phase, strength, real, imaginary, 0,
                   _originalData->ImageWidth(), y);
      _realBackground->SetValue(x, y, rfiR + rMean);
      _imaginaryBackground->SetValue(x, y, rfiI + iMean);
    } else {
      _realBackground->SetValue(x, y, rfiR);
      _imaginaryBackground->SetValue(x, y, rfiI);
    }
  }
}

void FringeStoppingFitter::PerformDynamicFrequencyFitOnOneChannel(
    unsigned y, unsigned windowSize) {
  SampleRow real = SampleRow::MakeFromRowWithMissings(
                _originalData->GetRealPart().get(), _originalMask.get(), y),
            imaginary = SampleRow::MakeFromRowWithMissings(
                _originalData->GetImaginaryPart().get(), _originalMask.get(),
                y);
  PerformDynamicFrequencyFitOnOneRow(real, imaginary, y, windowSize);
}

void FringeStoppingFitter::PerformDynamicFrequencyFitOnOneRow(
    const SampleRow &real, const SampleRow &imaginary, unsigned y,
    unsigned windowSize) {
  unsigned halfWindowSize = windowSize / 2;
  for (size_t x = 0; x < real.Size(); ++x) {
    size_t windowStart, windowEnd;
    if (x > halfWindowSize)
      windowStart = x - halfWindowSize;
    else
      windowStart = 0;
    if (x + halfWindowSize < real.Size())
      windowEnd = x + halfWindowSize;
    else
      windowEnd = real.Size();
    num_t windowPhase, windowStrength;
    MinimizeRFIFitError(windowPhase, windowStrength, real, imaginary,
                        windowStart, windowEnd, y);

    num_t rfiR, rfiI;
    GetRFIValue(rfiR, rfiI, x, y, windowPhase, windowStrength);
    if (_returnFittedValue) {
      num_t rMean, iMean;
      GetMeanValue(rMean, iMean, windowPhase, windowStrength, real, imaginary,
                   windowStart, windowEnd, y);
      _realBackground->SetValue(x, y, rfiR + rMean);
      _imaginaryBackground->SetValue(x, y, rfiI + iMean);
    } else {
      _realBackground->SetValue(x, y, rfiR);
      _imaginaryBackground->SetValue(x, y, rfiI);
    }
  }
}

void FringeStoppingFitter::PerformDynamicFrequencyFit() {
  for (size_t y = 0; y < _originalData->ImageHeight(); ++y) {
    PerformDynamicFrequencyFitOnOneChannel(y);
  }
}

void FringeStoppingFitter::PerformDynamicFrequencyFit(unsigned yStart,
                                                      unsigned yEnd,
                                                      unsigned windowSize) {
  SampleRow real = SampleRow::MakeFromRowSum(_originalData->GetRealPart().get(),
                                             yStart, yEnd),
            imaginary = SampleRow::MakeFromRowSum(
                _originalData->GetImaginaryPart().get(), yStart, yEnd);
  PerformDynamicFrequencyFitOnOneRow(real, imaginary, (yStart + yEnd) / 2,
                                     windowSize);
}

void FringeStoppingFitter::PerformDynamicFrequencyFit(unsigned yStart,
                                                      unsigned yEnd) {
  SampleRow real = SampleRow::MakeFromRowSum(_originalData->GetRealPart().get(),
                                             yStart, yEnd),
            imaginary = SampleRow::MakeFromRowSum(
                _originalData->GetImaginaryPart().get(), yStart, yEnd);
  PerformDynamicFrequencyFitOnOneRow(real, imaginary, (yStart + yEnd) / 2);
}

num_t FringeStoppingFitter::GetAmplitude(unsigned yStart, unsigned yEnd) {
  unsigned y = (yStart + yEnd) / 2;
  SampleRow real = SampleRow::MakeFromRowSum(_originalData->GetRealPart().get(),
                                             yStart, yEnd),
            imaginary = SampleRow::MakeFromRowSum(
                _originalData->GetImaginaryPart().get(), yStart, yEnd);
  num_t phase, amplitude;
  MinimizeRFIFitError(phase, amplitude, real, imaginary, 0,
                      _originalData->ImageWidth(), y);
  return amplitude;
}
