// Chemfiles, a modern library for chemistry file reading and writing
// Copyright (C) Guillaume Fraux and contributors -- BSD license

#include "chemfiles/Frame.hpp"
#include "chemfiles/Selection.hpp"
#include "chemfiles/ErrorFmt.hpp"

#include "chemfiles/selections/expr.hpp"
#include "chemfiles/selections/lexer.hpp"

using namespace chemfiles;
using namespace chemfiles::selections;

static const std::string EMPTY_STRING;

static std::string kind_as_string(Property::Kind kind) {
    switch (kind) {
    case Property::BOOL:
        return "bool";
    case Property::DOUBLE:
        return "double";
    case Property::STRING:
        return "string";
    case Property::VECTOR3D:
        return "Vector3D";
    }
    unreachable();
}

SubSelection::SubSelection(Variable variable): selection_(nullptr), variable_(variable) {}

SubSelection::SubSelection(std::string selection):
    selection_(new Selection(std::move(selection))), variable_(UINT8_MAX)
{
    if (selection_->size() != 1) {
        throw selection_error("sub-selection must have a size of 1");
    }
}

std::vector<size_t> SubSelection::eval(const Frame& frame, const Match& match) const {
    if (is_variable()) {
        return {match[variable_]};
    } else {
        return selection_->list(frame);
    }
}

std::string SubSelection::print() const {
    if (is_variable()) {
        return fmt::format("#{}", variable_ + 1);
    } else {
        return selection_->string();
    }
}

std::string And::print(unsigned delta) const {
    auto lhs = lhs_->print(7);
    auto rhs = rhs_->print(7);
    return fmt::format("and -> {}\n{:{}}    -> {}", lhs, "", delta, rhs);
}

bool And::is_match(const Frame& frame, const Match& match) const {
    return lhs_->is_match(frame, match) && rhs_->is_match(frame, match);
}

std::string Or::print(unsigned delta) const {
    auto lhs = lhs_->print(6);
    auto rhs = rhs_->print(6);
    return fmt::format("or -> {}\n{:{}}   -> {}", lhs, "", delta, rhs);
}

bool Or::is_match(const Frame& frame, const Match& match) const {
    return lhs_->is_match(frame, match) || rhs_->is_match(frame, match);
}

std::string Not::print(unsigned /*unused*/) const {
    return "not " + ast_->print(4);
}

bool Not::is_match(const Frame& frame, const Match& match) const {
    return !ast_->is_match(frame, match);
}

std::string All::print(unsigned /*unused*/) const {
    return "all";
}

bool All::is_match(const Frame& /*unused*/, const Match& /*unused*/) const {
    return true;
}

std::string None::print(unsigned /*unused*/) const {
    return "none";
}

bool None::is_match(const Frame& /*unused*/, const Match& /*unused*/) const {
    return false;
}

std::string BoolProperty::print(unsigned /*unused*/) const {
    if (is_ident(property_)) {
        return fmt::format("[{}](#{})", property_, argument_ + 1);
    } else {
        return fmt::format("[\"{}\"](#{})", property_, argument_ + 1);
    }
}

bool BoolProperty::is_match(const Frame& frame, const Match& match) const {
    const auto& property = frame[match[argument_]].get(property_);
    if (property) {
        if (property->kind() == Property::BOOL) {
            return property->as_bool();
        } else {
            throw selection_error(
                "invalid type for property [{}]: expected bool, got {}",
                property_, kind_as_string(property->kind())
            );
        }
    } else {
        // No property with the given name
        return false;
    }
}

std::string IsBonded::print(unsigned /*unused*/) const {
    return fmt::format("is_bonded({}, {})", i_.print(), j_ .print());
}

