#pragma once
#include <vector>
#include <memory>
#include <cmath>
#include <unordered_map>

namespace autodiff {

template<typename T, typename F>
auto central_difference(std::vector<T>& vec, F func, std::size_t arg, float epsilon = 1e-6) -> decltype(func(vec))
{
    std::vector<T> vec1=vec;
    std::vector<T> vec2=vec;
    vec1[arg]+=epsilon;
    vec2[arg]-=epsilon;
    return (func(vec1)-func(vec2))/(2.0*epsilon);
}

class ScalarFunction {
public:
    float data;
    float grad;
    int degree = 0;
public:
    ScalarFunction() {}
}; // class ScalarFunction

class ConstantScalar: public ScalarFunction {
public:
    ConstantScalar(float data): ScalarFunction() {
        this->data = data;
    }
}; // class ConstantScalar

class Add: public ScalarFunction {
public:
    std::shared_ptr<ScalarFunction> a;
    std::shared_ptr<ScalarFunction> b;
public:
    Add(std::shared_ptr<ScalarFunction> a, std::shared_ptr<ScalarFunction> b): a(a), b(b) {
        this->data = a->data + b->data;
        this->degree = 2;
    }
    float forward() {
        return a->data + b->data;
    }
    std::vector<float> backward(float d_input) {
        return {d_input, d_input};
    }
}; // class Add

class Log: public ScalarFunction {
public:
    std::shared_ptr<ScalarFunction> a;
public:
    Log(std::shared_ptr<ScalarFunction> a): a(a) {
        this->data = this->forward();
        this->degree = 1;
    }
    float forward() {
        return logf(a->data);
    }
    std::vector<float> backward(float d_input) {
        return {d_input / a->data};
    }
}; // class Log

class Mul: public ScalarFunction {
public:
    std::shared_ptr<ScalarFunction> a;
    std::shared_ptr<ScalarFunction> b;
public:
    Mul(std::shared_ptr<ScalarFunction> a, std::shared_ptr<ScalarFunction> b) : a(a), b(b) {
        this->data = this->forward();
        this->degree = 2;
    }
    float forward() {
        return a->data * b->data;
    }
    std::vector<float> backward(float d_input) {
        return {b->data * d_input, a->data * d_input};
    }
}; // class Mul

class Inv: public ScalarFunction {
public:
    std::shared_ptr<ScalarFunction> a;
public:
    Inv(std::shared_ptr<ScalarFunction> a): a(a) {
        this->data = this->forward();
        this->degree = 1;
    }
    float forward() {
        return 1.0f / a->data;
    }
    std::vector<float> backward(float d_input) {
        return {-d_input / (a->data * a->data)};
    }
}; // class Inv

class Sigmoid: public ScalarFunction {
public:
    std::shared_ptr<ScalarFunction> a;
public:
    Sigmoid(std::shared_ptr<ScalarFunction> a): a(a) {
        this->data = this->forward();
        this->degree = 1;
    }
    float forward() {
        float x = a->data;
        if (x >= 0) {
            return 1.0f / (1.0f + expf(-x));
        } else {
            float exp_x = expf(x);
            return exp_x / (1.0f + exp_x);
        }
    }
    std::vector<float> backward(float d_input) {
        float sig = this->data;
        return {d_input * sig * (1.0f - sig)};
    }
}; // class Sigmoid

// for testing
bool test_central_difference() {
    std::vector<float> x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
    auto func = [](const std::vector<float>& x) -> float {
        return x[0] + x[1] + x[2] + x[3] + x[4];
    };
    auto grad = central_difference(x, func, 2);
    if (abs(grad-1.0f) > 1e-4) {
        return false;
    }
    return true;
}

bool test_addscalar() {
    auto a = std::make_shared<ConstantScalar>(1.0f);
    auto b = std::make_shared<ConstantScalar>(2.0f);
    auto c = std::make_shared<Add>(a, b);
    if (c->data != 3.0f) {
        return false;
    }
    auto res = c->backward(2.0f);
    auto a_grad = res[0];
    auto b_grad = res[1];
    if (a_grad != 2.0f || b_grad != 2.0f) {
        return false;
    }
    return true;
}

bool test_mulscalar() {
    auto a = std::make_shared<ConstantScalar>(2.0f);
    auto b = std::make_shared<ConstantScalar>(3.0f);
    auto c = std::make_shared<Mul>(a, b);
    if (c->data != 6.0f) {
        return false;
    }
    auto res = c->backward(2.0f);
    auto a_grad = res[0];
    auto b_grad = res[1];
    if (a_grad != 6.0f || b_grad != 4.0f) {
        return false;
    }
    return true;
}

bool test_logscalar() {
    auto a = std::make_shared<ConstantScalar>(2.0f);
    auto b = std::make_shared<Log>(a);
    if (abs(b->data - logf(2.0f)) > 1e-4) {
        return false;
    }
    auto res = b->backward(2.0f);
    auto a_grad = res[0]; 
    if (abs(a_grad - 1.0f) > 1e-4) {
        return false;
    }
    return true;
}

bool test_invscalar() {
    auto a = std::make_shared<ConstantScalar>(2.0f);
    auto b = std::make_shared<Inv>(a);
    if (abs(b->data - 0.5f) > 1e-4) {
        return false;
    }
    auto res = b->backward(2.0f);
    auto a_grad = res[0];
    if (abs(a_grad + 0.5f) > 1e-4) {
        return false;
    }
    return true;
}

bool test_sigmoidscalar() {
    auto a = std::make_shared<ConstantScalar>(2.0f);
    auto b = std::make_shared<Sigmoid>(a);
    
    // 计算预期的sigmoid值
    float expected_data = 1.0f / (1.0f + expf(-2.0f));
    
    // 检查前向传播结果
    if (abs(b->data - expected_data) > 1e-4) {
        return false;
    }
    
    // 计算预期的导数
    float expected_grad = expected_data * (1.0f - expected_data);
    auto res = b->backward(2.0f);
    auto a_grad = res[0];
    
    // 检查反向传播结果
    if (abs(a_grad - 2.0f * expected_grad) > 1e-4) {
        return false;
    }
    
    return true;
}

}