/*
 * This file is part of the GIRAFFE Pipeline
 * Copyright (C) 2002-2019 European Southern Observatory
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

#ifdef HAVE_CONFIG_H
#  include <config.h>
#endif

#include <math.h>
#include <string.h>

#include <cxtypes.h>
#include <cxmemory.h>
#include <cxmessages.h>
#include <cxstrutils.h>

#include <cpl_error.h>

#include "gierror.h"
#include "gimodels.h"


/**
 * @defgroup gimodel Generic Model Interface
 *
 * TBD
 */

/**@{*/

inline static void
_giraffe_model_set_flag(GiModel *self, cxint idx, cxbool value)
{

    cx_assert(self != NULL);

    if (self->parameters.flags == NULL) {
        self->parameters.flags = cx_calloc(self->parameters.count,
                                           sizeof(cxint));
    }

    if (value == TRUE) {
        if (self->parameters.flags[idx] == 0) {
            self->parameters.flags[idx] = 1;
            self->fit.nfree += 1;
        }
    }
    else {
        if (self->parameters.flags[idx] == 1) {
            self->parameters.flags[idx] = 0;
            self->fit.nfree -= 1;
        }
    }

    return;

}


inline static cxbool
_giraffe_model_get_flag(const GiModel *self, cxint idx)
{

    cx_assert(self != NULL);

    if (self->parameters.flags == NULL) {
        return FALSE;
    }

    return self->parameters.flags[idx] == 0 ? FALSE : TRUE;

}


inline static cxdouble
_giraffe_compute_rsquare(cxdouble rss, cpl_matrix *y, cxint n)
{

    register cxint i;

    register cxdouble my = 0.;
    register cxdouble sy = 0.;
    register cxdouble ss = 0.;

    cxdouble r = 0.;
    cxdouble *_y = cpl_matrix_get_data(y);


    if (n < 1) {
        return 0.;
    }

    for (i = 0; i < n; i++) {
        my  += _y[i];
    }
    my /= n;


    for (i = 0; i < n; i++) {
        sy = _y[i] - my;
        ss += sy * sy;
    }

    r = rss / ss;

    if (isnan(r)) {
        return 0.;
    }

    return 1. - r;

}


inline static cxint
_giraffe_model_fit(GiModel *self, cpl_matrix *x, cpl_matrix *y,
                   cpl_matrix *sigma, cxint ndata, cxint start,
                   cxint stride)
{

    cxint status = 0;

    cxdouble _chisq = 0.;

    GiFitParams setup;


    if ((cpl_matrix_get_nrow(x) != cpl_matrix_get_nrow(y)) ||
        (cpl_matrix_get_nrow(x) != cpl_matrix_get_nrow(sigma))) {
            return -128;
        }

    if (cpl_matrix_get_ncol(x) != self->arguments.count) {
        return -128;
    }


    /*
     * Check data selection range
     */

    if (cpl_matrix_get_nrow(y) <= start + stride * (ndata - 1)) {
        return -255;
    }


    /*
     * Setup the fit
     */

    setup.iterations = self->fit.setup.iterations;
    setup.tests = self->fit.setup.tests;
    setup.dchisq = self->fit.setup.delta;

    if (self->fit.covariance != NULL) {
        cpl_matrix_set_size(self->fit.covariance, self->parameters.count,
                            self->parameters.count);
        cpl_matrix_fill(self->fit.covariance, 0.);
    }
    else {
        self->fit.covariance = cpl_matrix_new(self->parameters.count,
                                              self->parameters.count);
    }


    /*
     * Fit the data
     */

    giraffe_error_push();

    status = giraffe_nlfit(x, y, sigma, ndata, self->parameters.values,
                           self->parameters.limits, self->parameters.flags,
                           self->parameters.count, self->fit.covariance,
                           &_chisq, self->model, &setup);

    if (status < 0) {

        if (cpl_error_get_code() == CPL_ERROR_NONE) {
            giraffe_error_pop();
        }

        return status;

    }

    if (cpl_error_get_code() != CPL_ERROR_NONE) {
        return -255;
    }

    giraffe_error_pop();

    self->fit.df = ndata - self->fit.nfree;
    self->fit.iterations = status;
    self->fit.chisq = _chisq;
    self->fit.rsquare = _giraffe_compute_rsquare(self->fit.chisq, y, ndata);

    return 0;

}


