package inference.rewrite;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

import lombok.Getter;
import models.terms.DependencyTerm;
import models.terms.EvaluatableTerm;
import models.terms.PrimedTerm;
import models.terms.Resource;

public class ResourceTree {
	
	@Getter
	private EvaluatableTerm root;
	private Map<Position, List<Position>> tree;
	private Map<Position, EvaluatableTerm> resourceMap;
	
	public ResourceTree(EvaluatableTerm term) {
		tree = new HashMap<>();
		resourceMap = new HashMap<>();
		constructResourceTree(term, new Position(List.of(0)));
		root = resourceMap.get(new Position(List.of(0)));
	}
	
	public ResourceTree(Map<Position, List<Position>> tree, Map<Position, EvaluatableTerm> resourceMap) {
		this.tree = tree;
		this.resourceMap = resourceMap;
		root = resourceMap.get(new Position(List.of(0)));
	}
	
	
	public EvaluatableTerm getResource(Position pos) {
		if (pos == null) {
			return null;
		}
		return resourceMap.get(pos);
	}
	
	public List<Position> getChildren(Position pos) {
		if (tree.get(pos) == null) {
			return List.of();
		}
		return tree.get(pos);
	}
	
	
	
	private List<Position> constructResourceTree(EvaluatableTerm term, Position top) {
		if (term instanceof Resource resource) {
			resourceMap.put(top, resource);
			if (! tree.containsKey(top)) {
				tree.put(top, new ArrayList<>());
			}
			return List.of(top);
		} else if (term instanceof DependencyTerm depTerm) {
			EvaluatableTerm dependingTerm = depTerm.getDependingTerm();
			List<EvaluatableTerm> dependedResources = depTerm.getDependedTerms(); 
			List<EvaluatableTerm> argumentTerms = depTerm.getArgumentTerms();
			List<Position> dependingTermPositions = constructResourceTree(dependingTerm, top);
			List<Position> resultPositions = new ArrayList<>();
			for (Position pos: dependingTermPositions) {
				tree.put(pos, new ArrayList<>());
			}
			for (int i = 0; i < dependedResources.size(); i++) {
				for (Position pos : dependingTermPositions) {
					Position nextPos = pos.addPath(i);
					tree.get(pos).add(nextPos);
					Position dependedResourcePosition= constructResourceTree(dependedResources.get(i), nextPos).get(0);
					tree.put(dependedResourcePosition, new ArrayList<>());
					Position argumentTermPosition = dependedResourcePosition.addPath(0);
					tree.get(dependedResourcePosition).add(argumentTermPosition);
					tree.put(argumentTermPosition, new ArrayList<>());
					resultPositions.addAll(constructResourceTree(argumentTerms.get(i), argumentTermPosition));
				}
			}
			return resultPositions;
		} else if (term instanceof PrimedTerm pt) { 
			if (pt.isResource()) {
				resourceMap.put(top, pt);
				if (! tree.containsKey(top)) {
					tree.put(top, new ArrayList<>());
				}
				return List.of(top);
			} else {
				DependencyTerm depTerm = (DependencyTerm) (pt.getPrimedTerm());
				EvaluatableTerm dependingTerm = depTerm.getDependingTerm();
				List<EvaluatableTerm> dependedResources = depTerm.getDependedTerms(); 
				List<EvaluatableTerm> argumentTerms = depTerm.getArgumentTerms();
				List<Position> dependingTermPositions = constructResourceTree(dependingTerm, top);
				List<Position> resultPositions = new ArrayList<>();
				for (Position pos: dependingTermPositions) {
					tree.put(pos, new ArrayList<>());
				}
				for (int i = 0; i < dependedResources.size(); i++) {
					for (Position pos : dependingTermPositions) {
						Position nextPos = pos.addPath(i);
						tree.get(pos).add(nextPos);
						Position dependedResourcePosition= constructResourceTree(dependedResources.get(i), nextPos).get(0);
						tree.put(dependedResourcePosition, new ArrayList<>());
						Position argumentTermPosition = dependedResourcePosition.addPath(0);
						tree.get(dependedResourcePosition).add(argumentTermPosition);
						tree.put(argumentTermPosition, new ArrayList<>());
						resultPositions.addAll(constructResourceTree(argumentTerms.get(i), argumentTermPosition));
					}
				}
				return resultPositions;
			}
		} else {
			return null;
		}
	}
	
	@Override
	public String toString() {
		List<String> result = new ArrayList<>();
		toStringAllPath(new Position(), new ArrayList<>(), result);
		return "<" + result.stream().collect(Collectors.joining("\n")) + ">";
	}
	
	public void debug(Position pos) {
		System.out.println(pos + ", " + resourceMap.get(pos));
		if (tree.containsKey(pos)) {
			for (Position nextPos : tree.get(pos)) {
				debug(nextPos);
			}
		}
	}
	
	public void debugAllPath() {
		debugAllPath(new Position(), new ArrayList<>());
	}
	
	private void debugAllPath(Position pos, List<EvaluatableTerm> curPath) {
		curPath.add(resourceMap.get(pos));
		if (tree.get(pos).size() == 0) {
			System.out.println(curPath.stream().map(EvaluatableTerm::toString).collect(Collectors.joining("-")));
		} else {
			for (Position nextPos: tree.get(pos)) {
				debugAllPath(nextPos, curPath);
			}
		}
		curPath.remove(curPath.size() - 1);
	}
	
	private void toStringAllPath(Position pos, List<EvaluatableTerm> curPath, List<String> result) {
		curPath.add(resourceMap.get(pos));
		if (tree.get(pos).size() == 0) {
			result.add(curPath.stream().map(EvaluatableTerm::toString).collect(Collectors.joining("-")));
		} else {
			for (Position nextPos: tree.get(pos)) {
				toStringAllPath(nextPos, curPath, result);
			}
		}
		curPath.remove(curPath.size() - 1);
	}
	
	
	@Override
	public boolean equals(Object another) {
		if (another instanceof ResourceTree tree) {
			return this.tree.equals(tree.tree) && this.resourceMap.equals(tree.resourceMap);
		}
		return false;
	}
	
	@Override
	public int hashCode() {
		return Objects.hash(this.tree, this.resourceMap);
	}
	
}