import core.atomic;

import ae.sys.data;
import ae.sys.datamm;
import ae.sys.file;
import ae.utils.array;
import ae.utils.math : TypeForBits;
import ae.utils.meta;
import ae.utils.parallelism;

import std.algorithm.comparison;
import std.algorithm.iteration;
import std.algorithm.searching;
import std.algorithm.sorting;
import std.array;
import std.digest.sha;
import std.format;
import std.math.algebraic;
import std.math.exponential;
import std.parallelism;
import std.range;
import std.stdio;
import std.sumtype;
import std.typecons;

import gamelogic : Move, expand;
import driver;
import game;
import graph;
import packing;

/// Consider choices which affect score by less than this as inconsequential
enum Score scoreThreshold = 0.25;

template makeTree(Game.State state)
{
	static if (state == Game.State.bet)
		alias Decision = Bet;
	else
	static if (state == Game.State.playerTurn)
		alias Decision = Move;
	else
		static assert(false);

	alias Outcomes = Score[enumLength!Decision];
	alias SignificantDecision = Nullable!Decision;
	enum SignificantDecision inconsequential = SignificantDecision.init;

	SignificantDecision getSignificantDecision(ref const Outcomes outcomes)
	{
		auto minScore = outcomes[].reduce!min;
		auto maxScore = outcomes[].reduce!max;
		if (maxScore - minScore < scoreThreshold)
			return inconsequential;
		foreach (Decision decision, outcome; outcomes)
			if (outcome == maxScore)
				return SignificantDecision(decision);
		assert(false);
	}

	static struct GraphDriver
	{
		PackedScore[] scores;
		this(PackedScore[] scores) { this.scores = scores; }

		enum interactive = false;
		void log(scope string delegate() message) /*nothrow @nogc*/ {}
		Score recurse(const ref Game game) nothrow @nogc
		{
			return scores[game.packGame()].unpackScore;
		}

		alias GetScore(T) = Score delegate(T) /*nothrow @nogc*/;

		// alias choice = graphChoice!GetScore;

		bool topLevel = true;
		bool decisionRecorded = false;
		SignificantDecision bestDecision;

		Score choice(ChoiceMode mode, T)(
			scope string delegate() description,
			scope Score delegate(T) /*nothrow @nogc*/ getScore,
			scope string delegate(T) getDescription,
			scope const(T)[] items...
		) //nothrow @nogc
		{
			alias impl = graphChoice!GetScore;

			if (topLevel)
			{
				assert(!decisionRecorded);
				static if (mode == ChoiceMode.best && (is(T == Decision)))
				{
					topLevel = false; scope(exit) topLevel = true;

					Outcomes outcomes;
					foreach (item; items)
						outcomes[item] = getScore(item);
					bestDecision = getSignificantDecision(outcomes);

					decisionRecorded = true;
					return Score.nan; // We got what we needed
				}
				else
					assert(false);
			}
			else
				return impl!(mode, T)(description, getScore, getDescription, items);
		}

		alias gameOver = graphGameOver;
	}

	SignificantDecision getBestDecision(ref const Game game, PackedScore[] scores)
	{
		auto driver = GraphDriver(scores);
		expand(game, &driver);
		assert(driver.decisionRecorded);
		return driver.bestDecision;
	}

	enum Variable : ubyte
	{
		// All states
		deckSize,
		deck1s,
		deck2s,
		deck3s,
		deck4s,
		deck5s,
		deck6s,
		deck7s,
		deck8s,
		deck9s,
		deck10s,
		deck11s,
		deckTotalValue,

		// Player turn only
		bet,
		handNumCards,
		handTotalValue,
	}

	enum variableCard = state == Game.State.bet ? Variable.bet : enumLength!Variable;

	alias VariableValue = ubyte;
	enum variableValueCard = VariableValue.max + 1;

	VariableValue getVariable(ref const Game game, Variable v)
	{
		final switch (v)
		{
			case Variable.deckSize:
				return cast(VariableValue)game.deck[].sum();
			case Variable.deck1s:
				return game.deck[1];
			case Variable.deck2s:
				return game.deck[2];
			case Variable.deck3s:
				return game.deck[3];
			case Variable.deck4s:
				return game.deck[4];
			case Variable.deck5s:
				return game.deck[5];
			case Variable.deck6s:
				return game.deck[6];
			case Variable.deck7s:
				return game.deck[7];
			case Variable.deck8s:
				return game.deck[8];
			case Variable.deck9s:
				return game.deck[9];
			case Variable.deck10s:
				return game.deck[10];
			case Variable.deck11s:
				return game.deck[11];
			case Variable.deckTotalValue:
				uint totalValue;
				foreach (value, count; game.deck)
					totalValue += value * count;
				if (totalValue > VariableValue.max)
					return VariableValue.max;
				return cast(VariableValue)totalValue;
			case Variable.bet:
				return cast(VariableValue)game.bet;
			case Variable.handNumCards:
				return game.hand.numCards;
			case Variable.handTotalValue:
				return game.hand.totalValue;
		}
	}

	enum Operation : ubyte
	{
		equals,      // yes = variable value == threshold
		greaterThan, // yes = variable value >= threshold
	}

	struct Question
	{
		Variable variable;
		VariableValue threshold;
		Operation operation;
	}

	bool query(ref const Question question, VariableValue value)
	{
		final switch (question.operation)
		{
			case Operation.greaterThan:
				return value >= question.threshold;
			case Operation.equals:
				return value == question.threshold;
		}
	}

	struct QuestionNode
	{
		Question question;
		Node*[2] leaves;
	}

	struct LeafNode
	{
		Decision decision;
	}

	alias Node = SumType!(QuestionNode, LeafNode);

	enum maxDepth = 64;
	alias Selection = TypeForBits!maxDepth;

	void makeTree()
	{
		auto scoresData = mapFile(scoreFileName(), MmMode.read).asDataOf!PackedScore;
		auto scores = scoresData.unsafeContents;

		stderr.writeln("Enumerating states...");

		auto packedGamesFileName = "data/states-%s-%s.u%d".format(state, scoreThreshold, PackedGame.sizeof * 8);
		packedGamesFileName.cached!((string target) {
			PackedGame[] packedGames;
			{
				stderr.writefln(" > Allocating...");
				auto gameViableData = TData!bool(packedGameCard);
				auto gameViable = gameViableData.unsafeContents;
				stderr.writefln(" > Filtering...");
				foreach (PackedGame packedGame; packedGameCard.iota.parallel)
				{
					auto game = packedGame.unpackGame();
					gameViable[packedGame] =
						game.state == state
						&&
						!getBestDecision(game, scores).isNull
					;
				}
				stderr.writefln(" > Counting...");
				packedGames.length = gameViable.sum();
				stderr.writefln(" > Collecting...");
				size_t i;
				foreach (PackedGame packedGame, bool isViable; gameViable)
					if (isViable)
						packedGames[i++] = packedGame;
				assert(i == packedGames.length);
			}
			packedGames.toFile(target);
		})();
		auto packedGamesData = packedGamesFileName.mapFile(MmMode.read).asDataOf!PackedGame;
		PackedGame[] packedGames = packedGamesData.unsafeContents;
		stderr.writefln(" > %d / %d states.", packedGames.length, packedGameCard);

		stderr.writeln("Getting best decisions...");

		auto bestDecisions = new Decision[packedGames.length];
		foreach (i, packedGame; packedGames.parallel)
		{
			auto game = packedGame.unpackGame();
			bestDecisions[i] = getBestDecision(game, scores).get();
		}

		auto selections = new Selection[packedGames.length];

		Node* search(size_t depth, Selection selection)
		{
			if (depth == maxDepth)
				assert(false, "Too deep!");
			stderr.writefln("Searching @ %d / %b...", depth, selection);

			Selection mask = (Selection(1) << depth) - 1;

			// Check if we've reached a leaf node
			{
				bool[enumLength!Decision] sawDecision;
				foreach (i, decision; bestDecisions)
					if ((selections[i] & mask) == selection)
					{
						if (!sawDecision[decision])
						{
							sawDecision[decision] = true;
							if (sawDecision[].sum > 1)
								break;
						}
					}
				assert(sawDecision[].sum > 0, "Zero-state division");
				if (sawDecision[].sum == 1)
				{
					foreach (Decision decision, saw; sawDecision)
						if (saw)
							return new Node(LeafNode(decision));
					assert(false);
				}
			}

			double bestScore = -double.infinity;
			Question bestQuestion;

			alias GameCount = PackedGame;

			foreach (Variable variable; Variable.init .. variableCard)
			{
				GameCount[enumLength!Decision][variableValueCard] values =
					packedGames.length.iota.parallelChunks((typeof(packedGames.length.iota) chunk) {
						PackedGame[enumLength!Decision][variableValueCard] chunkValues;
						foreach (i; chunk)
							if ((selections[i] & mask) == selection)
							{
								auto packedGame = packedGames[i];
								auto game = packedGame.unpackGame();
								auto value = getVariable(game, variable);
								chunkValues[value][bestDecisions[i]]++;
							}
						return chunkValues;
					})
					.reduce!((ref a, ref b) {
						auto v = a;
						foreach (i, ref row; v)
							row[] += b[i][];
						return v;
					});

				// How many states with each outcome we have in total
				GameCount[enumLength!Decision] totalValues = values                .reduce!((ref a, ref b) { auto v = a; v[] += b[]; return v; });
				GameCount totalStates = totalValues[].sum;

				foreach (Operation operation; Operation.init .. enumLength!Operation)
					foreach (VariableValue threshold; VariableValue(1) .. variableValueCard)
					{
						auto question = Question(variable, threshold, operation);

						// How many states with each outcome we have on each side
						GameCount[enumLength!Decision] leftValues  = variableValueCard.iota.filter!(value => query(question, cast(VariableValue)value) == false).map!(value => values[value]).reduce!((ref a, ref b) { auto v = a; v[] += b[]; return v; });
						GameCount[enumLength!Decision] rightValues = variableValueCard.iota.filter!(value => query(question, cast(VariableValue)value) == true ).map!(value => values[value]).reduce!((ref a, ref b) { auto v = a; v[] += b[]; return v; });

						// Total states on both sides
						GameCount leftStates  = leftValues [].sum;
						GameCount rightStates = rightValues[].sum;
						if (leftStates == 0 || rightStates == 0)
							continue;

						double score;

						{
							double entropy(ref const GameCount[enumLength!Decision] values)
							{
								GameCount total = values[].sum();
								double entropy = 0;
								foreach (GameCount value; values)
								{
									if (value > 0)
									{
										double p = double(value) / total;
										entropy -= p * log2(p);
									}
								}
								return entropy;
							}

							double entropyBeforeSplit = entropy(totalValues);
							double entropyAfterSplit = (
								entropy( leftValues) *  leftStates +
								entropy(rightValues) * rightStates
							) / totalStates;

							score = entropyBeforeSplit - entropyAfterSplit;
						}

						if (score > bestScore)
						{
							bestScore = score;
							bestQuestion = question;
							stderr.writefln(" > %s (%s): [%s : %s, %s : %s]",
								bestQuestion, bestScore,
								leftStates , leftValues [].filter!(n => n > 0).walkLength,
								rightStates, rightValues[].filter!(n => n > 0).walkLength,
							);
						}
					}
			}

			assert(bestScore > -double.infinity);

			// Apply selection
			foreach (i, packedGame; packedGames.parallel)
				if ((selections[i] & mask) == selection)
				{
					auto game = packedGame.unpackGame();
					auto value = getVariable(game, bestQuestion.variable);
					auto answer = query(bestQuestion, value);
					auto bit = Selection(answer) << depth;
					selections[i] = (selections[i] & mask) | bit;
				}

			return new Node(QuestionNode(bestQuestion, [
				search(depth + 1, selection),
				search(depth + 1, selection | (Selection(1) << depth)),
			]));
		}

		auto rootNode = search(0, 0);

		{
			auto f = File(format("%s-%s.dot", state, scoreThreshold), "wb");
			f.writeln("digraph {");
			f.writeln("\tnode [shape=box];");
			void dump(Node* node)
			{
				string text;
				ubyte[] bytes;
				Node*[string] children;

				(*node).match!(
					(ref QuestionNode n) {
						text = format("%s %s %d ?",
							n.question.variable,
							["=", "≥"][n.question.operation],
							n.question.threshold,
						);
						bytes = n.question.asBytes;
						children = [
							"No"  : n.leaves[0],
							"Yes" : n.leaves[1],
						];
					},
					(ref LeafNode n) {
						text = format("%s", n.decision);
						bytes = n.decision.asBytes;
					},
				);

				f.writefln("\tnode_%s [label=%(%s%)]", cast(void*)node, [text]);
				foreach (edgeLabel, child; children)
					f.writefln("\tnode_%s -> node_%s [label=%(%s%)];", cast(void*)node, cast(void*)child, [edgeLabel]);

				foreach (edgeLabel, child; children)
					dump(child);
			}
			dump(rootNode);
			f.writeln("}");
		}
	}
}

void main()
{
	loadTables();

	makeTree!(Game.State.bet)();
	makeTree!(Game.State.playerTurn)();
}