// Fill out your copyright notice in the Description page of Project Settings.

#include "PDB_viewer.h"
#include "ProteinActor.h"

#include <set>
#include <utility>

#include "Protein.hpp"
#include "Parser.hpp"
#include <boost/optional.hpp>

#include "CPKColorScheme.h"
#include "ColorSchemeFromFile.h"

static FRotator getRotation(FVector actor, FVector target);

static const float BB_EXTENT_LENGTH = 50;
const std::string AProteinActor::RESIDUE_FILLER_DATA_PATH = TCHAR_TO_ANSI(*(FPaths::GameContentDir() + TEXT("ThirdPartyData/cachedAminoacids.txt")));

// Sets default values
AProteinActor::AProteinActor()
{
	// 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;

	//root component - everything else will be in relation to this. In editor it also shows the bounding box of the model
	rootBox = CreateDefaultSubobject<UBoxComponent>(TEXT("Root Component"));
	RootComponent = rootBox;
	rootBox->InitBoxExtent(FVector(BB_EXTENT_LENGTH, BB_EXTENT_LENGTH, BB_EXTENT_LENGTH));
	rootBox->SetCollisionEnabled(ECollisionEnabled::NoCollision);
	//protein parent - it will be moved and scaled in order to normalize the protein mesh
	proteinParent = CreateDefaultSubobject<USphereComponent>(TEXT("Protein Parent"));
	proteinParent->InitSphereRadius(1.);
	proteinParent->SetupAttachment(rootBox);
	proteinParent->SetCollisionEnabled(ECollisionEnabled::NoCollision);

	bCanBeDamaged = false;
	AutoReceiveInput = EAutoReceiveInput::Player0;
	material[MATERIAL_OPAQUE] = ConstructorHelpers::FObjectFinder<UMaterialInterface>(TEXT("Material'/Game/MCloudyColoredMaterial.MCloudyColoredMaterial'")).Object;
	material[MATERIAL_TRANSLUSCENT] = ConstructorHelpers::FObjectFinder<UMaterialInterface>(TEXT("Material'/Game/MTranslucentColoredMaterial.MTranslucentColoredMaterial'")).Object;
	sphereMesh = ConstructorHelpers::FObjectFinder<UStaticMesh>(TEXT("StaticMesh'/Game/UVsphere16.UVsphere16'")).Object;
	stickMesh = ConstructorHelpers::FObjectFinder<UStaticMesh>(TEXT("StaticMesh'/Game/cylinder.cylinder'")).Object;	
}

// Called when the game starts or when spawned
void AProteinActor::BeginPlay()
{																							
	Super::BeginPlay();
	InputComponent->BindAction("SpawnProtein", IE_Pressed, this, &AProteinActor::CustomInit);
	InputComponent->BindAction("ShowHydrogen", IE_Pressed, this, &AProteinActor::ToggleHydrogen);
}

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

static const TCHAR* ATOM_TAG = TEXT("ATOM");
static const TCHAR* STICK_TAG = TEXT("STICK");

void AProteinActor::Init(Protein _protein)
{
	protein = _protein;
	showHydrogen = false;
	ShowModel(0);
}

void AProteinActor::CustomInit()
{
	try
	{
		Parser parser("C:\\Users\\bartek\\Documents\\Unreal Projects\\PDB_viewer\\SamplePDBs\\DNA.pdb");
		parser.parseFile();
		Protein p = parser.getProtein();
		Init(p);
	}
	catch (std::string str)
	{
		printf("Error while loading file: %s\n", str.c_str());
	}
}

void AProteinActor::ToggleHydrogen()
{
	showHydrogen = !showHydrogen;
	UpdateVisibility();
}

void AProteinActor::LoadFromFile(FString path)
{
	try
	{
		ResidueConnection residueConnectionFiller = ResidueConnection(RESIDUE_FILLER_DATA_PATH);
		Parser parser((TCHAR_TO_UTF8(*path)));
		parser.parseFile();
		Protein p = parser.getProtein();
		p = residueConnectionFiller.addConnectionsWithinResidues(p);
		p = residueConnectionFiller.addConnectionsWithinChain(p);
		Init(p);
	}
	catch (std::string str)
	{
		printf("Error while loading file: %s\n", str.c_str());
	}
}

TArray<FString> AProteinActor::GetModelSerials()
{
	TArray<FString> result;
	for (Model& m : protein.models)
	{
		result.Add(std::to_string(m.serial).c_str());
	}
	return result;
}

