// RnHeightField.cpp : Implementation of CRnRect

#include "stdafx.h"
#include "RnHeightField.h"
#include "RpExamplesPlugin.h"


// {779A5BBE-B411-4883-90E2-684AA9E05917}
RS_DEFINE_GUID(CON_HF_WIDTH_SEGS,
0x779a5bbe, 0xb411, 0x4883, 0x90, 0xe2, 0x68, 0x4a, 0xa9, 0xe0, 0x59, 0x17);

// {ECDC30F7-F062-4f67-8474-C1230592381A}
RS_DEFINE_GUID(CON_HF_HEIGHT_SEGS,
0xecdc30f7, 0xf062, 0x4f67, 0x84, 0x74, 0xc1, 0x23, 0x5, 0x92, 0x38, 0x1a);

// {B1F0AE7B-BE97-44ae-A959-63AA6C031641}
RS_DEFINE_GUID(CON_HF_BITMAP, 
0xb1f0ae7b, 0xbe97, 0x44ae, 0xa9, 0x59, 0x63, 0xaa, 0x6c, 0x3, 0x16, 0x41);


// CRnHeightField

CRnHeightField::CRnHeightField()
{
}

HRESULT CRnHeightField::FinalConstruct()
{
    // get kernel package
    CComPtr<IRpKernel> spKernelPackage;
    HRESULT hr = CRpExamplesPlugin::GetKernelPackage(&spKernelPackage);
    
    // and create base node
    if (SUCCEEDED(hr))
        hr = spKernelPackage->CreateBaseNode(GetControllingUnknown(), &m_pUnkAgg);

    return hr;
}

STDMETHODIMP CRnHeightField::OnGetValue(CONID conID, IRsDataCache* pDataCache, VARIANT* pVtData, RtBool* pHandled)
{
    HRESULT hr = S_OK;

    if (conID == CON_MESH) {
		RtUlong segsx,segsy;
		CComPtr<IRsDataObject> spBitmapDo;

		//read input connectors
		{
			CComPtr<IRsDataCacheItem> spDi;
			pDataCache->GetDataCacheItem(CON_HF_WIDTH_SEGS,&spDi);
			if (spDi) spDi->GetDataUlong(&segsx);
			else segsx = 1;

			//params checking
			if (segsx > 1024) segsx = 1024;
		};

		{
			CComPtr<IRsDataCacheItem> spDi;
			pDataCache->GetDataCacheItem(CON_HF_HEIGHT_SEGS,&spDi);
			if (spDi) spDi->GetDataUlong(&segsy);
			else segsy = 1;

			//params checking
			if (segsy > 1024) segsy = 1024;
		};

		{
			CComPtr<IRsDataCacheItem> spDi;
			pDataCache->GetDataCacheItem(CON_HF_BITMAP,&spDi);
			if (spDi) spDi->GetDataObjectRef(&spBitmapDo);
		};

		//get bitmap interface from bitmap data object
		CComQIPtr<IRdBitmap> spBitmap = spBitmapDo;
        CComPtr<IRdMesh> spMesh;	
		GetMesh(segsx, segsy,spBitmap, &spMesh);
		
		//get output connector
        CComPtr<IRsDataCacheItem>    spConData;
        hr = pDataCache->GetDataCacheItem(CON_MESH, &spConData);
        if (FAILED(hr))
            return hr;

		//save value
        CComQIPtr<IRsDataObject>    spNewData = spMesh;
        hr = spConData->PutDataObject(spNewData);

		//validate output connector for caching
        spConData->Validate();
    }

    return hr;
}

STDMETHODIMP CRnHeightField::OnSetValue(CONID con_id, IRsDataCache* pDataCache, VARIANT* pvtData, RtBool* pHandled)
{
    return S_OK;
}

STDMETHODIMP CRnHeightField::OnDefaultValue(CONID con_id, VARIANT* pVtData, RtBool* pHandled)
{
    HRESULT hr = S_OK;

	if ((con_id == CON_HF_WIDTH_SEGS) || (con_id == CON_HF_HEIGHT_SEGS) ) {
		//default values for segment counts
        CComVariant defDiv((RtUlong)64);
        hr = defDiv.Detach(pVtData);
        *pHandled = TRUE;
    }
	else if (con_id == CON_HF_BITMAP) {
		//default bitmap
        CComVariant defDiv((IUnknown*)NULL);
        hr = defDiv.Detach(pVtData);
        *pHandled = TRUE;
    }


    return hr;
}
STDMETHODIMP CRnHeightField::OnInvalidate(CONID con_id, RtBool* pHandled)
{
    return S_OK;
}

