package dendrarium.trees.json;

import dendrarium.trees.Forest;
import dendrarium.trees.Node;
import dendrarium.trees.NonterminalNode;
import dendrarium.trees.TerminalNode;
import dendrarium.trees.NodeChildren;
import dendrarium.utils.Pair;
import java.util.Map;
import java.util.HashMap;
import java.util.List;
import java.util.Collection;
import java.util.ArrayList;
import java.util.Iterator;
import org.json.JSONObject;
import org.json.JSONArray;
import org.json.JSONException;

public class TreeOutputter {
	private static JSONObject startnodeToJSON(Forest forest)
			throws JSONException {
		JSONObject startnode = new JSONObject();
		startnode.put("label", forest.getStartnodeLabel());
		startnode.put("from", forest.getStartnodeFrom());
		startnode.put("to", forest.getStartnodeTo());
		if (!forest.isEmpty())
			startnode.put("nid", forest.getRoot().getId());
		return startnode;
	}

	private static JSONObject statsToJSON(Forest forest)
			throws JSONException {
		JSONObject stats = new JSONObject();
		stats.put("trees", forest.getTrees());
		stats.put("nodes", forest.getNodes());
		stats.put("inferences", forest.getInferences());
		stats.put("cputime", forest.getCputime());
		return stats;
	}

	private static int choiceTypeToJSON(NonterminalNode node) {
		NonterminalNode.ChoiceType type = node.getChoiceType();

		if (type == NonterminalNode.ChoiceType.NONE)
			return 1;
		if (type == NonterminalNode.ChoiceType.
				MANY_CHILDREN_ONE_SUGGESTED)
			return 2;
		if (type == NonterminalNode.ChoiceType.ONE_SUGGESTED_CHILD)
			return 3;
		if (type == NonterminalNode.ChoiceType.USER_SUGGESTED)
			return 4;
		if (type == NonterminalNode.ChoiceType.USER_NOT_SUGGESTED)
			return 5;

		return 0;
	}

	protected static class Segment {
		int from, to;

		public Segment(int from, int to) {
			this.from = from;
			this.to = to;
		}

		public boolean equals(Object arg0) {
			if (arg0 == null)
				return false;
			if (!Segment.class.isInstance(arg0))
				return false;

			Segment seg = (Segment) arg0;
			return from == seg.from && to == seg.to;
		}
	}

	protected static Pair<Collection, Collection> nodeOrth(Node node) {
		if (node instanceof TerminalNode) {
			Collection orth = new ArrayList();
			orth.add(((TerminalNode) node).getText());
			return new Pair<Collection, Collection>(orth,
					new ArrayList());
		}

		Collection orth = null;
		Collection allcentres = null;
		boolean orthready = false;

		for (NodeChildren ch : ((NonterminalNode) node).getChildren()) {
			List<Node> chlist = ch.getChildren();
			List<Integer> heads = ch.getHeads();
			int head = heads.size() == 1 ?
				heads.get(0).intValue() : -1;
			Collection centre = new ArrayList();

			for (Node chnode : chlist) {
				if (head --!= 0 && orthready)
					continue;
				Pair<Collection, Collection> sub =
					nodeOrth(chnode);
				if (orth == null)
					orth = sub.getF1();
				else if (!orthready)
					orth.addAll(sub.getF1());

				if (head != -1)
					continue;
				centre = sub.getF2();

				if (chlist.size() <= 1)
					continue;
				centre.add(new Segment(chnode.getFrom(),
							chnode.getTo()));
			}

			orthready = true;
			if (allcentres == null)
				allcentres = centre;
			else
				allcentres.retainAll(centre);
		}

		return new Pair<Collection, Collection>(orth, allcentres);
	}

	/* Ten sam wezel moze wystepowac w lesie w kilku miejscach i moze
	 * miec wtedy rozne zbiory atrybutow ale *powinien* miec zawsze
	 * jednakowe ich wartosci.  Dodajemy do Hasha wezlow nowy wezel
	 * jesli jeszcze go nie ma w Hashu, w przeciwnym wypadku dodajemy
	 * tylko nowe atrybuty.  Podobne do JSONObject.accumulate().
	 *
	 * TODO: wywalic wyjatek kiedy wartosci roznia sie (potrzebne
	 * "glebokie" porownanie).  */
	public static void unionToMap(JSONObject map,
			String key, JSONObject obj) throws JSONException {
		if (!map.has(key)) {
			map.put(key, obj);
			return;
		}

		map = map.getJSONObject(key);
		for (Iterator<String> i = obj.keys(); i.hasNext();) {
			key = i.next();
			map.put(key, obj.get(key));
		}
	}

