#ifndef HISTOGRAM_COLLECTION_H
#define HISTOGRAM_COLLECTION_H

#include "loghistogram.h"

#include <complex>
#include <map>
#include <vector>

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

#include "../util/serializable.h"

class HistogramCollection : public Serializable
{
	public:
		typedef std::pair<unsigned, unsigned> AntennaPair;
		
		HistogramCollection() : _polarizationCount(0), _totalHistograms(nullptr), _rfiHistograms(nullptr)
		{
		}
		
		explicit HistogramCollection(unsigned polarizationCount) : _polarizationCount(polarizationCount)
		{
			init();
		}
		
		HistogramCollection(const HistogramCollection &source) : _polarizationCount(source._polarizationCount)
		{
			init();
			for(unsigned i=0;i<_polarizationCount;++i)
			{
				copy(_totalHistograms[i], source._totalHistograms[i]);
				copy(_rfiHistograms[i], source._rfiHistograms[i]);
			}
		}
		
		~HistogramCollection()
		{
			destruct();
		}
		
		void SetPolarizationCount(unsigned polarizationCount)
		{
			destruct();
			_polarizationCount = polarizationCount;
			init();
		}
		
		void Add(const unsigned antenna1, const unsigned antenna2, const unsigned polarization, const std::complex<float> *values, const bool *isRFI, size_t sampleCount)
		{
			LogHistogram &totalHistogram = GetTotalHistogram(antenna1, antenna2, polarization);
			LogHistogram &rfiHistogram = GetRFIHistogram(antenna1, antenna2, polarization);
			
			for(size_t i=0;i<sampleCount;++i)
			{
				const double amplitude = sqrt(values[i].real()*values[i].real() + values[i].imag()*values[i].imag());
				totalHistogram.Add(amplitude);
				if(isRFI[i])
					rfiHistogram.Add(amplitude);
			}
		}
		
		void Add(const HistogramCollection &collection)
		{
			if(collection._polarizationCount != _polarizationCount)
				throw std::runtime_error("Polarization counts of histogram collections don't match");
			for(unsigned p=0;p<_polarizationCount;++p)
			{
				add(collection, p, p);
			}
		}
		
		bool Empty() const
		{
			if(_polarizationCount == 0)
				return true;
			for(unsigned p=0;p!=_polarizationCount;++p)
			{
				if(!_totalHistograms[p].empty() || !_rfiHistograms[p].empty())
					return false;
			}
			return true;
		}
		
		void Add(const unsigned antenna1, const unsigned antenna2, const unsigned polarization, Image2DCPtr image, Mask2DCPtr mask);
		
		void Add(const unsigned antenna1, const unsigned antenna2, const unsigned polarization, Image2DCPtr image, Mask2DCPtr flagMask, Mask2DCPtr correlatorMask);
		
		void Add(const unsigned antenna1, const unsigned antenna2, const unsigned polarization, Image2DCPtr real, Image2DCPtr imaginary, Mask2DCPtr flagMask, Mask2DCPtr correlatorMask);
		
		LogHistogram &GetTotalHistogram(const unsigned a1, const unsigned a2, const unsigned polarization)
		{
			return getHistogram(_totalHistograms, a1, a2, polarization);
		}
		
		LogHistogram &GetRFIHistogram(const unsigned a1, const unsigned a2, const unsigned polarization)
		{
			return getHistogram(_rfiHistograms, a1, a2, polarization);
		}
		
		const std::map<AntennaPair, LogHistogram*> &GetTotalHistogram(const unsigned polarization) const
		{
			return _totalHistograms[polarization];
		}
		
		const std::map<AntennaPair, LogHistogram*> &GetRFIHistogram(const unsigned polarization) const
		{
			return _rfiHistograms[polarization];
		}
		
		void GetTotalHistogramForCrossCorrelations(const unsigned polarization, LogHistogram &target) const
		{
			getHistogramForCrossCorrelations(_totalHistograms, polarization, target);
		}
		
		void GetRFIHistogramForCrossCorrelations(const unsigned polarization, LogHistogram &target) const
		{
			getHistogramForCrossCorrelations(_rfiHistograms, polarization, target);
		}
		
		void Clear()
		{
			destruct();
			init();
		}
		
		void Save(class HistogramTablesFormatter &histogramTables);
		
		void Load(class HistogramTablesFormatter &histogramTables);
		
		unsigned PolarizationCount() const { return _polarizationCount; }
		
		virtual void Serialize(std::ostream &stream) const final override
		{
			SerializeToUInt64(stream, _polarizationCount);
			serializeMapArray(stream, _totalHistograms);
			serializeMapArray(stream, _rfiHistograms);
		}
		
		virtual void Unserialize(std::istream &stream) final override
		{
			destruct();
			_polarizationCount = UnserializeUInt64(stream);
			init();
			unserializeMapArray(stream, _totalHistograms);
			unserializeMapArray(stream, _rfiHistograms);
		}
		
		HistogramCollection *CreateSummedPolarizationCollection() const
		{
			HistogramCollection *newCollection = new HistogramCollection(1);
			for(unsigned p=0;p<_polarizationCount;++p)
				newCollection->add(*this, p, 0);
			return newCollection;
		}
		
		void Rescale(double factor)
		{
			for(unsigned p=0;p<_polarizationCount;++p)
			{
				for(std::map<AntennaPair, LogHistogram*>::iterator i=_totalHistograms[p].begin(); i!=_totalHistograms[p].end(); ++i)
				{
					i->second->Rescale(factor);
				}
				for(std::map<AntennaPair, LogHistogram*>::iterator i=_rfiHistograms[p].begin(); i!=_rfiHistograms[p].end(); ++i)
				{
					i->second->Rescale(factor);
				}
			}
		}
		
