
//**************************************************************************/
// Copyright (c) 2008 Autodesk, Inc.
// All rights reserved.
//
// Use of this software is subject to the terms of the Autodesk license
// agreement provided at the time of installation or download, or which
// otherwise accompanies this software in either electronic or hard copy form.
//
//**************************************************************************/
// DESCRIPTION:
// CREATED: October 2008
//**************************************************************************/

#pragma warning( disable : 4311 4312 )

#include "displacer.h"
#include "math.h"
#include <QtCore/QFileInfo>
#include <QtCore/QDir>

IMPLEMENT_CLASS( DisplaceOperation, TreeNode, "Displace" );

// Register this plugin, use the DisplaceOperation::Initializer function as the initialized, this will be called when the plugin got loaded.
MB_PLUGIN( "MeshDisplacer", "Displacement operation for meshes", "Autodesk", "http://www.mudbox3d.com", DisplaceOperation::Initializer );

DisplaceOperation::DisplaceOperation( void ) : 
	Node( "Displacement" ), 
	m_aTiles( "displaceoperation tilelist" ),
//	m_aTiles( "displaceoperation tilelist" ), 
	m_pObject( this, "Target Mesh" ),
	m_iSubdivisionLevel( this, "Displace To Level" ),
	m_iMaskChannel( this, "Use Mask Channel" ),
	m_iUDim( this, "URange" ),
	m_iFirstTileIndex( this, "Initial Tile Number" ),
	m_sTileRange( this, "Use Tile Range" ),
	m_sDisplacementFileMask( this, "Displacement Map" ),
	m_sMaskFileMask( this, "Mask Map" ),
	m_iFinalFaceCount( this, "Final Face Count" ),
	m_eExecute( this, "Go" ),
	m_eDelete( this, "Delete this operation" ),
//	m_sName( "Node Name" ),
	m_iDisplacementChannel( this, "Use Displacement Channel" ),
	m_bSmoothTC( this, "Smooth Texture Coors" ),
	m_sNextCommand( this, "Next Command" ),
	m_eMapSpace( this, "Displacement Map Space" )
{
	m_eMapSpace.AddItem( "Normal (regular displacement map)" );
	m_eMapSpace.AddItem( "Tangent" );
	m_eMapSpace.AddItem( "Object" );
	m_eMapSpace.AddItem( "World" );
	m_eMapSpace = 0;
	m_eMapSpace.SetVisible( true );
	m_vDisplacementMax.Set( 1, 1, 1 );

	// Assign a unique name to this node.
	unsigned int i = 1;
	QString sName;
	bool b;
	do
	{
		b = false;
		sName = QString("Disp #%1").arg(i); // String().Format( "Disp #%d", i );
		for ( Node *p = First(); p; p = p->Next() )
			if ( p->IsKindOf( StaticClass() ) )
			{
				if ( sName == dynamic_cast<DisplaceOperation *>( p )->Name() )
				{
					i++;
					b = true;
				};
			};
	} while ( b );
	SetName( sName );
	m_bFloatMap = false;
	m_iBaseU = 0;
	m_iBaseV = 0;
	m_fGain = 1;

	m_pObject.m_sNullString = "- Select a Mesh -";

	LoadTemplate();

	m_sNextCommand.SetVisible( false );
	m_sNextCommand.Connect( Kernel()->m_sNextCommand );
};

void DisplaceOperation::Serialize( Stream &s )
{
	s == m_pObject == m_iSubdivisionLevel == m_iFinalFaceCount == m_iUDim == m_iFirstTileIndex == m_sTileRange == m_sDisplacementFileMask == m_sMaskFileMask == m_iMaskChannel;
	if( s.IsOlderThan( 300 ) )
	{
		QString s1;
		s == s1;
		SetName( s1 );
	};
	if ( !s.IsOlderThan( 155+1 ) )
		s == m_bSmoothTC;

	// TODO after merge
/*	QString sFN = m_sDisplacementFileMask;
	{
		m_bFloatMap = false;
		QFileInfo f( sFN );
		if ( f.suffix() == "tif" )
		{
			TIFF *t = TIFFOpen( sFN, "r" );
			if ( t )
			{
				unsigned short int iBitsPerSample;
				VERIFY( TIFFGetField( t, TIFFTAG_BITSPERSAMPLE, &iBitsPerSample ) );
				if ( iBitsPerSample == 32 )
					m_bFloatMap = true;
				TIFFClose( t );
			};
		};
		if ( m_bFloatMap )
		{
			m_fDisplacementMax.SetConst( true );
			m_fDisplacementMin.SetConst( true );
		}
		else
		{
			m_fDisplacementMax.SetConst( false );
			m_fDisplacementMin.SetConst( false );
		};
	};*/
	if ( s.IsNewerThan( 340 ) )
		TreeNode::Serialize( s );
};

