#include "spatialtimeloader.h"

#include <stdexcept>
#include <vector>

#include <casacore/tables/Tables/ArrayColumn.h>
#include <casacore/tables/Tables/ScalarColumn.h>

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

SpatialTimeLoader::SpatialTimeLoader(MSMetaData& msMetaData)
    : _msMetaData(msMetaData) {
  const casacore::Table rawTable(_msMetaData.Path());
  casacore::Block<casacore::String> names(4);
  names[0] = "DATA_DESC_ID";
  names[1] = "TIME";
  names[2] = "ANTENNA1";
  names[3] = "ANTENNA2";
  _sortedTable.reset(new casacore::Table(rawTable.sort(names)));

  _channelCount = _msMetaData.FrequencyCount(0);
  _timestepsCount = _msMetaData.TimestepCount();
  _antennaCount = _msMetaData.AntennaCount();
  _polarizationCount = _msMetaData.PolarizationCount();

  casacore::Block<casacore::String> selectionNames(1);
  selectionNames[0] = "DATA_DESC_ID";
  _tableIter.reset(new casacore::TableIterator(
      *_sortedTable, selectionNames, casacore::TableIterator::Ascending,
      casacore::TableIterator::NoSort));
}

SpatialTimeLoader::~SpatialTimeLoader() {}

TimeFrequencyData SpatialTimeLoader::Load(unsigned channelIndex,
                                          bool fringeStop) {
  const unsigned baselineCount = _antennaCount * (_antennaCount - 1) / 2;

  const casacore::Table table = _tableIter->table();
  const casacore::ScalarColumn<int> antenna1Column(table, "ANTENNA1");
  const casacore::ScalarColumn<int> antenna2Column(table, "ANTENNA2");
  const casacore::ScalarColumn<double> timeColumn(table, "TIME");
  const casacore::ArrayColumn<double> uvwColumn(table, "UVW");
  const casacore::ArrayColumn<bool> flagColumn(table, "FLAG");
  const casacore::ArrayColumn<casacore::Complex> dataColumn(table, "DATA");

  std::vector<Image2DPtr> realImages(_polarizationCount),
      imagImages(_polarizationCount);
  std::vector<Mask2DPtr> masks(_polarizationCount);
  for (unsigned p = 0; p < _polarizationCount; ++p) {
    realImages[p] =
        Image2D::CreateUnsetImagePtr(_timestepsCount, baselineCount);
    imagImages[p] =
        Image2D::CreateUnsetImagePtr(_timestepsCount, baselineCount);
    masks[p] = Mask2D::CreateUnsetMaskPtr(_timestepsCount, baselineCount);
  }

  const ChannelInfo channelInfo =
      _msMetaData.GetBandInfo(0).channels[channelIndex];

  unsigned timeIndex = 0;
  double lastTime = timeColumn(0);
  for (unsigned row = 0; row < table.nrow(); ++row) {
    const int a1 = antenna1Column(row), a2 = antenna2Column(row);
    const double time = timeColumn(row);
    if (time != lastTime) {
      timeIndex++;
      lastTime = time;
    }

    if (a1 != a2) {
      const casacore::Array<casacore::Complex> data = dataColumn(row);
      const casacore::Array<bool> flags = flagColumn(row);
      const casacore::Array<double> uvws = uvwColumn(row);

      casacore::Array<casacore::Complex>::const_iterator i = data.begin();
      casacore::Array<bool>::const_iterator fI = flags.begin();
      casacore::Array<double>::const_iterator uvwIter = uvws.begin();
      ++uvwIter;
      ++uvwIter;
      const double wRotation =
          -channelInfo.MetersToLambda(*uvwIter) * M_PI * 2.0;

      const unsigned baselineIndex =
          baselineCount - (_antennaCount - a1) * (_antennaCount - a1 - 1) / 2 +
          a2 - a1 - 1;

      for (unsigned c = 0; c < _channelCount; ++c) {
        if (c == channelIndex) {
          Logger::Debug << "Reading timeIndex=" << timeIndex
                        << ", baselineIndex=" << baselineIndex << ", a1=" << a1
                        << ", a2=" << a2 << ",w=" << wRotation << "\n";
          for (unsigned p = 0; p < _polarizationCount; ++p) {
            double realValue = i->real();
            double imagValue = i->imag();
            if (fringeStop) {
              const double newRealValue = realValue * std::cos(wRotation) -
                                          imagValue * std::sin(wRotation);
              imagValue = realValue * std::sin(wRotation) +
                          imagValue * std::cos(wRotation);
              realValue = newRealValue;
            }
            realImages[p]->SetValue(timeIndex, baselineIndex, realValue);
            imagImages[p]->SetValue(timeIndex, baselineIndex, imagValue);
            ++i;
            masks[p]->SetValue(timeIndex, baselineIndex, *fI);
            ++fI;
          }
        } else {
          for (unsigned p = 0; p < _polarizationCount; ++p) {
            ++i;
            ++fI;
          }
        }
      }
    }
  }
  const casacore::ROScalarColumn<int> bandColumn(table, "DATA_DESC_ID");
  const BandInfo band = _msMetaData.GetBandInfo(bandColumn(0));

  TimeFrequencyData data;
  if (_polarizationCount == 4) {
    data = TimeFrequencyData::FromLinear(
        realImages[0], imagImages[0], realImages[1], imagImages[1],
        realImages[2], imagImages[2], realImages[3], imagImages[3]);
    data.SetIndividualPolarizationMasks(masks[0], masks[1], masks[2], masks[3]);
  } else if (_polarizationCount == 2) {
    data = TimeFrequencyData(aocommon::Polarization::XX, realImages[0],
                             imagImages[0], aocommon::Polarization::YY,
                             realImages[1], imagImages[1]);
    data.SetIndividualPolarizationMasks(masks[0], masks[1]);
  } else if (_polarizationCount == 1) {
    data = TimeFrequencyData(aocommon::Polarization::StokesI, realImages[0],
                             imagImages[0]);
    data.SetGlobalMask(masks[0]);
  } else {
    throw std::runtime_error("Unknown number of polarizations!");
  }
  return data;
}
