#include "averagingmsreader.h"

#include "../structures/mask2d.h"
#include "../structures/msmetadata.h"
#include "../structures/image2d.h"

#include "../util/progress/progresslistener.h"

#include <aocommon/multibanddata.h>

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

using aocommon::Polarization;
using aocommon::PolarizationEnum;

AveragingMsReader::AveragingMsReader(const std::string& filename,
                                     const std::string& dataColumn)
    : _filename(filename), _dataColumn(dataColumn) {
  casacore::MeasurementSet ms(filename);

  _nBands = ms.spectralWindow().nrow();

  const casacore::ScalarColumn<int> fieldIdCol(
      ms,
      casacore::MeasurementSet::columnName(casacore::MeasurementSet::FIELD_ID));
  const casacore::ScalarColumn<double> timeCol(
      ms, casacore::MeasurementSet::columnName(casacore::MeasurementSet::TIME));

  int curField = -1;
  for (size_t row = 0; row != ms.nrow(); ++row) {
    const int field = fieldIdCol(row);
    if (field != curField) {
      _sequenceStart.emplace_back(row);
      curField = field;
    }
  }
  _sequenceStart.emplace_back(ms.nrow());

  _telescopeName = MSMetaData::GetTelescopeName(ms);
}

AveragingMsReader::Result AveragingMsReader::readSamples(
    BaselineIntegrationMode statType, size_t sequenceIndex, size_t bandIndex,
    bool includeAutos, bool includeFlags, bool freqDiff,
    ProgressListener& progress) {
  progress.OnStartTask("Read data & integrate baselines");
  casacore::MeasurementSet ms(_filename);
  const casacore::ScalarColumn<int> antenna1Col(
      ms,
      casacore::MeasurementSet::columnName(casacore::MeasurementSet::ANTENNA1));
  const casacore::ScalarColumn<int> antenna2Col(
      ms,
      casacore::MeasurementSet::columnName(casacore::MeasurementSet::ANTENNA2));
  const casacore::ScalarColumn<double> timeCol(
      ms, casacore::MeasurementSet::columnName(casacore::MeasurementSet::TIME));
  const casacore::ScalarColumn<int> dataDescIdCol(
      ms, casacore::MeasurementSet::columnName(
              casacore::MeasurementSet::DATA_DESC_ID));

  const casacore::ArrayColumn<std::complex<float>> dataColumn(ms, _dataColumn);
  const casacore::ArrayColumn<bool> flagColumn(
      ms, casacore::MeasurementSet::columnName(casacore::MeasurementSet::FLAG));

  casacore::IPosition dataShape = dataColumn.shape(0);
  size_t nPolarizations = dataShape[0], nChannels = dataShape[1],
         nTotal = freqDiff ? (nPolarizations * (nChannels - 1))
                           : nPolarizations * nChannels;
  casacore::Array<std::complex<float>> dataArray(dataShape);
  casacore::Array<bool> flagArray(dataShape);

  const aocommon::MultiBandData bands(ms.spectralWindow(),
                                      ms.dataDescription());

  const size_t startRow = _sequenceStart[sequenceIndex];
  const size_t endRow = _sequenceStart[sequenceIndex + 1];
  double time = timeCol(startRow) - 1;
  std::vector<Statistic> statsData;
  std::vector<double> times;
  size_t dataPos = 0;
  const size_t totalProgress = (endRow - startRow) * 104 / 100;
  for (size_t row = startRow; row != endRow; ++row) {
    size_t antenna1 = antenna1Col(row), antenna2 = antenna2Col(row);
    const bool baselineSelected = includeAutos || (antenna1 != antenna2);

    if (baselineSelected &&
        bands.GetBandIndex(dataDescIdCol(row)) == bandIndex) {
      if (timeCol(row) != time) {
        dataPos = statsData.size();
        statsData.resize(statsData.size() + nTotal);
        time = timeCol(row);
        times.emplace_back(time);
      }
      dataColumn.get(row, dataArray);
      flagColumn.get(row, flagArray);

      if (freqDiff) {
        for (size_t i = 0; i != nTotal; ++i) {
          const bool flag1 = flagArray.cbegin()[i];
          const bool flag2 = flagArray.cbegin()[i + nPolarizations];
          if ((!flag1 && !flag2) || includeFlags) {
            const std::complex<float> val =
                dataArray.cbegin()[i] - dataArray.cbegin()[i + nPolarizations];
            statsData[dataPos + i].Add(statType, val);
          }
        }
      } else {
        for (size_t i = 0; i != nTotal; ++i) {
          const bool flag = flagArray.cbegin()[i];
          if (!flag || includeFlags) {
            const std::complex<float> val = dataArray.cbegin()[i];
            statsData[dataPos + i].Add(statType, val);
          }
        }
      }
    }
    progress.OnProgress(row - startRow, totalProgress);
  }

  progress.OnStartTask("Combining statistics");
  Result result =
      makeResult(statType, statsData.data(), nPolarizations,
                 freqDiff ? (nChannels - 1) : nChannels, times.size());
  result.second->SetObservationTimes(times);
  fillBand(result, freqDiff, bands, bandIndex);
  progress.OnFinish();
  return result;
}

