#include <iostream>
#include <sstream>
#include <string>
#include <unordered_map>
#include <memory>
#include <cassert>
#include <cmath>
#include <variant>
#include <cctype>

/* converted form of the rayquake sdf parsing code I wrote with Caz, which was in Zig */

struct Expr;

struct BinArgs { std::unique_ptr<Expr> lhs, rhs; };
struct SqrtArg { std::unique_ptr<Expr> arg; };
struct Constant { float val; };
struct Input    { unsigned idx; };

using ExprVariant = std::variant<
    Constant, Input,
    struct Add, struct Sub, struct Mul, struct Div,
    struct Min, struct Max, struct Sqrt
>;

struct Add  { std::unique_ptr<Expr> lhs, rhs; };
struct Sub  { std::unique_ptr<Expr> lhs, rhs; };
struct Mul  { std::unique_ptr<Expr> lhs, rhs; };
struct Div  { std::unique_ptr<Expr> lhs, rhs; };
struct Min  { std::unique_ptr<Expr> lhs, rhs; };
struct Max  { std::unique_ptr<Expr> lhs, rhs; };
struct Sqrt { std::unique_ptr<Expr> arg; };

struct Expr {
    std::variant<Constant, Input, Add, Sub, Mul, Div, Min, Max, Sqrt> v;

    template<typename T>
    Expr(T &&t) : v(std::forward<T>(t)) {}

    Expr(Expr &&) = default;
    Expr &operator=(Expr &&) = default;
    Expr(const Expr &) = delete;
    Expr &operator=(const Expr &) = delete;
};

static std::unique_ptr<Expr> box(Expr e) {
    return std::make_unique<Expr>(std::move(e));
}


Expr make_constant(float f) { return Expr{ Constant{ f } }; }
Expr make_input(unsigned i)  { return Expr{ Input{ i } }; }

Expr expr_add(Expr l, Expr r) { return Expr{ Add{ box(std::move(l)), box(std::move(r)) } }; }
Expr expr_sub(Expr l, Expr r) { return Expr{ Sub{ box(std::move(l)), box(std::move(r)) } }; }
Expr expr_mul(Expr l, Expr r) { return Expr{ Mul{ box(std::move(l)), box(std::move(r)) } }; }
Expr expr_div(Expr l, Expr r) { return Expr{ Div{ box(std::move(l)), box(std::move(r)) } }; }
Expr expr_min(Expr l, Expr r) { return Expr{ Min{ box(std::move(l)), box(std::move(r)) } }; }
Expr expr_max(Expr l, Expr r) { return Expr{ Max{ box(std::move(l)), box(std::move(r)) } }; }
Expr expr_sqrt(Expr a)        { return Expr{ Sqrt{ box(std::move(a)) } }; }


/* deep copy */
Expr clone(const Expr &e);

template<typename T>
std::unique_ptr<Expr> clone_box(const std::unique_ptr<T> &p) {
    return box(clone(*p));
}

Expr clone(const Expr &e) {
    return std::visit([](const auto &node) -> Expr {
        using T = std::decay_t<decltype(node)>;
        /* filter 0-ary */
        if constexpr (std::is_same_v<T, Constant> || std::is_same_v<T, Input>) {
            return Expr{ node };
        } else if constexpr (std::is_same_v<T, Sqrt>) {
            /* only unary */
            return expr_sqrt(clone(*node.arg));
        } else {
            /* all binary nodes have lhs and rhs */
            using R = T;
            return Expr{ R{ box(clone(*node.lhs)), box(clone(*node.rhs)) } };
        }
    }, e.v);
}

