#include "starfit.h"
#include <gsl/gsl_blas.h>
#include <QDebug>
#include <iostream>

const int PARAM_AM = 0;
const int PARAM_X0 = 1;
const int PARAM_Y0 = 2;
const int PARAM_SX = 3;
const int PARAM_SY = 4;
const int PARAM_TH = 5;

const int MAX_ITER = 20;
const double TOL = 1.0e-3;

struct StarData
{
    size_t size;
    std::vector<double> val;
};

// a * exp(-0.5*((x-x0)/sx)^2 + ((y-y0)/sy)^2)
double gauss_model(double a, double x0, double y0, double sx, double sy, double x, double y)
{
    double _x = (x-x0)/sx;
    double _y = (y-y0)/sy;
    return a*exp(-0.5*(_x*_x + _y*_y));
}

int func_f(const gsl_vector *X, void *params, gsl_vector *f)
{
    StarData *d = static_cast<StarData*>(params);
    double am = gsl_vector_get(X, PARAM_AM);
    double x0 = gsl_vector_get(X, PARAM_X0);
    double y0 = gsl_vector_get(X, PARAM_Y0);
    double sx = gsl_vector_get(X, PARAM_SX);
    double sy = gsl_vector_get(X, PARAM_SY);

    int i = 0;
    for(size_t y=0;y<d->size;y++)
    {
        for(size_t x=0;x<d->size;x++)
        {
            double v = gauss_model(am, x0, y0, sx, sy, x, y);
            gsl_vector_set(f, i, d->val[i] - v);
            i++;
        }
    }

    return GSL_SUCCESS;
}

int func_df(const gsl_vector *X, void *params, gsl_matrix *J)
{
    StarData *d = static_cast<StarData*>(params);
    double am = gsl_vector_get(X, PARAM_AM);
    double x0 = gsl_vector_get(X, PARAM_X0);
    double y0 = gsl_vector_get(X, PARAM_Y0);
    double sx = gsl_vector_get(X, PARAM_SX);
    double sy = gsl_vector_get(X, PARAM_SY);

    int i = 0;
    for(size_t y=0;y<d->size;y++)
    {
        for(size_t x=0;x<d->size;x++)
        {
            double tx = x-x0;
            double ty = y-y0;
            double e = gauss_model(am, x0, y0, sx, sy, x, y);

            gsl_matrix_set(J, i, PARAM_AM, -e/am);
            gsl_matrix_set(J, i, PARAM_X0, -e*(tx/(sx*sx)));
            gsl_matrix_set(J, i, PARAM_Y0, -e*(ty/(sy*sy)));
            gsl_matrix_set(J, i, PARAM_SX, -e*(tx*tx/(sx*sx*sx)));
            gsl_matrix_set(J, i, PARAM_SY, -e*(ty*ty/(sy*sy*sy)));
            i++;
        }
    }

    return GSL_SUCCESS;
}

int func_f_an(const gsl_vector *X, void *params, gsl_vector *f)
{
    StarData *d = static_cast<StarData*>(params);
    double am = gsl_vector_get(X, PARAM_AM);
    double x0 = gsl_vector_get(X, PARAM_X0);
    double y0 = gsl_vector_get(X, PARAM_Y0);
    double sx = gsl_vector_get(X, PARAM_SX);
    double sy = gsl_vector_get(X, PARAM_SY);
    double th = gsl_vector_get(X, PARAM_TH);

    int i = 0;
    double a = sin(th);
    double b = cos(th);
    for(size_t y=0;y<d->size;y++)
    {
        for(size_t x=0;x<d->size;x++)
        {
            double v = gauss_model(am, x0, y0, sx, sy, x*b-y*a, x*a+y*b);
            gsl_vector_set(f, i, d->val[i] - v);
            i++;
        }
    }

    return GSL_SUCCESS;
}

int func_df_af(const gsl_vector *X, void *params, gsl_matrix *J)
{
    StarData *d = static_cast<StarData*>(params);
    double am = gsl_vector_get(X, PARAM_AM);
    double x0 = gsl_vector_get(X, PARAM_X0);
    double y0 = gsl_vector_get(X, PARAM_Y0);
    double sx = gsl_vector_get(X, PARAM_SX);
    double sy = gsl_vector_get(X, PARAM_SY);

    int i = 0;
    for(size_t y=0;y<d->size;y++)
    {
        for(size_t x=0;x<d->size;x++)
        {
            double tx = x-x0;
            double ty = y-y0;
            double e = gauss_model(am, x0, y0, sx, sy, x, y);

            gsl_matrix_set(J, i, PARAM_AM, -e/am);
            gsl_matrix_set(J, i, PARAM_X0, -e*(tx/(sx*sx)));
            gsl_matrix_set(J, i, PARAM_Y0, -e*(ty/(sy*sy)));
            gsl_matrix_set(J, i, PARAM_SX, -e*(tx*tx/(sx*sx*sx)));
            gsl_matrix_set(J, i, PARAM_SY, -e*(ty*ty/(sy*sy*sy)));
            i++;
        }
    }

    return GSL_SUCCESS;
}

//int func_fvv(const gsl_vector *x, const gsl_vector * v, void *params, gsl_vector *fvv)
//{
//    return GSL_SUCCESS;
//}

