// This file is a part of Framsticks SDK.  http://www.framsticks.com/
// Copyright (C) 1999-2020  Maciej Komosinski and Szymon Ulatowski.
// See LICENSE.txt for details.

#include "measure-hungarian.h"

const int SimilMeasureHungarian::iNOFactors = 4;

#define FIELDSTRUCT SimilMeasureHungarian

static ParamEntry simil_hungarian_paramtab[] = {
		{ "Creature: Similarity: Graph optimal", 1, 7, "SimilMeasureHungarian", "Evaluates morphological dissimilarity using the measure. More information:\nhttp://www.framsticks.com/bib/Komosinski-et-al-2001\nhttp://www.framsticks.com/bib/Komosinski-and-Kubiak-2011\nhttp://www.framsticks.com/bib/Komosinski-2016\nhttps://www.framsticks.com/bib/Komosinski-and-Mensfelt-2019", },
		{ "simil_parts", 0, 0, "Weight of parts count", "f 0 100 0", FIELD(m_adFactors[0]), "Differing number of parts is also handled by the 'part degree' similarity component.", },
		{ "simil_partdeg", 0, 0, "Weight of parts' degree", "f 0 100 1", FIELD(m_adFactors[1]), "", },
		{ "simil_neuro", 0, 0, "Weight of neurons count", "f 0 100 0.1", FIELD(m_adFactors[2]), "", },
		{ "simil_partgeom", 0, 0, "Weight of parts' geometric distances", "f 0 100 0", FIELD(m_adFactors[3]), "", },
		{ "simil_fixedZaxis", 0, 0, "Fix 'z' (vertical) axis?", "d 0 1 0", FIELD(fixedZaxis), "", },
		{ "simil_weightedMDS", 0, 0, "Should weighted MDS be used?", "d 0 1 0", FIELD(wMDS), "If activated, weighted MDS with vertex (i.e., Part) degrees as weights is used for 3D alignment of body structure.", },
		{ "evaluateDistance", 0, PARAM_DONTSAVE | PARAM_USERHIDDEN, "Evaluate model dissimilarity", "p f(oGeno,oGeno)", PROCEDURE(p_evaldistance), "Calculates dissimilarity between two models created from Geno objects.", },
		{ 0, },
};

#undef FIELDSTRUCT

SimilMeasureHungarian::SimilMeasureHungarian() : localpar(simil_hungarian_paramtab, this)
{
	localpar.setDefault();
	
	nSmaller = 0;
	nBigger = 0;
	
	for (int i = 0; i < 2; i++)
	{
		degrees[i] = nullptr;
		neurons[i] = nullptr;
		on_joint[i] = 0;
		anywhere[i] = 0;	
	}
	
	assignment = nullptr;
	parts_distances = nullptr;
	temp_parts_distances = nullptr;
	
	save_matching = false;
}

void SimilMeasureHungarian::prepareData()
{
	m_iSmaller = models[0]->getPartCount() <= models[1]->getPartCount() ? 0 : 1;
	nSmaller = models[m_iSmaller]->getPartCount();
	nBigger = models[1 - m_iSmaller]->getPartCount();

	for (int i = 0; i < 2; i++)
	{
		int size = models[i]->getPartCount();
		degrees[i] = new int[size]();
		neurons[i] = new int[size]();
	}
	
	countDegrees();
	countNeurons();
	
	parts_distances = new double[nBigger*nBigger]();
	fillPartsDistances(parts_distances, nBigger, nSmaller, false);
	assignment = new int[nBigger]();
	
	if (save_matching)
		for (int i = 0; i < nBigger; i++)
			min_assignment.push_back(0);
        
        if (m_adFactors[3] == 0)
            with_alignment = false;
}

void SimilMeasureHungarian::beforeTransformation()
{
	temp_parts_distances = new double[nBigger*nBigger]();
	std::copy(parts_distances, parts_distances + nBigger * nBigger, temp_parts_distances);
}

double SimilMeasureHungarian::distanceForTransformation()
{
	fillPartsDistances(temp_parts_distances, nBigger, nSmaller, true);
	std::fill_n(assignment, nBigger, 0);
	double distance = hungarian.Solve(temp_parts_distances, assignment, nBigger, nBigger);

	delete[] temp_parts_distances;
	return addNeuronsPartsDiff(distance);
}

