/* 
 * ionMapFile.cpp - Ion to Composition mapping file reader/writer 
 * Copyright (C) 2018, D Haley
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program 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 General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */
#include "ionMapFile.h"
#include "common/stringFuncs.h"

#include <stack>
#include <vector>
#include <map>
#include <utility>
#include <fstream>

using std::string;
using std::vector;
using std::pair;
using std::ofstream;
using std::endl;
using std::map;

enum
{
	IONMAP_ERR_NOPARSER=1,
	IONMAP_ERR_NOFILE,
	IONMAP_ERR_BADFILE,
	IONMAP_ERR_BADPROP,
	IONMAP_ERR_INCONSISTENT,
};
IonMapTable::IonMapTable()
{
}
bool IonMapTable::operator==(const IonMapTable &oth) const
{
	if(parentIons.size() !=oth.parentIons.size())
		return false;
	
	if(targetIons.size() !=oth.targetIons.size())
		return false;


	for(auto ui=0;ui<parentIons.size();ui++)
	{
		if(parentIons[ui] !=oth.parentIons[ui])
			return false;
	
	}

	for(auto ui=0;ui<targetIons.size();ui++)
	{
		if(targetIons[ui] !=oth.targetIons[ui])
			return false;
	
	}

	return true;
}

unsigned int IonMapTable::read(const char *filename) 
{
	targetIons.clear();
	parentIons.clear();

	xmlDocPtr doc;
	xmlParserCtxtPtr context;

	context =xmlNewParserCtxt();

	if(!context)
		return IONMAP_ERR_NOPARSER;

	//Open the XML file again, but without DTD validation
	doc = xmlCtxtReadFile(context, filename, NULL, XML_PARSE_NONET|XML_PARSE_NOENT);

	if(!doc)
		return IONMAP_ERR_NOFILE;

	//release the context
	xmlFreeParserCtxt(context);

	//retrieve root node	
	xmlNodePtr nodePtr = xmlDocGetRootElement(doc);
	
	try
	{
		std::stack<xmlNodePtr>  nodeStack;
		if(!nodePtr || !nodePtr->xmlChildrenNode)
			throw 1;
		//push root tag	
		nodeStack.push(nodePtr);
		
		//This *should* be an ionmap file
		if(xmlStrcmp(nodePtr->name, (const xmlChar *)"ionmap"))
			throw 1;



		//Spin through the entry elements
		nodePtr=nodePtr->xmlChildrenNode;	
		while(nodePtr)
		{
			if(XMLHelpFwdToElem(nodePtr,"entry"))
				break;	

			string molName;
			if(XMLHelpGetProp(molName, nodePtr,"name"))
				throw 1;
			parentIons.push_back(molName);

			vector<pair<string, unsigned int> > ionList;

			//Spin through the ion list.
			nodeStack.push(nodePtr);
			nodePtr = nodePtr->xmlChildrenNode;
			while(nodePtr)
			{
				if(XMLHelpFwdToElem(nodePtr,"ion"))
					break;	
				std::string ionName;
				unsigned int count;
					
				if(XMLHelpGetProp(ionName, nodePtr,"name"))
					throw 1;

				if(XMLHelpGetProp(count, nodePtr,"count"))
					count=1;

				ionList.push_back(make_pair(ionName,count));
			}
			
			targetIons.push_back(ionList);
			
			nodePtr = nodeStack.top();
			nodeStack.pop();

			nodePtr = nodePtr->next;
		}

	}
	catch (int)
	{
		parentIons.clear();
		targetIons.clear();
		//Code threw an error, just say "bad parse" and be done with it
		xmlFreeDoc(doc);
		return IONMAP_ERR_BADFILE;
	}
	
	xmlFreeDoc(doc);

	if(!isSelfConsistent())
		return IONMAP_ERR_INCONSISTENT;

	return 0;
}

unsigned int IonMapTable::write(const char *filename)
{
	ASSERT(parentIons.size() == targetIons.size() );
	ofstream f(filename);
	if(!f)
		return 1;


	f << "<ionmap>" << endl;
	for( auto ui=0;ui<parentIons.size();ui++)
	{
		f << tabs(1) << "<entry name=\"" << parentIons[ui]<< "\">" << endl;
		for( auto &p:  targetIons[ui])
		{
			f  << tabs(2) << "<ion name=\"" << p.first << 
				"\" count=\"" << p.second << "\"/>" << endl;

		}
		f<<  tabs(1) << "</entry>" << endl;
	}
	f<< "</ionmap> " << endl;

	

	return 0;
}


bool IonMapTable::isSelfConsistent() const
{
	return parentIons.size() == targetIons.size();
}