bool IsBonded::is_match(const Frame& frame, const Match& match) const {
    auto& bonds = frame.topology().bonds();
    for (auto i: i_.eval(frame, match)) {
        for (auto j: j_.eval(frame, match)) {
            if (i == j) {
                // This is not a valid bond
                continue;
            }

            auto candidate = chemfiles::Bond(i, j);
            auto it = std::lower_bound(bonds.begin(), bonds.end(), candidate);
            if (it != bonds.end() && *it == candidate) {
                return true;
            }
        }
    }
    return false;
}

std::string IsAngle::print(unsigned /*unused*/) const {
    return fmt::format("is_angle({}, {}, {})", i_.print(), j_.print(), k_.print());
}

bool IsAngle::is_match(const Frame& frame, const Match& match) const {
    auto& angles = frame.topology().angles();
    for (auto i: i_.eval(frame, match)) {
        for (auto j: j_.eval(frame, match)) {
            for (auto k: k_.eval(frame, match)) {
                if (i == j || j == k || i == k) {
                    // This is not a valid angle
                    continue;
                }

                auto candidate = chemfiles::Angle(i, j, k);
                auto it = std::lower_bound(angles.begin(), angles.end(), candidate);
                if (it != angles.end() && *it == candidate) {
                    return true;
                }
            }
        }
    }
    return false;
}

std::string IsDihedral::print(unsigned /*unused*/) const {
    return fmt::format("is_dihedral({}, {}, {}, {})", i_.print(), j_.print(), k_.print(), m_.print());
}

bool IsDihedral::is_match(const Frame& frame, const Match& match) const {
    auto& dihedrals = frame.topology().dihedrals();
    for (auto i: i_.eval(frame, match)) {
        for (auto j: j_.eval(frame, match)) {
            for (auto k: k_.eval(frame, match)) {
                for (auto m: m_.eval(frame, match)) {
                    if (i == j || j == k || k == m || i == k || j == m || i == m) {
                        // This is not a valid angle
                        continue;
                    }

                    auto candidate = chemfiles::Dihedral(i, j, k, m);
                    auto it = std::lower_bound(dihedrals.begin(), dihedrals.end(), candidate);
                    if (it != dihedrals.end() && *it == candidate) {
                        return true;
                    }
                }
            }
        }
    }
    return false;
}

std::string IsImproper::print(unsigned /*unused*/) const {
    return fmt::format("is_improper({}, {}, {}, {})", i_.print(), j_.print(), k_.print(), m_.print());
}

bool IsImproper::is_match(const Frame& frame, const Match& match) const {
    auto& impropers = frame.topology().impropers();
    for (auto i: i_.eval(frame, match)) {
        for (auto j: j_.eval(frame, match)) {
            for (auto k: k_.eval(frame, match)) {
                for (auto m: m_.eval(frame, match)) {
                    if (i == j || j == k || k == m || i == k || j == m || i == m) {
                        // This is not a valid angle
                        continue;
                    }

                    auto candidate = chemfiles::Improper(i, j, k, m);
                    auto it = std::lower_bound(impropers.begin(), impropers.end(), candidate);
                    if (it != impropers.end() && *it == candidate) {
                        return true;
                    }
                }
            }
        }
    }
    return false;
}

std::string StringSelector::print(unsigned /*unused*/) const {
    auto op = equals_ ? "==" : "!=";
    if (is_ident(value_)) {
        return fmt::format("{}(#{}) {} {}", name(), argument_ + 1, op, value_);
    } else {
        return fmt::format("{}(#{}) {} \"{}\"", name(), argument_ + 1, op, value_);
    }
}

bool StringSelector::is_match(const Frame& frame, const Match& match) const {
    return (this->value(frame, match[argument_]) == value_) == equals_;
}

std::string StringProperty::name() const {
    if (is_ident(property_)) {
        return "[" + property_ + "]";
    } else {
        return "[\"" + property_ + "\"]";
    }
}

