#include <iostream>
#include <fstream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
#include <string>
#include <map>
#include <set>
#include <queue>
#include <deque>
#include <stack>
#include "RerootTreeUtils.h"

using namespace std;




void split(string &content, vector<string> &elements) {
elements.clear();
string tmp;
for (int i=0;i<(int)content.size();++i) {
if (content[i]==',' || content[i]==':' || content[i]=='(' || content[i]==')') {
if (!tmp.empty())
elements.push_back(tmp);
char ch[2] ={content[i],0};
elements.push_back(string(ch));
tmp.clear();
continue;
} else if ((content[i]>='0' && content[i]<='9') ||
(content[i]>='A' && content[i]<='Z') ||
(content[i]>='a' && content[i]<='z') || content[i] =='.'){
const char ch[2] ={content[i],0};
tmp.append(ch);
}
}
}
struct Edge {
int a;
double weight;
Edge(int a, double weight) {
this->a =a;
this->weight =weight;
}
bool operator<(const Edge & edge) const {
if (a !=edge.a)
return a<edge.a;
return weight<edge.weight;
}
};


double stringToDouble(string & content) {
double ret =0;
int i =0;
for (;i<content.size() && content[i]!='.';++i) {
if (content[i] <'0' || content[i]>'9') {
printf("input tree string is not right\n");
exit(0);
}
ret =ret*10 +content[i] -'0';
}
double x =0;
if (content[i] =='.') {
for (int j=content.size()-1;j>i;--j) {
if (content[j] <'0' || content[j]>'9') {
printf("input tree string is not right\n");
exit(0);
}
x =x*0.1 +content[j] -'0';
}
}
x =x*0.1;
return ret+x;
}


int stringToInt(string & content) {
int ret =0;
for (int i=0;i<content.size();++i)
ret =ret*10 +content[i] -'0';
return ret;
}
void buildGraph(vector<string> & elements, map<int, map<int, double> > &graph, map<string, int> & leaf_to_label) {
graph.clear();
stack<char> s1;
stack<Edge> s2;
int a =-1;
int cc =0;
for (int i=0;i<elements.size();++i) {
if (elements[i].compare("(") ==0) {
s1.push('(');
} else if (elements[i].compare(",") ==0) {
s1.push(',');
} else if (elements[i].compare(":") ==0) {
s1.push(':');
} else if (elements[i].compare(")") ==0) {
if (s1.empty() || s1.top()!=',') {
printf("input tree string is not right\n");
exit(0);
}
s1.pop();
if (s1.empty() || s1.top()!='(') {
printf("input tree string is not right\n");
exit(0);
}
s1.pop();
a =cc;
     if ((int)s2.size()-2 <0) {
printf("input tree string is not right\n");
exit(0);
}
graph[a][s2.top().a] =s2.top().weight;
graph[s2.top().a][a] =s2.top().weight;
s2.pop();
graph[a][s2.top().a] =s2.top().weight;
graph[s2.top().a][a] =s2.top().weight;
s2.pop();
cc++;
} else {
if (s1.top()!=':') {
a =cc;
leaf_to_label[elements[i]] =cc;
cc++;
} else {
double xx =stringToDouble(elements[i]);
if (a ==-1) {
printf("input tree string is not right\n");
exit(0);
}
s1.pop();
s2.push(Edge(a, xx));
a =-1;
}
}
}
if (!s1.empty() || !s2.empty()) {
printf("input tree string is not right\n");
exit(0);
}
}




