/*
* @Author: namidame
* @Description: A Heatmap Generate Plugin, Supports Heightmap, Texture Coordinate Points And Geographic Location Data.
* @Date: 2023/03/24
*/

#include "HeatMapActor.h"
#include <vector>
#include <functional>

#include "GeoReferencingSystem.h"

using namespace std;

// Sets default values
AHeatMapActor::AHeatMapActor()
	: TextureSize(FVector2D(2048, 2048))
	, MapSize(FVector2D(100000, 100000))
	, MapSecment(FVector2D(500, 500))
	, LerpScale(2.5)
	, HeightScale(2500.0)
	, Opacity(0.7)
	, InfluenceSize(50)
{
 	// Set this actor to call Tick() every frame.  You can turn this off to improve performance if you don't need it.
	// PrimaryActorTick.bCanEverTick = true;

	RootComponent = CreateDefaultSubobject<USceneComponent>(TEXT("RootComponent"));
	
	MapMesh = CreateDefaultSubobject<UProceduralMeshComponent>(TEXT("MapMesh"));
	MapMesh->SetupAttachment(RootComponent);
}

// Called when the game starts or when spawned
void AHeatMapActor::BeginPlay()
{
	Super::BeginPlay();

}

// Called every frame
void AHeatMapActor::Tick(float DeltaTime)
{
	Super::Tick(DeltaTime);

}

void AHeatMapActor::CreateMapMesh(FVector2D mapSize)
{
	UE_LOG(LogTemp, Log, TEXT("Create Map Mesh"));
	MapSize = mapSize;
	float secmentX = MapSecment.X;
	float secmentY = MapSecment.Y;
	
	MapMesh->ClearAllMeshSections();

	TArray<FVector> Vertices;
	TArray<int32> Triangles;
	TArray<FVector> Normals;
	TArray<FVector2D> UV0;
	TArray<FColor> VertexColors;
	TArray<FProcMeshTangent> Tangents;

	int32 curX = 0;
	int32 curY = 0;
	FVector curLeftTop = FVector::ZeroVector;
	float xStep = MapSize.X / secmentX;
	float yStep = MapSize.Y / secmentY;

	float uStep = 1 / secmentX;
	float vStep = 1 / secmentY;
	int32 times = 0;
	while(curY < secmentY)
	{
		while(curX < secmentX)
		{
			FVector leftTop = curLeftTop;
			FVector rightTop = leftTop + GetActorForwardVector() * xStep;
			FVector leftBottom = curLeftTop + GetActorRightVector() * yStep;
			FVector rightBottom = leftBottom + GetActorForwardVector() * xStep;;

			int32 startIdx = Vertices.Num();
			Vertices += {leftTop, rightTop, leftBottom, rightBottom};
			Triangles += {startIdx+2, startIdx+1, startIdx, startIdx+3, startIdx+1, startIdx+2};
			Normals += {FVector::ZAxisVector, FVector::ZAxisVector, FVector::ZAxisVector, FVector::ZAxisVector};
			UV0 +=
				{
					FVector2D(curX*uStep, curY*vStep),
					FVector2D((curX+1)*uStep, curY*vStep),
					FVector2D(curX*uStep, (curY+1)*vStep),
					FVector2D((curX+1)*uStep, (curY+1)*vStep)
				};
			curLeftTop = rightTop;
			++curX;
			++times;
		}
		curX = 0;
		++curY;
		curLeftTop = GetActorRightVector() * curY * yStep;
	}
	MapMesh->CreateMeshSection(0, Vertices, Triangles, Normals, UV0, VertexColors, Tangents, false);
}

void AHeatMapActor::CreateWithHeightMap(UTexture2D* texture, FVector2D mapSize)
{
	m_type = E_HeatMapType::HEIGHT_MAP;
    int32 sizeX = texture->GetSizeX();
    int32 sizeY = texture->GetSizeY();

	FVector2D textureSize(sizeX, sizeY);
	TArray<FColor> colorArr = AHeatMapActor::GetColorDataFromTexture(texture, textureSize);

	this->CreateWithData(colorArr, textureSize, mapSize);
}

void AHeatMapActor::CreateWithData(const TArray<FColor>& colorArr, FVector2D textureSize, FVector2D mapSize)
{
	TextureSize = textureSize;
// UE_LOG(LogTemp, Log, TEXT("check===%d,%d"), colorArr[0].B, colorArr[TextureSize.X-1].B);
	UTexture2D* texture = UTexture2D::CreateTransient(textureSize.X, textureSize.Y);
	FTexture2DMipMap* MipMap = &texture->GetPlatformData()->Mips[0];
	FByteBulkData* ImageData = &MipMap->BulkData;
	uint8* RawImageData = (uint8*)ImageData->Lock(LOCK_READ_WRITE);
	FMemory::Memcpy(RawImageData, colorArr.GetData(), colorArr.Num()*4);
	ImageData->Unlock();
	texture->UpdateResource();
	
	// create map mesh
	this->CreateMapMesh(mapSize);
	
	// create material with texture
	if(!HeatMapActorMaterial) return;
	UMaterialInstanceDynamic* mat = UMaterialInstanceDynamic::Create(HeatMapActorMaterial, nullptr);
	mat->SetTextureParameterValue(TEXT("Texture"), texture);
	mat->SetScalarParameterValue(TEXT("ModelScale"), mapSize.X / 100000);
	
	Material = mat;
	
	this->setMaterialLerpScale(LerpScale);
	this->setMaterialHeightScale(HeightScale);
	this->setMaterialOpacity(Opacity);
	
	// set material into map mesh
	MapMesh->SetMaterial(0, mat);
}