float expr_eval(const Expr &e, const float *vals) {
    return std::visit([&](const auto &node) -> float {
        using T = std::decay_t<decltype(node)>;
        if constexpr (std::is_same_v<T, Constant>) return node.val;
        if constexpr (std::is_same_v<T, Input>)    return vals[node.idx];
        if constexpr (std::is_same_v<T, Add>)  return expr_eval(*node.lhs, vals) + expr_eval(*node.rhs, vals);
        if constexpr (std::is_same_v<T, Sub>)  return expr_eval(*node.lhs, vals) - expr_eval(*node.rhs, vals);
        if constexpr (std::is_same_v<T, Mul>)  return expr_eval(*node.lhs, vals) * expr_eval(*node.rhs, vals);
        if constexpr (std::is_same_v<T, Div>)  return expr_eval(*node.lhs, vals) / expr_eval(*node.rhs, vals);
        if constexpr (std::is_same_v<T, Min>)  return std::min(expr_eval(*node.lhs, vals), expr_eval(*node.rhs, vals));
        if constexpr (std::is_same_v<T, Max>)  return std::max(expr_eval(*node.lhs, vals), expr_eval(*node.rhs, vals));
        if constexpr (std::is_same_v<T, Sqrt>) return std::sqrt(expr_eval(*node.arg, vals));
        return 0.f;
    }, e.v);
}

template<class... Ts>
struct overloads : Ts... { using Ts::operator()...; };

// symbolic differentiation
Expr expr_diff(const Expr &e, unsigned i) {
    return std::visit([&](const auto &node) -> Expr {
        using T = std::decay_t<decltype(node)>;

        /* Rules for diff:
         *  - Variable/Input: same? 1 : 0
         *  - Constant: 0
         *  - Addition/Subtraction: a' +/- b'
         *  - Mul: f'g + g'f
         *  - Div: (fg' - gf') / g^2
         *  - Sqrt: (1/2)f' / sqrt(f)
         *  - Min/Max: Identity
         */

        if constexpr (std::is_same_v<T, Input>)
            return make_constant(node.idx == i ? 1.f : 0.f);

        if constexpr (std::is_same_v<T, Constant>)
            return make_constant(0.f);

        if constexpr (std::is_same_v<T, Add>)
            return expr_add(expr_diff(*node.lhs, i), expr_diff(*node.rhs, i));

        if constexpr (std::is_same_v<T, Sub>)
            return expr_sub(expr_diff(*node.lhs, i), expr_diff(*node.rhs, i));

        if constexpr (std::is_same_v<T, Mul>)
            return expr_add(
                expr_mul(expr_diff(*node.lhs, i), clone(*node.rhs)),
                expr_mul(expr_diff(*node.rhs, i), clone(*node.lhs))
            );

        if constexpr (std::is_same_v<T, Div>)
            return expr_div(
                expr_add(
                    expr_mul(clone(*node.lhs), expr_diff(*node.rhs, i)),
                    expr_mul(expr_mul(clone(*node.rhs), expr_diff(*node.lhs, i)),
                             make_constant(-1.f))
                ),
                expr_mul(clone(*node.rhs), clone(*node.rhs))
            );

        if constexpr (std::is_same_v<T, Sqrt>)
            return expr_mul(
                expr_diff(*node.arg, i),
                expr_mul(make_constant(0.5f),
                         expr_div(make_constant(1.f), expr_sqrt(clone(*node.arg))))
            );

        if constexpr (std::is_same_v<T, Min> || std::is_same_v<T, Max>)
            return clone(e);

        return make_constant(0.f);
    }, e.v);
}


struct Dual { float val, d; };


/* Autodiff */

