/**********************************************************************************

 Infomap software package for multi-level network clustering

 Copyright (c) 2013, 2014 Daniel Edler, Martin Rosvall
 
 For more information, see <http://www.mapequation.org>
 

 This file is part of Infomap software package.

 Infomap software package is free software: you can redistribute it and/or modify
 it under the terms of the GNU Affero General Public License as published by
 the Free Software Foundation, either version 3 of the License, or
 (at your option) any later version.

 Infomap software package is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU Affero General Public License for more details.

 You should have received a copy of the GNU Affero General Public License
 along with Infomap software package.  If not, see <http://www.gnu.org/licenses/>.

**********************************************************************************/


#include "FlowCalculator.h"
#include <iostream>
#include <cmath>
#include "../utils/Log.h"
#include "../core/StateNetwork.h"

namespace infomap {

void FlowCalculator::calculateFlow(StateNetwork& network, const Config& config)
{
	Log() << "Calculating global network flow using flow model '" << config.flowModel << "'... " << std::flush;

	// Prepare data in sequence containers for fast access of individual elements
    // Map to zero-based dense indexing
	unsigned int numNodes = network.numNodes();

	m_nodeFlow.assign(numNodes, 0.0);
	m_nodeTeleportRates.assign(numNodes, 0.0);

	std::vector<unsigned int> nodeOutDegree(numNodes, 0);
	std::vector<double> sumLinkOutWeight(numNodes, 0.0);
    
	unsigned int nodeIndex = 0;
	for (auto& nodeIt : network.nodes()) {
		auto& networkNode = nodeIt.second;
		m_nodeIndexMap[networkNode.id] = nodeIndex;
		++nodeIndex;
	}

	unsigned int numLinks = network.numLinks();
	m_flowLinks.reserve(numLinks);
	double sumLinkWeight = network.sumLinkWeight();
	double sumUndirLinkWeight = 2 * sumLinkWeight - network.sumSelfLinkWeight();

  for (auto& linkIt : network.nodeLinkMap()) {
		unsigned int linkSourceId = linkIt.first.id;
		unsigned int sourceIndex = m_nodeIndexMap[linkSourceId];

		const auto& subLinks = linkIt.second;
		for (auto& subIt : subLinks)
		{
			unsigned int linkTargetId = subIt.first.id;
			unsigned int targetIndex = m_nodeIndexMap[linkTargetId];
			double linkWeight = subIt.second.weight;

			++nodeOutDegree[sourceIndex];
			sumLinkOutWeight[sourceIndex] += linkWeight;
			m_nodeFlow[sourceIndex] += linkWeight / sumUndirLinkWeight;
			m_flowLinks.push_back(Link(sourceIndex, targetIndex, linkWeight));

			// Log() << "\n" << linkSourceId << " (" << sourceIndex << ") -> " << linkTargetId << " (" << targetIndex << ")"
			// 	<< ", weight: " << linkWeight << ". ";

			if (sourceIndex != targetIndex) {
				if (config.isUndirectedFlow()) {
					++nodeOutDegree[targetIndex];
					sumLinkOutWeight[targetIndex] += linkWeight;
				}
				if (config.flowModel != FlowModel::outdirdir)
					m_nodeFlow[targetIndex] += linkWeight / sumUndirLinkWeight;
			}
		}
	}

	if (config.flowModel == FlowModel::rawdir)
	{
		// Treat the link weights as flow (after global normalization) and
		// do one power iteration to set the node flow
		m_nodeFlow.assign(numNodes, 0.0);
		for (auto& link : m_flowLinks)
		{
			link.flow /= sumLinkWeight;
			m_nodeFlow[link.target] += link.flow;
		}
		Log() << "\n  -> Using directed links with raw flow.";
		Log() << "\n  -> Total link weight: " << sumLinkWeight << ".";
		finalize(network, config, true);
		return;
	}

	if (config.flowModel != FlowModel::directed)
	{
		if (config.flowModel == FlowModel::outdirdir)
			Log() << "\n  -> Counting only ingoing links.";
		else
			Log() << "\n  -> Using undirected links" << (config.undirdir? ", switching to directed after steady state." :
					".");

		if (config.flowModel == FlowModel::undirdir || config.flowModel == FlowModel::outdirdir)
		{
			//Take one last power iteration
			std::vector<double> nodeFlowSteadyState(m_nodeFlow);
			m_nodeFlow.assign(numNodes, 0.0);
			for (auto& link : m_flowLinks)
			{
				m_nodeFlow[link.target] += nodeFlowSteadyState[link.source] * link.flow / sumLinkOutWeight[link.source];
			}
			double sumNodeFlow = 0.0;
			for (unsigned int i = 0; i < m_nodeFlow.size(); ++i)
				sumNodeFlow += m_nodeFlow[i];
			// Update link data to represent flow instead of weight
			for (auto& link : m_flowLinks)
			{
				link.flow *= nodeFlowSteadyState[link.source] / sumLinkOutWeight[link.source] / sumNodeFlow;
			}
			finalize(network, config, true);
		}
		else // undirected
		{
			for (unsigned int i = 0; i < numLinks; ++i)
				m_flowLinks[i].flow /= sumUndirLinkWeight / 2;
			finalize(network, config);
		}
		return;
	}

	Log() << "\n  -> Using " << (config.recordedTeleportation ? "recorded" : "unrecorded") << " teleportation to " <<
			(config.teleportToNodes ? "nodes" : "links") << ". " << std::flush;

	// Log() << "\nnode flow: ";
	// double sum1 = 0.0;
	// for (auto& f : m_nodeFlow) {
	// 	Log() << f << ", ";
	// 	sum1 += f;
	// }
	// Log() << "\n => sum node flow: " << sum1;

	// Log() << "\nTeleport rates: ";
	// Calculate the teleport rate distribution
	if (config.teleportToNodes)
	{
		double sumNodeWeights = 0.0;
		for (auto& nodeIt : network.nodes()) {
			auto& networkNode = nodeIt.second;
			m_nodeTeleportRates[m_nodeIndexMap[networkNode.id]] = networkNode.weight;
			sumNodeWeights += networkNode.weight;
		}
		for (auto& weight : m_nodeTeleportRates) {
			weight /= sumNodeWeights;
		}
	}
	else // Teleport to links
	{
		// Teleport proportionally to out-degree, or in-degree if recorded teleportation.
		for (auto& link : m_flowLinks)
		{
			unsigned int toNode = config.recordedTeleportation ? link.target : link.source;
			m_nodeTeleportRates[toNode] += link.flow / sumLinkWeight;
			// Log() << "\nrate[" << toNode << "] += " << link.flow << " / " << sumLinkWeight << " = " << link.flow / sumLinkWeight << " => " << m_nodeTeleportRates[toNode];
		}
	}

	// double sumTeleportRates = 0.0;
	// Log() << "\n=> ";
	// for (auto& r : m_nodeTeleportRates) {
	// 	sumTeleportRates += r;
	// 	Log() << r << ", ";
	// }
	// Log() << "\n=> Sum: " << sumTeleportRates;

	double sumLinkFlow = 0.0;
	// Log() << "\nLinks:";
	// Normalize link weights with respect to its source nodes total out-link weight;
	for (auto& link : m_flowLinks)
	{
		// Log() << "\n" << linkIt->source << " -> " << linkIt->target << ": " <<  linkIt->flow;
		// Log() << ", sumLinkOutWeight[" << linkIt->source << "]: " << sumLinkOutWeight[linkIt->source];
		if (sumLinkOutWeight[link.source] > 0)
			link.flow /= sumLinkOutWeight[link.source];
		sumLinkFlow += link.flow;
		// Log() << " => sumLinkFlow += " << linkIt->flow << " = " << sumLinkFlow;
	}

	// Collect dangling nodes
	std::vector<unsigned int> danglings;
	for (unsigned int i = 0; i < numNodes; ++i)
	{
		if (nodeOutDegree[i] == 0)
			danglings.push_back(i);
	}

	// Log() << "\nsumLinkFlow: " << sumLinkFlow;

	// Log() << "\nDanglings: ";
	// for (auto& d : danglings) {
	// 	Log() << d << ", ";
	// }

	// Calculate PageRank
	std::vector<double> nodeFlowTmp(numNodes, 0.0);
	unsigned int numIterations = 0;
	double alpha = config.teleportationProbability;
	double beta = 1.0 - alpha;
	double sqdiff = 1.0;
	double danglingRank = 0.0;
	do
	{
		// Calculate dangling rank
		danglingRank = 0.0;
		for (auto& danglingIndex : danglings)
		{
			danglingRank += m_nodeFlow[danglingIndex];
		}

		// Flow from teleportation
		for (unsigned int i = 0; i < numNodes; ++i)
		{
			nodeFlowTmp[i] = (alpha + beta*danglingRank) * m_nodeTeleportRates[i];
		}

		// Flow from links
		for (auto& link : m_flowLinks)
		{
			nodeFlowTmp[link.target] += beta * link.flow * m_nodeFlow[link.source];
		}

		// Update node flow from the power iteration above and check if converged
		double sum = 0.0;
		double sqdiff_old = sqdiff;
		sqdiff = 0.0;
		for (unsigned int i = 0; i < numNodes; ++i)
		{
			sum += nodeFlowTmp[i];
			sqdiff += std::abs(nodeFlowTmp[i] - m_nodeFlow[i]);
			m_nodeFlow[i] = nodeFlowTmp[i];
		}

		// Normalize if needed
		if (std::abs(sum - 1.0) > 1.0e-10)
		{
			Log() << "(Normalizing ranks after " <<	numIterations << " power iterations with error " << (sum-1.0) << ") ";
			for (unsigned int i = 0; i < numNodes; ++i)
			{
				m_nodeFlow[i] /= sum;
			}
			// break;
		}

		// Perturb the system if equilibrium
		if(sqdiff == sqdiff_old)
		{
			alpha += 1.0e-10;
			beta = 1.0 - alpha;
		}

		numIterations++;
	}  while((numIterations < 200) && (sqdiff > 1.0e-15 || numIterations < 50));

	double sumNodeRank = 1.0;

	if (!config.recordedTeleportation)
	{
		//Take one last power iteration excluding the teleportation (and normalize node flow to sum 1.0)
		sumNodeRank = 1.0 - danglingRank;
		m_nodeFlow.assign(numNodes, 0.0);
		for (auto& link : m_flowLinks)
		{
			m_nodeFlow[link.target] += link.flow * nodeFlowTmp[link.source] / sumNodeRank;
		}
		beta = 1.0;
	}


	// Update the links with their global flow from the PageRank values. (Note: beta is set to 1 if unrec)
	for (auto& link : m_flowLinks)
	{
		link.flow *= beta * nodeFlowTmp[link.source] / sumNodeRank;
	}

	Log() << "\n  -> PageRank calculation done in " << numIterations << " iterations." << std::endl;
	finalize(network, config);
}

void FlowCalculator::finalize(StateNetwork& network, const Config& config, bool normalizeNodeFlow)
{
	// TODO: Skip bipartite flow adjustment for directed / rawdir / .. ?
	if (network.isBipartite() && !config.skipAdjustBipartiteFlow)
	{
		Log() << "\n  -> Using bipartite links.";
		// Only links between ordinary nodes and feature nodes in bipartite network
		// Don't code feature nodes -> distribute all flow from those to ordinary nodes
		unsigned int bipartiteStartId = network.bipartiteStartId();

		for (auto& link : m_flowLinks)
		{
			unsigned int sourceIsFeature = link.source >= bipartiteStartId;

			if (sourceIsFeature) {
				m_nodeFlow[link.target] += link.flow;
				m_nodeFlow[link.source] = 0.0; // Doesn't matter if done multiple times on each node.
			}
			else { // if (config.parseAsUndirected()) {
				m_nodeFlow[link.source] += link.flow;
				m_nodeFlow[link.target] = 0.0; // Doesn't matter if done multiple times on each node.
			}
			//TODO: Should flow double before moving to nodes, does it cancel out in normalization?
			link.flow *= 2; // Markov time 2 on the full network will correspond to markov time 1 between the real nodes.
		}

		normalizeNodeFlow = true;
	}

	if (normalizeNodeFlow)
	{
		double sumNodeFlow = 0.0;
		for (unsigned int i = 0; i < m_nodeFlow.size(); ++i)
			sumNodeFlow += m_nodeFlow[i];
		for (unsigned int i = 0; i < m_nodeFlow.size(); ++i)
			m_nodeFlow[i] /= sumNodeFlow;
	}

	// Write back flow to network
	//TODO: Add enter/exit flow
	double sumNodeFlow = 0.0;
	unsigned int nodeIndex = 0;
	for (auto& nodeIt : network.m_nodes) {
		auto& networkNode = nodeIt.second;
		networkNode.flow = m_nodeFlow[nodeIndex];
		sumNodeFlow += networkNode.flow;
		++nodeIndex;
	}

	double sumLinkFlow = 0.0;
	unsigned int linkIndex = 0;
	for (auto& linkIt : network.m_nodeLinkMap) {
		for (auto& subIt : linkIt.second)
		{
			auto& linkData = subIt.second;
			linkData.flow = m_flowLinks[linkIndex].flow;
			sumLinkFlow += linkData.flow;
			++linkIndex;
		}
	}
	Log() << "\n  => Sum node flow: " << sumNodeFlow << ", sum link flow: " << sumLinkFlow << "\n";
}

}