	public static void nodesToJSON(JSONObject map, Node node,
			int disambNode, Map<Integer, Integer> links,
			boolean strip, boolean full, boolean leaf)
			throws JSONException {
		int id = node.getId();
		boolean doCentre = !full &&
			(node.isDrawLeaf() || node.isLeaf() || leaf);

		/* TODO: aktualnie musimy przejsc kazdy wezel calego widocznego
		 * lasu tyle razy ile razy moze byc widoczny, byc moze da
		 * sie to zoptymalizowac?  */
		if (map.has(new Integer(id).toString()) && full)
			return;

		JSONObject obj = new JSONObject();
		obj.put("from", node.getFrom());
		obj.put("to", node.getTo());
		obj.put("nid", id);
		obj.put("subtrees", node.getSubtrees());
		obj.put("leaf", (!full && (node.isDrawLeaf() ||
						node.isLeaf())) ||
				(full && node instanceof TerminalNode));

		strip |= !full && node.isDrawLeaf();

		List<String> order = new ArrayList();
		for (Pair<String, String> attr : node.getAttributes())
			order.add(attr.getF1());
		obj.put("attrs", node.getAttributesMap());
		obj.put("attrs_order", order);

		if (links.containsKey(new Integer(id)))
			obj.put("target", links.get(new Integer(id)));

		if (id == disambNode)
			obj.put("selected", true);

		JSONObject properties = new JSONObject();
		if (node instanceof TerminalNode) {
			properties.put("base",
					((TerminalNode) node).getBase());
			properties.put("orth",
					((TerminalNode) node).getOrth());
			properties.put("interp_id",
					((TerminalNode) node).getInterp_id());
			obj.put("terminal", properties);

			map.put(new Integer(id).toString(), obj);
			return;
		}

		int chtype = choiceTypeToJSON((NonterminalNode) node);
		if (chtype > 0)
			obj.put("chtype", chtype);

		properties.put("category",
				((NonterminalNode) node).getCategory());
		obj.put("nonterminal", properties);

		if (doCentre) {
			Pair<Collection, Collection> orth = nodeOrth(node);
			obj.put("orth", orth.getF1());

			JSONArray centre = new JSONArray();
			for (Segment seg : (List<Segment>) orth.getF2())
				centre.put(new int[] { seg.from, seg.to });
			obj.put("centre", centre);

			if ((node.isDrawLeaf() || node.isLeaf()) &&
					!obj.has("selected")) {
				map.put(new Integer(id).toString(), obj);
				return;
			}
		}

		Integer chnum = ((NonterminalNode) node).getChosenChildrenNo();
		List<JSONObject> children = new ArrayList();
		for (NodeChildren ch : ((NonterminalNode) node).getChildren()) {
			if (!obj.has("selected") && strip && (
						(chnum == null &&
						 ch.getIndex() != 0) ||
						(chnum != null &&
						 ch.getIndex() !=
						 chnum.intValue())))
				continue;

			List<JSONObject> child = new ArrayList();
			for (Node chnode : ch.getChildren()) {
				JSONObject chobj = new JSONObject();
				chobj.put("nid", chnode.getId());
				child.add(chobj);

				nodesToJSON(map, chnode, disambNode,
						links, strip, full,
						obj.has("selected"));
			}

			JSONObject choice = new JSONObject();
			choice.put("rule", ch.getRule());
			if (ch.getHeads().size() > 0)
				choice.put("heads", ch.getHeads());
			choice.put("child", child);
			children.add(choice);
		}
		obj.put("children", children);

		if ((obj.has("selected") || !strip) && chnum != null)
			obj.put("current", chnum.intValue());

		unionToMap(map, new Integer(id).toString(), obj);
	}

	public static JSONObject forestToJSON(Forest forest,
			Node disambNode, Map<Integer, Integer> links,
			boolean strip, boolean full) throws JSONException {
		JSONObject obj = new JSONObject();

		if (disambNode != null)
			obj.put("disamb_node_id", disambNode.getId());

		obj.put("grammar_no", forest.getGrammar_no());

		obj.put("startnode", startnodeToJSON(forest));

		obj.put("text", forest.getText());

		obj.put("stats", statsToJSON(forest));

		if (forest.isEmpty()) {
			obj.put("empty", true);
			obj.put("node", new JSONObject());
			return obj;
		}

		JSONObject node = new JSONObject();
		nodesToJSON(node, forest.getRoot(), disambNode == null ? -1 :
				disambNode.getId(), links, strip, full, false);
		obj.put("node", node);

		return obj;
	}
}