void callback(const size_t iter, void *, const gsl_multifit_nlinear_workspace *w)
{
    double rcond;
    gsl_vector *x = gsl_multifit_nlinear_position(w);
    gsl_multifit_nlinear_rcond(&rcond, w);
    QString r = "Iter: " + QString::number(iter)
                + " Am: " + QString::number(gsl_vector_get(x, PARAM_AM))
                + " X0: " + QString::number(gsl_vector_get(x, PARAM_X0))
                + " Y0: " + QString::number(gsl_vector_get(x, PARAM_Y0))
                + " SX: " + QString::number(gsl_vector_get(x, PARAM_SX))
                + " SY: " + QString::number(gsl_vector_get(x, PARAM_SY))
                + " J(X) :" + QString::number(1.0/rcond)
                + " av: " + QString::number(gsl_multifit_nlinear_avratio(w));
    std::cout << r.toStdString() << std::endl;
}

void callback_an(const size_t iter, void *, const gsl_multifit_nlinear_workspace *w)
{
    double rcond;
    gsl_vector *x = gsl_multifit_nlinear_position(w);
    gsl_multifit_nlinear_rcond(&rcond, w);
    qDebug() << "Iter:" << iter << "Am:" << gsl_vector_get(x, PARAM_AM)
                << "X0:" << gsl_vector_get(x, PARAM_X0)
                << "Y0:" << gsl_vector_get(x, PARAM_Y0)
                << "SX:" << gsl_vector_get(x, PARAM_SX)
                << "SY:" << gsl_vector_get(x, PARAM_SY)
                << "TH:" << gsl_vector_get(x, PARAM_TH)
                << "J(X):" << 1.0/rcond
                << "av:" << gsl_multifit_nlinear_avratio(w);
}

Star::Star()
{
    m_am = m_x = m_y = m_sx = m_sy = NAN;
}

bool Star::valid() const
{
    return !isnan(m_am);
}

//half width at half maximum = sqrt(2*ln(2))
double Star::hwhmX() const
{
    return 1.177410023*m_sx;
}

double Star::hwhmY() const
{
    return 1.177410023*m_sy;
}

// half width at 1/20 maximum
double Star::hw20X() const
{
    return 2.447746831*m_sx;
}

double Star::hw20Y() const
{
    return 2.447746831*m_sy;
}

// full width at half maximum
double Star::fwhmX() const
{
    return 2.354820045*m_sx;
}

double Star::fwhmY() const
{
    return 2.354820045*m_sy;
}

bool Star::operator<(const Star &d) const
{
    return m_am < d.m_am;
}

StarFit::StarFit(int size)
{
    m_size = size;
    m_fdf_params = gsl_multifit_nlinear_default_parameters();
    m_fdf_params.trs = gsl_multifit_nlinear_trs_lmaccel;

    m_fdf.f = func_f;
    m_fdf.df = func_df;
    m_fdf.fvv = nullptr;
    m_fdf.n = size*size;
    m_fdf.p = 5;//number of model parameters amplitude, x, y, fwhm_x, fwhm_y

    m_fdf_an.f = func_f_an;
    m_fdf_an.df = nullptr;
    m_fdf_an.fvv = nullptr;
    m_fdf_an.n = size*size;
    m_fdf_an.p = 6;//number of model parameters amplitude, x, y, sigma_x, sigma_y, angle
}

StarFit::~StarFit()
{
}

Star StarFit::fitStar(const std::vector<double> &data, bool angle)
{
    gsl_multifit_nlinear_fdf *fdf = angle ? &m_fdf_an : &m_fdf;
    Star star;
    StarData d;
    d.val = data;
    d.size = m_size;
    d.val = data;
    fdf->params = &d;
    int info;

    double min = *std::min_element(data.begin(), data.end());
    double max = *std::max_element(data.begin(), data.end()) - min;
    for(double &v : d.val)
    {
        v -= min;
    }

    gsl_vector *start = gsl_vector_alloc(fdf->p);
    gsl_vector_set(start, PARAM_AM, max);
    gsl_vector_set(start, PARAM_X0, m_size/2);
    gsl_vector_set(start, PARAM_Y0, m_size/2);
    gsl_vector_set(start, PARAM_SX, 1.0);
    gsl_vector_set(start, PARAM_SY, 1.0);
    if(angle)
        gsl_vector_set(start, PARAM_TH, 0.0);

    gsl_multifit_nlinear_workspace *workspace = gsl_multifit_nlinear_alloc(gsl_multifit_nlinear_trust, &m_fdf_params, fdf->n, fdf->p);

    gsl_multifit_nlinear_init(start, fdf, workspace);
    gsl_vector *f = gsl_multifit_nlinear_residual(workspace);

    double cost, cost0;
    gsl_blas_ddot(f, f, &cost0);
    int ret = gsl_multifit_nlinear_driver(MAX_ITER, TOL, TOL, TOL, nullptr, nullptr, &info, workspace);

    gsl_blas_ddot(f, f, &cost);

    if(ret==0)
    {
        gsl_vector *y = gsl_multifit_nlinear_position(workspace);
        star.m_am = gsl_vector_get(y, PARAM_AM);
        star.m_x = gsl_vector_get(y, PARAM_X0);
        star.m_y = gsl_vector_get(y, PARAM_Y0);
        star.m_sx = gsl_vector_get(y, PARAM_SX);
        star.m_sy = gsl_vector_get(y, PARAM_SY);
        if(angle)
            star.m_theta = gsl_vector_get(y, PARAM_TH);
        //qDebug() << "finished" << star.m_am << star.m_sx << star.m_sy;
    }

    gsl_vector_free(start);
    gsl_multifit_nlinear_free(workspace);

    return star;
}