double SimilMeasureHungarian::distanceWithoutAlignment()
{
	double distance = hungarian.Solve(parts_distances, assignment, nBigger, nBigger);
	if (save_matching)
		copyMatching();
	return addNeuronsPartsDiff(distance);
}

double SimilMeasureHungarian::addNeuronsPartsDiff(double dist)
{
	//add difference in anywhere and onJoint neurons
	dist += m_adFactors[2] * (abs(on_joint[0] - on_joint[1]) + abs(anywhere[0] - anywhere[1]));
	//add difference in part numbers
	dist += (nBigger - nSmaller) * m_adFactors[0];
	return dist;
}

void SimilMeasureHungarian::copyMatching()
{
	min_assignment.clear();
	min_assignment.insert(min_assignment.begin(), assignment, assignment + nBigger);
}

void SimilMeasureHungarian::cleanData()
{
	for (int i = 0; i < 2; i++)
	{
		// delete degree and position arrays
		SAFEDELETEARRAY(degrees[i]);
		SAFEDELETEARRAY(neurons[i]);	

		on_joint[i] = 0;
		anywhere[i] = 0;	
	}

	delete[] assignment;
	delete[] parts_distances;
	
	if (save_matching)
		min_assignment.clear();
        
        with_alignment = true; //restore default value
}

void SimilMeasureHungarian::countDegrees()
{
	Part *P1, *P2;
	int i, j, i1, i2;

	for (i = 0; i < 2; i++)
	{
		for (j = 0; j < models[i]->getJointCount(); j++)
		{
			Joint *J = models[i]->getJoint(j);

			P1 = J->part1;
			P2 = J->part2;

			i1 = models[i]->findPart(P1);
			i2 = models[i]->findPart(P2);

			degrees[i][i1]++;
			degrees[i][i2]++;
		}
	}
}

void SimilMeasureHungarian::countNeurons()
{
	Part *P1;
	Joint *J1;
	int i, j, i2;

	for (i = 0; i < 2; i++)
	{
		for (j = 0; j < models[i]->getNeuroCount(); j++)
		{
			Neuro *N = models[i]->getNeuro(j);
			// count parts attached to neurons
			P1 = N->getPart();
			if (P1)
			{
				i2 = models[i]->findPart(P1);
				neurons[i][i2]++;
			}
			else
			// count unattached neurons
			{
				J1 = N->getJoint();
				if (J1)
					on_joint[i]++;
				else
					anywhere[i]++;
			}
		}
	}
}

void SimilMeasureHungarian::fillPartsDistances(double*& dist, int bigger, int smaller, bool geo)
{
	for (int i = 0; i < bigger; i++)
	{
		for (int j = 0; j < bigger; j++)
		{
			// assign penalty for unassignment for vertex from bigger model
			if (j >= smaller)
			{
				if (geo)
					dist[i*bigger + j] += m_adFactors[3] * coordinates[1 - m_iSmaller][i].length();
				else
					dist[i*bigger + j] = m_adFactors[1] * degrees[1 - m_iSmaller][i] + m_adFactors[2] * neurons[1 - m_iSmaller][i];
			}
			// compute distance between parts
			else
			{
				if (geo){
					dist[i*bigger + j] += m_adFactors[3] * coordinates[1 - m_iSmaller][i].distanceTo(coordinates[m_iSmaller][j]);
	}
				else
					dist[i*bigger + j] = m_adFactors[1] * abs(degrees[1 - m_iSmaller][i] - degrees[m_iSmaller][j])
					+ m_adFactors[2] * abs(neurons[1 - m_iSmaller][i] - neurons[m_iSmaller][j]);
			}
		}
	}
}

/** Returns number of factors involved in final distance computation.
		These factors include differences in numbers of parts, degrees,
		number of neurons.
		*/
int SimilMeasureHungarian::getNOFactors()
{
	return SimilMeasureHungarian::iNOFactors;
}

int SimilMeasureHungarian::setParams(std::vector<double> params)
{
	int i = 0;
	for (i = 0; i < SimilMeasureHungarian::iNOFactors; i++)
		m_adFactors[i] = params.at(i);
	fixedZaxis = params.at(i);
	return 0;
}