void AHeatMapActor::setMaterialLerpScale(float scale)
{
	LerpScale = scale;
	if(Material)
	{
		Material->SetScalarParameterValue(TEXT("LerpScale"), scale);
	}
}

void AHeatMapActor::setMaterialHeightScale(float scale)
{
	HeightScale = scale;
	if(Material)
	{
		Material->SetScalarParameterValue(TEXT("HeightScale"), scale);
	}
}

void AHeatMapActor::setMaterialOpacity(float opacity)
{
	Opacity = opacity;
	if(Material)
	{
		Material->SetScalarParameterValue(TEXT("Opacity"), opacity);
	}
}

TArray<FColor> AHeatMapActor::GetColorDataFromTexture(UTexture2D* texture, FVector2D textureSize)
{
	TArray<FColor> rlt;
	
	TextureCompressionSettings oldSetting = texture->CompressionSettings;
#if WITH_EDITORONLY_DATA
	TextureMipGenSettings mipSetting = texture->MipGenSettings;
#endif
	bool isSRGB = texture->SRGB;
	
	texture->CompressionSettings = TextureCompressionSettings::TC_VectorDisplacementmap;
#if WITH_EDITORONLY_DATA
	texture->MipGenSettings = TextureMipGenSettings::TMGS_NoMipmaps;
#endif
		
	texture->SRGB = false;
	texture->UpdateResource();
	
	const FColor* colorArr = (FColor*)(texture->GetPlatformData()->Mips[0].BulkData.Lock(LOCK_READ_ONLY));
	
	for(int32 y = 0; y < textureSize.Y; ++y)
	{
		for(int32 x = 0; x < textureSize.X; ++x)
		{
			FColor color = colorArr[y * (int32)textureSize.X + x];
			rlt.Push(color);
		}
	}
	texture->GetPlatformData()->Mips[0].BulkData.Unlock();

	texture->CompressionSettings = oldSetting;
#if WITH_EDITORONLY_DATA
	texture->MipGenSettings = mipSetting;
#endif
	texture->SRGB = isSRGB;
	texture->UpdateResource();

	return rlt;
}

TArray<FColor> AHeatMapActor::GetColorDataFromPointHeightMap(const TMap<FVector2D, float>& map, FVector2D textureSize, int32 influenceSize)
{
	// 根据零散的点集生成高度图数据

	int32 sizeX = textureSize.X;
	int32 sizeY = textureSize.Y;

	std::vector<float> heightVec(sizeX * sizeY, 0);
	// 遍历所有数据点，对纹理数据增加权重
	for(auto it : map)
	{
		int32 px = it.Key.X;
		int32 py = it.Key.Y;
		float height = it.Value;

		int32 startX = FMath::Max(px - influenceSize, 0);
		int32 endX = FMath::Min(px + influenceSize, sizeX - 1);
		int32 startY = FMath::Max(py - influenceSize, 0);
		int32 endY = FMath::Min(py + influenceSize, sizeY - 1);
		for(int32 j = startY; j <= endY; ++j)
		{
			for(int32 i = startX; i <= endX; ++i)
			{
				int32 distance = (i - px)*(i - px) + (j - py)*(j - py);

				if(distance <= influenceSize * influenceSize)
				{
					// 在影响圈内 增加权重
					float percent = 1.0f * distance / (influenceSize * influenceSize);
					float influence = 1 - percent;
					if(InfluenceCurve)
					{
						influence = InfluenceCurve->GetFloatValue(percent);
					}
					float newHeight = influence * height;
					heightVec[j * sizeX + i] = FMath::Max(heightVec[j * sizeX + i], newHeight);
				}
			}
		}
	}

	
	TArray<FColor> rlt;
	for (int32 i = 0; i < heightVec.size(); ++i)
	{
		float height = heightVec[i];
		int32 bValue = FMath::CeilToInt(height * 255.0f);
		bValue = FMath::Min(bValue, 255);
		FColor color(0, 0, bValue);
		rlt.Add(color);
	}
	return rlt;
}

void AHeatMapActor::setMapSecment(FVector2D secment)
{
	MapSecment = FVector2D(round(secment.X), round(secment.Y));
	this->ReGenerateMapAndMaterial();
}

void AHeatMapActor::setMapSize(FVector2D size)
{
	MapSize = size;
	this->ReGenerateMapAndMaterial();
}

