package jnpf.elnFormula;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public final class FormulaLatexRenderer {

    private final FormulaFunctionRegistry functions = new FormulaFunctionRegistry();

    public String render(FormulaAst.Node node, Map<String, String> parameterDisplayNames) {
        if (node == null || parameterDisplayNames == null) {
            throw new IllegalArgumentException("node and parameterDisplayNames must not be null");
        }
        return renderNode(node, parameterDisplayNames, 0, false);
    }

    private String renderNode(FormulaAst.Node node, Map<String, String> names,
                              int parentPrecedence, boolean rightChild) {
        int precedence = precedence(node);
        String rendered;
        if (node instanceof FormulaAst.Literal) {
            rendered = literal((FormulaAst.Literal) node);
        } else if (node instanceof FormulaAst.Parameter) {
            FormulaAst.Parameter parameter = (FormulaAst.Parameter) node;
            String displayName = names.get(parameter.getCode());
            rendered = "\\text{" + escapeText(displayName == null ? parameter.getCode() : displayName) + "}";
        } else if (node instanceof FormulaAst.Unary) {
            FormulaAst.Unary unary = (FormulaAst.Unary) node;
            String operator = unaryOperator(unary.getOperator());
            rendered = operator + renderNode(unary.getOperand(), names, precedence, true);
        } else if (node instanceof FormulaAst.Binary) {
            rendered = binary((FormulaAst.Binary) node, names);
        } else if (node instanceof FormulaAst.FunctionCall) {
            rendered = function((FormulaAst.FunctionCall) node, names);
        } else {
            throw new IllegalArgumentException("unsupported formula node");
        }

        boolean parentheses = precedence < parentPrecedence
                || (rightChild && precedence == parentPrecedence
                && node instanceof FormulaAst.Binary);
        return parentheses ? "\\left(" + rendered + "\\right)" : rendered;
    }

    private String binary(FormulaAst.Binary binary, Map<String, String> names) {
        String operator = binary.getOperator();
        int precedence = precedence(binary);
        if ("/".equals(operator)) {
            return "\\frac{" + renderNode(binary.getLeft(), names, 0, false) + "}{"
                    + renderNode(binary.getRight(), names, 0, false) + "}";
        }
        if ("^".equals(operator)) {
            String left = renderNode(binary.getLeft(), names, precedence, false);
            if (binary.getLeft() instanceof FormulaAst.Binary
                    && "^".equals(((FormulaAst.Binary) binary.getLeft()).getOperator())) {
                left = "\\left(" + left + "\\right)";
            }
            return left + "^{" + renderNode(binary.getRight(), names, 0, false) + "}";
        }
        String left = renderNode(binary.getLeft(), names, precedence, false);
        if (precedence == 3 && binary.getLeft() instanceof FormulaAst.Binary
                && precedence(binary.getLeft()) == precedence) {
            left = "\\left(" + left + "\\right)";
        }
        return left + latexOperator(operator)
                + renderNode(binary.getRight(), names, precedence, true);
    }

    private String function(FormulaAst.FunctionCall call, Map<String, String> names) {
        if (!functions.contains(call.getName())) {
            throw new IllegalArgumentException("unsupported formula function");
        }
        List<String> arguments = new ArrayList<>();
        for (FormulaAst.Node argument : call.getArguments()) {
            arguments.add(renderNode(argument, names, 0, false));
        }
        String name = call.getName();
        if ("ABS".equals(name)) {
            return "\\left|" + arguments.get(0) + "\\right|";
        }
        if ("SQRT".equals(name)) {
            return "\\sqrt{" + arguments.get(0) + "}";
        }
        if ("POW".equals(name)) {
            return powerBase(call.getArguments().get(0), arguments.get(0))
                    + "^{" + arguments.get(1) + "}";
        }
        if ("NTH_ROOT".equals(name)) {
            return "\\sqrt[" + arguments.get(1) + "]{" + arguments.get(0) + "}";
        }
        if ("LN".equals(name)) {
            return "\\ln\\left(" + arguments.get(0) + "\\right)";
        }
        if ("LOG10".equals(name)) {
            return "\\log_{10}\\left(" + arguments.get(0) + "\\right)";
        }
        if ("LOG".equals(name)) {
            return "\\log_{" + arguments.get(1) + "}\\left("
                    + arguments.get(0) + "\\right)";
        }
        if ("EXP".equals(name)) {
            return "e^{" + arguments.get(0) + "}";
        }
        if ("MIN".equals(name) || "MAX".equals(name)) {
            return "\\" + name.toLowerCase() + "\\left(" + join(arguments) + "\\right)";
        }
        if ("SUM".equals(name)) {
            return "\\sum\\left(" + arguments.get(0) + "\\right)";
        }
        if ("MEAN".equals(name)) {
            return "\\overline{" + arguments.get(0) + "}";
        }
        if ("STDDEV_S".equals(name)) {
            return "s\\left(" + arguments.get(0) + "\\right)";
        }
        if ("STDDEV_P".equals(name)) {
            return "\\sigma\\left(" + arguments.get(0) + "\\right)";
        }
        if ("VAR_S".equals(name)) {
            return "s^{2}\\left(" + arguments.get(0) + "\\right)";
        }
        if ("VAR_P".equals(name)) {
            return "\\sigma^{2}\\left(" + arguments.get(0) + "\\right)";
        }
        return "\\operatorname{" + escapeFunctionName(name) + "}\\left("
                + join(arguments) + "\\right)";
    }

    private static String powerBase(FormulaAst.Node node, String rendered) {
        return precedence(node) < 9 ? "\\left(" + rendered + "\\right)" : rendered;
    }

    private static String escapeFunctionName(String name) {
        return name.replace("_", "\\_");
    }

    private static String literal(FormulaAst.Literal literal) {
        switch (literal.getKind()) {
            case DECIMAL:
                return ((BigDecimal) literal.getValue()).toPlainString();
            case TEXT:
                return "\\text{\"" + escapeText((String) literal.getValue()) + "\"}";
            case BOOLEAN:
                return "\\mathrm{" + literal.getValue().toString() + "}";
            case NULL:
                return "\\mathrm{null}";
            default:
                throw new IllegalArgumentException("unsupported literal kind");
        }
    }

    private static String latexOperator(String operator) {
        if ("+".equals(operator) || "-".equals(operator)
                || ">".equals(operator) || "<".equals(operator)) {
            return " " + operator + " ";
        }
        if ("*".equals(operator)) {
            return " \\cdot ";
        }
        if ("%".equals(operator)) {
            return " \\bmod ";
        }
        if (">=".equals(operator)) {
            return " \\ge ";
        }
        if ("<=".equals(operator)) {
            return " \\le ";
        }
        if ("!=".equals(operator)) {
            return " \\ne ";
        }
        if ("==".equals(operator)) {
            return " = ";
        }
        if ("&&".equals(operator)) {
            return " \\land ";
        }
        if ("||".equals(operator)) {
            return " \\lor ";
        }
        throw new IllegalArgumentException("unsupported binary operator");
    }

    private static String unaryOperator(String operator) {
        if ("+".equals(operator) || "-".equals(operator)) {
            return operator;
        }
        if ("!".equals(operator)) {
            return "\\lnot ";
        }
        throw new IllegalArgumentException("unsupported unary operator");
    }

    private static int precedence(FormulaAst.Node node) {
        if (node instanceof FormulaAst.Binary) {
            String operator = ((FormulaAst.Binary) node).getOperator();
            if ("||".equals(operator)) {
                return 1;
            }
            if ("&&".equals(operator)) {
                return 2;
            }
            if ("==".equals(operator) || "!=".equals(operator)) {
                return 3;
            }
            if (">".equals(operator) || ">=".equals(operator)
                    || "<".equals(operator) || "<=".equals(operator)) {
                return 4;
            }
            if ("+".equals(operator) || "-".equals(operator)) {
                return 5;
            }
            if ("*".equals(operator) || "/".equals(operator) || "%".equals(operator)) {
                return 6;
            }
            if ("^".equals(operator)) {
                return 8;
            }
            throw new IllegalArgumentException("unsupported binary operator");
        }
        if (node instanceof FormulaAst.Unary) {
            return 7;
        }
        return 9;
    }

    private static String escapeText(String value) {
        StringBuilder escaped = new StringBuilder(value.length() + 16);
        for (int i = 0; i < value.length(); i++) {
            char current = value.charAt(i);
            switch (current) {
                case '\\':
                    escaped.append("\\textbackslash{}");
                    break;
                case '{':
                    escaped.append("\\{");
                    break;
                case '}':
                    escaped.append("\\}");
                    break;
                case '_':
                case '%':
                case '#':
                case '&':
                case '$':
                    escaped.append('\\').append(current);
                    break;
                case '<':
                    escaped.append("\\textless{}");
                    break;
                case '>':
                    escaped.append("\\textgreater{}");
                    break;
                case '^':
                    escaped.append("\\textasciicircum{}");
                    break;
                case '~':
                    escaped.append("\\textasciitilde{}");
                    break;
                case '\n':
                case '\r':
                case '\t':
                    escaped.append(' ');
                    break;
                default:
                    escaped.append(current);
            }
        }
        return escaped.toString();
    }

    private static String join(List<String> values) {
        StringBuilder joined = new StringBuilder();
        for (int i = 0; i < values.size(); i++) {
            if (i > 0) {
                joined.append(", ");
            }
            joined.append(values.get(i));
        }
        return joined.toString();
    }
}