string convert(char *content, char *new_root) {
    string strRes;
if (content ==NULL || new_root ==NULL)
return strRes;
string tree_str(content);
vector<string> elements;
split(tree_str, elements);
map<int, map<int, double> > graph;
map<string, int> leaf_to_label;
buildGraph(elements, graph, leaf_to_label);


string new_root_str(new_root);
if (leaf_to_label.find(new_root_str)==leaf_to_label.end()) {
printf("No such root %s\n", new_root);
exit(0);
}
int nr =graph.size();
int xx =-1;
double yy =0;


//modify graph, add new root
int nl =leaf_to_label[new_root_str];
for (map<int, map<int, double> >::iterator iter =graph.begin();iter!=graph.end();++iter) {
if (iter->second.find(nl)!=iter->second.end()) {
yy =iter->second[nl];
xx =iter->first;
iter->second[nr] =yy/2;
iter->second.erase(nl);
break;
}
}
graph[nr][xx] =yy/2;
graph[nr][nl] =yy/2;
graph[nl].clear();
graph[nl][nr] =yy/2;
#if 0
printf("graph\n");
for (map<int, map<int, double> >::iterator iter1 =graph.begin();iter1!=graph.end();++iter1) {
printf("%d:", iter1->first);
for (map<int, double>::iterator iter2 =iter1->second.begin();iter2!=iter1->second.end();++iter2) {
printf("(%d,%lf) ", iter2->first, iter2->second);
}
printf("\n");
}
#endif
// bfs, get new weight
int n =graph.size();
vector<double> wei;
vector<bool> flag;
wei.reserve(n);
flag.reserve(n);
for (int i=0;i<n;i++) {
flag.push_back(false);
wei.push_back(0);
}
queue<int> qu;
qu.push(nr);
flag[nr] =true;
map<int, set<int> > tree;
map<int, int> parent;
while (!qu.empty()) {
int t =qu.front();
qu.pop();
if (graph.find(t)==graph.end())
continue;
for (map<int, double>::iterator iter =graph[t].begin();iter!=graph[t].end();++iter) {
if (flag[iter->first])
continue;
flag[iter->first] =true;
qu.push(iter->first);
wei[iter->first] =wei[t] +(iter->second);
tree[t].insert(iter->first);
parent[iter->first] =t;
}
}
#if 0
printf("tree\n");
for (map<int, set<int> >::iterator iter1 =tree.begin();iter1!=tree.end();++iter1) {
printf("%d:", iter1->first);
for (set<int>::iterator iter2 =iter1->second.begin();iter2!=iter1->second.end();++iter2) {
printf("%d ", *iter2);
}
printf("\n");
}
printf("parent\n");
for (map<int, int>::iterator iter =parent.begin();iter!=parent.end();++iter) {
printf("%d %d\n", iter->first, iter->second);
}


printf("weight\n");
for (int i=0;i<wei.size();++i) {
printf("%d %lf\n", i, wei[i]);
}
#endif
// eliminate old root
int old =n-2;
xx =*(tree[old].begin());
for (map<int, set<int> >::iterator iter =tree.begin();iter!=tree.end();++iter) {
if (iter->second.find(old)!=iter->second.end()) {
iter->second.erase(old);
iter->second.insert(xx);
parent[xx] =iter->first;
break;
}
}
tree.erase(old);
parent.erase(old);


#if 0
printf("tree\n");
for (map<int, set<int> >::iterator iter1 =tree.begin();iter1!=tree.end();++iter1) {
printf("%d:", iter1->first);
for (set<int>::iterator iter2 =iter1->second.begin();iter2!=iter1->second.end();++iter2) {
printf("%d ", *iter2);
}
printf("\n");
}
printf("parent\n");
for (map<int, int>::iterator iter =parent.begin();iter!=parent.end();++iter) {
printf("%d %d\n", iter->first, iter->second);
}


printf("weight\n");
for (int i=0;i<wei.size();++i) {
printf("%d %lf\n", i, wei[i]);
}
#endif
// print new tree
map<int, string> nts;


deque<pair<int, int> > de;
for (map<int, set<int> >::iterator iter =tree.begin();iter!=tree.end();++iter) {
if (iter->second.size() ==2) {
int a[3] ={0,0,0};
for (set<int>::iterator iter2 =iter->second.begin();iter2!=iter->second.end();++iter2) {
if (tree.find(*iter2) ==tree.end())
a[++a[0]] =*iter2;
}
if (a[0] ==2) {
de.push_back(pair<int, int>(a[1], a[2]));
}
}
}


for (map<string, int>::iterator iter =leaf_to_label.begin();iter!=leaf_to_label.end();++iter) {
char tmp[100];
double tt =0;
if (parent.find(iter->second)!=parent.end())
tt =wei[parent[iter->second]];
snprintf(tmp, 100,"%f", wei[iter->second] -tt);
if (iter->second !=nr)
nts[iter->second] =iter->first+':'+string(tmp);
else
nts[iter->second] =iter->first;
}
#if 0
printf("node to string\n");
for (map<int, string>::iterator iter =nts.begin();iter!=nts.end();++iter) {
printf("%d: %s\n", iter->first, iter->second.c_str());
}
#endif
while (!de.empty()){
pair<int, int> a =de.front();
de.pop_front();
int pa =parent[a.first];
char tmp[100];
double tt =0;
if (parent.find(pa)!=parent.end())
tt =wei[parent[pa]];
snprintf(tmp,100,"%f", wei[pa] -tt);
if (pa!=nr)
nts[pa] ='('+nts[a.first]+','+nts[a.second]+')'+':'+string(tmp);
else
nts[pa] ='('+nts[a.first]+','+nts[a.second]+')';
tree.erase(pa);
if (parent.find(pa)!=parent.end()) {
int ppa =parent[pa];
int sibling;
for (set<int>::iterator iter =tree[ppa].begin();iter!=tree[ppa].end();++iter) {
if ((*iter)!=pa)
sibling =*iter;
}
if (tree.find(sibling) ==tree.end()) {
de.push_back(pair<int, int>(pa, sibling));
}
}
}
//printf("%s\n", nts[nr].c_str());
    strRes = nts[nr];
    return strRes;
}


void Test_split() {
string a("(((1:1.0,2:2.0):1.2,(3:1.0,4:2.0):1.6):1.5,5:1.0)");
string b("( ( ( 1 : 1.0 , 2 : 2.0 ) : 1.2 , ( 3 : 1.0 , 4 : 2.0 ) : 1.6 ) : 1.5 , 5 : 1.0 ) ");
vector<string> elements;
split(a, elements);
for (int i=0;i<elements.size();++i) {
printf("%s ", elements[i].c_str());
}
printf("\n");
split(b, elements);
for (int i=0;i<elements.size();++i) {
printf("%s ", elements[i].c_str());
}
printf("\n");
}


void Test_buildGraph() {
string a("(((1:1.0,2:2.0):1.2,(3:1.0,4:2.0):1.6):1.5,5:1.0)");
vector<string> elements;
split(a, elements);
map<int, map<int, double> > graph;
map<string, int> leaf_to_label;
buildGraph(elements, graph, leaf_to_label);
printf("leaf to label\n");
for (map<string, int>::iterator iter =leaf_to_label.begin();iter!=leaf_to_label.end();++iter) {
printf("%s:%d\n", iter->first.c_str(), iter->second);
}
printf("Graph\n");
for (map<int, map<int, double> >::iterator iter1 =graph.begin();iter1!=graph.end();++iter1) {
printf("%d:", iter1->first);
for (map<int, double>::iterator iter2 =iter1->second.begin();iter2!=iter1->second.end();++iter2) {
printf("(%d,%lf) ", iter2->first, iter2->second);
}
printf("\n");
}
}


string ReRootTreeNewick( char *nwFile, char *taxaNewRoot ) {
//char * a="(((1:1.0,2:2.0):1.2,(3:1.0,4:2.0):1.6):1.5,5:1.0)";
//char *b ="3";
// usage for converting
return convert(nwFile, taxaNewRoot);
}
