#include "core/error/error.h"
#include "core/math/math.h"
#include "core/strings/string.inl"
#include "resource/expression_language.inl"
#include <alloca.h>
#include <limits.h> // UINT_MAX
#include <stdlib.h> // strtof
#include <string.h> // memmove

namespace crown
{
namespace expression_language
{
	/// Byte code constants.
	///
	/// If the upper 12 bits of the byte code do not match one of these values, the operation is
	/// BC_PUSH_FLOAT and the byte code specify the 32 bit float to push. If the upper 12 bits
	/// match one of these values (which are all NaNs, so they should never appear as regular
	/// floats), the operation will instead be the one matching.
	///
	/// The remaining 20 bits of the byte code are used for the id of functions and variables.
	enum ByteCode { BC_FUNCTION = 0x7f800000, BC_PUSH_VAR = 0x7f900000, BC_END = 0x7fa00000 };

	/// Returns the byte code operation part of the byte code word.
	static inline unsigned bc_mask(unsigned i)
	{
		return i & 0xfff00000;
	}

	/// Returns the id part of the byte code word.
	static inline unsigned id_mask(unsigned i)
	{
		return i & 0x000fffff;
	}

	/// Opcodes for functions
	enum OpCode
	{
		OP_ADD,
		OP_SUB,
		OP_MUL,
		OP_DIV,
		OP_UNARY_MINUS,
		OP_NOP,
		OP_SIN,
		OP_COS,
		OP_ABS,
		OP_MATCH,
		OP_MATCH_2D
	};

	inline float fmax(float a, float b)
	{
		return a > b ? a : b;
	}

	inline float length(float a, float b)
	{
		return fsqrt((b - a) * (b - a));
	}

	inline float match(float a, float b)
	{
		return fmax(1.0f - length(a, b), 0.0f);
	}

	inline float match2d(float a, float b, float c, float d)
	{
		return match(a, b) * match(c, d);
	}

	/// Computes the function specified by @a op_code on the @a stack.
	static inline void default_compute_function(int op_code, Stack &stack)
	{
#define POP() pop(stack)
#define PUSH(f) push(stack, f)

		float a, b, c, d;

		switch (op_code) {
		case OP_ADD: b = POP(); a = POP(); PUSH(a + b); break;
		case OP_SUB: b = POP(); a = POP(); PUSH(a - b); break;
		case OP_MUL: b = POP(); a = POP(); PUSH(a * b); break;
		case OP_DIV: b = POP(); a = POP(); PUSH(a / b); break;
		case OP_UNARY_MINUS: PUSH(-POP()); break;
		case OP_SIN: PUSH(fsin(POP())); break;
		case OP_COS: PUSH(fcos(POP())); break;
		case OP_ABS: a = POP(); PUSH(fabs(a)); break;
		case OP_MATCH: b = POP(); a = POP(); PUSH(match(a, b)); break;
		case OP_MATCH_2D: d = POP(); c = POP(); b = POP(); a = POP(); PUSH(match2d(a, b, c, d)); break;
		case OP_NOP: break;
		default:
			CE_FATAL("Unknown opcode");
		}

#undef POP
#undef PUSH
	}

	/// Union to cast through to convert between float and unsigned.
	union FloatAndUnsigned
	{
		float f;
		unsigned u;
	};

	static inline float unsigned_to_float(unsigned u)
	{
		FloatAndUnsigned fu;
		fu.u = u;
		return fu.f;
	}

	bool run(const unsigned *byte_code, const float *variables, Stack &stack, ComputeFunction compute_function)
	{
		if (compute_function == NULL)
			compute_function = default_compute_function;

		const unsigned *p = byte_code;
		while (true) {
			unsigned bc = *p++;
			unsigned op = bc_mask(bc);
			unsigned id = id_mask(bc);
			switch (op) {
			case BC_PUSH_VAR:
				if (stack.size == stack.capacity) return false;
				stack.data[stack.size++] = variables[id];
				break;
			case BC_FUNCTION:
				compute_function((OpCode)id, stack);
				break;
			case BC_END:
				return true;
			default: // BC_PUSH_FLOAT
				if (stack.size == stack.capacity) return false;
				stack.data[stack.size++] = unsigned_to_float(bc);
				break;
			}
		}
	}

} // namespace expression_language

namespace expression_language
{
	static inline unsigned float_to_unsigned(float f)
	{
		FloatAndUnsigned fu;
		fu.f = f;
		return fu.u;
	}

#ifdef WIN32
	float strtof(const char *nptr, char **endptr)
	{
		return (float)strtod(nptr, endptr);
	}
#endif