void DisplaceOperation::OnNodeEvent( const Attribute &a, NodeEventType eType )
{
	if ( a == m_pObject && eType == etValueChanged )
	{
		// Disable subdivision if the object already has more than one subdivision levels.
		//if ( m_pObject != 0 && m_pObject->LevelCount() > 1 )
		//{
		//	m_iSubdivisionLevel.SetConst( true );
		//	m_bSmoothTC.SetConst( true );
		//	m_iSubdivisionLevel = m_pObject->ActiveLevel()->Index();
		//}
		//else
		//{
		//	m_iSubdivisionLevel.SetConst( m_pObject == 0 );
		//	m_bSmoothTC.SetConst( m_pObject == 0 );
		//};
		if ( m_pObject )
			m_iSubdivisionLevel = m_pObject->LevelCount()-1;
	};
	// Recalculate the face count if the object or level is changed.
	if ( (a == m_iSubdivisionLevel || a == m_pObject) && eType == etValueChanged )
	{
		if ( m_pObject.Value() && m_iSubdivisionLevel < int(m_pObject->LevelCount()-1) )
		{
			Kernel()->MessageBox( Kernel::msgError, "Error", QString("Subdivision level cannot be less than %1").arg( m_pObject->LevelCount()-1 ) );
			m_iSubdivisionLevel = m_pObject->LevelCount()-1;
			return;
		};
		if ( m_pObject != 0 && m_iSubdivisionLevel != -1)
			m_iFinalFaceCount = int(m_pObject->LowestLevel()->TotalFaceCount()*pow(float(4),m_iSubdivisionLevel));
		else
			m_iFinalFaceCount = 0;
	};

	if ( eType == etValueChanged && a == m_sNextCommand )
	{
		QString sOp = m_sNextCommand.Value().section( ' ', 0, 0 );
		if ( sOp == "meshdisplacement" )
		{
			QString sName = m_sNextCommand.Value().section( '"', 1, 1 );
			if ( sName == Name() )
				Do();
		};
	};

	// Displacement filename is changed
	if ( a == m_sDisplacementFileMask && eType == etValueChanged )
	{
		QString sFN = m_sDisplacementFileMask.Value();
		m_sDisplacementFileMask = FilterFileName( sFN );
	};
	// Mask filename is changed, replace 0001 with %i since it is possibly the tile index.
	if ( a == m_sMaskFileMask && eType == etValueChanged )
	{
		QString sFN = m_sMaskFileMask.Value();
		m_sMaskFileMask = FilterFileName( sFN );
	};
	// Start operation was pressed
	if ( a == m_eExecute && eType == etEventTriggered )
	{
		Kernel()->RecordCommand( "meshdisplacement \"%s\"", (const char *) Name().toAscii() );
		Do();
	};
	// Delete pressed, delete the operation.
	if ( a == m_eDelete && eType == etEventTriggered )
	{
		// Seems unsafe to delete this here.  I have no idea who may have
		// references to 'this'.  Adding the RemoveChild to at least remove
		// the node from the graph.  Otherwise children will have invalid
		// sibling ptrs.
		Kernel()->Scene()->RemoveChild(this);
		Kernel()->RefreshUI();

		delete this;
	}
}

