#pragma once

#ifndef levenberg_marquardt_h
#define levenberg_marquardt_h

#include <Eigen/Dense>
#include <iostream>

class LevenbergMarquardt {
public:
    LevenbergMarquardt(int L_, double* parma_, double (*Func_)(double, double*),
                       double* (*Gunc_)(double, double*), int type_ = 0);
    ~LevenbergMarquardt();

public:
    void addData(double x, double y);
    double* solve();

public:
    double mu = 1., eps = 1e-6;
    double* parma;
    double (*Func)(double, double*);
    double* (*Gunc)(double, double*);
    int L, type_, maxIter = 300;
    std::vector<Eigen::Vector2d> data;

private:
    void calmJ_vF();
    void calmH_vG();
    double calMse(double* parma_);

private:
    Eigen::MatrixXd mJ;  // 雅克比矩阵
    Eigen::MatrixXd mH;  // H矩阵
    Eigen::VectorXd vF;  // 误差向量
    Eigen::VectorXd vG;  // 左侧
};

LevenbergMarquardt::LevenbergMarquardt(int L_, double* parma_, double (*Func_)(double, double*),
                                       double* (*Gunc_)(double, double*), int type_) {
    L = L_;
    parma = parma_;
    Func = Func_;
    Gunc = Gunc_;
    this->type_ = type_;
}

LevenbergMarquardt::~LevenbergMarquardt() {}

void LevenbergMarquardt::addData(double x, double y) { data.push_back(Eigen::Vector2d(x, y)); }

void LevenbergMarquardt::calmJ_vF() {
    double x, y;
    double* resJ;

    mJ.resize(data.size(), L);
    vF.resize(data.size());

    for (int i = 0; i < data.size(); i++) {
        Eigen::Vector2d& point = data.at(i);
        x = point(0), y = point(1);
        resJ = (*Gunc)(x, parma);
        for (int j = 0; j < L; j++) mJ(i, j) = resJ[j];
        vF(i) = y - (*Func)(x, parma);
    }
}

void LevenbergMarquardt::calmH_vG() {
    mH = mJ.transpose() * mJ;
    vG = -mJ.transpose() * vF;
}

double LevenbergMarquardt::calMse(double* parma_) {
    double x, y, mse = 0;

    for (int i = 0; i < data.size(); i++) {
        Eigen::Vector2d& point = data.at(i);
        x = point(0), y = point(1);
        mse += std::pow(y - (*Func)(x, parma_), 2);
    }

    return mse;
}

double* LevenbergMarquardt::solve() {
    double v = 2, cost;
    double* parmaNew = new double[L];

    Eigen::VectorXd vX(L);
    Eigen::MatrixXd mT, mD(L, L);
    for (int i = 0; i < L; i++) {
        for (int j = 0; j < L; j++) mD(i, j) = 0;
        mD(i, i) = 1;
    }

    for (int k = 0; k < maxIter; k++) {
        calmJ_vF();
        calmH_vG();

        if (type_ == 1) {
            mT = mJ * mJ.transpose();
            for (int i = 0; i < L; i++) mD(i, i) = sqrt(mT(i, i));
        }

        mH = mH + mu * mD;
        vX = mH.ldlt().solve(vG);
        if (vX.norm() <= eps) return parma;

        for (int i = 0; i < L; i++) parmaNew[i] = parma[i] + vX[i];
        cost = (calMse(parma) - calMse(parmaNew)) / (vX.transpose() * (mu * vX + vG));

        if (cost > 0) {
            for (int i = 0; i < L; i++) parma[i] = parmaNew[i];
            mu = mu * std::max(1. / 3, 1 - std::pow(2 * cost - 1, 3));
            v = 2;
        } else {
            mu *= v;
            v *= 2;
        }
    }

    return parma;
}

#endif