Dual expr_auto_diff(const Expr &e, const Dual *pos) {
    return std::visit([&](const auto &node) -> Dual {
        using T = std::decay_t<decltype(node)>;

        if constexpr (std::is_same_v<T, Input>)    return pos[node.idx];
        if constexpr (std::is_same_v<T, Constant>) return { node.val, 0.f };

        if constexpr (std::is_same_v<T, Add>) {
            auto l = expr_auto_diff(*node.lhs, pos), r = expr_auto_diff(*node.rhs, pos);
            return { l.val + r.val, l.d + r.d };
        }
        if constexpr (std::is_same_v<T, Sub>) {
            auto l = expr_auto_diff(*node.lhs, pos), r = expr_auto_diff(*node.rhs, pos);
            return { l.val - r.val, l.d - r.d };
        }
        if constexpr (std::is_same_v<T, Mul>) {
            auto l = expr_auto_diff(*node.lhs, pos), r = expr_auto_diff(*node.rhs, pos);
            return { l.val * r.val, l.val * r.d + l.d * r.val };
        }
        if constexpr (std::is_same_v<T, Div>) {
            auto l = expr_auto_diff(*node.lhs, pos), r = expr_auto_diff(*node.rhs, pos);
            return { l.val / r.val, (l.d * r.val - l.val * r.d) / (r.val * r.val) };
        }
        if constexpr (std::is_same_v<T, Sqrt>) {
            auto v = expr_auto_diff(*node.arg, pos);
            float sv = std::sqrt(v.val);
            return { sv, v.d / (2.f * sv) };
        }
        if constexpr (std::is_same_v<T, Min>) {
            auto l = expr_auto_diff(*node.lhs, pos), r = expr_auto_diff(*node.rhs, pos);
            return l.val < r.val ? l : r;
        }
        if constexpr (std::is_same_v<T, Max>) {
            auto l = expr_auto_diff(*node.lhs, pos), r = expr_auto_diff(*node.rhs, pos);
            return l.val > r.val ? l : r;
        }
        return { 0.f, 0.f };
    }, e.v);
}


void expr_ndiff(const Expr &e, const float *vals, float eps, float out[3]) {
    for (int k = 0; k < 3; k++) {
        float vp[3] = { vals[0], vals[1], vals[2] };
        float vm[3] = { vals[0], vals[1], vals[2] };
        vp[k] += eps; vm[k] -= eps;
        out[k] = (expr_eval(e, vp) - expr_eval(e, vm)) / (2.f * eps);
    }
}

void expr_print(const Expr &e, std::ostream &os = std::cout) {
    const auto visitor = overloads {
        [&](const Input& node) { os << "x" << node.idx; },
        [&](const Constant& node) { os << node.val; },
        [&](const Sqrt& node)  { os << "(sqrt "; expr_print(*node.arg, os); os << ")"; },
        [&](const auto& node) {
            using T = std::decay_t<decltype(node)>;
            if constexpr (std::is_same_v<T, Input>)    { os << "x" << node.idx; return; }
            if constexpr (std::is_same_v<T, Constant>) { os << node.val; return; }
            if constexpr (std::is_same_v<T, Sqrt>) {
                os << "(sqrt "; expr_print(*node.arg, os); os << ")"; return;
            }

            const char *op =
                std::is_same_v<T, Add> ? "+" :
                std::is_same_v<T, Sub> ? "-" :
                std::is_same_v<T, Mul> ? "*" :
                std::is_same_v<T, Div> ? "/" :
                std::is_same_v<T, Min> ? "min" : "max";

            os << "(" << op << " ";
            expr_print(*node.lhs, os);
            os << " ";
            expr_print(*node.rhs, os);
            os << ")";
        }
    };

    std::visit(visitor, e.v);
}