unsigned int IonMapTable::remapCounts(const vector<string> &ionNames, const vector<unsigned int> &inCounts,
				vector<string> &mappedNames, vector<unsigned int> &mappedCounts, bool errorOnUnmappable) const
{
	//This is brute force, but should be OK as these will normally be small.
	// TODO : replace with better lookup structure (eg sort + search)

	map<string,unsigned int> countMap;
	vector<unsigned int> unMapped;

	for(auto ui=0;ui<ionNames.size();ui++)
	{

		unsigned int target;
		auto it=std::find(parentIons.begin(),parentIons.end(),ionNames[ui]);
		if(it == parentIons.end())
		{
			if(errorOnUnmappable)
				return 1;

			unMapped.push_back(ui);
			
			continue;
		}

		target =std::distance(parentIons.begin(),it);
	
		unsigned int sourceCounts;
		sourceCounts=inCounts[ui];

		//Loop through each target in the the target ions
		for(auto p : targetIons[target])
		{
			auto itMap =countMap.find(p.first);
			if(itMap == countMap.end())
				countMap[p.first ] =p.second*sourceCounts;
			else
				itMap->second+=p.second*sourceCounts;
		}

	}



	//Append the ions that could not be broken apart
	for(auto &i : unMapped)
	{
		mappedNames.push_back(ionNames[i]);
		mappedCounts.push_back(inCounts[i]);
	}

	//Now flatten the map to produce ountput
	for( auto &p : countMap)
	{
		//The input ions could contain, e.g. "Ti" , but this is not mapped
		// to any other units, so it does not appear in our mapping.
		// we need to merge these with any fragments in the mapping output, as needed
		auto it=std::find(mappedNames.begin(),mappedNames.end(),p.first);
		if(it == mappedNames.end())
		{
			mappedNames.push_back(p.first);
			mappedCounts.push_back(p.second);
		}
		else
		{
			unsigned int offset = std::distance(mappedNames.begin(),it);
			mappedCounts[offset]+=p.second;
		}

	}

	return 0;
}


unsigned int IonMapTable::remapCounts(const std::vector<std::string> &ionNames, 
	const std::vector<std::vector<unsigned int> > &inCounts, std::vector<std::string> &mappedNames, 
	std::vector<std::vector<unsigned int> > &mappedCount, bool errorOnUnmappable) const
{
	vector<vector<string> > mappedNamesVec;
	vector<vector<unsigned int> > mappedCountsVec;

	//Transpose the counts so we can treat the outer vector as independent
	// and the inner one as dependant
	vector<vector<unsigned int> > transposedCounts;
	transposedCounts = inCounts;
	vectorTranspose(inCounts,transposedCounts);

	mappedNamesVec.resize(transposedCounts.size());
	mappedCountsVec.resize(transposedCounts.size());
	//Run the remap using the single function implementation
	for(auto ui=0; ui<transposedCounts.size();ui++)
	{
		ASSERT(ionNames.size() == transposedCounts[ui].size());
		unsigned int errCode;
		errCode=remapCounts(ionNames,transposedCounts[ui],
				mappedNamesVec[ui],mappedCountsVec[ui],errorOnUnmappable);
		if(errCode)
			return errCode;
	}

	
	//Now merge the grouping by name, such that we end up with one grouping
	// per element

	//Flatten names of ions
	vector<string> totalNames;
	for(auto v : mappedNamesVec)
	{
		for(auto s : v)
		{
			if(std::find(totalNames.begin(),totalNames.end(),s) == totalNames.end())
				totalNames.push_back(s);
		}
	}

	vectorTranspose(mappedCountsVec,mappedCount);	
	mappedNames=totalNames;

	ASSERT(mappedCount.size() == mappedNames.size());

	return 0;
}



#ifdef DEBUG
#include "wx/wxcommon.h"

bool testIonMapFile()
{
	IonMapTable mapTable, mapTableB;

	TEST(!mapTable.read("../test/ionMapExample.xml"),"ion map read");

	string s;
	s=createTmpFilename();
	mapTable.write(s.c_str());
	mapTableB.read(s.c_str());

	ASSERT((mapTable == mapTableB));

	vector<string> ionNames,mappedNames;
	vector<vector<unsigned int> > inCounts,mappedCounts;

	//In our example case, TiO is decomped to TiO, but not NiO
	ionNames.push_back("Ti");
	ionNames.push_back("TiO");
	ionNames.push_back("NiO");

	inCounts.resize(3);
	for(auto i=0;i<inCounts.size();i++)
	{
		for(auto j=0;j<4;j++)
			inCounts[i].push_back(j);
	}


	TEST(!mapTable.remapCounts(ionNames,inCounts,mappedNames,mappedCounts,false),"Vector remap");

	TEST(mappedNames.size() == mappedCounts.size(),"Mapped names for each counts");
	TEST(mappedNames.size() ==3,"Name count"); //Ti, O, NiO
	
	TEST(std::find(mappedNames.begin(),mappedNames.end(),"O") != mappedNames.end(),"Check decomp");

	return true;
}

#endif