STDMETHODIMP CRnHeightField::OnInitialize(void)
{
    return S_OK;
}

STDMETHODIMP CRnHeightField::OnLoadDefault(void)
{
    CComQIPtr<IRsBaseNode> spBase = m_pUnkAgg;   // Get the BaseNode interface
    ATLASSERT(spBase);

    RtConDesc   sSegsWidth  = { &CON_HF_WIDTH_SEGS,  L"Width segments",  L"Count of width segments",         VT_UI4, NULL, CONFLG_IN};
    RtConDesc   sSegsHeight = { &CON_HF_HEIGHT_SEGS, L"Height segments", L"Count of height segments",        VT_UI4, NULL, CONFLG_IN};
    RtConDesc   sBitmap	= { &CON_HF_BITMAP, L"Height field",  L"Height field bitmap",   VT_UNKNOWN, DEF_GUIDNAME_PTR(IRdBitmap), CONFLG_IN};

    CComPtr<IRsConnector>   spBitmap, spSegsWidth, spSegsHeight, spMeshCon;

    HRESULT hr = S_OK;

	//add registered mesh connector
    if(FAILED(hr = spBase->AddRegisteredConnector(CON_MESH, &spMeshCon)))
        return hr;

    if (FAILED(hr = spBase->AddNewConnector(&sBitmap, &spBitmap)))
        return hr;
    if (FAILED(hr = spBase->AddNewConnector(&sSegsWidth, &spSegsWidth)))
        return hr;
    if (FAILED(hr = spBase->AddNewConnector(&sSegsHeight, &spSegsHeight)))
        return hr;

	//set dependencies
    spBitmap->AddDependency(spMeshCon);
    spSegsWidth->AddDependency(spMeshCon);
    spSegsHeight->AddDependency(spMeshCon);
    
    return hr;
}

STDMETHODIMP CRnHeightField::OnDestroy(void)
{
    return S_OK;
}