static size_t gen_glsl_step(
    const Expr &e,
    std::unordered_map<const Expr *, size_t> &visited,
    std::ostringstream &ss)
{
    auto it = visited.find(&e);
    if (it != visited.end()) return it->second;

    size_t idx = visited.size();
    visited[&e] = idx;

    std::ostringstream tmp;

    const auto visitor = overloads {
        [&](const Constant& node) -> std::string  {
            tmp << "vec2(" << node.val << ", 0.0)";
            return tmp.str();
        },
        [&](const Input& node)-> std::string  {
            switch (node.idx) {
                case 0: return "vec2(v.x, vp.x)";
                case 1: return "vec2(v.y, vp.y)";
                case 2: return "vec2(v.z, vp.z)";
                default: return "vec2(0.0, 0.0)";
            }
        },
        [&](const Sqrt& node) -> std::string  {
            size_t a = gen_glsl_step(*node.arg, visited, ss);
            tmp << "nsqrt(_" << a << ")";
            return tmp.str();
        },
        [&](const auto& node) -> std::string {
            /* binops */
            size_t l = gen_glsl_step(*node.lhs, visited, ss);
            size_t r = gen_glsl_step(*node.rhs, visited, ss);

            using T = std::decay_t<decltype(node)>;
        
            const char *fn =
                std::is_same_v<T, Add> ? nullptr :
                std::is_same_v<T, Sub> ? nullptr :
                std::is_same_v<T, Mul> ? "nmul" :
                std::is_same_v<T, Div> ? "ndiv" :
                std::is_same_v<T, Min> ? "nmin" : "nmax";

            if constexpr (std::is_same_v<T, Add>)
                tmp << "(_" << l << " + _" << r << ")";
            else if constexpr (std::is_same_v<T, Sub>)
                tmp << "(_" << l << " - _" << r << ")";
            else
                tmp << fn << "(_" << l << ", _" << r << ")";

            return tmp.str();
        }
    };
    
    std::string rhs = std::visit(visitor, e.v);

    ss << "vec2 _" << idx << " = " << rhs << ";\n";
    return idx;
}

std::string gen_glsl(const Expr &e) {
    std::unordered_map<const Expr *, size_t> visited;
    std::ostringstream ss;
    size_t last = gen_glsl_step(e, visited, ss);
    ss << "return _" << last << ";";
    return ss.str();
}


static void skip_ws(const std::string &s, size_t &i) {
    while (i < s.size() && std::isspace((unsigned char)s[i])) i++;
}

static bool is_op_char(char c) {
    return std::isalpha((unsigned char)c) || std::string("+-/*").find(c) != std::string::npos;
}

static std::string parse_op_str(const std::string &s, size_t &i) {
    size_t start = i;
    while (i < s.size() && is_op_char(s[i])) i++;
    return s.substr(start, i - start);
}

static Expr parse_arg(const std::string &s, size_t &i,
                      std::unordered_map<std::string, unsigned> &vm);

static Expr parse_iden(const std::string &s, size_t &i,
                       std::unordered_map<std::string, unsigned> &vm) {
    size_t start = i;
    while (i < s.size() && std::isalnum((unsigned char)s[i])) i++;
    std::string name = s.substr(start, i - start);
    auto [it, inserted] = vm.emplace(name, (unsigned)vm.size());
    return make_input(it->second);
}

static Expr parse_num(const std::string &s, size_t &i) {
    bool neg = false;
    if (i < s.size() && s[i] == '-') { neg = true; ++i; }
    float num = 0.f;
    int frac_digits = 0;
    size_t start = i;
    bool saw_digit = false;


    while (i < s.size() && std::isdigit((unsigned char)s[i])) {
        saw_digit = true;
        num = num * 10.f + (s[i] - '0');
        ++i;
    }

    if (i < s.size() && s[i] == '.') {
        ++i;
        while (i < s.size() && std::isdigit((unsigned char)s[i])) {
            saw_digit = true;
            num = num * 10.f + (s[i] - '0');
            ++i;
            ++frac_digits;
        }
    }

    if (!saw_digit) {
        return make_constant(0.0f);
    }

    if (frac_digits > 0) {
        num /= std::pow(10.f, (float)frac_digits);
    }

    if (neg) num = -num;
    return make_constant(num);
}



