#pragma once

#ifndef utils_h
#define utils_h

#include "LevenbergMarquardt.h"

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

double Gaussian(double x, double* p) {
    return p[0] * exp(-(x - p[1]) * (x - p[1]) / (2 * p[2] * p[2]));
}

double* GaussianJacobian(double x, double* p) {
    double* resJ = new double[3];
    resJ[0] = -Gaussian(x, p) / p[0];
    resJ[1] = -(x - p[1]) * Gaussian(x, p) / (p[2] * p[2]);
    resJ[2] = -(x - p[1]) * (x - p[1]) * Gaussian(x, p) / (p[2] * p[2] * p[2]);
    return resJ;
}

class GaussFit {
public:
    GaussFit();
    ~GaussFit();

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

public:
    double* parma = new double[3];
    std::vector<Eigen::Vector2d> data;
};

GaussFit::GaussFit() {}
GaussFit::~GaussFit() {}

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

double* GaussFit::fit() {
    int n = data.size();
    double x, y, sigma;

    Eigen::VectorXd mY(n);
    Eigen::MatrixXd mX(n, 3);
    Eigen::MatrixXd mB(3, 1);

    for (int i = 0; i < n; i++) {
        Eigen::Vector2d& point = data.at(i);
        x = point(0), y = point(1);
        mY(i, 0) = log(y);
        mX(i, 0) = 1, mX(i, 1) = x, mX(i, 2) = x * x;
    }

    mB = (mX.transpose() * mX).inverse() * mX.transpose() * mY;
    sigma = -1 / mB(2, 0);
    parma[1] = mB(1, 0) * sigma / 2;
    parma[0] = exp(mB(0, 0) + parma[1] * parma[1] / sigma);
    parma[2] = sqrt(sigma);

    LevenbergMarquardt LM(3, parma, Gaussian, GaussianJacobian);
    LM.data = data;
    parma = LM.solve();

    return parma;
}

#endif