#include "Reduction.h"

#include "IR.h"
#include "IREquality.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "IRVisitor.h"
#include "Simplify.h"
#include "Var.h"
#include <utility>

namespace Halide {
namespace Internal {

namespace {

void check(const Expr &pred, std::vector<Expr> &expected) {
    std::vector<Expr> result;
    split_into_ands(pred, result);
    bool is_equal = true;

    if (result.size() != expected.size()) {
        is_equal = false;
    } else {
        for (size_t i = 0; i < expected.size(); ++i) {
            if (!equal(simplify(result[i]), simplify(expected[i]))) {
                is_equal = false;
                break;
            }
        }
    }

    if (!is_equal) {
        std::cout << "Expect predicate " << pred << " to be split into:\n";
        for (const auto &e : expected) {
            std::cout << "  " << e << "\n";
        }
        std::cout << "Got:\n";
        for (const auto &e : result) {
            std::cout << "  " << e << "\n";
        }
        internal_error << "\n";
    }
}

}  // namespace

void split_predicate_test() {
    Expr x = Var("x"), y = Var("y"), z = Var("z"), w = Var("w");

    {
        std::vector<Expr> expected;
        expected.push_back(z < 10);
        check(z < 10, expected);
    }

    {
        std::vector<Expr> expected;
        expected.push_back((x < y) || (x == 10));
        check((x < y) || (x == 10), expected);
    }

    {
        std::vector<Expr> expected;
        expected.push_back(x < y);
        expected.push_back(x == 10);
        check((x < y) && (x == 10), expected);
    }

    {
        std::vector<Expr> expected;
        expected.push_back(x < y);
        expected.push_back(x == 10);
        expected.push_back(y == z);
        check((x < y) && (x == 10) && (y == z), expected);
    }

    {
        std::vector<Expr> expected;
        expected.push_back((w == 1) || ((x == 10) && (y == z)));
        check((w == 1) || ((x == 10) && (y == z)), expected);
    }

    {
        std::vector<Expr> expected;
        expected.push_back(x < y);
        expected.push_back((w == 1) || ((x == 10) && (y == z)));
        check((x < y) && ((w == 1) || ((x == 10) && (y == z))), expected);
    }

    std::cout << "Split predicate test passed" << std::endl;
}

struct ReductionDomainContents {
    mutable RefCount ref_count;
    std::vector<ReductionVariable> domain;
    Expr predicate;
    bool frozen = false;

    ReductionDomainContents()
        : predicate(const_true()) {
    }

    // Pass an IRVisitor through to all Exprs referenced in the ReductionDomainContents
    void accept(IRVisitor *visitor) {
        for (const ReductionVariable &rvar : domain) {
            if (rvar.min.defined()) {
                rvar.min.accept(visitor);
            }
            if (rvar.extent.defined()) {
                rvar.extent.accept(visitor);
            }
        }
        if (predicate.defined()) {
            predicate.accept(visitor);
        }
    }

    // Pass an IRMutator through to all Exprs referenced in the ReductionDomainContents
    void mutate(IRMutator *mutator) {
        for (ReductionVariable &rvar : domain) {
            if (rvar.min.defined()) {
                rvar.min = mutator->mutate(rvar.min);
            }
            if (rvar.extent.defined()) {
                rvar.extent = mutator->mutate(rvar.extent);
            }
        }
        if (predicate.defined()) {
            predicate = mutator->mutate(predicate);
        }
    }
};

template<>
RefCount &ref_count<Halide::Internal::ReductionDomainContents>(const ReductionDomainContents *p) noexcept {
    return p->ref_count;
}

template<>
void destroy<Halide::Internal::ReductionDomainContents>(const ReductionDomainContents *p) {
    delete p;
}

ReductionDomain::ReductionDomain(const std::vector<ReductionVariable> &domain)
    : contents(new ReductionDomainContents) {
    contents->domain = domain;
}

ReductionDomain ReductionDomain::deep_copy() const {
    if (!contents.defined()) {
        return ReductionDomain();
    }
    ReductionDomain copy(contents->domain);
    copy.contents->predicate = contents->predicate;
    copy.contents->frozen = contents->frozen;
    return copy;
}

const std::vector<ReductionVariable> &ReductionDomain::domain() const {
    return contents->domain;
}

namespace {
class DropSelfReferences : public IRMutator {
    using IRMutator::visit;

    Expr visit(const Variable *op) override {
        if (op->reduction_domain.defined()) {
            user_assert(op->reduction_domain.same_as(domain))
                << "An RDom's predicate may only refer to its own RVars, "
                << " not the RVars of some other RDom. "
                << "Cannot set the predicate to : " << predicate << "\n";
            return Variable::make(op->type, op->name);
        } else {
            return op;
        }
    }

public:
    Expr predicate;
    const ReductionDomain &domain;
    DropSelfReferences(Expr p, const ReductionDomain &d)
        : predicate(std::move(p)), domain(d) {
    }
};
}  // namespace

void ReductionDomain::set_predicate(const Expr &p) {
    // The predicate can refer back to the RDom. We need to break
    // those cycles to prevent a leak.
    contents->predicate = DropSelfReferences(p, *this).mutate(p);
}

void ReductionDomain::where(Expr predicate) {
    set_predicate(simplify(contents->predicate && std::move(predicate)));
}

Expr ReductionDomain::predicate() const {
    return contents->predicate;
}

std::vector<Expr> ReductionDomain::split_predicate() const {
    std::vector<Expr> predicates;
    split_into_ands(contents->predicate, predicates);
    return predicates;
}

void ReductionDomain::freeze() {
    contents->frozen = true;
}

bool ReductionDomain::frozen() const {
    return contents->frozen;
}

void ReductionDomain::accept(IRVisitor *visitor) const {
    if (contents.defined()) {
        contents->accept(visitor);
    }
}

void ReductionDomain::mutate(IRMutator *mutator) {
    if (contents.defined()) {
        contents->mutate(mutator);
    }
}

}  // namespace Internal
}  // namespace Halide