void DisplaceOperation::Do( void )
{
	if ( m_pObject == 0 )
		throw new Error( "You must select a mesh in the scene first" );

	// Calculate UV bounds, and align m_iUDim if it is not yet set to a proper value.
	AxisAlignedBoundingBox b;
	for ( unsigned int i = 0; i < m_pObject->LowestLevel()->TCCount(); i++ )
		b.Extend( Vector( m_pObject->LowestLevel()->m_pTCs[i].m_fU, m_pObject->LowestLevel()->m_pTCs[i].m_fV ) );
	if ( float(m_iUDim) < ceilf(b.m_vEnd.x) )
		m_iUDim = int(ceilf(b.m_vEnd.x));

	m_iBaseU = int(floorf(b.m_vStart.x));
	m_iBaseV = int(floorf(b.m_vStart.y));

	// Calculate the tile list based on m_sTileRange.
	m_aTiles.Clear();
	QString sTileRange = m_sTileRange.Value();
	if ( sTileRange.size() )
	{
		int p = 0, i[2] = { 0, 0 }, k = 0;
		while ( p < sTileRange.size() )
		{
			if ( p == sTileRange.size() || sTileRange[p] == ',' )
			{
				i[1] = Max( i[1], i[0] );
				for ( int h = i[0]; h <= i[1]; h++ )
				{
					Tile t;
					t.m_iIndex = h;
					t.m_bZero = false;
					m_aTiles.Add( t );
				};
				i[0] = i[1] = 0;
				k = 0;
			};
			if ( sTileRange[p] >= '0' && sTileRange[p] <= '9')
				i[k] = i[k]*10+sTileRange.toAscii()[p]-'0';
			else if ( sTileRange[p] == '-' )
				k = 1;
			p++;
		};
		i[1] = Max( i[1], i[0] );
		for ( int h = i[0]; h <= i[1]; h++ )
		{
			Tile t;
			t.m_iIndex = h;
			t.m_bZero = false;
			m_aTiles.Add( t );
		};

		// now we put those tiles into the array which the user want to skip with zero factor. the reason for this is that at the edge of tiles
		// (if the edge vertex has different UV one in each tile) we want to result an average value of the zero and user selected tile. this
		// way if the user applies the displacement operation to a multiple tile mesh separately for each tile, the result will be the same as
		// applying the tiles in one operation.
		int iWidth = int(ceilf(b.m_vEnd.x)-floorf(b.m_vStart.x));
		int iHeight = int(ceilf(b.m_vEnd.y)-floorf(b.m_vStart.y));
		int iStart = m_iFirstTileIndex;
		if ( iWidth > m_iUDim )
			m_iUDim = iWidth;
		int iEnd = m_iFirstTileIndex+m_iUDim*iHeight;
		for ( int i = iStart; i < iEnd; i++ )
		{
			unsigned int j = 0;
			while ( j < m_aTiles.ItemCount() && m_aTiles[j].m_iIndex != i )
				j++;
			if ( j < m_aTiles.ItemCount() )
				continue;
			Tile t;
			t.m_iIndex = i;
			t.m_bZero = true;
			m_aTiles.Add( t );
		};
		MB_ASSERT( m_aTiles.ItemCount() == iEnd-iStart );
	}
	else
	{
		int iWidth = int(ceilf(b.m_vEnd.x)-floorf(b.m_vStart.x));
		int iHeight = int(ceilf(b.m_vEnd.y)-floorf(b.m_vStart.y));
		int iStart = m_iFirstTileIndex;
		if ( iWidth > m_iUDim )
			m_iUDim = iWidth;
		int iEnd = m_iFirstTileIndex+m_iUDim*iHeight;
		for ( int i = iStart; i < iEnd; i++ )
		{
			Tile t;
			t.m_iIndex = i;
			t.m_bZero = false;
			m_aTiles.Add( t );
		};
	};

	// Execute the operation.
	Execute();

	// Redraw the scene.
	Kernel()->Redraw();
};