GiModel *
giraffe_model_new(const cxchar *name)
{

    register cxint i = 0;

    GiModel *self = NULL;


    if (name == NULL) {
        return NULL;
    }

    while (giraffe_models[i].name != NULL) {

        if (strcmp(name, giraffe_models[i].name) == 0) {

            self = cx_calloc(1, sizeof(GiModel));

            giraffe_error_push();

            giraffe_models[i].ctor(self, &giraffe_models[i]);

            if (cpl_error_get_code() != CPL_ERROR_NONE) {

                giraffe_model_delete(self);
                self = NULL;

                break;

            }

            break;

        }

        ++i;

    }

    self->fit.setup.iterations = 0;
    self->fit.setup.tests = 0;
    self->fit.setup.delta = 0.;

    self->fit.iterations = 0;
    self->fit.nfree = 0;
    self->fit.df = 0;
    self->fit.covariance = NULL;

    return self;

}


GiModel *
giraffe_model_clone(const GiModel *other)
{

    GiModel *self = NULL;


    if (other != NULL) {

        self = giraffe_model_new(other->name);

        /*
         * The model constructor creates already an a property list and
         * a matrix as container for the argument and parameter names and
         * values. We get rid of them here before we create the copies.
         * This way is simpler than copying each argument and parameter
         * individually.
         */

        cpl_propertylist_delete(self->arguments.names);
        self->arguments.names =
            cpl_propertylist_duplicate(other->arguments.names);

        cpl_matrix_delete(self->arguments.values);
        self->arguments.values =
            cpl_matrix_duplicate(other->arguments.values);

        self->arguments.count = other->arguments.count;

        cx_assert(cpl_propertylist_get_size(self->arguments.names) ==
                  self->arguments.count);
        cx_assert(cpl_matrix_get_nrow(self->arguments.values) *
                  cpl_matrix_get_ncol(self->arguments.values) ==
                  self->arguments.count);


        cpl_propertylist_delete(self->parameters.names);
        self->parameters.names =
            cpl_propertylist_duplicate(other->parameters.names);

        cpl_matrix_delete(self->parameters.values);
        self->parameters.values =
            cpl_matrix_duplicate(other->parameters.values);

        self->parameters.count = other->parameters.count;

        cx_assert(cpl_propertylist_get_size(self->parameters.names) ==
                  self->parameters.count);
        cx_assert(cpl_matrix_get_nrow(self->parameters.values) *
                  cpl_matrix_get_ncol(self->parameters.values) ==
                  self->parameters.count);

        self->fit.setup = other->fit.setup;
        self->fit.iterations = other->fit.iterations;
        self->fit.nfree = other->fit.nfree;
        self->fit.df = other->fit.df;

        if (other->fit.covariance == NULL) {
            self->fit.covariance = NULL;
        }
        else {
            self->fit.covariance =
                cpl_matrix_duplicate(other->fit.covariance);
        }

    }

    return self;

}

void
giraffe_model_delete(GiModel *self)
{

    if (self) {

        register cxint i = 0;

        while (giraffe_models[i].name != NULL) {

            if (strcmp(self->name, giraffe_models[i].name) == 0) {
                giraffe_models[i].dtor(self);
                cx_free(self);

                break;
            }

            ++i;

        }

    }

    return;

}


const cxchar *
giraffe_model_get_name(const GiModel *self)
{

    cx_assert(self != NULL);
    return self->name;

}


GiModelType
giraffe_model_get_type(const GiModel *self)
{

    cx_assert(self != NULL);
    return self->type;

}


cxsize
giraffe_model_count_arguments(const GiModel *self)
{

    cx_assert(self != NULL);
    return self->arguments.count;

}


cxsize
giraffe_model_count_parameters(const GiModel *self)
{

    cx_assert(self != NULL);
    return self->parameters.count;

}


const cxchar*
giraffe_model_argument_name(const GiModel* self, cxsize position)
{

    const cpl_property* p = NULL;


    cx_assert(self != NULL);

    p = cpl_propertylist_get(self->arguments.names, position);
    if (p == NULL) {
        return NULL;
    }

    return cpl_property_get_name(p);

}


const cxchar*
giraffe_model_parameter_name(const GiModel* self, cxsize position)
{

    const cpl_property* p = NULL;


    cx_assert(self != NULL);

    p = cpl_propertylist_get(self->parameters.names, position);
    if (p == NULL) {
        return NULL;
    }

    return cpl_property_get_name(p);

}


cxint
giraffe_model_set_argument(GiModel *self, const cxchar *name, cxdouble value)
{

    const cxchar *const fctid = "giraffe_model_set_argument";


    cx_assert(self != NULL);

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return 1;
    }

    if (!cpl_propertylist_has(self->arguments.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return 1;
    }
    else {

        register cxint idx = cpl_propertylist_get_int(self->arguments.names,
                                                      name);

        cpl_matrix_set(self->arguments.values, idx, 0, value);

    }

    return 0;

}