	/// Represents a token in the expression language. The tokens are used
	/// both during the tokenization phase and as a representation of the
	/// program. (A list of tokens in reverse polish notation form.)
	struct Token
	{
		enum TokenType { EMPTY, NUMBER, FUNCTION, VARIABLE, LEFT_PARENTHESIS, RIGHT_PARENTHESIS };

		Token()
			: type(EMPTY)
			, id(0)
		{
		}

		explicit Token(TokenType type)
			: type(type)
		{
		}

		Token(TokenType type, unsigned id)
			: type(type)
			, id(id)
		{
		}

		Token(TokenType type, float value)
			: type(type)
			, value(value)
		{
		}

		TokenType type; ///< Identifies the type of the token.
		union
		{
			unsigned id; ///< Id for FUNCTION and VARIABLE tokens.
			float value; ///< Numeric value for NUMBER tokens.
		};
	};

	namespace compile_env
	{
		/// Finds a string in @a strings matching @a s of length @a len and returns its index.
		/// Returns UINT_MAX if no such string is found.
		static unsigned find_string(const char *s, unsigned len, unsigned num_strings, const char **strings)
		{
			for (unsigned i = 0; i < num_strings; ++i)
				if (strncmp(s, strings[i], len) == 0 && strlen32(strings[i]) == len)
					return i;
			return UINT_MAX;
		}

		/// Finds a token representing the identifier in the environment.
		Token token_for_identifier(const CompileEnvironment &env, const char *identifier, unsigned len, Token default_token = Token(Token::NUMBER, 0.0f))
		{
			unsigned i;
			if ((i = find_string(identifier, len, env.num_variables, env.variable_names)) != UINT_MAX) {
				return Token(Token::VARIABLE, i);
			} else if ((i = find_string(identifier, len, env.num_constants, env.constant_names)) != UINT_MAX) {
				return Token(Token::NUMBER, env.constant_values[i]);
			} else if ((i = find_string(identifier, len, env.num_functions, env.function_names)) != UINT_MAX) {
				return Token(Token::FUNCTION, i);
			} else {
				// CE_FATAL("Unknown identifier: %s", identifier);
				return default_token;
			}
		}

		/// Finds a token representing the identifier in the environment.
		Token token_for_identifier(const CompileEnvironment &env, const char *identifier)
		{
			return token_for_identifier(env, identifier, strlen32(identifier));
		}

		/// True if there is a function matching the specified identifier.
		bool has_function(const CompileEnvironment &env, char *identifier)
		{
			return find_string(identifier, strlen32(identifier), env.num_functions, env.function_names) != UINT_MAX;
		}

	} // namespace compile_env

	/// Tokenizes the source code @a p into a sequence of tokens. The environment @a env
	/// is used for looking up source code identifiers.
	/// Returns the total number of tokens. If the returned number is greater than the @a capacity, only
	/// the first @a capacity items will be converted.
	static unsigned tokenize(const char *p, const CompileEnvironment &env, Token *tokens, unsigned capacity)
	{
		// Determines if the next + or - is a binary or unary operator.
		bool binary = false;

		unsigned num_tokens = 0;
		unsigned overflow_tokens = 0;

		while (*p != 0) {
			Token token(Token::EMPTY);

			// Numbers
			if (*p >= '0' && *p <= '9') {
				char *out;
				token = Token(Token::NUMBER, strtof(p, &out));
				p = out;
				binary = true;
				// Identifiers
			} else if ((*p >= 'a' && *p <= 'z') || (*p >= 'A' && *p <= 'Z') || (*p == '_')) {
				const char *identifier = p;
				while ((*p >= 'a' && *p <= 'z') || (*p >= 'A' && *p <= 'Z') || (*p == '_') || (*p >= '0' && *p <= '9'))
					p++;
				token = compile_env::token_for_identifier(env, identifier, u32(p - identifier));
				binary = true;
				// Operators
			} else {
				switch (*p) {
				case '(': token = Token(Token::LEFT_PARENTHESIS); binary = false; break;
				case ')': token = Token(Token::RIGHT_PARENTHESIS); binary = true; break;
				case ' ': case '\t': case '\n': case '\r': break;
				case '-': token = compile_env::token_for_identifier(env, binary ? "-" : "u-"); binary = false; break;
				case '+': token = compile_env::token_for_identifier(env, binary ? "+" : "u+"); binary = false; break;

				default: {
					char s2[3] = {*p, *(p + 1), 0};

					if (s2[1] && compile_env::has_function(env, s2)) {
						token = compile_env::token_for_identifier(env, s2);
						++p;
					} else {
						char s1[2] = {*p, 0};
						token = compile_env::token_for_identifier(env, s1);
					}

					binary = false;
					break;
				}
				}
				++p;
			}

			if (token.type != Token::EMPTY) {
				if (num_tokens == capacity)
					++overflow_tokens;
				else
					tokens[num_tokens++] = token;
			}
		}

		return num_tokens + overflow_tokens;
	}

