/*======================================================================================

	Program:		Range Imaging Toolkit (RITK)
	Module:			ritkRImageActorPipeline.cxx
	Version:		$Revision: 1761 $
	Date:			$Date: 2011-09-23 17:25:41 +0200 (Fr, 23 Sep 2011) $

	Copyright (c) 2010-2011
	Pattern Recognition Lab, 
    University Erlangen-Nuremberg, Germany
    All rights reserved.

	This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. 
	See RITKCopyright.txt for more information.

	http://www5.informatik.uni-erlangen.de/research/software/ritk

========================================================================================*/

#include "ritkRImageActorPipeline.h"
#include "ritkRGBRImage.h"
#include "ritkDebugManager.h"

#include <itkRescaleIntensityImageFilter.h>
#include <itkImageRegionIterator.h>

#include <vtkObjectFactory.h>
#include <vtkPointData.h>
#include <vtkFloatArray.h>
#include <vtkTextProperty.h>

#include <algorithm>

namespace ritk
{

// Implements static RImageActorPipeline *New();
vtkStandardNewMacro(RImageActorPipeline);


//----------------------------------------------------------------------------
RImageActorPipeline::RImageActorPipeline()
{
	// Default visualization mode is monochrome
	m_VisualizationMode = MONOCHROME;

	// By default, no LUT
	m_UseLUT = false;

	// Initialize the LUT
	m_LUT = vtkSmartPointer<vtkLookupTable>::New();
	unsigned int JetLength = 200;
	m_LUT->SetNumberOfColors(LUT_Jet_Length);
	for( unsigned int i = 0; i < JetLength; i++ ) 
		m_LUT->SetTableValue(i, LUT_Jet[i][0], LUT_Jet[i][1], LUT_Jet[i][2], 1.0); 
		//m_LUT->SetTableValue(JetLength-1-i, LUT_Jet[i][0], LUT_Jet[i][1], LUT_Jet[i][2], 1.0); 
	m_RangeInterval[0] = 0;
	m_RangeInterval[1] = 1000;
	m_RGBInterval[0] = 0;
	m_RGBInterval[1] = 600;

	// By default the poly data is a point cloud
	m_PolyDataType = RImageVTKData::POINTCLOUD;

	// Initialize the RImageVTK data
	m_RImageVTKData = vtkSmartPointer<RImageVTKData>::New();

	// Initialize the scalar bar actor
	m_ScalarBarVisibility = false;
	m_ScalarBarActor = vtkSmartPointer<vtkScalarBarActor>::New();
	m_ScalarBarActor->SetTitle("Depth [mm]");
	m_ScalarBarActor->SetTextPositionToPrecedeScalarBar();
	m_ScalarBarActor->SetLookupTable(m_Mapper->GetLookupTable());
	m_ScalarBarActor->SetVisibility(m_ScalarBarVisibility);
	m_ScalarBarActor->SetMaximumNumberOfColors(1024);
	//m_ScalarBarActor->SetMaximumHeightInPixels(600);
	m_ScalarBarActor->SetMaximumWidthInPixels(80);
	vtkSmartPointer<vtkTextProperty> SATextProperty = m_ScalarBarActor->GetTitleTextProperty();
	SATextProperty->SetFontSize(10);
	SATextProperty->SetJustificationToCentered();
	m_ScalarBarActor->SetTitleTextProperty(SATextProperty);
	m_ScalarBarActor->SetPosition(0.9, 0.1);
	m_ScalarBarActor->SetPosition2(0.9, 0.8);
}


//----------------------------------------------------------------------------
RImageActorPipeline::~RImageActorPipeline()
{
}


//----------------------------------------------------------------------------
void
RImageActorPipeline::SetVisualizationMode(VisualizationMode Mode)
{
	// We're changing the pipeline so lock!
	Lock();

	// Update the mode
	m_VisualizationMode = Mode;

	// Done here
	Unlock();
}


//----------------------------------------------------------------------------
void
RImageActorPipeline::UseLUT(bool UseLUT)
{
	// We're changing the pipeline so lock!
	Lock();

	// Update the flag
	m_UseLUT = UseLUT;

	// Done here
	Unlock();
}


//----------------------------------------------------------------------------
void 
RImageActorPipeline::SetPolyDataType(PolyDataType Type)
{
	// We're changing the pipeline so lock!
	Lock();

	// Update the type
	m_PolyDataType = Type;

	// Done here
	Unlock();
}


//----------------------------------------------------------------------------
void
RImageActorPipeline::SetData(PolyDataPointer Data, bool DeepCopy)
{
	// Note: Don't lock here as this will be done in the superclass. Locking here
	// would cause a dead lock 
	Superclass::SetData(Data,DeepCopy);	

	/*if(Data->GetPointData()->GetScalars())
	{
		if(Data->GetPointData()->GetScalars()->GetNumberOfTuples()>0)
		{
			// scalar data available
			
			// Compute a poly data structure with the correct colors. This will cause a deep copy of the input
			PolyDataPointer ColoredData = ComputeColoredPolyData(Data);

			// Pass the colored data to the superclass. We don't need a deep copy anymore
			Superclass::SetData(ColoredData,false);
		}
	}*/

	// Update the mapper
	UpdateMapper(m_VisualizationMode, m_UseLUT);
}


//----------------------------------------------------------------------------
void
RImageActorPipeline::SetData(PolyDataPointer PolyData, const RImageType *Image)
{
	// If no image is set we can pass the data straight to the superclass. This will cause a deep copy of the input
	if ( !Image )
	{
		Superclass::SetData(PolyData);
	}
	else
	{
		// Compute a poly data structure with the correct colors. This will cause a deep copy of the input
		PolyDataPointer ColoredData = ComputeColoredPolyData(PolyData, Image, m_VisualizationMode, m_UseLUT);

		// Pass the colored data to the superclass. We don't need a deep copy anymore
		Superclass::SetData(ColoredData,false);
	}

	// Update the mapper
	UpdateMapper(m_VisualizationMode, m_UseLUT);
}


//----------------------------------------------------------------------------
void
RImageActorPipeline::SetData(const RImageType *Image)
{
	// Create the poly data from the image
	m_RImageVTKData->SetPolyDataType(m_PolyDataType);
	m_RImageVTKData->GeneratePolyData(Image);
	PolyDataPointer Data = m_RImageVTKData->GetPolyData();

	// Compute a poly data structure with the correct colors. This will cause a deep copy of the input
	PolyDataPointer ColoredData = ComputeColoredPolyData(Data, Image, m_VisualizationMode, m_UseLUT);

	// Pass the colored data to the superclass. We don't need a deep copy anymore due to our smartpointer
	Superclass::SetData(ColoredData,false);

	// Update the mapper
	UpdateMapper(m_VisualizationMode, m_UseLUT);
}


//----------------------------------------------------------------------------
void
RImageActorPipeline::SetRangeInterval(float min, float max)
{
	Lock();

	m_RangeInterval[0] = min;
	m_RangeInterval[1] = max;

	Unlock();
}


//----------------------------------------------------------------------------
void
RImageActorPipeline::SetRGBInterval(float min, float max)
{
	Lock();

	m_RGBInterval[0] = min;
	m_RGBInterval[1] = max;

	Unlock();
}


//----------------------------------------------------------------------------
void 
RImageActorPipeline::SetScalarBarVisibility(bool b)
{
	Lock();

	m_ScalarBarVisibility = b;

	Unlock();
}

float ComputeMedian(std::vector<float>& vec) 
{ 
  std::nth_element( vec.begin(), vec.begin() + vec.size() / 2, vec.end() ); 
  return *( vec.begin() + vec.size() / 2 ); 
}

//----------------------------------------------------------------------------
RImageActorPipeline::PolyDataPointer
RImageActorPipeline::ComputeColoredPolyData(PolyDataConstPointer Data)
{
	// Deep copy of the input
	Data->Update();
	PolyDataPointer NewData = PolyDataPointer::New();
	NewData->DeepCopy(Data);
	
	// This will hold the scalar values (colors)
	vtkSmartPointer<vtkDataArray> Colors;

	if(Data->GetPointData()->GetScalars()->GetNumberOfComponents() == 1)
	{
		Colors = vtkSmartPointer<vtkUnsignedCharArray>::New();
		Colors->SetNumberOfComponents(4);

		double range[2];
		// Compute scalar data range
		//Data->GetPointData()->GetScalars()->GetRange(range);

		// Compute median
		std::vector<float> ScalarFieldVector;
		for(unsigned int i = 0; i < Data->GetPointData()->GetScalars()->GetNumberOfTuples(); i++)
		{
			ScalarFieldVector.push_back(Data->GetPointData()->GetScalars()->GetTuple1(i));
		}
		float Median = ComputeMedian(ScalarFieldVector);

		// Compute deviation
		std::vector<float> DeviationVector;
		for(unsigned int i = 0; i < Data->GetPointData()->GetScalars()->GetNumberOfTuples(); i++)
		{
			DeviationVector.push_back(fabs(Data->GetPointData()->GetScalars()->GetTuple1(i) - Median));
		}
		std::sort(DeviationVector.begin(), DeviationVector.end());

		// Set range for scaling
		range[0] = Median - DeviationVector[Data->GetPointData()->GetScalars()->GetNumberOfTuples()*0.95];
		range[1] = Median + DeviationVector[Data->GetPointData()->GetScalars()->GetNumberOfTuples()*0.95];

		for(unsigned int i = 0; i < Data->GetPointData()->GetScalars()->GetNumberOfTuples(); i++)
		{
			if(Data->GetPointData()->GetScalars()->GetTuple1(i)<range[0])
			{
				unsigned char c = 0;
				Colors->InsertNextTuple4(c,c,c,255);
			}
			else if(Data->GetPointData()->GetScalars()->GetTuple1(i)>range[1])
			{
				unsigned char c = 255;
				Colors->InsertNextTuple4(c,c,c,255);
			}
			else
			{
				// Value lies within range
				unsigned char c = (unsigned char) (255.0f*(Data->GetPointData()->GetScalars()->GetTuple1(i) - range[0])/(range[1]-range[0]));
				Colors->InsertNextTuple4(c,c,c,255);
			}
		}
	}
	else if(Data->GetPointData()->GetScalars()->GetNumberOfComponents() == 3)
	{
		Colors = vtkSmartPointer<vtkUnsignedCharArray>::New();
		Colors->SetNumberOfComponents(4);
		
		double * ColorTuple3 = new double[3];

		for(unsigned int i = 0; i < Data->GetPointData()->GetScalars()->GetNumberOfTuples(); i++)
		{
			// xyzrgb data: colors must be inverted
			ColorTuple3 = Data->GetPointData()->GetScalars()->GetTuple3(i);
			unsigned char r = (unsigned char) (255.0f*(1.0f-ColorTuple3[0]));
			unsigned char g = (unsigned char) (255.0f*(1.0f-ColorTuple3[1]));
			unsigned char b = (unsigned char) (255.0f*(1.0f-ColorTuple3[2]));
			Colors->InsertNextTuple4(r,g,b,255);
		}		
	}

	// Set the new scalars
	NewData->GetPointData()->SetScalars(Colors);

	// Return the new data
	return NewData;
}


//----------------------------------------------------------------------------
RImageActorPipeline::PolyDataPointer
RImageActorPipeline::ComputeColoredPolyData(PolyDataConstPointer Data, const RImageType *Image, VisualizationMode Mode, bool UseLUT)
{
	// Deep copy of the input
	Data->Update();
	PolyDataPointer NewData = PolyDataPointer::New();
	NewData->DeepCopy(Data);

	// This will hold the scalar values (colors)
	vtkSmartPointer<vtkDataArray> Colors;

	// Depending on the color modes
	if ( Mode == MONOCHROME )
	{
		// Simple unsigned chars that define the colors
		Colors = vtkSmartPointer<vtkUnsignedCharArray>::New();
		Colors->SetNumberOfComponents(4);
		for ( unsigned int i = 0; i < Image->GetRequestedRegion().GetNumberOfPixels(); i++ )
			Colors->InsertNextTuple4(255,255,255,255);	
			//Colors->InsertNextTuple4(255,255,255,(Image->GetFlagImage()->GetBufferPointer()[i]&RImageType::FOREGROUND)*255);		
	}
	else if ( Mode == RGB )
	{
		RGBRImageUCF2::ConstPointer CurrentFrameRGBP = dynamic_cast<const ritk::RGBRImageUCF2*>((const ritk::RImageF2*) Image);
		if ( !CurrentFrameRGBP )
		{
			LOG_DEB("RImageActorPipeline cannot provide RGB visualization mode for non-RGBRImage data.");
			return NULL;
		}

		// If a LUT is requested we have to use floats, otherwise unsigned chars
		if ( UseLUT )
		{
			// Copy the scalars 
			Colors = vtkSmartPointer<vtkFloatArray>::New();
			Colors->SetNumberOfComponents(4);
			typedef itk::ImageRegionConstIterator<RGBRImageUCF2::RGBImageType> IteratorType;
			IteratorType it(CurrentFrameRGBP->GetRGBImage(), CurrentFrameRGBP->GetRGBImage()->GetRequestedRegion());
			for ( it.GoToBegin(); !it.IsAtEnd(); ++it )
			{
				float r = it.Value()[0];
				float g = it.Value()[1];
				float b = it.Value()[2];
				
				float gray = 0.3*r + 0.59*g + 0.11*b; // convert the rgb to gray -> vtk mapper uses LUT for the first channel
				Colors->InsertNextTuple4(gray,gray,gray,255);	
				//Colors->InsertNextTuple4(f,f,f,(Image->GetFlagImage()->GetBufferPointer()[i]&RImageType::FOREGROUND)*255);	
			}
		}
		else
		{
			typedef itk::ImageRegionConstIterator<RGBRImageUCF2::RGBImageType> IteratorType;
			IteratorType it(CurrentFrameRGBP->GetRGBImage(), CurrentFrameRGBP->GetRGBImage()->GetRequestedRegion());
			Colors = vtkSmartPointer<vtkUnsignedCharArray>::New();
			Colors->SetNumberOfComponents(4);
			for ( it.GoToBegin(); !it.IsAtEnd(); ++it )
			{
				float r = it.Value()[0];
				float g = it.Value()[1];
				float b = it.Value()[2];
        
				Colors->InsertNextTuple4(r,g,b,255);	
				//Colors->InsertNextTuple4(c,c,c,(Image->GetFlagImage()->GetBufferPointer()[i]&RImageType::FOREGROUND)*255);
			}
		}
	}
	else if ( Mode == ZBUFFER )
	{
		// Iterator operating on the world coords
		typedef itk::ImageRegionConstIterator<RImageType::WorldCoordImageType> WCIteratorType;
		WCIteratorType itWC(Image->GetWorldCoordImage(),Image->GetWorldCoordImage()->GetRequestedRegion());

		// If a LUT is requested we have to use floats, otherwise unsigned chars
		if ( UseLUT )
		{
			// float z-Coords in mm 
			Colors = vtkSmartPointer<vtkFloatArray>::New();
			Colors->SetNumberOfComponents(4);
			typedef itk::ImageRegionConstIterator<RImageType::WorldCoordImageType> IteratorType;
			IteratorType it(Image->GetWorldCoordImage(),Image->GetWorldCoordImage()->GetRequestedRegion());
			for ( it.GoToBegin(); !it.IsAtEnd(); ++it )
			{
				float z = it.Value()[2];
				Colors->InsertNextTuple4(z,z,z,255);
				//Colors->InsertNextTuple4(z,z,z,(Image->GetFlagImage()->GetBufferPointer()[i]&RImageType::FOREGROUND)*255);
			}
		}
		else
		{
			typedef itk::RescaleIntensityImageFilter<RImageType::RangeImageType,itk::Image<unsigned char,2> > RescalerType;
			RescalerType::Pointer Rescaler = RescalerType::New();
			Rescaler->SetOutputMinimum(0);
			Rescaler->SetOutputMaximum(255);
			Rescaler->SetInput(Image->GetRangeImage());
			Rescaler->GetOutput()->SetRequestedRegion(Image->GetRangeImage()->GetRequestedRegion());
			Rescaler->Update();

			// Copy the scalars into an unsigned char array that defines the colors
			Colors = vtkSmartPointer<vtkUnsignedCharArray>::New();
			Colors->SetNumberOfComponents(4);
			typedef itk::ImageRegionConstIterator<RescalerType::OutputImageType> IteratorType;
			IteratorType it(Rescaler->GetOutput(),Rescaler->GetOutput()->GetRequestedRegion());
			for ( it.GoToBegin(); !it.IsAtEnd(); ++it )
			{
				unsigned char c = it.Value();
				Colors->InsertNextTuple4(c,c,c,255);
				//Colors->InsertNextTuple4(c,c,c,(Image->GetFlagImage()->GetBufferPointer()[i]&RImageType::FOREGROUND)*255);
			}
		}
	}
	else if ( Mode == FLAGS )
	{
		// Set the colors
		Colors = vtkSmartPointer<vtkUnsignedCharArray>::New();
		Colors->SetNumberOfComponents(4);
		typedef itk::ImageRegionConstIterator<RImageType::FlagImageType> FlagIteratorType;
		FlagIteratorType itFlag(Image->GetFlagImage(),Image->GetFlagImage()->GetRequestedRegion());
		for ( itFlag.GoToBegin(); !itFlag.IsAtEnd(); ++itFlag )
		{
			// The color
			unsigned char Color[4];

			// Depending on the flag
			if ( !(itFlag.Value()&RImageType::VALID) )
			{
				Color[0] = 255;
				Color[1] = 0;
				Color[2] = 0;
				Color[3] = 255;
			}
			else
			{
				Color[0] = 0;
				Color[1] = 255;
				Color[2] = 0;
				Color[3] = 255;
			}

			// Insert color
			Colors->InsertNextTuple4(Color[0],Color[1],Color[2],Color[3]);
		}
	}
	else
	{
		// This should never happen...
		std::cout << "RImageActorPipeline: Unknown color type..." << std::endl;
		return PolyDataPointer::New();
	}	

	// Set the new scalars
	NewData->GetPointData()->SetScalars(Colors);

	// Return the new data
	return NewData;
}


//----------------------------------------------------------------------------
void
RImageActorPipeline::UpdateMapper(VisualizationMode Mode, bool UseLUT)
{
	// We're changing the pipeline so lock!
	Lock();

	// If we should use a LUT update the mapper
	if ( UseLUT )
	{
		m_Mapper->SetColorModeToMapScalars();
		m_Mapper->SetLookupTable(m_LUT);
		m_Mapper->ScalarVisibilityOn();
		if ( Mode == ZBUFFER )
			m_Mapper->SetScalarRange(m_RangeInterval[0],m_RangeInterval[1]);
		else if ( Mode == RGB ) // Right now this only takes the first channel into account
			m_Mapper->SetScalarRange(m_RGBInterval[0],m_RGBInterval[1]);

    m_ScalarBarActor->SetLookupTable(m_Mapper->GetLookupTable());
		m_ScalarBarActor->SetVisibility(m_ScalarBarVisibility);
	}
	else
	{
		m_Mapper->SetColorModeToDefault();
		m_Mapper->ScalarVisibilityOn();
		m_ScalarBarActor->SetVisibility(false);
	}

	// Done here
	Unlock();
}

} // end namespace ritk