const std::string& StringProperty::value(const Frame& frame, size_t i) const {
    const auto& property = frame[i].get(property_);
    if (property) {
        if (property->kind() == Property::STRING) {
            return property->as_string();
        } else {
            throw selection_error(
                "invalid type for property [{}]: expected string, got {}",
                property_, kind_as_string(property->kind())
            );
        }
    } else {
        // No property with the given name
        return EMPTY_STRING;
    }
}

std::string Type::name() const {
    return "type";
}

const std::string& Type::value(const Frame& frame, size_t i) const {
    return frame[i].type();
}

std::string Name::name() const {
    return "name";
}

const std::string& Name::value(const Frame& frame, size_t i) const {
    return frame[i].name();
}

std::string Resname::name() const {
    return "resname";
}

const std::string& Resname::value(const Frame& frame, size_t i) const {
    auto residue = frame.topology().residue_for_atom(i);
    if (residue) {
        return residue->name();
    } else {
        return EMPTY_STRING;
    }
}

bool Math::is_match(const Frame& frame, const Match& match) const {
    auto lhs = lhs_->eval(frame, match);
    auto rhs = rhs_->eval(frame, match);
    switch (op_) {
    case Math::Operator::EQUAL:
        return lhs == rhs;
    case Math::Operator::NOT_EQUAL:
        return lhs != rhs;
    case Math::Operator::LESS:
        return lhs < rhs;
    case Math::Operator::LESS_EQUAL:
        return lhs <= rhs;
    case Math::Operator::GREATER:
        return lhs > rhs;
    case Math::Operator::GREATER_EQUAL:
        return lhs >= rhs;
    }
    unreachable();
}

std::string Math::print(unsigned /*unused*/) const {
    std::string op;
    switch (op_) {
    case Math::Operator::EQUAL:
        op = " == ";
        break;
    case Math::Operator::NOT_EQUAL:
        op = " != ";
        break;
    case Math::Operator::LESS:
        op = " < ";
        break;
    case Math::Operator::LESS_EQUAL:
        op = " <= ";
        break;
    case Math::Operator::GREATER:
        op = " > ";
        break;
    case Math::Operator::GREATER_EQUAL:
        op = " >= ";
        break;
    default:
        unreachable();
    }

    return lhs_->print() + op + rhs_->print();
}

void Math::optimize() {
    auto lhs_opt = lhs_->optimize();
    if (lhs_opt) {
        lhs_ = MathAst(new Number(lhs_opt.value()));
    }

    auto rhs_opt = rhs_->optimize();
    if (rhs_opt) {
        rhs_ = MathAst(new Number(rhs_opt.value()));
    }
}

double Add::eval(const Frame& frame, const Match& match) const {
    return lhs_->eval(frame, match) + rhs_->eval(frame, match);
}

optional<double> Add::optimize() {
    auto lhs_opt = lhs_->optimize();
    auto rhs_opt = rhs_->optimize();
    if (lhs_opt && lhs_opt) {
        return lhs_opt.value() + rhs_opt.value();
    } else if (lhs_opt) {
        lhs_ = MathAst(new Number(lhs_opt.value()));
    } else if (rhs_opt) {
        rhs_ = MathAst(new Number(rhs_opt.value()));
    }
    return nullopt;
}

std::string Add::print() const {
    return fmt::format("({} + {})", lhs_->print(), rhs_->print());
}

double Sub::eval(const Frame& frame, const Match& match) const {
    return lhs_->eval(frame, match) - rhs_->eval(frame, match);
}

optional<double> Sub::optimize() {
    auto lhs_opt = lhs_->optimize();
    auto rhs_opt = rhs_->optimize();
    if (lhs_opt && lhs_opt) {
        return lhs_opt.value() - rhs_opt.value();
    } else if (lhs_opt) {
        lhs_ = MathAst(new Number(lhs_opt.value()));
    } else if (rhs_opt) {
        rhs_ = MathAst(new Number(rhs_opt.value()));
    }
    return nullopt;
}