	/// Performs constant folding on the program represented by @a rpl which is a
	/// sequence of tokens in reverse polish notation. Any function found in the
	/// token stream which only takes constant arguments is replaced by the
	/// result of evaluating the function over the constant arguments.
	static void fold_constants(Token *rpl, unsigned &num_tokens, const CompileEnvironment &env)
	{
		static const int MAX_ARITY = 4;
		float stack_data[MAX_ARITY];

		for (unsigned i = 0; i < num_tokens; ++i) {
			if (rpl[i].type != Token::FUNCTION)
				continue;

			Stack stack(stack_data, MAX_ARITY);
			bool constant_arguments = true;
			Function f = env.function_values[rpl[i].id];
			unsigned arity = f.arity;
			CE_ASSERT(arity <= MAX_ARITY, "MAX_ARITY too small");
			CE_ASSERT(i >= arity, "Too few arguments to function");
			unsigned arg_start = i - arity;
			for (unsigned j = 0; j < arity && constant_arguments; ++j) {
				constant_arguments = constant_arguments && rpl[i - j - 1].type == Token::NUMBER;
				stack.data[j] = rpl[arg_start + j].value;
			}
			if (!constant_arguments)
				continue;

			stack.size = arity;
			env.compute_function(f.op_code, stack);
			unsigned results = stack.size;
			int to_remove = int(arity + 1) - int(results);
			if (to_remove > 0) {
				memmove(&rpl[arg_start], &rpl[arg_start + to_remove], sizeof(Token)*(num_tokens - arg_start - to_remove));
				num_tokens -= to_remove;
			}
			for (unsigned res = 0; res < stack.size; ++res)
				rpl[arg_start + res] = Token(Token::NUMBER, stack.data[res]);
			i = arg_start - 1;
		}
	}

	/// Generates bytecode from a program in RPL token stream form.
	/// Returns the number of byte_code tokens generated. If the returned number is > capacity, only the first
	/// capacity items are generated.
	static unsigned generate_bytecode(const Token *rpl, unsigned num_tokens, const CompileEnvironment &env, unsigned *byte_code, unsigned capacity)
	{
		unsigned size = 0;
		unsigned overflow = 0;

		for (unsigned i = 0; i < num_tokens; ++i) {
			Function f;
			Token t = rpl[i];
			unsigned op;
			switch (t.type) {
			case Token::NUMBER:
				op = float_to_unsigned(t.value);
				break;
			case Token::VARIABLE:
				op = BC_PUSH_VAR + t.id;
				break;
			case Token::FUNCTION:
				f = env.function_values[t.id];
				op = BC_FUNCTION + f.op_code;
				break;
			default:
				op = 0xdeadbeef;
				CE_FATAL("Unknown token");
				break;
			}
			if (size < capacity)
				byte_code[size++] = op;
			else
				++overflow;
		}

		unsigned op = BC_END;
		if (size < capacity)
			byte_code[size++] = op;
		else
			++overflow;

		return size + overflow;
	}

	/// Represents an item on the function call stack.
	/// This object is comparable so that functions that have higher precedence are executed
	/// before others.
	struct FunctionStackItem
	{
		FunctionStackItem()
		{
		}

		FunctionStackItem(Token t, int p, int pl)
			: token(t)
			, precedence(p)
			, par_level(pl)
		{
		}

		inline int cmp(const FunctionStackItem &f) const
		{
			if (par_level != f.par_level) return par_level - f.par_level;
			return precedence - f.precedence;
		}

		inline bool operator<(const FunctionStackItem &other) const
		{
			return cmp(other) < 0;
		}

		inline bool operator<=(const FunctionStackItem &other) const
		{
			return cmp(other) <= 0;
		}

		inline bool operator==(const FunctionStackItem &other) const
		{
			return cmp(other) == 0;
		}