AveragingMsReader::Result AveragingMsReader::readTimeDiff(
    BaselineIntegrationMode statType, size_t sequenceIndex, size_t bandIndex,
    bool includeAutos, bool includeFlags, ProgressListener& progress) {
  progress.OnStartTask("Read data & integrate baselines");
  casacore::MeasurementSet ms(_filename);
  const casacore::ScalarColumn<int> antenna1Col(
      ms,
      casacore::MeasurementSet::columnName(casacore::MeasurementSet::ANTENNA1));
  const casacore::ScalarColumn<int> antenna2Col(
      ms,
      casacore::MeasurementSet::columnName(casacore::MeasurementSet::ANTENNA2));
  const casacore::ScalarColumn<double> timeCol(
      ms, casacore::MeasurementSet::columnName(casacore::MeasurementSet::TIME));
  const casacore::ScalarColumn<int> dataDescIdCol(
      ms, casacore::MeasurementSet::columnName(
              casacore::MeasurementSet::DATA_DESC_ID));

  const casacore::ArrayColumn<std::complex<float>> dataColumn(ms, _dataColumn);
  const casacore::ArrayColumn<bool> flagColumn(
      ms, casacore::MeasurementSet::columnName(casacore::MeasurementSet::FLAG));

  casacore::IPosition dataShape = dataColumn.shape(0);
  size_t nPolarizations = dataShape[0], nChannels = dataShape[1],
         nTotal = nPolarizations * nChannels;
  casacore::Array<std::complex<float>> dataArray(dataShape);
  casacore::Array<bool> flagArray(dataShape);

  const aocommon::MultiBandData bands(ms.spectralWindow(),
                                      ms.dataDescription());

  const size_t startRow = _sequenceStart[sequenceIndex];
  const size_t endRow = _sequenceStart[sequenceIndex + 1];
  double time = timeCol(startRow) - 1;
  std::vector<Statistic> statsData;
  std::vector<double> times;
  size_t dataPos = 0;
  const size_t totalProgress = (endRow - startRow) * 104 / 100;
  using Baseline = std::pair<size_t, size_t>;
  using TimeData = std::vector<std::pair<std::complex<float>, bool>>;
  std::map<Baseline, TimeData> previousTimeData, currentTimeData;
  for (size_t row = startRow; row != endRow; ++row) {
    size_t antenna1 = antenna1Col(row), antenna2 = antenna2Col(row);
    const bool baselineSelected = includeAutos || (antenna1 != antenna2);

    if (baselineSelected &&
        bands.GetBandIndex(dataDescIdCol(row)) == bandIndex) {
      const std::pair<size_t, size_t> baseline(antenna1, antenna2);
      if (timeCol(row) != time) {
        time = timeCol(row);
        times.emplace_back(time);

        if (times.size() > 1) {
          dataPos = statsData.size();
          statsData.resize(statsData.size() + nTotal);
          previousTimeData = std::move(currentTimeData);
          currentTimeData.clear();
        }
      }

      dataColumn.get(row, dataArray);
      flagColumn.get(row, flagArray);

      auto& curBaseline = currentTimeData[baseline];
      curBaseline.resize(nTotal);
      for (size_t i = 0; i != nTotal; ++i) {
        const bool curFlag = flagArray.cbegin()[i];
        const std::complex<float> curVal = dataArray.cbegin()[i];
        curBaseline[i].first = curVal;
        curBaseline[i].second = curFlag;
      }

      auto prevBaselinePtr = previousTimeData.find(baseline);
      if (times.size() > 1 && prevBaselinePtr != previousTimeData.end()) {
        for (size_t i = 0; i != nTotal; ++i) {
          const bool curFlag = flagArray.cbegin()[i];
          const bool prevFlag = prevBaselinePtr->second[i].second;
          if ((!curFlag && !prevFlag) || includeFlags) {
            const std::complex<float> val =
                prevBaselinePtr->second[i].first - dataArray.cbegin()[i];
            statsData[dataPos + i].Add(statType, val);
          }
        }
      }
    }
    progress.OnProgress(row - startRow, totalProgress);
  }

  progress.OnStartTask("Combining statistics");
  Result result = makeResult(statType, statsData.data(), nPolarizations,
                             nChannels, times.size() - 1);
  fillBand(result, false, bands, bandIndex);
  for (size_t i = 0; i != times.size() - 1; ++i)
    times[i] = 0.5 * (times[i] + times[i + 1]);
  times.resize(times.size() - 1);
  result.second->SetObservationTimes(times);
  progress.OnFinish();
  return result;
}

