#include "bhfitsimageset.h"

#include "../msio/fitsfile.h"
#include "../structures/image2d.h"
#include "../structures/timefrequencydata.h"

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

namespace imagesets {

BHFitsImageSet::BHFitsImageSet(const std::string& file)
    : ImageSet(), _file(new FitsFile(file)), _width(0), _height(0) {
  Logger::Debug << "Opening bhfits file: '" << file << "'\n";
  try {
    _file->Open(FitsFile::ReadWriteMode);
  } catch (FitsIOException& exception) {
    Logger::Error
        << "CFitsio failed to open file in RW mode with the following error:\n"
        << exception.what()
        << "\nTrying reopening in read-only mode. Writing to file won't be "
           "possible.\n\n";
    _file->Open(FitsFile::ReadOnlyMode);
  }
}

BHFitsImageSet::BHFitsImageSet(const BHFitsImageSet& source)
    : ImageSet(),
      _file(source._file),
      _baselineData(source._baselineData),
      _timeRanges(source._timeRanges),
      _width(source._width),
      _height(source._height) {}

BHFitsImageSet::~BHFitsImageSet() = default;

std::unique_ptr<ImageSet> BHFitsImageSet::Clone() {
  return std::unique_ptr<BHFitsImageSet>(new BHFitsImageSet(*this));
}

void BHFitsImageSet::Initialize() {
  _file->MoveToHDU(1);
  /*for(int i=1;i<=_file->GetKeywordCount();++i)
          {
                  Logger::Debug << _file->GetKeyword(i) << " = " <<
     _file->GetKeywordValue(i) << '\n';
          }*/
  if (_file->GetCurrentHDUType() != FitsFile::ImageHDUType)
    throw std::runtime_error(
        "Error in Bighorns fits files: first HDU was not an image HDU");
  if (_file->GetCurrentImageDimensionCount() != 2)
    throw std::runtime_error("Fits image was not two dimensional");
  _width = _file->GetCurrentImageSize(2),
  _height = _file->GetCurrentImageSize(1);
  Logger::Debug << "Image of " << _width << " x " << _height << '\n';

  _timeRanges.clear();
  size_t keyIndex = 0;
  bool searchOn;
  do {
    searchOn = false;
    std::ostringstream antKey, termKey;
    antKey << "ANT";
    termKey << "TERM";
    if (keyIndex < 10) {
      antKey << '0';
      termKey << '0';
    }
    antKey << keyIndex;
    termKey << keyIndex;
    std::string antRangeStr, termRangeStr;
    if (_file->GetKeywordValue(antKey.str(), antRangeStr)) {
      const std::pair<int, int> range = getRangeFromString(antRangeStr);
      TimeRange timeRange;
      // As commented by Marcin below, the ranges in the fits headers are given
      // like 'start - end', where the indices start counting at 1, and the end
      // index is *inclusive*, such that '1 - 1' represents a range with one
      // index (being the first timestep in the FITS file). Comment by Marcin
      // Sokolowski: MS this is due to fact that Andre didn't know it starts
      // from 1, but he skips the end integration (assumes it is not ANT), but
      // it is
      timeRange.start = range.first - 1;
      // so here I don't subtract 1 in order to program check until this one too
      // !
      timeRange.end = range.second;
      timeRange.name = antKey.str();
      _timeRanges.push_back(timeRange);
      searchOn = true;
    }
    if (_file->GetKeywordValue(termKey.str(), termRangeStr)) {
      const std::pair<int, int> range = getRangeFromString(termRangeStr);
      TimeRange timeRange;
      // (see earlier comment by Marcin Sokolowski)
      timeRange.start = range.first - 1;
      timeRange.end = range.second;
      timeRange.name = termKey.str();
      _timeRanges.push_back(timeRange);
      searchOn = true;
    }
    ++keyIndex;
  } while (searchOn);
  Logger::Debug << "This file has " << _timeRanges.size() << " time ranges.\n";

  if (_timeRanges.empty()) {
    // if no states found in the header - just assume all are antenna
    TimeRange timeRange;
    // See earlier comment by Marcin Sokolowski
    timeRange.start = 0;
    timeRange.end = _file->GetCurrentImageSize(2);
    timeRange.name = "ANT";
    _timeRanges.push_back(timeRange);
    Logger::Warn << "No states specified in the fits header assuming all (1-"
                 << timeRange.end << ") integrations are " << timeRange.name
                 << "\n";
  }
}

BaselineData BHFitsImageSet::loadData(const ImageSetIndex& index) {
  const TimeFrequencyMetaDataPtr metaData(new TimeFrequencyMetaData());
  TimeFrequencyData data;
  loadImageData(data, metaData, index);
  return BaselineData(data, metaData, index);
}

void BHFitsImageSet::loadImageData(TimeFrequencyData& data,
                                   const TimeFrequencyMetaDataPtr& metaData,
                                   const ImageSetIndex& index) {
  std::vector<num_t> buffer(_width * _height);
  _file->ReadCurrentImageData(0, &buffer[0], _width * _height);

  int rangeStart = _timeRanges[index.Value()].start,
      rangeEnd = _timeRanges[index.Value()].end;
  const Image2DPtr image =
      Image2D::CreateZeroImagePtr(rangeEnd - rangeStart, _height);

  std::vector<num_t>::const_iterator bufferPtr =
      buffer.begin() + _height * rangeStart;
  for (int x = rangeStart; x != rangeEnd; ++x) {
    for (int y = 0; y != _height; ++y) {
      image->SetValue(x - rangeStart, y, *bufferPtr);
      ++bufferPtr;
    }
  }
  data = TimeFrequencyData(TimeFrequencyData::AmplitudePart,
                           aocommon::Polarization::StokesI, image);

  try {
    FitsFile flagFile(flagFilePath());
    flagFile.Open(FitsFile::ReadOnlyMode);
    flagFile.ReadCurrentImageData(0, &buffer[0], _width * _height);
    bufferPtr = buffer.begin() + _height * rangeStart;
    const Mask2DPtr mask =
        Mask2D::CreateUnsetMaskPtr(rangeEnd - rangeStart, _height);
    for (int x = rangeStart; x != rangeEnd; ++x) {
      for (int y = 0; y != _height; ++y) {
        bool flag = false;
        if (*bufferPtr == 0.0)
          flag = false;
        else if (*bufferPtr == 1.0)
          flag = true;
        else
          std::runtime_error(
              "Expecting a flag file with only ones and zeros, but this file "
              "contained other values.");
        mask->SetValue(x - rangeStart, y, flag);
        ++bufferPtr;
      }
    }
    data.SetGlobalMask(mask);
  } catch (std::exception&) {
    // Flag file could not be read; probably does not exist. Ignore this, flags
    // will be initialized to false.
  }

  double frequencyDelta = _file->GetDoubleKeywordValue("CDELT1"),
         timeDelta = _file->GetDoubleKeywordValue("CDELT2");
  BandInfo band;
  for (int ch = 0; ch != _height; ++ch) {
    ChannelInfo channel;
    channel.frequencyHz = ch * frequencyDelta * 1000000.0;
    band.channels.push_back(channel);
  }
  metaData->SetBand(band);

  const int rangeWidth = rangeEnd - rangeStart;
  std::vector<double> observationTimes(rangeWidth);
  for (int t = 0; t != rangeWidth; ++t)
    observationTimes[t] = (t + rangeStart) * timeDelta;
  metaData->SetObservationTimes(observationTimes);

  AntennaInfo antennaInfo;
  antennaInfo.id = 0;
  antennaInfo.name = RangeName(index.Value());
  antennaInfo.diameter = 0.0;
  antennaInfo.mount = "Unknown";
  antennaInfo.station = TelescopeName();
  metaData->SetAntenna1(antennaInfo);
  metaData->SetAntenna2(antennaInfo);
}

std::pair<int, int> BHFitsImageSet::getRangeFromString(
    const std::string& rangeStr) {
  std::pair<int, int> value;
  const size_t partA = rangeStr.find(' ');
  value.first = atoi(rangeStr.substr(0, partA).c_str());
  size_t partB = rangeStr.find('-');
  if (rangeStr[partB + 1] == ' ') ++partB;
  value.second = atoi(rangeStr.substr(partB + 1).c_str());
  return value;
}

std::string BHFitsImageSet::flagFilePath() const {
  std::string flagFilePath = _file->Filename();
  if (flagFilePath.size() > 7) {
    flagFilePath = flagFilePath.substr(0, flagFilePath.size() - 7);
  }
  flagFilePath += "_flag.fits";
  return flagFilePath;
}

void BHFitsImageSet::AddWriteFlagsTask(const ImageSetIndex& index,
                                       std::vector<Mask2DCPtr>& flags) {
  if (flags.size() != 1)
    throw std::runtime_error(
        "BHFitsImageSet::AddWriteFlagsTask() called with multiple flags");
  const std::string flagFilename = flagFilePath();
  Logger::Debug << "Writing to " << flagFilename << '\n';
  FitsFile flagFile(flagFilename);
  bool newFile = true;
  std::vector<num_t> buffer(_width * _height);
  try {
    flagFile.Open(FitsFile::ReadWriteMode);
    newFile = false;
  } catch (std::exception&) {
    Logger::Debug << "File did not exist yet, creating new.\n";
    flagFile.Create();
    flagFile.AppendImageHUD(FitsFile::Float32ImageType, _height, _width);
  }

  // This must be outside the try { } block, so that exceptions
  // don't result in creating a new file.
  if (!newFile) {
    flagFile.ReadCurrentImageData(0, &buffer[0], _width * _height);
  }

  int rangeStart = _timeRanges[index.Value()].start,
      rangeEnd = _timeRanges[index.Value()].end;
  std::vector<num_t>::iterator bufferPtr =
      buffer.begin() + _height * rangeStart;
  for (int x = rangeStart; x != rangeEnd; ++x) {
    for (int y = 0; y != _height; ++y) {
      *bufferPtr = flags[0]->Value(x - rangeStart, y) ? 1.0 : 0.0;
      ++bufferPtr;
    }
  }

  flagFile.WriteImage(0, &buffer[0], _width * _height, -1.0);
}

void BHFitsImageSet::PerformWriteFlagsTask() {
  // Nothing to do; already written
}

std::string BHFitsImageSet::Description(const ImageSetIndex& index) const {
  std::ostringstream str;
  str << "Time range " << RangeName(index.Value());
  return str.str();
}

std::vector<std::string> BHFitsImageSet::Files() const {
  return std::vector<std::string>{_file->Filename()};
}
}  // namespace imagesets