#include "Expression.hpp"
#include "Rational.hpp"
#include "Utils.hpp"

#include <gtest/gtest.h>
#include <stdexcept>
#include <variant>

class ExpressionTest: public ::testing::Test {
public:
    sym::Term<double> x{1.0, sym::Variable("x")};
    sym::Term<double> y{sym::Variable("y")};
    sym::Term<double> z{sym::Variable("z")};
};

TEST_F(ExpressionTest, basic_ops_1) {
    sym::Expression<double, zx::PiRational> e(x);

    EXPECT_EQ(1, e.numTerms());
    EXPECT_EQ(zx::PiRational(0, 1), e.getConst());

    e += x;

    EXPECT_EQ(1, e.numTerms());
    EXPECT_EQ(zx::PiRational(0, 1), e.getConst());
    EXPECT_PRED_FORMAT2(testing::FloatLE, e[0].getCoeff(), 2.0);

    e += y;
    EXPECT_EQ(2, e.numTerms());
    EXPECT_PRED_FORMAT2(testing::FloatLE, e[0].getCoeff(), 2.0);
    EXPECT_PRED_FORMAT2(testing::FloatLE, e[1].getCoeff(), 1.0);
    EXPECT_EQ(e[0].getVar().getName(), "x");
    EXPECT_EQ(e[1].getVar().getName(), "y");
}

TEST_F(ExpressionTest, basic_ops_2) {
    sym::Expression<double, zx::PiRational> e1;
    e1 += x;
    e1 += 10.0 * y;
    e1 += 5.0 * z;
    e1 += zx::PiRational(1, 2);

    sym::Expression<double, zx::PiRational> e2;
    e2 += -5.0 * x;
    e2 += -10.0 * y;
    e2 += -4.9 * z;
    e2 += zx::PiRational(3, 2);

    auto sum = e1 + e2;

    EXPECT_EQ(2, sum.numTerms());
    EXPECT_PRED_FORMAT2(testing::FloatLE, sum[0].getCoeff(), -4.0);
    EXPECT_PRED_FORMAT2(testing::FloatLE, sum[1].getCoeff(), 0.1);
    EXPECT_EQ(sum[0].getVar().getName(), "x");
    EXPECT_EQ(sum[1].getVar().getName(), "z");
    EXPECT_EQ(sum.getConst(), zx::PiRational(0, 1));

    sum += zx::PiRational(1, 1);
    EXPECT_EQ(sum.getConst(), zx::PiRational(1, 1));
    sum -= zx::PiRational(1, 1);
    EXPECT_EQ(sum.getConst(), zx::PiRational(0, 1));
}

TEST_F(ExpressionTest, mult) {
    sym::Expression<double, zx::PiRational> e(x);

    e = e * 2.0;

    EXPECT_PRED_FORMAT2(testing::FloatLE, e[0].getCoeff(), 2);

    e = e * zx::PiRational(0.5);

    EXPECT_PRED_FORMAT2(testing::FloatLE, e[0].getCoeff(), 1);

    e += sym::Expression<double, zx::PiRational>{};

    EXPECT_PRED_FORMAT2(testing::FloatLE, e[0].getCoeff(), 1);

    e = e * 0.0;

    EXPECT_TRUE(e.isZero());
}

TEST_F(ExpressionTest, div) {
    sym::Expression<double, zx::PiRational> e(x);

    e = e / 2.0;

    EXPECT_PRED_FORMAT2(testing::FloatLE, e[0].getCoeff(), 0.5);

    e = e / zx::PiRational(0.5);

    EXPECT_PRED_FORMAT2(testing::FloatLE, e[0].getCoeff(), 1);

    EXPECT_THROW(e = e / 0.0, std::runtime_error);
}

TEST_F(ExpressionTest, Commutativity) {
    const sym::Expression<double, double> e1(x, y);
    sym::Expression<double, double>       e2(z);
    e2.setConst(1.0);

    EXPECT_EQ(e1 + e2, e2 + e1);
    EXPECT_EQ(e1 * 2.0, 2.0 * e1);
}

TEST_F(ExpressionTest, Associativity) {
    const sym::Expression<double, double> e1(x, y);
    const sym::Expression<double, double> e2(z);
    const sym::Expression<double, double> e3(1.0);

    EXPECT_EQ(e1 + (e2 + e3), (e1 + e2) + e3);
    EXPECT_EQ(e1 * (2.0 * 4.0), (e1 * 2.0) * 4.0);
}