cxdouble
giraffe_model_get_argument(const GiModel *self, const cxchar *name)
{

    const cxchar *const fctid = "giraffe_model_get_argument";

    register cxint idx;


    cx_assert(self != NULL);

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return 0.;
    }

    if (!cpl_propertylist_has(self->arguments.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return 0.;
    }


    idx = cpl_propertylist_get_int(self->arguments.names, name);

    return cpl_matrix_get(self->arguments.values, idx, 0);

}


cxint
giraffe_model_set_parameter(GiModel *self, const cxchar *name,
                            cxdouble value)
{

    const cxchar *const fctid = "giraffe_model_set_parameter";


    cx_assert(self != NULL);

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return 1;
    }

    if (!cpl_propertylist_has(self->parameters.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return 1;
    }
    else {

        register cxint idx = cpl_propertylist_get_int(self->parameters.names,
                                                      name);

        cpl_matrix_set(self->parameters.values, idx, 0, value);

    }

    return 0;

}


cxdouble
giraffe_model_get_parameter(const GiModel *self, const cxchar *name)
{

    const cxchar *const fctid = "giraffe_model_get_parameter";

    register cxint idx;


    cx_assert(self != NULL);

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return 0.;
    }

    if (!cpl_propertylist_has(self->parameters.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return 0.;
    }


    idx = cpl_propertylist_get_int(self->parameters.names, name);

    return cpl_matrix_get(self->parameters.values, idx, 0);

}


cxint
giraffe_model_freeze_parameter(GiModel *self, const cxchar *name)
{

    const cxchar *const fctid = "giraffe_model_freeze_parameter";


    if (self == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return 1;
    }

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return 1;
    }

    if (!cpl_propertylist_has(self->parameters.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return 1;
    }
    else {

        register cxint idx = cpl_propertylist_get_int(self->parameters.names,
                                                      name);

        _giraffe_model_set_flag(self, idx, FALSE);

    }

    return 0;

}


cxint
giraffe_model_thaw_parameter(GiModel *self, const cxchar *name)
{

    const cxchar *const fctid = "giraffe_model_thaw_parameter";


    cx_assert(self != NULL);

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return 1;
    }

    if (!cpl_propertylist_has(self->parameters.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return 1;
    }
    else {

        register cxint idx = cpl_propertylist_get_int(self->parameters.names,
                                                      name);

        _giraffe_model_set_flag(self, idx, TRUE);

    }

    return 0;

}


cxbool
giraffe_model_frozen_parameter(const GiModel *self, const cxchar *name)
{

    const cxchar *const fctid = "giraffe_model_frozen_parameter";

    register cxint idx;


    cx_assert(self != NULL);

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return FALSE;
    }

    if (!cpl_propertylist_has(self->parameters.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return FALSE;
    }


    idx = cpl_propertylist_get_int(self->parameters.names, name);

    return _giraffe_model_get_flag(self, idx) == FALSE;

}


cxint
giraffe_model_freeze(GiModel *self)
{

    register cxint i;


    cx_assert(self != NULL);

    for (i = 0; i < cpl_propertylist_get_size(self->parameters.names); i++) {

        const cpl_property *p = cpl_propertylist_get(self->parameters.names,
                                                     i);

        if (p == NULL) {
            return 1;
        }

        _giraffe_model_set_flag(self, cpl_property_get_int(p), FALSE);

    }

    return 0;

}


cxint
giraffe_model_thaw(GiModel *self)
{

    register cxint i;


    cx_assert(self != NULL);

    for (i = 0; i < cpl_propertylist_get_size(self->parameters.names); i++) {

        const cpl_property *p = cpl_propertylist_get(self->parameters.names,
                                                     i);

        if (p == NULL) {
            return 1;
        }

        _giraffe_model_set_flag(self, cpl_property_get_int(p), TRUE);

    }

    return 0;

}


cxint
giraffe_model_evaluate(const GiModel *self, cxdouble *result, cxint *status)
{

    const cxchar *const fctid = "giraffe_model_evaluate";

    cxdouble _result = 0.;
    cxdouble *arg = NULL;
    cxdouble *par = NULL;


    cx_assert(self != NULL);

    arg = cpl_matrix_get_data(self->arguments.values);

    if (arg == NULL) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return 2;
    }

    par = cpl_matrix_get_data(self->parameters.values);

    if (par == NULL) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return 3;
    }

    giraffe_error_push();

    self->model(&_result, arg, par, self->parameters.count, NULL, NULL);

    if (cpl_error_get_code() != CPL_ERROR_NONE) {

        if (status) {
            *status = 1;
        }

        return 4;

    }

    giraffe_error_pop();

    *result = _result;
    *status = 0;

    return 0;

}