void CRnHeightField::GetMesh(RtUlong segsx,RtUlong segsy, IRdBitmap *pHeightField, IRdMesh **ppMesh)
{
	//generate height field mesh from bitmap

	CComPtr<IRdMesh> spMesh;
	CRpExamplesPlugin::CreateDataObject(&spMesh);
	
	//get height map resolution. 
	RtDWORD resx,resy;
	if (pHeightField) {
		pHeightField->GetSizeX(&resx);
		pHeightField->GetSizeY(&resy);
		resx--;
		resy--;
	};

	//calculate number of output triangles and vertices

	RtUlong dwNumTriangles = segsx * segsy * 2;
	RtUlong dwNumVertices = (segsx + 1) * (segsy + 1);


	//create triangle indices stream
	CComPtr<IRdTriangleStream> spTriangles;
	CRpExamplesPlugin::CreateDataObject(&spTriangles);
	CComQIPtr<IRsGenStream> spTrianglesGen = spTriangles;

	
	//create vertex stream
	CComPtr<IRdVertexStream> spVertices;
	CRpExamplesPlugin::CreateDataObject(&spVertices);
	CComQIPtr<IRsGenStream> spVerticesGen = spVertices;

	//create UV indices stream
	CComPtr<IRdUVTriangleStream> spUvTris;
	CRpExamplesPlugin::CreateDataObject(&spUvTris);
	CComQIPtr<IRsGenStream> spUvTrisGen = spUvTris;

	//create UV coordinates stream
	CComPtr<IRdUVCoordStream> spUvCoords;
	CRpExamplesPlugin::CreateDataObject(&spUvCoords);
	CComQIPtr<IRsGenStream> spUvCoordsGen = spUvCoords;

	//create edge flags stream
	CComPtr<IRdTriangleEdgesStream> spEdges;
	CRpExamplesPlugin::CreateDataObject(&spEdges);
	CComQIPtr<IRsGenStream> spEdgesGen = spEdges;

	//set number of items for streams and lock them for writing
	
	RtTripleIndex *pTriangles, *pUvIndices;
	RtPoint *pVertices;
	RtUVcoord *pUvCoords;
	RtDWORD *pEdgeFlags;

	spTrianglesGen->SetNumItems(dwNumTriangles);
	spTrianglesGen->LockStream(MFLG_WILLOVERWRITE, (void**)&pTriangles);

	spUvTrisGen->SetNumItems(dwNumTriangles);
	spUvTrisGen->LockStream(MFLG_WILLOVERWRITE, (void**)&pUvIndices);

	spEdgesGen->SetNumItems(dwNumTriangles);
	spEdgesGen->LockStream(MFLG_WILLOVERWRITE, (void**)&pEdgeFlags);

	spVerticesGen->SetNumItems(dwNumVertices);
	spVerticesGen->LockStream(MFLG_WILLOVERWRITE, (void**)&pVertices);

	spUvCoordsGen->SetNumItems(dwNumVertices);
	spUvCoordsGen->LockStream(MFLG_WILLOVERWRITE, (void**)&pUvCoords);



	float u,v;
	RtUlong x,y;
	RtUlong ptr = 0;
	//generate vertices and UV coordinates
	for (y = 0; y<=segsy; y++) {
		v = float(y) / float(segsy);

		for (x = 0; x<=segsx;x++) {
			u = float(x) / float(segsx);

			//calculate position
			pVertices[ptr].x = 2*u - 1;
			pVertices[ptr].y = 2*v - 1;
			if (pHeightField) {
				COLORREF c;
				pHeightField->GetPixelCOLORREF(RtUlong(resx*u), RtUlong(resy*v),&c);

				//average rgb to get height
				pVertices[ptr].z = float(((c & 0xFF0000) >> 16) + ((c & 0x00FF00) >> 8) + ((c & 0x0000FF))) / 768.0f;
			}
			else pVertices[ptr].z = 0;


			//calculate UV coordinates
			pUvCoords[ptr].u = u;
			pUvCoords[ptr].v = v;
			ptr++;
		};
	};


	//calculate triangles
	ptr = 0;
	for (y = 0; y<segsy; y++) {
		for (x = 0; x<segsx;x++) {
			pUvIndices->i = pTriangles->i = ptr;
			pUvIndices->j = pTriangles->j = ptr+1;
			pUvIndices->k = pTriangles->k = ptr+segsx+2;
			pEdgeFlags[0] = 1+2;
			pTriangles++;
			pUvIndices++;
			pEdgeFlags++;

			pUvIndices->i = pTriangles->i = ptr;
			pUvIndices->j = pTriangles->j = ptr+segsx+2;
			pUvIndices->k = pTriangles->k = ptr+segsx+1;
			pEdgeFlags[0] = 2+4;
			pTriangles++;
			pUvIndices++;
			pEdgeFlags++;

			ptr++;
		};
		ptr++;
	};


	//unlock streams
	spTrianglesGen->UnLockStream();
	spUvTrisGen->UnLockStream();
	spVerticesGen->UnLockStream();
	spUvCoordsGen->UnLockStream();
	spEdgesGen->UnLockStream();
	

	//attach streams to mesh

	//UV indices, vertex indices and triangle edge flags go to triangle stream group
	spMesh->AttachTrianglesStream(spTrianglesGen);
	spMesh->AttachTrianglesStream(spUvTrisGen);
	spMesh->AttachTrianglesStream(spEdgesGen);

	//vertex coordinates belong to vertex group
	spMesh->AttachVerticesStream(spVerticesGen);

	//uv coordinates go to custom stream group
	spMesh->AttachCustStream(spUvCoordsGen);

	//use space3d mesh fn set to generate smooth normals
	CComPtr<IRpSpace3DPck> spSpace3d;
	CRpExamplesPlugin::FindPackage(DEF_GUIDNAME_PTR(CRpSpace3DPck), &spSpace3d);
	CComPtr<IRfMeshSet> spMeshSet;
	spSpace3d->CreateMeshFnSet(&spMeshSet);

	//attach mesh
	spMeshSet->Attach(spMesh);

	//calculate vertex normals (smooth)
	spMeshSet->CreateVertexNormals();


	//copy mesh to output...
	spMesh.CopyTo(ppMesh);	
};