		inline bool operator>=(const FunctionStackItem &other) const
		{
			return cmp(other) >= 0;
		}

		inline bool operator>(const FunctionStackItem &other) const
		{
			return cmp(other) > 0;
		}

		Token token;
		int precedence;
		int par_level;
	};

	const int NUM_DEFAULT_FUNCTIONS = 12;

	/// Sets up the functions that should be usable in the language.
	static unsigned setup_functions(const char **names, Function *functions, unsigned capacity)
	{
		CE_ASSERT(capacity >= NUM_DEFAULT_FUNCTIONS, "Not enough space for default functions");
		CE_UNUSED(capacity);
		names[0] = ","; functions[0] = Function(OP_NOP, 1, 0);
		names[1] = "+"; functions[1] = Function(OP_ADD, 12, 2);
		names[2] = "-"; functions[2] = Function(OP_SUB, 12, 2);
		names[3] = "*"; functions[3] = Function(OP_MUL, 13, 2);
		names[4] = "/"; functions[4] = Function(OP_DIV, 13, 2);
		names[5] = "u-"; functions[5] = Function(OP_UNARY_MINUS, 16, 1);
		names[6] = "u+"; functions[6] = Function(OP_NOP, 16, 0);
		names[7] = "sin"; functions[7] = Function(OP_SIN, 17, 1);
		names[8] = "cos"; functions[8] = Function(OP_COS, 17, 1);
		names[9] = "abs"; functions[9] = Function(OP_ABS, 17, 1);
		names[10] = "match"; functions[10] = Function(OP_MATCH, 17, 2);
		names[11] = "match_2d"; functions[11] = Function(OP_MATCH_2D, 17, 4);
		return NUM_DEFAULT_FUNCTIONS;
	}

	unsigned compile(const char *source
		, CompileEnvironment &env
		, unsigned *byte_code
		, unsigned capacity
		)
	{
		unsigned num_tokens = tokenize(source, env, NULL, 0);

		// Change alloca to some other temp memory allocator if you want to
		Token *tokens = (Token *)alloca(sizeof(Token) * num_tokens);
		tokenize(source, env, tokens, num_tokens);

		Token *rpl = (Token *)alloca(sizeof(Token) * num_tokens);
		unsigned num_rpl = 0;

		FunctionStackItem *function_stack = (FunctionStackItem *)alloca(sizeof(FunctionStackItem)*num_tokens);
		unsigned num_function_stack = 0;

		int par_level = 0;
		for (unsigned i = 0; i < num_tokens; ++i) {
			Token &token = tokens[i];
			switch (token.type) {
			case Token::NUMBER:
			case Token::VARIABLE:
				rpl[num_rpl++] = token;
				break;
			case Token::LEFT_PARENTHESIS:
				++par_level;
				break;
			case Token::RIGHT_PARENTHESIS:
				--par_level;
				break;
			case Token::FUNCTION: {
				FunctionStackItem f(token, env.function_values[token.id].precedence, par_level);
				while (num_function_stack > 0 && function_stack[num_function_stack - 1] >= f)
					rpl[num_rpl++] = function_stack[--num_function_stack].token;
				function_stack[num_function_stack++] = f;
				break;
			}
			default:
				CE_FATAL("Unknown token");
				break;
			}
		}

		while (num_function_stack > 0)
			rpl[num_rpl++] = function_stack[--num_function_stack].token;

		fold_constants(rpl, num_rpl, env);
		return generate_bytecode(rpl, num_rpl, env, byte_code, capacity);
	}

	unsigned compile(const char *source
		, unsigned num_variables
		, const char **variables
		, unsigned num_constants
		, const char **constant_names
		, const float *constant_values
		, unsigned *byte_code
		, unsigned capacity
		)
	{
		const char *function_names[NUM_DEFAULT_FUNCTIONS];
		Function functions[NUM_DEFAULT_FUNCTIONS];
		unsigned num_functions = setup_functions(function_names, functions, NUM_DEFAULT_FUNCTIONS);

		CompileEnvironment env;
		env.num_variables = num_variables;
		env.variable_names = variables;
		env.num_constants = num_constants;
		env.constant_names = constant_names;
		env.constant_values = constant_values;
		env.num_functions = num_functions;
		env.function_names = function_names;
		env.function_values = functions;
		env.compute_function = default_compute_function;

		return compile(source, env, byte_code, capacity);
	}

} // namespace expression_language

} // namespace crown