// converts the internal filename into a real one by inserting tile data into it, and looking for the file in the directory
QString DisplaceOperation::GetFileName( const QString &sMask, unsigned int iTile, unsigned int iUPos, unsigned int iVPos ) const
{	
	QString sFileName = sMask;

	// get the directory of the file
	QFileInfo i( sFileName );
	QDir d = i.absoluteDir();
	
	// insert the tile data after the filename, and the * mark before searching for the file
	QString g;
	if ( m_aTiles.ItemCount() == 1 )
		g = i.completeBaseName()+"*."+i.suffix();
	else
	{
		int iU = iUPos, iV = iVPos;
		if ( iU >= 0 )
			iU++;
		if ( iV >= 0 )
			iV++;
		g = i.completeBaseName()+QString("_u%1_v%2*.").arg(iU).arg(iV)+i.suffix();
	};

	// set the mask for the files we are looking for
	QStringList f;
	f << g;
	d.setNameFilters( f );

	// lets see how many files we found
	if ( d.count() )
	{
		if ( d.count() > 1 )
			Kernel()->MessageBox( Kernel::msgWarning, "Warning", QString("More than one files beginning with the mask %1 is found. The result of the mesh displacement might be incorrect. Delete the unneeded ones to solve the problem.").arg(g) );
		QString f = d[0];
		QFileInfo k( f );

		// get the string which replaced the "*" in the search pattern, this is probably the gain information
		QString a = k.completeBaseName().mid( g.length()-i.suffix().length()-2 );
		if ( a.left( 2 ) == "_g" )
		{
			// obtain the gain value from the string "a". the first two caracter is "_g" if everything goes right.
			m_fGain = a.mid( 2 ).toFloat();
			m_vDisplacementMin.Set( -m_fGain/2, -m_fGain/2, -m_fGain/2 );
			m_vDisplacementMax.Set( m_fGain/2, m_fGain/2, m_fGain/2 );
		};

		// return the name of the file we found
		return QString(d.filePath( f ));
	};

	// the file not found, return an empty string.
	return "";
};	

// convert the filename from a real existing one to an internal one, by removing the _u, _v and _g parts from the end of the filename.
QString DisplaceOperation::FilterFileName( const QString &sFileName ) const
{
	QString s = sFileName;
	QString e;

	// search for the extension in the filename
	int k = sFileName.size();
	while ( k > 0 )
	{
		k--;
		if ( sFileName[k] == '/' || sFileName[k] == '\\' )
			break;
		if ( sFileName[k] == '.' )
		{
			e = s.right( s.size()-k-1 );
			s = s.left( k );
			break;
		};
	};

	// start from the end of the filename, look for the character '_'. that indicates a section which must be removed. later this string will be added back before looking for the file.
	for ( int j = s.size()-1; j >= 0; j-- )
	{
		// if its a '_' character, we skip everything beginning with it
		if ( s[j] == '_' )
		{
			s = s.left( j );
			continue;
		};

		// if the character is not a valid addition character, then we already finished removing the added part, so finish the loop
		if ( s[j] != 	'u' && s[j] != 'v' && s[j] != 'g' && s[j] != 'e' && s[j] != '-' && s[j] != '+' && s[j] != '.' && (s[j] < '0' || s[j] > '9') )
			break;
	};

	// return the naked filename with the original path and extension
	return s+"."+e;
};