void AProteinActor::ShowModel(int modelIndex)
{
	if (modelIndex < 0 || modelIndex >= protein.models.size())
		return; //fail silently
	selectedModel = modelIndex;
	ClearComponents();
	GenerateComponentsForModel(protein.models[selectedModel]);
	UpdateVisibility();
	NormalizeMesh();
}

FTransform AProteinActor::NormalizeMesh()
{
	if (atomVisualVector.empty())
		return FTransform();
	FVector BB_min(+BIG_NUMBER), BB_max(-BIG_NUMBER);
	for (UStaticMeshComponent* atomMeshComponent : atomVisualVector)
	{
		FVector relativeLocation = atomMeshComponent->GetRelativeTransform().GetLocation();
		BB_min = BB_min.ComponentMin(relativeLocation);
		BB_max = BB_max.ComponentMax(relativeLocation);
	}
	FVector center = (BB_max + BB_min) / 2;
	//the added constant alleviates the 0.5 radius of the spheres used
	float scale_factor = (BB_EXTENT_LENGTH * 2)/((BB_max - BB_min).GetMax() + 1);
	const FTransform loc(-center);
	const FTransform rotScale(FQuat(0.5, 0.5, 0.5, 0.5), FVector(0), FVector(scale_factor));
	FTransform finalNormalizedTransform = FTransform();
	FTransform::Multiply(&finalNormalizedTransform, &loc, &rotScale);
	proteinParent->SetRelativeTransform(finalNormalizedTransform);
	return finalNormalizedTransform;
}

void AProteinActor::UpdateVisibility()
{
	for (UStaticMeshComponent* atomMeshComponent : atomVisualVector)
	{
		int serial = FCString::Atoi(*atomMeshComponent->ComponentTags[1].ToString());
		boost::optional<Atom&> atom = protein.models[selectedModel].findAtom(serial);
		if (atom && atom->element == "H")
		{
			atomMeshComponent->SetHiddenInGame(!showHydrogen);
		}
	}
	for (UStaticMeshComponent* stickMeshComponent : stickVisualVector)
	{
		bool connectsWithHydrogen = false;
		for (int i = 0; i < 2; ++i)
		{
			int serial = FCString::Atoi(*stickMeshComponent->ComponentTags[1+i].ToString());
			boost::optional<Atom&> atom = protein.models[selectedModel].findAtom(serial);
			if (atom && atom->element == "H")
			{
				connectsWithHydrogen = true;
				break;
			}
		}
		if (connectsWithHydrogen)
		{
			stickMeshComponent->SetHiddenInGame(!showHydrogen);
		}
	}
}

void AProteinActor::ClearComponents()
{
	for (UStaticMeshComponent* thing : atomVisualVector)
	{
		thing->DetachFromComponent(FDetachmentTransformRules::KeepRelativeTransform);
		thing->UnregisterComponent();
		thing->DestroyComponent();
	}
	atomVisualVector.clear();
	for (UStaticMeshComponent* thing : stickVisualVector)
	{
		thing->DetachFromComponent(FDetachmentTransformRules::KeepRelativeTransform);
		thing->UnregisterComponent();
		thing->DestroyComponent();
	}
	stickVisualVector.clear();
	//Have to clear cached materials - they are trashed by GC anyway when not used in a mesh for 1 minute
	for (int i = 0; i < MATERIAL_COUNT; ++i)
		cachedMaterials[i].clear();
}

