#include "DeviceArgument.h"
#include "CodeGen_GPU_Dev.h"
#include "IRPrinter.h"

namespace Halide {
namespace Internal {

std::vector<DeviceArgument> HostClosure::arguments() {
    if (debug::debug_level() >= 2) {
        debug(2) << *this;
    }

    std::vector<DeviceArgument> res;
    for (const auto &v : vars) {
        res.emplace_back(v.first, false, MemoryType::Auto, v.second, 0);
    }
    for (const auto &b : buffers) {
        DeviceArgument arg(b.first, true, b.second.memory_type, b.second.type, b.second.dimensions, b.second.size);
        arg.read = b.second.read;
        arg.write = b.second.write;
        res.push_back(arg);
    }
    return res;
}

void HostClosure::visit(const Call *op) {
    if (op->is_intrinsic(Call::image_load) ||
        op->is_intrinsic(Call::image_store)) {

        // The argument to the call is either a StringImm or a broadcasted
        // StringImm if this is part of a vectorized expression

        const StringImm *string_imm = op->args[0].as<StringImm>();
        if (!string_imm) {
            internal_assert(op->args[0].as<Broadcast>());
            string_imm = op->args[0].as<Broadcast>()->value.as<StringImm>();
        }

        internal_assert(string_imm);

        std::string bufname = string_imm->value;
        Buffer &ref = buffers[bufname];
        ref.type = op->type;
        ref.memory_type = op->is_intrinsic(Call::image_load) ||
                                  op->is_intrinsic(Call::image_store) ?
                              MemoryType::GPUTexture :
                              MemoryType::Auto;

        if (op->is_intrinsic(Call::image_load)) {
            ref.read = true;
            ref.dimensions = (op->args.size() - 2) / 2;
        } else if (op->is_intrinsic(Call::image_store)) {
            ref.write = true;
            ref.dimensions = op->args.size() - 3;
        }

        // The Func's name and the associated .buffer are mentioned in the
        // argument lists, but don't treat them as free variables.
        ScopedBinding<> p1(ignore, bufname);
        ScopedBinding<> p2(ignore, bufname + ".buffer");
        Internal::Closure::visit(op);
    } else {
        Internal::Closure::visit(op);
    }
}

void HostClosure::visit(const For *loop) {
    if (CodeGen_GPU_Dev::is_gpu_var(loop->name)) {
        // The size of the threads and blocks is not part of the closure
        ScopedBinding<> p(ignore, loop->name);
        loop->body.accept(this);
    } else {
        Internal::Closure::visit(loop);
    }
}

}  // namespace Internal
}  // namespace Halide