AveragingMsReader::Result AveragingMsReader::makeResult(
    BaselineIntegrationMode statType, const Statistic* statsData,
    size_t nPolarizations, size_t nChannels, size_t nTimes) {
  const size_t nTotal = nPolarizations * nChannels;
  std::vector<Image2DPtr> images(nPolarizations);
  std::vector<Mask2DPtr> masks(nPolarizations);
  for (size_t p = 0; p != nPolarizations; ++p) {
    images[p] = Image2D::CreateUnsetImagePtr(nTimes, nChannels);
    masks[p] = Mask2D::CreateUnsetMaskPtr(nTimes, nChannels);
  }
  for (size_t y = 0; y != nChannels; ++y) {
    for (size_t x = 0; x != nTimes; ++x) {
      for (size_t p = 0; p != nPolarizations; ++p) {
        const Statistic& stat = statsData[x * nTotal + y * nPolarizations + p];
        if (stat.count == 0) {
          images[p]->SetValue(x, y, 0.0);
          masks[p]->SetValue(x, y, true);
        } else {
          const double val = stat.Calculate(statType);
          images[p]->SetValue(x, y, val);
          masks[p]->SetValue(x, y, false);
        }
      }
    }
  }
  std::vector<PolarizationEnum> pols;
  if (nPolarizations == 4)
    pols = std::vector<PolarizationEnum>{Polarization::XX, Polarization::XY,
                                         Polarization ::YX, Polarization::YY};
  else if (nPolarizations == 2)
    pols = std::vector<PolarizationEnum>{Polarization::XX, Polarization::YY};
  else
    pols.emplace_back(Polarization::StokesI);

  std::vector<TimeFrequencyData> tfs;
  for (size_t p = 0; p != nPolarizations; ++p) {
    tfs.emplace_back(TimeFrequencyData::AmplitudePart, pols[p], images[p]);
    tfs.back().SetGlobalMask(masks[p]);
  }
  Result result;
  if (nPolarizations == 4)
    result.first = TimeFrequencyData::MakeFromPolarizationCombination(
        tfs[0], tfs[1], tfs[2], tfs[3]);
  else if (nPolarizations == 2)
    result.first =
        TimeFrequencyData::MakeFromPolarizationCombination(tfs[0], tfs[1]);
  else
    result.first = std::move(tfs[0]);

  result.second.reset(new TimeFrequencyMetaData());

  return result;
}

void AveragingMsReader::fillBand(AveragingMsReader::Result& result,
                                 bool freqDiff,
                                 const aocommon::MultiBandData& bands,
                                 size_t bandIndex) {
  BandInfo bandInfo;
  bandInfo.windowIndex = 0;
  if (freqDiff) {
    bandInfo.channels.resize(bands[bandIndex].ChannelCount() - 1);
    for (size_t i = 0; i != bands[bandIndex].ChannelCount() - 1; ++i) {
      bandInfo.channels[i].frequencyIndex = i;
      bandInfo.channels[i].frequencyHz =
          0.5 * (bands[bandIndex].Channel(i).Frequency() +
                 bands[bandIndex].Channel(i + 1).Frequency());
      bandInfo.channels[i].effectiveBandWidthHz =
          bands[bandIndex].Channel(i).Width() +
          bands[bandIndex].Channel(i + 1).Width();
    }
  } else {
    bandInfo.channels.resize(bands[bandIndex].ChannelCount());
    for (size_t i = 0; i != bands[bandIndex].ChannelCount(); ++i) {
      bandInfo.channels[i].frequencyIndex = i;
      bandInfo.channels[i].frequencyHz =
          bands[bandIndex].Channel(i).Frequency();
      bandInfo.channels[i].effectiveBandWidthHz =
          bands[bandIndex].Channel(i).Width();
    }
  }
  result.second->SetBand(bandInfo);
}

