#include "timefrequencydata.h"
#include "stokesimager.h"

#include "../util/ffttools.h"

Image2DCPtr TimeFrequencyData::GetAbsoluteFromComplex(const Image2DCPtr &real, const Image2DCPtr &imag) const
{
	return Image2DPtr(FFTTools::CreateAbsoluteImage(*real, *imag));
}
			
Image2DCPtr TimeFrequencyData::GetSum(const Image2DCPtr &left, const Image2DCPtr &right) const
{
	return StokesImager::CreateSum(left, right);
}

Image2DCPtr TimeFrequencyData::GetNegatedSum(const Image2DCPtr &left, const Image2DCPtr &right) const
{
	return StokesImager::CreateNegatedSum(left, right);
}

Image2DCPtr TimeFrequencyData::GetDifference(const Image2DCPtr &left, const Image2DCPtr &right) const
{
	return StokesImager::CreateDifference(left, right);
}

Image2DCPtr TimeFrequencyData::getSinglePhaseFromTwoPolPhase(size_t polA, size_t polB) const
{
	return StokesImager::CreateAvgPhase(_data[polA]._images[0], _data[polB]._images[0]);
}

Mask2DCPtr TimeFrequencyData::GetCombinedMask() const
{
	if(MaskCount() == 0)
		return GetSetMask<false>();
	else if(MaskCount() == 1)
		return GetMask(0);
	else
	{
		Mask2DPtr mask(new Mask2D(*GetMask(0)));
		size_t i = 0;
		while(i!= MaskCount())
		{
			const Mask2DCPtr& curMask = GetMask(i);
			for(unsigned y=0;y<mask->Height();++y) {
				for(unsigned x=0;x<mask->Width();++x) {
					bool v = curMask->Value(x, y);
					if(v)
						mask->SetValue(x, y, true);
				}
			}
			++i;
		}
		return mask;
	}
}

TimeFrequencyData TimeFrequencyData::MakeZeroLinearData(size_t width, size_t height)
{
  Image2DPtr zero = Image2D::CreateSetImagePtr(width, height, 0.0);
  TimeFrequencyData data;
  data._data.resize(4);
  data._data[0] = PolarizedTimeFrequencyData(aocommon::Polarization::XX, zero, zero);
  data._data[1] = PolarizedTimeFrequencyData(aocommon::Polarization::XY, zero, zero);
  data._data[2] = PolarizedTimeFrequencyData(aocommon::Polarization::YX, zero, zero);
  data._data[3] = PolarizedTimeFrequencyData(aocommon::Polarization::YY, zero, zero);
  return data;
}
	
TimeFrequencyData TimeFrequencyData::Make(enum ComplexRepresentation representation) const
{
	if(representation == _complexRepresentation)
		return TimeFrequencyData(*this);
	else if(_complexRepresentation == ComplexParts)
	{
		TimeFrequencyData data;
		data._complexRepresentation = representation;
		data._data.resize(_data.size());
		for(size_t i=0; i!=_data.size(); ++i)
		{
			const PolarizedTimeFrequencyData& source = _data[i];
			PolarizedTimeFrequencyData& dest = data._data[i];
			dest._polarization = source._polarization;
			dest._flagging = source._flagging;
			switch(representation)
			{
				case RealPart:
					dest._images[0] = source._images[0];
					break;
				case ImaginaryPart:
					dest._images[0] = source._images[1];
					break;
				case AmplitudePart:
					dest._images[0] = GetAbsoluteFromComplex(source._images[0], source._images[1]);
					break;
				case PhasePart:
					dest._images[0] = StokesImager::CreateAvgPhase(source._images[0], source._images[1]);
					break;
				case ComplexParts:
					break; // already handled above.
			}
		}
		return data;
	} else if(representation == ComplexParts && _complexRepresentation == AmplitudePart)
	{
		return MakeFromComplexCombination(*this, *this);
	}
	else throw std::runtime_error("Request for time/frequency data with a phase representation that can not be extracted from the source (source is not complex)");

}

