#include "UniquifyVariableNames.h"
#include "IRMutator.h"
#include <sstream>

namespace Halide {
namespace Internal {

using std::map;
using std::string;
using std::vector;

class UniquifyVariableNames : public IRMutator2 {

    using IRMutator2::visit;

    map<string, int> vars;

    void push_name(string s) {
        if (vars.find(s) == vars.end()) {
            vars[s] = 0;
        } else {
            vars[s]++;
        }
    }

    string get_name(string s) {
        if (vars.find(s) == vars.end()) {
            return s;
        } else if (vars[s] == 0) {
            return s;
        } else {
            std::ostringstream oss;
            oss << s << "_" << vars[s];
            return oss.str();
        }
    }

    void pop_name(string s) {
        vars[s]--;
    }

    template<typename LetOrLetStmt>
    auto visit_let(const LetOrLetStmt *op) -> decltype(op->body) {
        struct Frame {
            const LetOrLetStmt *op;
            Expr value;
            string new_name;
        };

        vector<Frame> frames;
        decltype(op->body) result;
        while (op) {
            Expr val = mutate(op->value);
            push_name(op->name);
            frames.push_back({op, val, get_name(op->name)});
            result = op->body;
            op = result.template as<LetOrLetStmt>();
        }

        result = mutate(result);

        for (auto it = frames.rbegin(); it != frames.rend(); it++) {
            pop_name(it->op->name);
            if (it->new_name == it->op->name &&
                result.same_as(it->op->body) &&
                it->op->value.same_as(it->value)) {
                result = it->op;
            } else {
                result = LetOrLetStmt::make(it->new_name, it->value, result);
            }
        }

        return result;
    }

    Stmt visit(const LetStmt *op) override {
        return visit_let(op);
    }

    Expr visit(const Let *op) override {
        return visit_let(op);
    }

    Stmt visit(const For *op) override {
        Expr min = mutate(op->min);
        Expr extent = mutate(op->extent);
        push_name(op->name);
        string new_name = get_name(op->name);
        Stmt body = mutate(op->body);
        pop_name(op->name);

        if (new_name == op->name &&
            body.same_as(op->body) &&
            min.same_as(op->min) &&
            extent.same_as(op->extent)) {
            return op;
        } else {
            return For::make(new_name, min, extent, op->for_type, op->device_api, body);
        }
    }

    Expr visit(const Variable *op) override {
        string new_name = get_name(op->name);
        if (op->name != new_name) {
            return Variable::make(op->type, new_name);
        } else {
            return op;
        }
    }
};

Stmt uniquify_variable_names(Stmt s) {
    UniquifyVariableNames u;
    return u.mutate(s);
}

}  // namespace Internal
}  // namespace Halide