cxint
giraffe_model_set_iterations(GiModel *self, cxint iterations)
{

    cx_assert(self != NULL);

    if (iterations < 1) {
        return 1;
    }

    self->fit.setup.iterations = iterations;

    return 0;

}


cxint
giraffe_model_get_iterations(const GiModel *self)
{

    cx_assert(self != NULL);

    return self->fit.setup.iterations;

}


cxint
giraffe_model_set_tests(GiModel *self, cxint tests)
{

    cx_assert(self != NULL);

    if (tests < 1) {
        return 1;
    }

    self->fit.setup.tests = tests;

    return 0;

}


cxint
giraffe_model_get_tests(const GiModel *self)
{

    cx_assert(self != NULL);

    return self->fit.setup.tests;

}


cxint
giraffe_model_set_delta(GiModel *self, cxdouble delta)
{

    cx_assert(self != NULL);

    if (delta < 0.) {
        return 1;
    }

    self->fit.setup.delta = delta;

    return 0;

}


cxdouble
giraffe_model_get_delta(const GiModel *self)
{

    cx_assert(self != NULL);

    return self->fit.setup.delta;

}


cxint
giraffe_model_get_position(const GiModel *self)
{

    cx_assert(self != NULL);

    if (self->fit.iterations <= 0) {
        return -1;
    }

    return self->fit.iterations;

}


cxdouble
giraffe_model_get_variance(const GiModel *self, const cxchar *name)
{

    const cxchar *const fctid = "giraffe_model_get_variance";


    cxdouble variance = 0.;


    cx_assert(self != NULL);

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return variance;
    }

    if (!cpl_propertylist_has(self->parameters.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return variance;
    }
    else {

        if (self->fit.covariance == NULL) {
            cpl_error_set(fctid, CPL_ERROR_DATA_NOT_FOUND);
            return variance;
        }
        else {

            register cxint idx =
                cpl_propertylist_get_int(self->parameters.names, name);

            variance = cpl_matrix_get(self->fit.covariance, idx, idx);

        }

    }

    return variance;

}


cxdouble
giraffe_model_get_sigma(const GiModel *self, const cxchar *name)
{

    const cxchar *const fctid = "giraffe_model_get_sigma";


    cxdouble sigma = 0.;


    cx_assert(self != NULL);

    if (name == NULL) {
        cpl_error_set(fctid, CPL_ERROR_NULL_INPUT);
        return sigma;
    }

    if (!cpl_propertylist_has(self->parameters.names, name)) {
        cpl_error_set(fctid, CPL_ERROR_ILLEGAL_INPUT);
        return sigma;
    }
    else {

        if (self->fit.covariance == NULL) {
            cpl_error_set(fctid, CPL_ERROR_DATA_NOT_FOUND);
            return sigma;
        }
        else {

            register cxint idx =
                cpl_propertylist_get_int(self->parameters.names, name);

            sigma = cpl_matrix_get(self->fit.covariance, idx, idx);

            if (isnan(sigma) || sigma < 0.) {
                sigma = 0.;
            }
            else {
                sigma = sqrt(sigma);
            }

        }

    }

    return sigma;

}


cxint
giraffe_model_get_df(const GiModel *self)
{

    cx_assert(self != NULL);

    return self->fit.df;

}


cxdouble
giraffe_model_get_chisq(const GiModel *self)
{

    cx_assert(self != NULL);

    return self->fit.chisq;

}


cxdouble
giraffe_model_get_rsquare(const GiModel *self)
{

    cx_assert(self != NULL);

    return self->fit.rsquare;

}


cxint
giraffe_model_fit(GiModel *self, cpl_matrix *x, cpl_matrix *y,
                  cpl_matrix *sigma)
{

    cxint ndata = 0;


    cx_assert(self != NULL);

    if (x == NULL || y == NULL) {
        return -128;
    }

    if (sigma == NULL) {
        return -128;
    }

    ndata = cpl_matrix_get_nrow(y);

    return _giraffe_model_fit(self, x, y, sigma, ndata, 0, 1);

}


cxint
giraffe_model_fit_sequence(GiModel *self, cpl_matrix *x, cpl_matrix *y,
                           cpl_matrix *sigma, cxint ndata, cxint start,
                           cxint stride)
{

    cx_assert(self != NULL);

    /* FIXME: Currently only (start == 0 && stride == 1) is supported! */
    cx_assert((start == 0) || (stride == 1));


    if (x == NULL || y == NULL) {
        return -128;
    }

    if (sigma == NULL) {
        return -128;
    }

    if ((start < 0) || (stride < 0)) {
        return -128;
    }

    return _giraffe_model_fit(self, x, y, sigma, ndata, 0, 1);

}
/**@}*/