TimeFrequencyData TimeFrequencyData::MakeFromComplexCombination(const TimeFrequencyData& real, const TimeFrequencyData& imaginary)
{
	if(real.ComplexRepresentation() == ComplexParts ||
		imaginary.ComplexRepresentation() == ComplexParts)
		throw std::runtime_error("Trying to create complex TF data from incorrect phase representations");
	if(real.Polarizations() != imaginary.Polarizations())
		throw std::runtime_error("Combining real/imaginary time frequency data from different polarisations");
	TimeFrequencyData data;
	data._data.resize(real._data.size());
	data._complexRepresentation = ComplexParts;
	for(size_t i=0; i!=real._data.size(); ++i)
	{
		data._data[i]._polarization = real._data[i]._polarization;
		data._data[i]._images[0] = real._data[i]._images[0];
		data._data[i]._images[1] = imaginary._data[i]._images[0];
		data._data[i]._flagging = real._data[i]._flagging;
	}
	return data;
}

TimeFrequencyData TimeFrequencyData::MakeFromPolarizationCombination(const TimeFrequencyData &xx, const TimeFrequencyData &xy, const TimeFrequencyData &yx, const TimeFrequencyData &yy)
{
	if(xx.ComplexRepresentation() != xy.ComplexRepresentation() ||
		xx.ComplexRepresentation() != yx.ComplexRepresentation() ||
		xx.ComplexRepresentation() != yy.ComplexRepresentation())
		throw std::runtime_error("Trying to create dipole time frequency combination from data with different phase representations!");

	TimeFrequencyData data;
	data._data.resize(4);
	data._complexRepresentation = xx._complexRepresentation;
	for(size_t i=0; i!=xx._data.size(); ++i)
	{
		data._data[0] = xx._data[0];
		data._data[1] = xy._data[0];
		data._data[2] = yx._data[0];
		data._data[3] = yy._data[0];
	}
	return data;
}

TimeFrequencyData TimeFrequencyData::MakeFromPolarizationCombination(const TimeFrequencyData &first, const TimeFrequencyData &second)
{
	if(first.ComplexRepresentation() != second.ComplexRepresentation())
		throw std::runtime_error("Trying to create polarization combination from data with different phase representations!");

	TimeFrequencyData data;
	data._data = first._data;
	data._complexRepresentation = first._complexRepresentation;
	data._data.insert(data._data.end(), second._data.begin(), second._data.end());
	return data;
}

void TimeFrequencyData::SetImagesToZero()
{
	if(!IsEmpty())
	{
		Image2DPtr zeroImage = Image2D::CreateZeroImagePtr(ImageWidth(), ImageHeight());
		Mask2DPtr mask = Mask2D::CreateSetMaskPtr<false>(ImageWidth(), ImageHeight());
		for(PolarizedTimeFrequencyData& data : _data)
		{
			data._images[0] = zeroImage;
			if(data._images[1])
				data._images[1] = zeroImage;
			data._flagging = mask;
		}
	}
}

void TimeFrequencyData::MultiplyImages(long double factor)
{
	for(PolarizedTimeFrequencyData& data : _data)
	{
		if(data._images[0])
		{
			Image2DPtr newImage(new Image2D(*data._images[0]));
			newImage->MultiplyValues(factor);
			data._images[0] = newImage;
		}
		if(data._images[1])
		{
			Image2DPtr newImage(new Image2D(*data._images[1]));
			newImage->MultiplyValues(factor);
			data._images[1] = newImage;
		}
	}
}

void TimeFrequencyData::JoinMask(const TimeFrequencyData& other)
{
	if(other.MaskCount() == 0)
	{
		// Nothing to be done; other has no flags
	} else if(other.MaskCount() == MaskCount())
	{
		for(size_t i=0;i<MaskCount();++i)
		{
			Mask2D mask(*GetMask(i));
			mask.Join(*other.GetMask(i));
			SetMask(i, Mask2DPtr(new Mask2D(mask)));
		}
	} else if(other.MaskCount() == 1)
	{
		if(MaskCount() == 0)
		{
			for(size_t i=0; i!=_data.size(); ++i)
				_data[i]._flagging = other._data[0]._flagging;
		}	
		else {
			for(size_t i=0;i<MaskCount();++i)
			{
				Mask2D mask(*GetMask(i));
				mask.Join(*other.GetMask(0));
				SetMask(i, Mask2DPtr(new Mask2D(mask)));
			}
		}
	} else if(MaskCount() == 1)
	{
		Mask2D mask(*GetMask(0));
		mask.Join(*other.GetSingleMask());
		SetMask(0, Mask2DPtr(new Mask2D(mask)));
	}	else if(MaskCount() == 0 && _data.size() == other._data.size())
	{
		for(size_t i=0; i!=_data.size(); ++i)
			_data[i]._flagging = other._data[i]._flagging;
	}
	else
		throw std::runtime_error("Joining time frequency flagging with incompatible structures");
}

