#include <stdexcept>
#include <string>
#include "environment.h"
#include "error.h"

using namespace std;

#define DEBUG


void Environment::load(const Environment &other){
	for(const auto &p : other.defs){
		defs[p.first]=p.second;
	}
	for(const auto &p : other.hooks){
		hooks[p.first]=p.second;
	}
}

void Environment::define(const Name &name,const AST &ast){
	defs[name]=ast;
}

void Environment::define(const Name &name,const Hook &hook){
	hooks[name]=hook;
}

AST Environment::get(const Name &name){
	auto it=defs.find(name);
	if(it==defs.end()){
		auto it=hooks.find(name);
		if(it==hooks.end()){
			throw NameError(name);
		}
		Hook hook=it->second;
		return AST::makeNative([this,hook](const AST &ast){
			return hook(*this,ast);
		});
	}
	return it->second;
}


// ------------------------------------------------------------


static void indexReplace(AST &ast,const Name &name,Index index){
	if(ast.quoted){
		return;
	}

	switch(ast.type){
		case AST::Type::number:
		case AST::Type::string:
		case AST::Type::index:
		case AST::Type::native:
			break;

		case AST::Type::name:
			if(ast.nameval==name){
				ast.type=AST::Type::index;
				ast.indexval=index;
			}
			break;

		case AST::Type::tuple:
			for(AST &term : ast.terms){
				indexReplace(term,name,index);
			}
			break;

		case AST::Type::lambda:
			if(ast.lambdaval.arg!=name){
				indexReplace(*ast.lambdaval.body,name,index+1);
			}
			break;
	}
}


static void indexify(AST &ast){
	if(ast.quoted){
		return;
	}

	switch(ast.type){
		case AST::Type::number:
		case AST::Type::string:
		case AST::Type::name:
		case AST::Type::index:
		case AST::Type::native:
			break;

		case AST::Type::tuple:
			for(AST &term : ast.terms){
				indexify(term);
			}
			break;

		case AST::Type::lambda:
			indexReplace(*ast.lambdaval.body,ast.lambdaval.arg,1);
			ast.lambdaval.arg="";
			indexify(*ast.lambdaval.body);
			break;
	}
}

static void singlify(AST &ast){
	if(ast.quoted){
		return;
	}

	switch(ast.type){
		case AST::Type::number:
		case AST::Type::string:
		case AST::Type::name:
		case AST::Type::index:
		case AST::Type::native:
			break;

		case AST::Type::lambda:
			singlify(*ast.lambdaval.body);
			break;

		case AST::Type::tuple:
			if(ast.terms.size()==0){
				break;
			}
			if(ast.terms.size()==1){
				AST newast=ast.terms[0];
				ast=newast;
				singlify(ast);
				break;
			}
			while(ast.terms.size()>2){
				AST two=AST::makeTuple({ast.terms[0],ast.terms[1]});
				ast.terms.erase(ast.terms.begin()+1);
				ast.terms[0]=move(two);
			}
			if(ast.terms.size()==2){
				singlify(ast.terms[0]);
				singlify(ast.terms[1]);
			}
			break;
	}
}

bool Environment::resolve(AST &ast){
	unordered_set<Name> avoid;
	return resolve(ast,avoid);
}

bool Environment::resolve(AST &ast,unordered_set<Name> &avoid){
	if(ast.quoted){
		return false;
	}

	bool ret=false;

	switch(ast.type){
		case AST::Type::number:
		case AST::Type::string:
		case AST::Type::index:
		case AST::Type::native:
			break;

		case AST::Type::name:
			if(avoid.find(ast.nameval)==avoid.end()){
				try {
					avoid.insert(ast.nameval);
					ast=get(ast.nameval);
					resolve(ast,avoid);
					ret=true;
				} catch(NameError){
					// just leave an unknown name as it is
				}
			}
			break;

		case AST::Type::tuple:
			for(AST &term : ast.terms){
				ret=resolve(term,avoid)||ret;
			}
			break;

		case AST::Type::lambda:
			ret=resolve(*ast.lambdaval.body,avoid)||ret;
			break;
	}

	return ret;
}


AST Environment::run(const AST &astinput){
	AST ast(astinput);
	indexify(ast);
#ifdef DEBUG
	cerr<<"indexify gave "<<ast<<endl;
#endif
	reduce(ast);
	return ast;
}


static void recursiveFindLevel(AST &ast,Index index,vector<pair<AST*,Index>> &nodes){
	if(ast.quoted){
		return;
	}

	switch(ast.type){
		case AST::Type::number:
		case AST::Type::string:
		case AST::Type::name:
		case AST::Type::native:
			break;

		case AST::Type::index:
			if(ast.indexval==index){
				nodes.emplace_back(&ast,index);
			}
			break;

		case AST::Type::tuple:
			for(AST &term : ast.terms){
				recursiveFindLevel(term,index,nodes);
			}
			break;

		case AST::Type::lambda:
			recursiveFindLevel(*ast.lambdaval.body,index+1,nodes);
			break;
	}
}