static Expr parse_expr(const std::string &s, size_t &i,
                       std::unordered_map<std::string, unsigned> &vm) {
    skip_ws(s, i);
    assert(s[i] == '('); i++;
    skip_ws(s, i);

    std::string op = parse_op_str(s, i);

    Expr e = [&]() -> Expr {
        if (op == "sqrt") return expr_sqrt(parse_arg(s, i, vm));
        Expr lhs = parse_arg(s, i, vm);
        Expr rhs = parse_arg(s, i, vm);
        if (op == "+")   return expr_add(std::move(lhs), std::move(rhs));
        if (op == "-")   return expr_sub(std::move(lhs), std::move(rhs));
        if (op == "*")   return expr_mul(std::move(lhs), std::move(rhs));
        if (op == "/")   return expr_div(std::move(lhs), std::move(rhs));
        if (op == "min") return expr_min(std::move(lhs), std::move(rhs));
        if (op == "max") return expr_max(std::move(lhs), std::move(rhs));
        std::cerr << "Unknown op: " << op << "\n"; std::exit(1);
    }();

    skip_ws(s, i);
    assert(s[i] == ')'); i++;
    return e;
}

static Expr parse_arg(const std::string &s, size_t &i,
                      std::unordered_map<std::string, unsigned> &vm) {
    skip_ws(s, i);
    char c = s[i];
    if (c == '(')                           return parse_expr(s, i, vm);
    if (std::isalpha((unsigned char)c))     return parse_iden(s, i, vm);
    if (std::isdigit((unsigned char)c) || c == '-' || c == '.')
                                            return parse_num(s, i);
    std::cerr << "Unexpected char '" << c << "'\n"; std::exit(1);
}

Expr expr_parse(const std::string &src) {
    std::unordered_map<std::string, unsigned> vm;
    size_t i = 0;
    skip_ws(src, i);
    return parse_expr(src, i, vm);
}


Expr sdf_length(Expr pos[3]) {
    return expr_sqrt(
        expr_add(
            expr_mul(clone(pos[0]), clone(pos[0])),
            expr_add(
                expr_mul(clone(pos[1]), clone(pos[1])),
                expr_mul(clone(pos[2]), clone(pos[2]))
            )
        )
    );
}

Expr sdf_sphere(Expr pos[3], float radius) {
    return expr_sub(sdf_length(pos), make_constant(radius));
}

Expr sdf_box(Expr p[3], float b[3]) {
    Expr q[3] = {
        expr_sub(expr_sqrt(expr_mul(clone(p[0]), clone(p[0]))), make_constant(b[0])),
        expr_sub(expr_sqrt(expr_mul(clone(p[1]), clone(p[1]))), make_constant(b[1])),
        expr_sub(expr_sqrt(expr_mul(clone(p[2]), clone(p[2]))), make_constant(b[2])),
    };
    Expr inner[3] = {
        expr_max(clone(q[0]), make_constant(0.f)),
        expr_max(clone(q[1]), make_constant(0.f)),
        expr_max(clone(q[2]), make_constant(0.f)),
    };
    return expr_add(
        sdf_length(inner),
        expr_min(
            expr_max(clone(q[0]), expr_max(clone(q[1]), clone(q[2]))),
            make_constant(0.f)
        )
    );
}

/*
int main() {
    Expr x = make_input(0), y = make_input(1), z = make_input(2);
    Expr x2 = expr_mul(clone(x), clone(x));
    Expr y2 = expr_mul(clone(y), clone(y));
    Expr z2 = expr_mul(clone(z), clone(z));
    Expr l  = expr_sqrt(expr_add(expr_add(std::move(x2), std::move(y2)), std::move(z2)));
    Expr dl = expr_diff(l, 0);

    expr_print(l);  std::cout << "\n";
    expr_print(dl); std::cout << "\n";

    float v[3] = { 0.5f, 0.3f, 0.2f };
    std::cout << expr_eval(l, v) << "\n";

    std::string src =
        "(min"
        "  (- y 1)"
        "  (max"
        "      (max x (+ -9 (* -1 x)))"
        "      (* -1 (- (sqrt (+ (+ (* x x) (* y y)) (* z z))) 15))"
        "   ))";

    Expr parsed = expr_parse(src);
    expr_print(parsed);
    std::cout << "\n";

    std::cout << "\nGLSL:\n" << gen_glsl(l) << "\n";

    return 0;
}
*/