std::string Sub::print() const {
    return fmt::format("({} - {})", lhs_->print(), rhs_->print());
}

double Mul::eval(const Frame& frame, const Match& match) const {
    return lhs_->eval(frame, match) * rhs_->eval(frame, match);
}

optional<double> Mul::optimize() {
    auto lhs_opt = lhs_->optimize();
    auto rhs_opt = rhs_->optimize();
    if (lhs_opt && lhs_opt) {
        return lhs_opt.value() * rhs_opt.value();
    } else if (lhs_opt) {
        lhs_ = MathAst(new Number(lhs_opt.value()));
    } else if (rhs_opt) {
        rhs_ = MathAst(new Number(rhs_opt.value()));
    }
    return nullopt;
}

std::string Mul::print() const {
    return fmt::format("({} * {})", lhs_->print(), rhs_->print());
}

double Div::eval(const Frame& frame, const Match& match) const {
    return lhs_->eval(frame, match) / rhs_->eval(frame, match);
}

optional<double> Div::optimize() {
    auto lhs_opt = lhs_->optimize();
    auto rhs_opt = rhs_->optimize();
    if (lhs_opt && lhs_opt) {
        return lhs_opt.value() / rhs_opt.value();
    } else if (lhs_opt) {
        lhs_ = MathAst(new Number(lhs_opt.value()));
    } else if (rhs_opt) {
        rhs_ = MathAst(new Number(rhs_opt.value()));
    }
    return nullopt;
}

std::string Div::print() const {
    return fmt::format("({} / {})", lhs_->print(), rhs_->print());
}

double Pow::eval(const Frame& frame, const Match& match) const {
    return pow(lhs_->eval(frame, match), rhs_->eval(frame, match));
}

optional<double> Pow::optimize() {
    auto lhs_opt = lhs_->optimize();
    auto rhs_opt = rhs_->optimize();
    if (lhs_opt && lhs_opt) {
        return pow(lhs_opt.value(), rhs_opt.value());
    } else if (lhs_opt) {
        lhs_ = MathAst(new Number(lhs_opt.value()));
    } else if (rhs_opt) {
        rhs_ = MathAst(new Number(rhs_opt.value()));
    }
    return nullopt;
}

std::string Pow::print() const {
    return fmt::format("{} ^({})", lhs_->print(), rhs_->print());
}

double Neg::eval(const Frame& frame, const Match& match) const {
    return - ast_->eval(frame, match);
}

optional<double> Neg::optimize() {
    auto optimized = ast_->optimize();
    if (optimized) {
        return - optimized.value();
    } else {
        return nullopt;
    }
}

std::string Neg::print() const {
    return fmt::format("(-{})", ast_->print());
}

double Mod::eval(const Frame& frame, const Match& match) const {
    return fmod(lhs_->eval(frame, match), rhs_->eval(frame, match));
}

optional<double> Mod::optimize() {
    auto lhs_opt = lhs_->optimize();
    auto rhs_opt = rhs_->optimize();
    if (lhs_opt && lhs_opt) {
        return fmod(lhs_opt.value(), rhs_opt.value());
    } else if (lhs_opt) {
        lhs_ = MathAst(new Number(lhs_opt.value()));
    } else if (rhs_opt) {
        rhs_ = MathAst(new Number(rhs_opt.value()));
    }
    return nullopt;
}

std::string Mod::print() const {
    return fmt::format("({} % {})", lhs_->print(), rhs_->print());
}

double Function::eval(const Frame& frame, const Match& match) const {
    return fn_(ast_->eval(frame, match));
}

optional<double> Function::optimize() {
    auto optimized = ast_->optimize();
    if (optimized) {
        return fn_(optimized.value());
    } else {
        return nullopt;
    }
}

std::string Function::print() const {
    return fmt::format("{}({})", name_, ast_->print());
}

double Number::eval(const Frame& /*unused*/, const Match& /*unused*/) const {
    return value_;
}