void DisplaceOperation::Execute( void )
{
	Instance<Image> pDisplacement, pMask;

	// If no object specified, throw an error message.
	if ( !m_pObject )
		throw new Error( "No target specified");

	// Subdivision
	SubdivisionLevel::SetUVCreation( true );
	SubdivisionLevel *pM = m_pObject->LowestLevel();

	// Check if the object has texture coordinates
	if ( pM->TCCount() == 0 )
		throw new Error("Target mesh has no texture coordinates");

	// If the object has less subdivision level as we want, subdivide it further
	while ( int(m_pObject->LevelCount()-1) < m_iSubdivisionLevel )
	{
		m_pObject->HighestLevel()->Subdivide();
		if ( m_bSmoothTC )
			m_pObject->HighestLevel()->SmoothTextureCoordinates( 1.0f );
	};

	pM = m_pObject->Level( m_iSubdivisionLevel );
	if ( pM == 0 )
		MB_ERROR( "Corrupt subdivision structure" );
	m_pObject->SetActiveLevel( pM );
	pM->AddFaceComponent( pM->fcTCIndex );

	// Initialize the subdivision level to be ready for global vertex index lookups
	Kernel()->Log( "Displacement operation on %s\n", (const char *)pM->Name().toAscii() );

	Kernel()->ProgressStart( "applying maps...", m_aTiles.ItemCount()*pM->FaceCount() );

	// This array will hold the collected information.
	Store<VertexInfo> aVertices( "displaceoperation vertex tmp" );

	TangentGenerator *pTG = NULL;
	if ( m_eMapSpace == spaceTangent )
		pTG = pM->ChildByClass<TangentGenerator>( true );

	// Enumerate the possible tiles.
	for ( unsigned int i = 0; i < m_aTiles.ItemCount(); i++ )
	{
		unsigned int iTile = m_aTiles[i].m_iIndex;
		bool bZero = m_aTiles[i].m_bZero;
		int iUPos = (iTile-m_iFirstTileIndex)%m_iUDim+m_iBaseU, iVPos = (iTile-m_iFirstTileIndex)/m_iUDim+m_iBaseV;
		Kernel()->Log( "Displacing tile %d\n", iTile );

		bool bMask = false;
		if ( !bZero )
		{
			Kernel()->Log( "\tLoading %s for dislacement map", (const char*)GetFileName( (const char*) m_sDisplacementFileMask.Value().toAscii(), iTile, iUPos, iVPos ).toAscii() );
			try 
			{ 
				if ( m_eMapSpace.Value() )
					pDisplacement->Load( GetFileName( m_sDisplacementFileMask.Value(), iTile, iUPos, iVPos ) ); 
				else
					pDisplacement->LoadChannel( GetFileName( m_sDisplacementFileMask.Value(), iTile, iUPos, iVPos ), m_iDisplacementChannel ); 
			}
			catch ( Error *pError )
			{
				Kernel()->Log( " - map not found, skipping\n" );
				pError->Discard();
				continue;
			};
			Kernel()->Log( " - OK\n" );

			// Load the mask map.
			if ( m_sMaskFileMask == "" )
				bMask = false;
			else
			{
				Kernel()->Log( "\tLoading %s for mask map", (const char*)GetFileName( (const char*) m_sMaskFileMask.Value().toAscii(), iTile, iUPos, iVPos ).toAscii() );
				bool bMask = true;
				try { 	pMask->LoadChannel( GetFileName( m_sMaskFileMask.Value(), iTile, iUPos, iVPos ), m_iMaskChannel ); }
				catch ( Error *pError ) { pError->Discard(); bMask = false; };
				if ( bMask )
					Kernel()->Log( " - OK\n" );
				else
					Kernel()->Log( " - not found\n" );
			};
		};

		unsigned int iBefore = aVertices.ItemCount();

		// Use an array to store which vertex got processed yet.
		Store<int> aProcessedTC( pM->TCCount(), "displace - processedtc" );
		for ( unsigned int j = 0; j < pM->TCCount(); j++ )
			aProcessedTC[j] = -1;

		// Enumerate all faces on the surface.
		for ( unsigned int v = 0; v < pM->FaceCount(); v++ )
		{
			Kernel()->ProgressSet( i*pM->FaceCount()+v );

			// Enumerate all the corners of the face.
			for ( int c = 0; c < pM->SideCount(); c++ )
			{
				// Get position and texture coordinate indices for the current vertex.
				unsigned int iVertexIndex, iTCIndex;
				if ( pM->Type() == Mesh::typeTriangular )
				{
					iVertexIndex = pM->TriangleIndex( v, c );
					iTCIndex = pM->TriangleTCI( v, c );
				}
				else
				{
					iVertexIndex = pM->QuadIndex( v, c );
					iTCIndex = pM->QuadTCI( v, c );
				};

				MB_ONBUG( iTCIndex >= pM->TCCount() )
					continue;

				// If this vertex was touched by us then skip it.
				if ( pM->VertexStrokeID( iVertexIndex ) == pM->CollectionID() )
					continue;

				// Check if the texture coordinate 
				float fU = pM->VertexTC( iTCIndex ).m_fU-float(iUPos), fV = 1-(pM->VertexTC( iTCIndex ).m_fV-float(iVPos));
				if ( fU < 0.0f || fU > 1.0f || fV < 0.0f || fV > 1.0f )
					continue;

				// Set this vertex as touched.
				pM->SetVertexStrokeID( iVertexIndex, pM->CollectionID() );

				// Get the displacement map value for this vertex from the map.
				if ( bZero )
				{
					// this tile is not listed by the user, so we skip it, except when a vertex is affected by another processed tile.
					// in that case we treat this tile as a black zero tile. this is important because the user may want to apply
					// the tiles separately in different displacement operations and the result must be the same as with a single
					// operation

					if ( aProcessedTC[iTCIndex] != -1 )
						aVertices[aProcessedTC[iTCIndex]].Refine( Vector(), 1 );
					continue;
				};

				float fMask = 1;
				Vector vDisp;
				vDisp.x = pDisplacement->ValueAt( fU, fV, 0 );
				if ( m_eMapSpace.Value() )
				{
					vDisp.y = pDisplacement->ValueAt( fU, fV, 1 );
					vDisp.z = pDisplacement->ValueAt( fU, fV, 2 );
				};

				// Transform the displacement map value into the [m_fDisplacementMin:m_fDisplacementMax] range. 8 and 16 bit maps can hold values only in the range of [0:1]
				if( !m_bFloatMap )
					vDisp = m_vDisplacementMin+vDisp*(m_vDisplacementMax-m_vDisplacementMin);

				// If there is a mask texture, get the value from it also, otherwise use 1.
				if ( bMask )
					fMask = pMask->ValueAt( fU, fV, 0 );

				// If the displacement value is not zero, store it for later processing
				if ( vDisp )
				{
					// Check if the current texture coordinate was processed before or not.
					if ( aProcessedTC[iTCIndex] != -1 )
					{
						// Refine the previously stored data for this vertex.
						aVertices[aProcessedTC[iTCIndex]].Refine( vDisp, fMask );
					}
					else
					{
						// Its a new vertex, add it to the aVertices array.
						unsigned int iIndex;
						iIndex = aVertices.Add( VertexInfo( iVertexIndex, iVertexIndex, vDisp, fMask, v, c ) );

						// If the returned index is 0xffffffff, the Add operation was not successful.

						if ( iIndex == 0xffffffff )
						{
							Kernel()->Log( "\tOut of memory\n" );
							throw &Error::s_cBadAlloc;
						};

						// Store the index for later usage.
						aProcessedTC[iTCIndex] = iIndex;
					};
				};
			};
		};
		Kernel()->Log( "\tProcessed %d vertices\n", aVertices.ItemCount()-iBefore );
	};

	// Increase the collection counter since we used the current one for some vertices.
	pM->IncreaseCollectionID();

	// Creating a new layer for the results.
	LayerMeshData *pL = pL = pM->AddLayer();
	if ( pL == 0 ) 
		throw Error( "Unable to create new layer" );
	pL->SetName( "Displacement" );

	// Fill the layer with the collected data
	pL->SetVertexCount( aVertices.ItemCount() );
	aVertices.Sort();
	for ( unsigned int k = 0; k < aVertices.ItemCount(); k++ )
	{
		Vector vDelta;
		if ( m_eMapSpace.Value() )
		{
			vDelta = aVertices[k].m_vDisplacement;
			if ( m_eMapSpace == spaceTangent )
			{
				MB_SAFELY( pTG )
				{
					Base b = pTG->LocalBase( aVertices[k].m_iFaceIndex, aVertices[k].m_iCornerIndex );
					// we are scaling the normal vector to have a length which is the geometrical average of the tangent and binormal
					float al = b.a.Length();
					float bl = b.b.Length();
					float cl = b.c.Length();
					b.b *= sqrtf(al*cl)/bl;
					vDelta = b.TransformFrom( vDelta );
				};
			};
			if ( m_eMapSpace == spaceWorld )
				vDelta = pM->Geometry()->Transformation()->TransformToLocal( vDelta, 0 );
		}
		else
			vDelta = Vector(pM->VertexNormal(  aVertices[k].m_iVertexIndex  ))*aVertices[k].m_vDisplacement.x;
		pL->SetVertexData( k, aVertices[k].m_iGlobalVertexIndex, aVertices[k].m_fMask );
		pL->SetVertexDelta( k, aVertices[k].m_iVertexIndex, vDelta );
	};

	// HACK recreate render buffers
	pM->ContentChanged();

	// Since vertex positions changed, recalculate normal vectors.
	pM->RecalculateNormals();

	Kernel()->ProgressEnd();
};

// This function will be called then the plugin got loaded.
void DisplaceOperation::Initializer( void )
{
	// Add a menu entry for the DisplaceOperation class/
	Kernel()->AddClassMenuItem( "Maps", "Sculpt Model Using Displacement Map", StaticClass(), "New Operation..." );
};
