#include "filterbankset.h"

#include "../../structures/date.h"
#include "../../structures/system.h"

#include "../../lua/telescopefile.h"

#include <fstream>

namespace rfiStrategy {

FilterBankSet::FilterBankSet(const std::string& location)
    : _location(location),
      _timeOfSample(0.0),
      _timeStart(0.0),
      _fch1(0.0),
      _foff(0.0),
      _channelCount(0),
      _ifCount(0),
      _bitCount(0),
      _sampleCount(0),
      _nBeams(0),
      _iBeam(0),
      _machineId(0),
      _intervalCount(0),
      _headerEnd(0) {
  std::ifstream file(_location.c_str());
  if (!file.good())
    throw std::runtime_error(std::string("Error opening filterbank file ") +
                             _location);

  std::string keyword = readString(file);
  if (keyword != "HEADER_START")
    throw std::runtime_error(
        "Filterbank file does not start with 'HEADER_START' -- corrupt file?");

  while (file.good() && keyword != "HEADER_END") {
    keyword = readString(file);

    if (keyword == "tsamp")
      _timeOfSample = readDouble(file);
    else if (keyword == "tstart")
      _timeStart = readDouble(file);
    else if (keyword == "fch1")
      _fch1 = readDouble(file);
    else if (keyword == "foff")
      _foff = readDouble(file);
    else if (keyword == "nchans")
      _channelCount = readInt(file);
    else if (keyword == "nifs")
      _ifCount = readInt(file);
    else if (keyword == "nbits")
      _bitCount = readInt(file);
    else if (keyword == "nsamples")
      _sampleCount = readInt(file);
    else if (keyword == "machine_id")
      _machineId = readInt(file);
    else if (keyword == "telescope_id")
      _telescopeId = readInt(file);
    else if (keyword == "nbeams")
      _nBeams = readInt(file);
    else if (keyword == "ibeam")
      _iBeam = readInt(file);
    else if (keyword == "src_raj" || keyword == "src_dej" ||
             keyword == "az_start" || keyword == "za_start" ||
             keyword == "refdm" || keyword == "period")
      readDouble(file);
    else if (keyword == "data_type" || keyword == "barycentric" ||
             keyword == "pulsarcentric")
      readInt(file);
  }
  _headerEnd = file.tellg();
  if (_sampleCount == 0) {
    file.seekg(0, std::ios::end);
    std::streampos endPos = file.tellg();
    size_t dataSize = endPos - _headerEnd;
    _sampleCount = (dataSize * 8) / _channelCount / _bitCount;
  }
  Logger::Debug << "tsamp=" << _timeOfSample << ", tstart=" << _timeStart
                << ", fch1=" << _fch1 << ", foff=" << _foff << '\n'
                << "nChans=" << _channelCount << ", nIFs=" << _ifCount
                << ", nBits=" << _bitCount << ", nSamples=" << _sampleCount
                << "\nmachine_ID=" << _machineId
                << ", telescope_ID=" << _telescopeId << '\n';

  _timeStart = Date::MJDToAipsMJD(_timeStart);

  double sizeOfImage = double(_channelCount) * _sampleCount * _bitCount / 8.0;
  double memSize = System::TotalMemory();
  _intervalCount = ceil(sizeOfImage / (memSize / 16.0));
  if (_intervalCount < 1) _intervalCount = 1;
  if (_intervalCount * 8 > _sampleCount) _intervalCount = _sampleCount / 8;
  Logger::Debug << round(sizeOfImage * 1e-8) * 0.1
                << " GB/image required of total of "
                << round(memSize * 1e-8) * 0.1 << " GB of mem, splitting in "
                << _intervalCount << " intervals\n";
}

void FilterBankSet::AddReadRequest(const ImageSetIndex& index) {
  _requests.push_back(new BaselineData(index));
}

std::unique_ptr<BaselineData> FilterBankSet::GetNextRequested() {
  if (_bitCount != 32)
    throw std::runtime_error(
        "Only support for 32-bit filterbank sets has been added as of yet");

  std::unique_ptr<BaselineData> baseline(std::move(_requests.front()));
  _requests.pop_front();
  const size_t intervalIndex = baseline->Index().Value();

  const size_t startIndex = (_sampleCount * intervalIndex) / _intervalCount,
               endIndex = (_sampleCount * (intervalIndex + 1)) / _intervalCount;

  std::ifstream file(_location.c_str());
  file.seekg(_headerEnd +
             std::streampos(startIndex * sizeof(float) * _channelCount));

  Image2DPtr image =
      Image2D::CreateUnsetImagePtr(endIndex - startIndex, _channelCount);
  Mask2DPtr mask =
      Mask2D::CreateUnsetMaskPtr(endIndex - startIndex, _channelCount);
  std::vector<float> buffer(_channelCount);
  for (size_t x = 0; x != endIndex - startIndex; ++x) {
    file.read(reinterpret_cast<char*>(&buffer[0]),
              _channelCount * sizeof(float));
    for (size_t y = 0; y != _channelCount; ++y) {
      image->SetValue(x, y, buffer[y]);
      mask->SetValue(x, y, !std::isfinite(buffer[y]));
    }
  }
  TimeFrequencyData tfData(TimeFrequencyData::AmplitudePart,
                           aocommon::Polarization::StokesI, image);
  tfData.SetGlobalMask(mask);
  TimeFrequencyMetaDataPtr metaData(new TimeFrequencyMetaData());
  AntennaInfo antenna;
  antenna.diameter = 0;
  antenna.id = 0;
  antenna.mount = "unknown";
  antenna.name = "unknown";
  antenna.position = EarthPosition();
  antenna.station = "unknown";
  metaData->SetAntenna1(antenna);
  metaData->SetAntenna2(antenna);
  BandInfo band;
  band.windowIndex = 0;
  for (size_t ch = 0; ch != _channelCount; ++ch) {
    ChannelInfo channel;
    channel.frequencyHz = 1e6 * (_fch1 + _foff * ch);
    channel.effectiveBandWidthHz = 1e6 * std::fabs(_foff);
    channel.frequencyIndex = ch;
    channel.channelWidthHz = 1e6 * std::fabs(_foff);
    channel.resolutionHz = 1e6 * std::fabs(_foff);
    band.channels.push_back(channel);
  }
  metaData->SetBand(band);
  std::vector<double> observationTimes(endIndex - startIndex);
  for (size_t t = startIndex; t != endIndex; ++t)
    observationTimes[t - startIndex] = (_timeStart + _timeOfSample * t);
  metaData->SetObservationTimes(observationTimes);
  metaData->SetValueDescription("Power");

  baseline->SetData(tfData);
  baseline->SetMetaData(metaData);

  return baseline;
}

void FilterBankSet::AddWriteFlagsTask(const ImageSetIndex& index,
                                      std::vector<Mask2DCPtr>& flags) {
  if (_bitCount != 32)
    throw std::runtime_error(
        "Only support for 32-bit filterbank sets has been added as of yet");

  const size_t intervalIndex = index.Value();

  const size_t startIndex = (_sampleCount * intervalIndex) / _intervalCount,
               endIndex = (_sampleCount * (intervalIndex + 1)) / _intervalCount;

  std::fstream file(_location.c_str());
  file.seekg(_headerEnd +
             std::streampos(startIndex * sizeof(float) * _channelCount));

  std::vector<float> buffer(_channelCount);
  for (size_t x = 0; x != endIndex - startIndex; ++x) {
    std::streampos pos = file.tellg();
    file.read(reinterpret_cast<char*>(&buffer[0]),
              _channelCount * sizeof(float));
    for (size_t y = 0; y != _channelCount; ++y) {
      if (flags[0]->Value(x, y))
        buffer[y] = std::numeric_limits<float>::quiet_NaN();
    }
    file.seekp(pos);
    file.write(reinterpret_cast<char*>(&buffer[0]),
               _channelCount * sizeof(float));
  }
}

std::string FilterBankSet::TelescopeName() {
  return TelescopeFile::TelescopeName(TelescopeFile::GENERIC_TELESCOPE);
}

void FilterBankSet::Initialize() {}

void FilterBankSet::PerformReadRequests(class ProgressListener&) {}

void FilterBankSet::PerformWriteDataTask(
    const ImageSetIndex& index, std::vector<Image2DCPtr> realImages,
    std::vector<Image2DCPtr> imaginaryImages) {
  throw std::runtime_error(
      "Can't write data back to filter bank format: not implemented");
}

std::string FilterBankSet::Description(const ImageSetIndex& index) const {
  std::ostringstream str;
  str << "Filterbank set -- interval " << (index.Value() + 1) << '/'
      << _intervalCount;
  return str.str();
}

}  // namespace rfiStrategy