		void CreateMissingBins()
		{
			for(unsigned p=0;p<_polarizationCount;++p)
			{
				for(std::map<AntennaPair, LogHistogram*>::iterator i=_totalHistograms[p].begin(); i!=_totalHistograms[p].end(); ++i)
				{
					i->second->CreateMissingBins();
				}
				for(std::map<AntennaPair, LogHistogram*>::iterator i=_rfiHistograms[p].begin(); i!=_rfiHistograms[p].end(); ++i)
				{
					i->second->CreateMissingBins();
				}
			}
		}
	private:
		unsigned _polarizationCount;
		std::map<AntennaPair, LogHistogram*> *_totalHistograms;
		std::map<AntennaPair, LogHistogram*> *_rfiHistograms;
		
		void init()
		{
			if(_polarizationCount != 0)
			{
				_totalHistograms = new std::map<AntennaPair, LogHistogram*>[_polarizationCount];
				_rfiHistograms = new std::map<AntennaPair, LogHistogram*>[_polarizationCount];
			} else {
				_totalHistograms = 0;
				_rfiHistograms = 0;
			}
		}
		
		void destruct()
		{
			if(_polarizationCount != 0)
			{
				for(unsigned p=0;p<_polarizationCount;++p)
				{
					for(std::map<AntennaPair, LogHistogram*>::iterator i=_totalHistograms[p].begin(); i!=_totalHistograms[p].end(); ++i)
					{
						delete i->second;
					}
					for(std::map<AntennaPair, LogHistogram*>::iterator i=_rfiHistograms[p].begin(); i!=_rfiHistograms[p].end(); ++i)
					{
						delete i->second;
					}
				}
				delete[] _totalHistograms;
				delete[] _rfiHistograms;
			}
		}
		
		void serializeMapArray(std::ostream &stream, const std::map<AntennaPair, LogHistogram*> *map) const
		{
			for(unsigned p=0;p<_polarizationCount;++p)
				serializeMap(stream, map[p]);
		}
		
		void unserializeMapArray(std::istream &stream, std::map<AntennaPair, LogHistogram*> *map)
		{
			for(unsigned p=0;p<_polarizationCount;++p)
				unserializeMap(stream, map[p]);
		}
		
		void serializeMap(std::ostream &stream, const std::map<AntennaPair, LogHistogram*> &map) const
		{
			SerializeToUInt64(stream, map.size());
			for(std::map<AntennaPair, LogHistogram*>::const_iterator i=map.begin();i!=map.end();++i)
			{
				const AntennaPair &antennae = i->first;
				const LogHistogram *histogram = i->second;
				SerializeToUInt32(stream, antennae.first);
				SerializeToUInt32(stream, antennae.second);
				histogram->Serialize(stream);
			}
		}
		
		void unserializeMap(std::istream &stream, std::map<AntennaPair, LogHistogram*> &map)
		{
			map.clear();
			size_t mapSize = UnserializeUInt64(stream);
			std::map<AntennaPair, LogHistogram*>::iterator insertPos = map.begin();
			for(size_t i=0;i!=mapSize;++i)
			{
				std::pair<AntennaPair, LogHistogram*> p;
				p.first.first = UnserializeUInt32(stream);
				p.first.second = UnserializeUInt32(stream);
				p.second = new LogHistogram();
				p.second->Unserialize(stream);
				insertPos = map.insert(insertPos, p);
			}
		}
		
		void copy(std::map<AntennaPair, LogHistogram*> &destination, const std::map<AntennaPair, LogHistogram*> &source)
		{
			for(std::map<AntennaPair, LogHistogram*>::const_iterator i=source.begin();i!=source.end();++i)
			{
				destination.insert(std::pair<AntennaPair, LogHistogram*>(i->first, new LogHistogram(*i->second)));
			}
		}

		void add(const HistogramCollection &collection, unsigned fromPolarization, unsigned toPolarization)
		{
			for(std::map<AntennaPair, LogHistogram*>::const_iterator i=collection._totalHistograms[fromPolarization].begin(); i!=collection._totalHistograms[fromPolarization].end(); ++i)
			{
				LogHistogram &histogram = GetTotalHistogram(i->first.first, i->first.second, toPolarization);
				histogram.Add(*i->second);
			}
			
			for(std::map<AntennaPair, LogHistogram*>::const_iterator i=collection._rfiHistograms[fromPolarization].begin(); i!=collection._rfiHistograms[fromPolarization].end(); ++i)
			{
				LogHistogram &histogram = GetRFIHistogram(i->first.first, i->first.second, toPolarization);
				histogram.Add(*i->second);
			}
		}
		
		LogHistogram &getHistogram(std::map<AntennaPair, LogHistogram*> *histograms, const unsigned a1, const unsigned a2, const unsigned polarization)
		{
			const AntennaPair antennae(a1, a2);
			std::map<AntennaPair, LogHistogram*>::iterator i = histograms[polarization].find(antennae);
			if(i == histograms[polarization].end())
			{
				i = histograms[polarization].insert(std::pair<AntennaPair, LogHistogram*>(antennae, new LogHistogram())).first;
			}
			return *i->second;
		}
		
		void getHistogramForCrossCorrelations(std::map<AntennaPair, LogHistogram*> *histograms, const unsigned polarization, LogHistogram &target) const
		{
			for(std::map<AntennaPair, LogHistogram*>::const_iterator i=histograms[polarization].begin(); i!=histograms[polarization].end(); ++i)
			{
				if(i->first.first != i->first.second)
					target.Add(*i->second);
			}
		}
};

#endif