void AHeatMapActor::CreateWithPointHeightValue(const TMap<FVector2D, float>& map, FVector2D textureSize, FVector2D mapSize, bool isGeoData)
{
	if(isGeoData)
	{
		m_type = E_HeatMapType::GEO_DATA;
	}
	else
	{
		m_type = E_HeatMapType::POINT_HEIGHT_DATA;
	}
	PointHeightValueMap = map;

	TArray<FColor> colorArr = this->GetColorDataFromPointHeightMap(map, textureSize, InfluenceSize);

	this->CreateWithData(colorArr, textureSize, mapSize);
}

void AHeatMapActor::CreateWithGeoInfoMap(const TMap<FVector, float>& map)
{
	GeoInfoMap = map;
	TMap<FVector2D, float> rlt = this->ConvertGeoValueMapToUnrealValueMap(map);
	this->CreateWithPointHeightValue(rlt, TextureSize, MapSize, true);
}

void AHeatMapActor::AddPointHeightValue(const FVector2D& pos, float value)
{
	if(m_type == E_HeatMapType::POINT_HEIGHT_DATA)
	{
		if(PointHeightValueMap.Contains(pos))
		{
			PointHeightValueMap[pos] = value;
		}
		else
		{
			PointHeightValueMap.Add(pos, value);
		}
		this->ReGenerateMapAndMaterial();
	}
}

void AHeatMapActor::DeletePointHeightValue(const FVector2D& pos)
{
	if(m_type == E_HeatMapType::POINT_HEIGHT_DATA)
	{
		if(PointHeightValueMap.Contains(pos))
		{
			PointHeightValueMap.Remove(pos);
		}
		this->ReGenerateMapAndMaterial();
	}
}

void AHeatMapActor::setPointInfluenceSize(int32 size)
{
	InfluenceSize = size;
	this->ReGenerateMapAndMaterial();
}

void AHeatMapActor::ReGenerateMapAndMaterial()
{
	if(m_type == E_HeatMapType::HEIGHT_MAP)
	{
		this->CreateWithHeightMap(HeightMapTexture, MapSize);	
	}
	else if(m_type == E_HeatMapType::POINT_HEIGHT_DATA)
	{
		this->CreateWithPointHeightValue(PointHeightValueMap, TextureSize, MapSize, false);
	}
	else if(m_type == E_HeatMapType::GEO_DATA)
	{
		this->CreateWithGeoInfoMap(GeoInfoMap);
	}
}

void AHeatMapActor::setMapGeoLocation(FVector location)
{
	FVector newPos = this->ConvertGeoGraphicToEngine(location);	
	SetActorLocation(newPos);

	this->ReGenerateMapAndMaterial();
}

TMap<FVector2D, float> AHeatMapActor::ConvertGeoValueMapToUnrealValueMap(const TMap<FVector, float>& map)
{
	TMap<FVector2D, float> rlt1;
	float minX = 0, maxX = 0, minY = 0, maxY = 0;
	
	for(auto it : map)
	{
		FVector loc = it.Key;
		
		FVector newLoc = this->ConvertGeoGraphicToEngine(loc);
		
		FVector2D loc2(newLoc.X, newLoc.Y);
		rlt1.Add(loc2, it.Value);

		// 找出最大最小坐标 从而确定平面的位置
		if(loc2.X < minX) minX = loc2.X;
		if(loc2.X > maxX) maxX = loc2.X;
		if(loc2.Y < minY) minY = loc2.Y;
		if(loc2.Y > maxY) maxY = loc2.Y;
	}

	SetActorLocation(FVector(minX, minY, 0));
	
	float xLen = maxX - minX;
	float yLen = maxY - minY;
	TMap<FVector2D, float> rlt2;
	
	// 计算textureSize
	int32 defaultSizeBase = 2048;
	if(xLen > yLen)
	{
		TextureSize = FVector2D((int32)(defaultSizeBase * xLen / yLen), defaultSizeBase);
	}
	else
	{
		TextureSize = FVector2D(defaultSizeBase, (int32)(defaultSizeBase * yLen / xLen));
	}
	
	// 将点的坐标映射到tex.X, tex.Y范围内
	for(auto it : rlt1)
	{
		FVector2D loc = it.Key;
		float xPercent = (loc.X - minX) / xLen;
		float x = xPercent * TextureSize.X;

		float yPercent = (loc.Y - minY) / yLen;
		float y = yPercent * TextureSize.Y;

		rlt2.Add(FVector2D(x, y), it.Value);
	}
	// 计算mapSize
	MapSize = FVector2D((int32)(xLen), (int32)(yLen));
	
	return rlt2;
}

FVector AHeatMapActor::ConvertGeoGraphicToEngine(FVector location)
{
	FGeographicCoordinates geo(location.X, location.Y, location.Z);
	FVector newPos;
	AGeoReferencingSystem::GetGeoReferencingSystem(GetWorld())->GeographicToEngine(geo, newPos);
	return newPos;
}