#ifndef MS_STAT_SET_H
#define MS_STAT_SET_H

#include <memory>
#include <queue>
#include <string>
#include <utility>
#include <vector>

#include "../msio/averagingmsreader.h"

#include "imageset.h"

namespace imagesets {

class AveragingMsSet final : public ImageSet {
 public:
  AveragingMsSet(const std::string& path, const std::string& dataColumn,
                 BaselineIntegrationMode mode,
                 BaselineIntegrationDifferencing diffType, bool includeAutos,
                 bool includeFlags)
      : path_(path),
        mode_(mode),
        differencing_type_(diffType),
        include_autos_(includeAutos),
        include_flags_(includeFlags),
        reader_(std::make_shared<AveragingMsReader>(
            path_, dataColumn.empty() ? "DATA" : dataColumn)) {}

  std::unique_ptr<ImageSet> Clone() final {
    return std::make_unique<AveragingMsSet>(*this);
  }

  size_t Size() const final {
    return reader_->NBands() * reader_->NSequences();
  }

  std::string Description(const ImageSetIndex& index) const final {
    return Name() + " (over all baselines, " + ToString(mode_) + ", " +
           ToString(differencing_type_) + ")";
  }

  void Initialize() final {}

  std::string Name() const final { return path_; }

  std::vector<std::string> Files() const final {
    return std::vector<std::string>{path_};
  }

  std::string TelescopeName() final { return reader_->TelescopeName(); }

  void AddReadRequest(const ImageSetIndex& index) final {
    _requests.emplace_back(index);
  }

  void PerformReadRequests(class ProgressListener& progress) final {
    for (const ImageSetIndex& request : _requests) {
      size_t bandIndex = request.Value() % reader_->NBands(),
             seqIndex = request.Value() / reader_->NBands();
      AveragingMsReader::Result result =
          reader_->Read(mode_, differencing_type_, seqIndex, bandIndex,
                        include_autos_, include_flags_, progress);
      _results.emplace(result.first, result.second, request);
    }
    _requests.clear();
  }

  std::unique_ptr<BaselineData> GetNextRequested() final {
    std::unique_ptr<BaselineData> result(
        new BaselineData(std::move(_results.front())));
    _results.pop();
    return result;
  }

  void AddWriteFlagsTask(const ImageSetIndex& index,
                         std::vector<Mask2DCPtr>& flags) final {
    size_t bandIndex = index.Value() % reader_->NBands(),
           seqIndex = index.Value() / reader_->NBands();
    reader_->StoreFlags(flags, differencing_type_, seqIndex, bandIndex,
                        include_autos_);
  }

  void PerformWriteFlagsTask() final {}

  bool HasCrossCorrelations() const final { return false; }

 private:
  std::string path_;
  BaselineIntegrationMode mode_;
  BaselineIntegrationDifferencing differencing_type_;
  bool include_autos_, include_flags_;
  std::shared_ptr<AveragingMsReader> reader_;
  std::vector<ImageSetIndex> _requests;
  std::queue<BaselineData> _results;
};

}  // namespace imagesets

#endif