void AveragingMsReader::StoreFlags(const std::vector<Mask2DCPtr>& flags,
                                   BaselineIntegrationDifferencing diffType,
                                   size_t sequenceIndex, size_t bandIndex,
                                   bool includeAutos) {
  std::vector<Mask2DCPtr> appliedFlags;
  switch (diffType) {
    case BaselineIntegrationDifferencing::NoDifference:
      appliedFlags = flags;
      break;
    case BaselineIntegrationDifferencing::TimeDifference:
      for (size_t i = 0; i != flags.size(); ++i) {
        const Mask2DCPtr input = flags[i];
        Mask2DPtr mask =
            Mask2D::CreateUnsetMaskPtr(input->Width() + 1, input->Height());
        for (size_t y = 0; y != input->Height(); ++y) {
          mask->SetValue(0, y, input->Value(0, y));
          for (size_t x = 1; x != input->Width(); ++x) {
            mask->SetValue(x, y, input->Value(x - 1, y) || input->Value(x, y));
          }
          mask->SetValue(input->Width(), y,
                         input->Value(input->Width() - 1, y));
        }
        appliedFlags.emplace_back(std::move(mask));
      }
      break;
    case BaselineIntegrationDifferencing::FrequencyDifference:
      for (size_t i = 0; i != flags.size(); ++i) {
        const Mask2DCPtr input = flags[i];
        Mask2DPtr mask =
            Mask2D::CreateUnsetMaskPtr(input->Width(), input->Height() + 1);
        for (size_t x = 0; x != input->Width(); ++x)
          mask->SetValue(x, 0, input->Value(x, 0));
        for (size_t y = 1; y != input->Height(); ++y) {
          for (size_t x = 0; x != input->Width(); ++x) {
            mask->SetValue(x, y, input->Value(x, y - 1) || input->Value(x, y));
          }
        }
        for (size_t x = 0; x != input->Width(); ++x)
          mask->SetValue(x, input->Height(),
                         input->Value(x, input->Height() - 1));
        appliedFlags.emplace_back(std::move(mask));
      }
      break;
  }
  storeFlags(appliedFlags, sequenceIndex, bandIndex, includeAutos);
}

void AveragingMsReader::storeFlags(const std::vector<Mask2DCPtr>& flags,
                                   size_t sequenceIndex, size_t bandIndex,
                                   bool includeAutos) {
  casacore::MeasurementSet ms(_filename,
                              casacore::MeasurementSet::TableOption::Update);
  const casacore::ScalarColumn<int> antenna1Col(
      ms,
      casacore::MeasurementSet::columnName(casacore::MeasurementSet::ANTENNA1));
  const casacore::ScalarColumn<int> antenna2Col(
      ms,
      casacore::MeasurementSet::columnName(casacore::MeasurementSet::ANTENNA2));
  const casacore::ScalarColumn<double> timeCol(
      ms, casacore::MeasurementSet::columnName(casacore::MeasurementSet::TIME));
  const casacore::ScalarColumn<int> dataDescIdCol(
      ms, casacore::MeasurementSet::columnName(
              casacore::MeasurementSet::DATA_DESC_ID));

  casacore::ArrayColumn<bool> flagColumn(
      ms, casacore::MeasurementSet::columnName(casacore::MeasurementSet::FLAG));

  casacore::IPosition dataShape = flagColumn.shape(0);
  size_t nPolarizations = dataShape[0], nChannels = dataShape[1];
  if (nPolarizations != flags.size())
    throw std::runtime_error(
        "Invalid nr of masks specified in call to MSStatReader::storeFlags()");
  casacore::Array<bool> flagArray(dataShape);

  const aocommon::MultiBandData bands(ms.spectralWindow(),
                                      ms.dataDescription());

  size_t startRow = _sequenceStart[sequenceIndex],
         endRow = _sequenceStart[sequenceIndex + 1];
  double time = timeCol(startRow) - 1;
  size_t timeIndex = 0;
  for (size_t row = startRow; row != endRow; ++row) {
    size_t antenna1 = antenna1Col(row), antenna2 = antenna2Col(row);
    const bool baselineSelected = includeAutos || (antenna1 != antenna2);

    if (baselineSelected &&
        bands.GetBandIndex(dataDescIdCol(row)) == bandIndex) {
      const std::pair<size_t, size_t> baseline(antenna1, antenna2);
      if (timeCol(row) != time) {
        time = timeCol(row);
        ++timeIndex;
      }

      flagColumn.get(row, flagArray);
      auto iter = flagArray.cbegin();
      for (size_t ch = 0; ch != nChannels; ++ch) {
        for (size_t p = 0; p != nPolarizations; ++p) {
          *iter = *iter || flags[p]->Value(timeIndex - 1, ch);
          ++iter;
        }
      }
      flagColumn.put(row, flagArray);
    }
  }
}