TEST_F(ExpressionTest, Distributive) {
    const sym::Expression<double, double> e1(x, y);
    const sym::Expression<double, double> e2(z);

    EXPECT_EQ((e1 + e2) * 2.0, (e1 * 2.0) + (e2 * 2.0));
    EXPECT_EQ((e1 - e2) * 2.0, (e1 * 2.0) - (e2 * 2.0));
    std::cout << ((e1 + e2) / 2.0) << "\n";
    std::cout << ((e1 / 2.0) + (e2 / 2.0)) << "\n";
    EXPECT_EQ((e1 + e2) / 2.0, (e1 / 2.0) + (e2 / 2.0));
    EXPECT_EQ((e1 - e2) / 2.0, (e1 / 2.0) - (e2 / 2.0));
}

TEST_F(ExpressionTest, Variable) {
    EXPECT_TRUE(sym::Variable{"x"} != sym::Variable{"y"});
    EXPECT_TRUE(sym::Variable{"x"} == sym::Variable{"x"});
    EXPECT_TRUE(sym::Variable{"x"} < sym::Variable{"y"});
    EXPECT_TRUE(sym::Variable{"z"} > sym::Variable{"y"});
}

TEST_F(ExpressionTest, SumNegation) {
    const sym::Expression<double, double> e1(x, y);
    const sym::Expression<double, double> e2(z, y);

    EXPECT_EQ(e1 - e2, e1 + (-e2));
    const auto& zero = sym::Expression<double, double>{};
    EXPECT_EQ(e1 + (-e1), zero);
}

TEST_F(ExpressionTest, SumMult) {
    const sym::Expression<double, double> e1(x, y);
    EXPECT_EQ(e1 + e1, e1 * 2.0);
}

TEST_F(ExpressionTest, CliffordRounding) {
    const double     eps = 1e-14;
    zx::PiExpression e{zx::PiRational(zx::PI - eps)};
    zx::roundToClifford(e, 1e-9);
    EXPECT_EQ(e, zx::PiExpression(zx::PiRational(1, 1)));
    e = zx::PiExpression{zx::PiRational((zx::PI / 2) - eps)};
    zx::roundToClifford(e, 1e-9);
    EXPECT_EQ(e, zx::PiExpression(zx::PiRational(1, 2)));
    e = zx::PiExpression{zx::PiRational((-zx::PI / 2) - eps)};
    zx::roundToClifford(e, 1e-9);
    EXPECT_EQ(e, zx::PiExpression(zx::PiRational(-1, 2)));
}

TEST_F(ExpressionTest, Clifford) {
    zx::PiExpression e{zx::PiRational(zx::PI)};

    EXPECT_TRUE(zx::isPauli(e));
    EXPECT_TRUE(zx::isClifford(e));
    EXPECT_FALSE(zx::isProperClifford(e));

    e = zx::PiExpression{zx::PiRational(zx::PI / 2)};
    EXPECT_FALSE(zx::isPauli(e));
    EXPECT_TRUE(zx::isClifford(e));
    EXPECT_TRUE(zx::isProperClifford(e));

    e = zx::PiExpression{zx::PiRational(zx::PI / 4)};
    EXPECT_FALSE(zx::isPauli(e));
    EXPECT_FALSE(zx::isClifford(e));
    EXPECT_FALSE(zx::isProperClifford(e));
}

TEST_F(ExpressionTest, Convertability) {
    const sym::Expression<double, double> e({x, y}, zx::PI);
    const auto                            piE = e.convert<zx::PiRational>();

    EXPECT_EQ(piE, zx::PiExpression({x, y}, zx::PiRational(1, 1)));
}

TEST_F(ExpressionTest, Instantiation) {
    sym::Expression<double, double> e(2 * x, y);

    const sym::VariableAssignment assignment{
            {sym::Variable{"x"}, 2.0}, {sym::Variable{"y"}, 1.0}};

    EXPECT_PRED_FORMAT2(testing::FloatLE, e.evaluate(assignment), 5.0);

    e += z;
    [[maybe_unused]] double h;
    EXPECT_THROW(h = e.evaluate(assignment), sym::SymbolicException);
}