optional<double> Number::optimize() {
    return value_;
}

std::string Number::print() const {
    if (std::round(value_) == value_) {
        return std::to_string(std::lround(value_));
    } else {
        return std::to_string(value_);
    }
}

double Distance::eval(const Frame& frame, const Match& match) const {
    return frame.distance(match[i_], match[j_]);
}

std::string Distance::print() const {
    return fmt::format("distance(#{}, #{})", i_ + 1, j_ + 1);
}

double selections::Angle::eval(const Frame& frame, const Match& match) const {
    return frame.angle(match[i_], match[j_], match[k_]);
}

std::string selections::Angle::print() const {
    return fmt::format("angle(#{}, #{}, #{})", i_ + 1, j_ + 1, k_ + 1);
}

double selections::Dihedral::eval(const Frame& frame, const Match& match) const {
    return frame.dihedral(match[i_], match[j_], match[k_], match[m_]);
}

std::string selections::Dihedral::print() const {
    return fmt::format("dihedral(#{}, #{}, #{}, #{})", i_ + 1, j_ + 1, k_ + 1, m_ + 1);
}

double OutOfPlane::eval(const Frame& frame, const Match& match) const {
    return frame.out_of_plane(match[i_], match[j_], match[k_], match[m_]);
}

std::string OutOfPlane::print() const {
    return fmt::format("out_of_plane(#{}, #{}, #{}, #{})", i_ + 1, j_ + 1, k_ + 1, m_ + 1);
}

double NumericSelector::eval(const Frame& frame, const Match& match) const {
    return this->value(frame, match[argument_]);
}

optional<double> NumericSelector::optimize() {
    return nullopt;
}

std::string NumericSelector::print() const {
    return fmt::format("{}(#{})", name(), argument_ + 1);
}

std::string NumericProperty::name() const {
    if (is_ident(property_)) {
        return "[" + property_ + "]";
    } else {
        return "[\"" + property_ + "\"]";
    }
}

double NumericProperty::value(const Frame& frame, size_t i) const {
    const auto& property = frame[i].get(property_);
    if (property) {
        if (property->kind() == Property::DOUBLE) {
            return property->as_double();
        } else {
            throw selection_error(
                "invalid type for property [{}]: expected double, got {}",
                property_, kind_as_string(property->kind())
            );
        }
    } else {
        // No property with the given name
        // return nan so that all comparaison down the line evaluate to false
        return nan("0");
    }
}

std::string Index::name() const {
    return "index";
}

double Index::value(const Frame& /*unused*/, size_t i) const {
    return static_cast<double>(i);
}

std::string Resid::name() const {
    return "resid";
}

double Resid::value(const Frame& frame, size_t i) const {
    auto residue = frame.topology().residue_for_atom(i);
    if (residue && residue->id()) {
        return static_cast<double>(*residue->id());
    } else {
        return -1;
    }
}

std::string Mass::name() const {
    return "mass";
}

double Mass::value(const Frame& frame, size_t i) const {
    return frame[i].mass();
}

std::string Position::name() const {
    switch (coordinate_) {
    case Coordinate::X:
        return "x";
    case Coordinate::Y:
        return "y";
    case Coordinate::Z:
        return "z";
    }
    unreachable();
}

double Position::value(const Frame& frame, size_t i) const {
    return frame.positions()[i][static_cast<size_t>(coordinate_)];
}

std::string Velocity::name() const {
    switch (coordinate_) {
    case Coordinate::X:
        return "vx";
    case Coordinate::Y:
        return "vy";
    case Coordinate::Z:
        return "vz";
    }
    unreachable();
}

double Velocity::value(const Frame& frame, size_t i) const {
    if (frame.velocities()) {
        auto& velocities = *frame.velocities();
        return velocities[i][static_cast<size_t>(coordinate_)];
    } else {
        return 0.0;
    }
}
