#include <iostream>
#include <fstream>
#include <cstdlib>
#include <cmath>
#include <cstring>
#include "mst.h"

using namespace std;

int createMST(char *filename, float alpha)
{
	int x,y,z,j,k,l;
	float f1,f2,f3;
	MST_Tree tree;

	/* Load Multivariate Data into n nodes with l attributes
	 * First line of file contains two integers, n l
	 *	n - number of datasets
	 *	l - number is attributes for each dataset
	 * The next n lines are the datasets, 
	 *  each with l floats describing the attributes */

	//Open data
	ifstream fptr(filename);
	if (!fptr.is_open())
	{
		cerr << "MST: Unable to open " << filename << endl;
		return 1;
	}
	fptr >> tree.numNodes >> tree.numAttributes;

	cout << "Nodes: " << tree.numNodes << " Attributes: " << tree.numAttributes << endl;
	//Initialize Tree Structure
	if(tree.numNodes > 0 && tree.numAttributes > 0)
	{
		tree.Nodes = (MST_Node *) malloc(sizeof(MST_Node) * tree.numNodes);
		if(tree.Nodes == NULL)
		{
			perror("malloc");
			exit(1);
		}
		for(x=0;x<tree.numNodes;x++)
		{
			tree.Nodes[x].Attributes = (float *) malloc(sizeof(float) * (tree.numAttributes));
			if(tree.Nodes[x].Attributes == NULL)
			{
				perror("malloc");
				exit(1);
			}
			tree.Nodes[x].numEdges = 0;
			tree.Nodes[x].Edges = NULL;
		}
	}
	else
	{
		cerr << "MST: Invalid multivariate data file" << endl;
		return 1;
	}

	cout << "Reading MST Data\n";
	//Read Data into Tree
	for(x=0;x<tree.numNodes;x++)
	{
		for(y=0;y<tree.numAttributes;y++)
		{
			if(!fptr.eof())
				fptr >> tree.Nodes[x].Attributes[y];
			else
			{
				cerr << "MST: Invalid multivariate data file" << endl;
				return 1;
			}
			/* Assign Fragment Sizes to each Node */
			tree.Nodes[x].FragmentSize = 1;
			tree.Nodes[x].FragmentNumber = x+1;
		}
	}

	fptr.close();

	cout << "Creating MST...\n";
	/* Connect the smallest fragment to the closest different fragment */

	f3 = -1.f; //All time shortest edge
	while(tree.Nodes[0].FragmentSize < tree.numNodes)
	{
		//Find Smallest Fragment
		x=0;y=tree.numNodes+1;
		for(z=0;z<tree.numNodes;z++)
		{
			if(tree.Nodes[z].FragmentSize < y)
			{
				x = tree.Nodes[z].FragmentNumber;
				y = tree.Nodes[z].FragmentSize;
			}
		}
		
		//Find Closest Different Fragment
		f1 = -1.f;
		y = 0;
		l = 0;
		for(z=0;z<tree.numNodes;z++)
		{
			if(tree.Nodes[z].FragmentNumber == x)
			{
				for(j=0;j<tree.numNodes;j++)
				{
					if(tree.Nodes[j].FragmentNumber != x)
					{
						f2 = 0.f;
						for(k=0;k<tree.numAttributes;k++)
							f2 = f2 + ((tree.Nodes[z].Attributes[k] - tree.Nodes[j].Attributes[k]) * 
									   (tree.Nodes[z].Attributes[k] - tree.Nodes[j].Attributes[k]));
						f2 = (float)sqrt(f2);
						if((f1 == -1.f) || (f2 < f1))
						{
							f1 = f2;
							y = j;
							l = z;
							if((f3 == -1.f) || (f1 < f3))
								f3 = f1;
						}
					}
				}
			}
		}

		//Join Disjoint Fragments - Edge is between nodes y and l
		tree.Nodes[y].numEdges++;
		tree.Nodes[l].numEdges++;
		tree.Nodes[y].Edges = (MST_Node **) realloc(tree.Nodes[y].Edges,
				sizeof(MST_Node *) * tree.Nodes[y].numEdges);
		tree.Nodes[l].Edges = (MST_Node **) realloc(tree.Nodes[l].Edges,
				sizeof(MST_Node *) * tree.Nodes[l].numEdges);
		if(tree.Nodes[y].Edges == NULL || tree.Nodes[l].Edges == NULL)
		{
			perror("malloc");
			exit(1);
		}
		tree.Nodes[y].Edges[tree.Nodes[y].numEdges-1] = &(tree.Nodes[l]);
		tree.Nodes[l].Edges[tree.Nodes[l].numEdges-1] = &(tree.Nodes[y]);
		x = tree.Nodes[y].FragmentSize + tree.Nodes[l].FragmentSize;
		y = tree.Nodes[y].FragmentNumber;
		j = tree.Nodes[l].FragmentNumber;
		for(z=0;z<tree.numNodes;z++)
		{
			if(tree.Nodes[z].FragmentNumber == y)
				tree.Nodes[z].FragmentSize = x;
			else if(tree.Nodes[z].FragmentNumber == j)
			{
				tree.Nodes[z].FragmentSize = x;
				tree.Nodes[z].FragmentNumber = y;
			}
		}
		cout << x << " ";
	}

	cout << endl;

	/* Find Shortest Edge Length -- Done f3 contains it */
	f3 = f3 * alpha;

	/* Add edges between all nodes without edges where the distance
	 *  between them is less than or equal to (alpha x shortest edge length) */
	for(x=0;x<tree.numNodes;x++)
	{
		for(y=0;y<tree.numNodes;y++)
		{
			j = 0;
			for(z=0;z<tree.Nodes[y].numEdges;z++)
			{
				if(tree.Nodes[y].Edges[z] == &(tree.Nodes[x]))
				{
					j = 1;
					break;
				}
			}
			if(j == 0)
			{
				f1 = 0.f;
				for(z=0;z<tree.numAttributes;z++)
					f1 = f1 + ((tree.Nodes[x].Attributes[z] - tree.Nodes[y].Attributes[z]) *
							   (tree.Nodes[x].Attributes[z] - tree.Nodes[y].Attributes[z]));
				f1 = (float)sqrt(f1);
				if(f1 <= f3)
				{
					tree.Nodes[x].numEdges++;
					tree.Nodes[y].numEdges++;
					tree.Nodes[x].Edges = (MST_Node **) realloc(tree.Nodes[x].Edges,
							sizeof(MST_Node *) * tree.Nodes[x].numEdges);
					tree.Nodes[y].Edges = (MST_Node **) realloc(tree.Nodes[y].Edges,
							sizeof(MST_Node *) * tree.Nodes[y].numEdges);
					if(tree.Nodes[x].Edges == NULL || tree.Nodes[y].Edges == NULL)
					{
						perror("malloc");
						exit(1);
					}
					tree.Nodes[x].Edges[tree.Nodes[x].numEdges-1] = &(tree.Nodes[y]);
					tree.Nodes[y].Edges[tree.Nodes[y].numEdges-1] = &(tree.Nodes[x]);
				}
			}
		}
	}

	/* Create .wel file */

	//Count Edges
	k=0;
	f2 = 0.f;
	for(x=1;x<tree.numNodes;x++)
	{
		for(y=0;y<x;y++)
		{
			for(z=0;z<tree.Nodes[x].numEdges;z++)
			{
				if(tree.Nodes[x].Edges[z] == &(tree.Nodes[y]))
				{
					//Get Max Weight
					f1 = 0.f;
					for(j=0;j<tree.numAttributes;j++)
						f1 = f1 + ((tree.Nodes[x].Attributes[j] - tree.Nodes[y].Attributes[j]) *
								   (tree.Nodes[x].Attributes[j] - tree.Nodes[y].Attributes[j]));
					f1 = (float)sqrt(f1);
					if(f1 > f2)
						f2 = f1;
					k++;
					break;
				}
			}
		}
	}

	ofstream welptr;
	char *welfilename;
	welfilename = (char *) malloc(sizeof(char) * strlen(filename));
	strcpy(welfilename,filename);
	strcpy(welfilename + strlen(filename) - 4,".wel");

	welptr.open(welfilename);
	if(!welptr.is_open())
	{
		cerr << "Unable to open wel file\n";
		return 1;
	}
	welptr << tree.numNodes << "\t" << k << endl;

	for(x=1;x<tree.numNodes;x++)
	{
		for(y=0;y<x;y++)
		{
			for(z=0;z<tree.Nodes[x].numEdges;z++)
			{
				if(tree.Nodes[x].Edges[z] == &(tree.Nodes[y]))
				{
					//Get Weight
					f1 = 0.f;
					for(j=0;j<tree.numAttributes;j++)
						f1 = f1 + ((tree.Nodes[x].Attributes[j] - tree.Nodes[y].Attributes[j]) *
								   (tree.Nodes[x].Attributes[j] - tree.Nodes[y].Attributes[j]));
					f1 = (float)sqrt(f1);
					f1 = f1 / f2;
					welptr << x << "\t" << y << "\t" << f1 << endl;
					break;
				}
			}
		}
	}
	welptr.close();
	free(welfilename);
	for(x = 0;x < tree.numNodes;x++)
	{
		free(tree.Nodes[x].Attributes);
		free(tree.Nodes[x].Edges);
	}
	free(tree.Nodes);
	return 0;
}