void AProteinActor::GenerateComponentsForModel(Model& model)
{
	ColorSchemeFromFile fileScheme(TCHAR_TO_ANSI(*(FPaths::GameContentDir() + TEXT("ThirdPartyData/CPK_Jmol_color_scheme.tsv"))));
	CPKColorScheme CPKScheme;
	AtomColorScheme& colorScheme = (fileScheme.isInitialized()) ? ((AtomColorScheme&)fileScheme) : ((AtomColorScheme&)CPKScheme);
	std::set<std::pair<int, int>> sticks;
	for (Chain chain : model.chains)
	{
		for (Residue residue : chain.residues)
		{
			for (Atom atom : residue.atoms)
			{
				FVector atomPos(atom.x, atom.y, atom.z);
				UStaticMeshComponent* newComponent = NewObject<UStaticMeshComponent>(this);

				if (sphereMesh)
				{
					newComponent->SetStaticMesh(sphereMesh);
				}
				newComponent->ComponentTags.Add(ATOM_TAG);
				newComponent->ComponentTags.Add(*FString::Printf(TEXT("%d"), atom.serial));

				for (int i = 0; i < newComponent->GetNumMaterials(); ++i)
				{
					int material_type = (newComponent->GetMaterial(i)->GetBlendMode() == BLEND_Translucent)
						? MATERIAL_TRANSLUSCENT : MATERIAL_OPAQUE;
					if (cachedMaterials[material_type].find(atom.element) == cachedMaterials[material_type].end())
					{
						UMaterialInstanceDynamic* newMaterial = UMaterialInstanceDynamic::Create(material[material_type], this);
						newMaterial->SetVectorParameterValue(TEXT("InputColor"), colorScheme.getColorForElement(atom.element));
						cachedMaterials[material_type][atom.element] = newMaterial;
					}

					UMaterialInstanceDynamic* materialToUse = cachedMaterials[material_type][atom.element];
					newComponent->SetMaterial(i, materialToUse);
				}

				newComponent->RegisterComponent();
				newComponent->AttachToComponent(proteinParent, FAttachmentTransformRules::KeepRelativeTransform);
				newComponent->bGenerateOverlapEvents = true;
				newComponent->SetCollisionEnabled(ECollisionEnabled::QueryOnly);
				newComponent->SetCollisionResponseToChannel(ECC_WorldStatic, ECR_Overlap);
				newComponent->SetCollisionResponseToChannel(ECC_WorldDynamic, ECR_Overlap);
				newComponent->SetCollisionResponseToChannel(ECC_Pawn, ECR_Overlap);
				newComponent->SetRelativeLocation(atomPos);
				newComponent->SetRelativeScale3D(FVector(0.5));
				atomVisualVector.push_back(newComponent);

				for (int atomID : atom.bondedAtomIds)
				{
					int atom1 = atom.serial;
					int atom2 = atomID;
					if (atom1 > atom2)
						std::swap(atom1, atom2);
					std::pair<int, int> stick(atom1, atom2);
					if (sticks.find(stick) == sticks.end())
						sticks.insert(stick);
					else
						continue;

					boost::optional<Atom&> otherAtom = model.findAtom(atomID);
					if (otherAtom)
					{
						FVector otherAtomPos(otherAtom.get().x, otherAtom.get().y, otherAtom.get().z);
						FVector diff = otherAtomPos - atomPos;

						newComponent = NewObject<UStaticMeshComponent>(this);
						if (stickMesh)
						{
							newComponent->SetStaticMesh(stickMesh);
						}

						newComponent->ComponentTags.Add(STICK_TAG);
						newComponent->ComponentTags.Add(*FString::Printf(TEXT("%d"), atom1));
						newComponent->ComponentTags.Add(*FString::Printf(TEXT("%d"), atom2));

						newComponent->RegisterComponent();
						newComponent->AttachToComponent(proteinParent, FAttachmentTransformRules::KeepRelativeTransform);
						newComponent->SetCollisionEnabled(ECollisionEnabled::NoCollision);

						float dist = (atomPos - otherAtomPos).Size();
						newComponent->SetRelativeScale3D(FVector(0.0008, 0.0008, dist * 0.005));
						newComponent->SetRelativeLocation((atomPos + otherAtomPos)/2);
						FRotator rotator = getRotation(atomPos, otherAtomPos);
						newComponent->SetRelativeRotation(rotator);
						stickVisualVector.push_back(newComponent);
					}
				}
			}
		}
	}

}

FRotator getRotation(FVector host, FVector target)
{
	FVector diff = target - host;
	FVector normalized = diff.GetSafeNormal();

	FVector XY = FVector(diff.X, diff.Y, 0);
	FVector XYnormalized = XY.GetSafeNormal();
	
	float Yaw  = 0;
	float Roll = 0;
	float Pitch = 0;

	Roll = FVector::DotProduct(FVector::UpVector, normalized);
	Roll = acosf(Roll);

	Yaw = FVector::DotProduct(FVector::RightVector, XYnormalized);
	Yaw = acosf(Yaw);
	if (diff.X > 0)
		Yaw *= -1;

	Roll *= 180.0f / PI;
	Yaw  *= 180.0f / PI;

	return FRotator(Pitch, Yaw, Roll);
}
