#pragma once

#ifndef gauss_fit_h
#define gauss_fit_h

#include "GaussNewton.h"
#include "LevenbergMarquardt.h"
#include "clip.h"
#include "utils.h"
#include <TCanvas.h>
#include <TH1F.h>

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

double Gaussian(double x, double* p) {
    return p[0] * std::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:
    double* parma = new double[3];
    std::vector<Eigen::Vector2d> data;

public:
    void addData(double x, double y);
    double* fit(int type_ = 0);
    double RSquare();
    int getTotal();
    void draw(std::string title = "./Figure.png");
};

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

double* GaussFit::fit(int type_) {
    double x, y;

    SigmaClip* SC2 = new SigmaClip();
    data = SC2->clip(data);

    parma[0] = dataMax2D(data);
    parma[1] = dataAvg2D(data);
    parma[2] = dataStd2D(data);

    if (DEBUG > 1)
        for (int i = 0; i < data.size(); i++) {
            Eigen::Vector2d& point = data.at(i);
            x = point(0), y = point(1);
            std::cout << x << " " << y << std::endl;
        }

    if (type_ == 0) {
        LevenbergMarquardt LM(3, parma, Gaussian, GaussianJacobian);
        LM.data = data;
        parma = LM.solve();
    } else if (type_ == 1) {
        LevenbergMarquardt LM(3, parma, Gaussian, GaussianJacobian, type_ = type_);
        LM.data = data;
        parma = LM.solve();
    } else {
        GaussNewton GN(3, parma, Gaussian, GaussianJacobian);
        GN.data = data;
        parma = GN.solve();
    }

    if (DEBUG) std::cout << parma[0] << " " << parma[1] << ", " << parma[2] << std::endl;
    if (DEBUG) std::cout << RSquare() << std::endl;

    if (RSquare() < 0.6)
        for (int i = 0; i < 3; i++) parma[i] = 0;

    return parma;
}

double GaussFit::RSquare() {
    double x, y, mu;
    double RSS = 0, TSS = 0;

    std::vector<double> yData;
    for (int i = 0; i < data.size(); i++) yData.push_back(data.at(i)(1));
    mu = dataAvg(yData);

    for (int i = 0; i < data.size(); i++) {
        Eigen::Vector2d& point = data.at(i);
        x = point(0), y = point(1);
        RSS += std::pow(y - Gaussian(x, parma), 2);
        TSS += std::pow(y - mu, 2);
    }

    return 1 - RSS / TSS;
}

int GaussFit::getTotal() {
    int sum = 0;
    for (int i = 0; i < data.size(); i++) {
        Eigen::Vector2d& point = data.at(i);
        sum += point(1);
    }
    return sum;
}

void GaussFit::draw(std::string title) {
    TCanvas* c1 = new TCanvas("Gauss Fit", "Gauss Fit", 0, 0, 1600, 1200);
    TH1F* h1 = new TH1F("", "Raw Data", CHANNEL_NUMBER, 0, CHANNEL_NUMBER);
    TH1F* h2 = new TH1F("", "Fit Data", CHANNEL_NUMBER, 0, CHANNEL_NUMBER);

    for (int i = 0; i < data.size(); i++) {
        Eigen::Vector2d& point = data.at(i);
        h1->SetBinContent(point(0), point(1));
    }
    for (int i = 0; i < CHANNEL_NUMBER; i++) h2->SetBinContent(i, Gaussian(i, parma));

    c1->cd(1);
    h1->Draw("L");
    h2->SetLineColor(kRed);
    h2->Draw("SAME");
    c1->SaveAs(title.c_str());
}

#endif