static void increaseFree(AST &ast,Index amount,Index fromIndex){
	if(ast.quoted){
		return;
	}

	switch(ast.type){
		case AST::Type::number:
		case AST::Type::string:
		case AST::Type::name:
		case AST::Type::native:
			break;

		case AST::Type::index:
			if(ast.indexval>=fromIndex){
				ast.indexval+=amount;
			}
			break;

		case AST::Type::tuple:
			for(AST &term : ast.terms){
				increaseFree(term,amount,fromIndex);
			}
			break;

		case AST::Type::lambda:
			increaseFree(*ast.lambdaval.body,amount,fromIndex+1);
			break;
	}
}


static bool hasFree(AST &ast,Index fromIndex){
	if(ast.quoted){
		return false;
	}

	switch(ast.type){
		case AST::Type::number:
		case AST::Type::string:
		case AST::Type::native:
			return false;

		case AST::Type::name:
			return false;

		case AST::Type::index:
			return ast.indexval>=fromIndex;

		case AST::Type::tuple:
			for(AST &term : ast.terms){
				if(hasFree(term,fromIndex)){
					return true;
				}
			}
			return false;

		case AST::Type::lambda:
			return hasFree(*ast.lambdaval.body,fromIndex+1);

		default:
			throw logic_error("Unknown AST::Type?");
	}
}


#ifdef DEBUG
static string indent(i64 depth){
	string s(2*depth,' ');
	for(i64 i=0;i<depth;i++){
		s[2*i]='|';
	}
	return s;
}
#endif

bool Environment::betareduce(AST &ast,i64 depth){
	if(ast.quoted){
		return false;
	}

	if(ast.type!=AST::Type::tuple){
		return false;
	}
	if(ast.terms.size()!=2){
		return false;
	}
	if(ast.terms[0].type!=AST::Type::lambda){
		return false;
	}

#ifdef DEBUG
	cerr<<indent(depth)<<"Betareducing "<<ast<<endl;
#endif

	bool success;

	// cerr<<"=== β-REDUCE LAMBDA ==="<<endl;
	// cerr<<"ast = "<<ast<<endl;
	AST newterm=*ast.terms[0].lambdaval.body;
	ast.terms[0]=newterm;
	// cerr<<"ast = "<<ast<<endl;

	vector<pair<AST*,Index>> repl;
	recursiveFindLevel(ast.terms[0],1,repl);
	// cerr<<"Level 2:"; for(i64 i=0;i<(i64)repl.size();i++)cerr<<" {"<<*repl[i].first<<','<<repl[i].second<<'}'; cerr<<endl;

	increaseFree(ast.terms[0],-1,2);

	for(const pair<AST*,Index> &p : repl){
		*p.first=ast.terms[1];
		increaseFree(*p.first,p.second-1,1);
	}
	newterm=ast.terms[0];
	ast=newterm;
	success=true;

#ifdef DEBUG
	cerr<<indent(depth)<<"'=> "<<ast<<endl;
#endif
	return success;
}


static bool etareduce(AST &ast){
	if(ast.quoted){
		return false;
	}

	// do nothing yet
	return false;
}


bool Environment::reduce(AST &ast,i64 depth){
#ifdef DEBUG
	cerr<<indent(depth)<<"Reducing "<<ast<<endl;
#endif
	if(ast.quoted){
		return false;
	}

	bool ret=false;

	switch(ast.type){
		case AST::Type::number:
		case AST::Type::string:
		case AST::Type::index:
		case AST::Type::native:
			break;

		case AST::Type::name:
			ret=resolve(ast);
			break;

		case AST::Type::lambda:
			ret=reduce(*ast.lambdaval.body,depth+1);
			ret=etareduce(ast)||ret;
			break;

		case AST::Type::tuple:
			if(ast.terms.size()==0){
				break;
			}
			if(ast.terms.size()==1){
				AST newast=ast.terms[0];
				ast=newast;
				ret=true;
				break;
			}
			if(ast.terms.size()>2){
				singlify(ast);
			}
			ret=reduce(ast.terms[0],depth+1);
			if(ast.terms[0].type==AST::Type::lambda){
				bool brret=betareduce(ast,depth+1);
				ret=ret||brret;
				if(brret){
					ret=reduce(ast,depth+1)||ret;
				}
			} else if(ast.terms[0].type==AST::Type::native){
#ifdef DEBUG
				cerr<<indent(depth+1)<<"Preparation for native call, reducing argument"<<endl;
#endif
				ret=reduce(ast.terms[1],depth+1)||ret;
				if(!hasFree(ast.terms[1],1)){
					ast=ast.terms[0].nativeval(ast.terms[1]);
					reduce(ast,depth+1);
					ret=true;
				} else {
					cerr<<indent(depth+1)<<"Argument contained free indices, not calling"<<endl;
				}
			}
			break;
	}

#ifdef DEBUG
	cerr<<indent(depth)<<"'-> "<<ast<<" (ret="<<ret<<")"<<endl;
#endif
	return ret;
}