// Copyright (c) 2021 Stig Rune Sellevag
//
// This file is distributed under the MIT License. See the accompanying file
// LICENSE.txt or http://www.opensource.org/licenses/mit-license.php for terms
// and conditions.

#ifndef SCILIB_LINALG_MATRIX_NORM_H
#define SCILIB_LINALG_MATRIX_NORM_H

#ifdef USE_MKL
#include <mkl.h>
#else
#include <lapacke.h>
#endif

#include <scilib/mdarray.h>
#include <scilib/linalg_impl/lapack_types.h>
#include <type_traits>

namespace Sci {
namespace Linalg {

namespace stdex = std::experimental;

// Matrix norm of a general rectangular matrix:
//
// Types of matrix norms:
// - M, m:       largest absolute value of the matrix
// - 1, O, o:    1-norm of the matrix (maximum column sum)
// - I, i:       infinity norm of the matrix (maximum row sum)
// - F, f, E, e: Frobenius norm of the matrix (square root of sum of squares)
//
template <class T, class Layout>
    requires(std::is_same_v<std::remove_cv_t<T>, double>)
inline auto matrix_norm(Sci::Matrix_view<T, Layout> a, char norm)
{
    static_assert(a.is_contiguous());

    assert(norm == 'M' || norm == 'm' || norm == '1' || norm == 'O' ||
           norm == 'o' || norm == 'I' || norm == 'i' || norm == 'F' ||
           norm == 'f' || norm == 'E' || norm == 'e');

    auto matrix_layout = LAPACK_ROW_MAJOR;
    BLAS_INT m = static_cast<BLAS_INT>(a.extent(0));
    BLAS_INT n = static_cast<BLAS_INT>(a.extent(1));
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = LAPACK_COL_MAJOR;
        lda = m;
    }
    return LAPACKE_dlange(matrix_layout, norm, m, n, a.data(), lda);
}

template <class Layout, class Allocator>
inline double matrix_norm(const Sci::Matrix<double, Layout, Allocator>& a,
                          char norm)
{
    return matrix_norm(a.view(), norm);
}

} // namespace Linalg
} // namespace Sci

#endif // SCILIB_LINALG_MATRIX_NORM_H


/* permutation/permutation.c
 *
 * Copyright (C) 2001, 2002 Nicolas Darnis
 *
 * 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 Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/* Modified for GSL by Brian Gough.
 * Use in-place algorithms, no need for workspace
 * Use conventions for canonical form given in Knuth (opposite of Sedgewick)
 */

#include <config.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_permutation.h>

int
gsl_permutation_linear_to_canonical (gsl_permutation * q,
                                     const gsl_permutation * p)
{
  const size_t n = p->size;
  size_t i, k, s;
  size_t t = n;

  const size_t *const pp = p->data;
  size_t *const qq = q->data;

  if (q->size != p->size)
    {
      GSL_ERROR ("size of q does not match size of p", GSL_EINVAL);
    }

  for (i = 0; i < n; i++)
    {

      k = pp[i];
      s = 1;

      while (k > i)
        {
          k = pp[k];
          s++;
        }

      if (k < i)
        continue;

      /* Now have k == i, i.e the least in its cycle, and s == cycle length */

      t -= s;

      qq[t] = i;

      k = pp[i];
      s = 1;

      while (k > i)
        {
          qq[t + s] = k;
          k = pp[k];
          s++;
        }

      if (t == 0)
        break;
    }

  return GSL_SUCCESS;
}

int
gsl_permutation_canonical_to_linear (gsl_permutation * p,
                                     const gsl_permutation * q)
{
  size_t i, k, kk, first;
  const size_t n = p->size;

  size_t *const pp = p->data;
  const size_t *const qq = q->data;

  if (q->size != p->size)
    {
      GSL_ERROR ("size of q does not match size of p", GSL_EINVAL);
    }

  for (i = 0; i < n; i++)
    {
      pp[i] = i;
    }

  k = qq[0];
  first = pp[k];

  for (i = 1; i < n; i++)
    {
      kk = qq[i];

      if (kk > first)
        {
          pp[k] = pp[kk];
          k = kk;
        }
      else
        {
          pp[k] = first;
          k = kk;
          first = pp[kk];
        }
    }

  pp[k] = first;

  return GSL_SUCCESS;
}

size_t
gsl_permutation_inversions (const gsl_permutation * p)
{
  size_t count = 0;
  size_t i, j;
  const size_t size = p->size;

  for (i = 0; i < size - 1; i++)
    {
      for (j = i + 1; j < size; j++)
        {
          if (p->data[i] > p->data[j])
            {
              count++;
            }
        }
    }

  return count;
}

size_t
gsl_permutation_linear_cycles (const gsl_permutation * p)
{
  size_t i, k;
  size_t count = 0;
  const size_t size = p->size;

  for (i = 0; i < size; i++)
    {

      k = p->data[i];

      while (k > i)
        {
          k = p->data[k];
        }

      if (k < i)
        continue;

      count++;
    }

  return count;
}

size_t
gsl_permutation_canonical_cycles (const gsl_permutation * p)
{
  size_t i;
  size_t count = 1;
  size_t min = p->data[0];

  for (i = 0; i < p->size; i++)
    {
      if (p->data[i] < min)
        {
          min = p->data[i];
          count++;
        }
    }

  return count;
}


/*
 # This file is part of the Astrometry.net suite.
 # Licensed under a 3-clause BSD style license - see LICENSE
 */

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

#include <gsl/gsl_matrix_double.h>
#include <gsl/gsl_vector_double.h>

#include "sip-utils.h"
#include "tweak.h"
#include "healpix.h"
#include "dualtree_rangesearch.h"
#include "kdtree_fits_io.h"
#include "mathutil.h"
#include "log.h"
#include "permutedsort.h"
#include "gslutils.h"
#include "errors.h"
#include "fit-wcs.h"

// TODO:
//
//  1. Write the document which explains every step of tweak in detail with
//     comments relating exactly to the code.
//  2. Implement polynomial terms - use order parameter to zero out A matrix
//  3. Make it robust to outliers
//     - Make jitter evolve as fit evolves
//     - Sigma clipping
//  4. Need test image with non-trivial rotation to test CD transpose problem
//
//
//  - Put CD inverse into its own function
//
//  BUG? transpose of CD matrix is similar to CD matrix!
//  BUG? inverse when computing sx/sy (i.e. same transpose issue)
//  Ability to fit without re-doing correspondences
//  Split fit x/y (i.e. two fits one for x one for y)

#define KERNEL_SIZE 5
#define KERNEL_MARG ((KERNEL_SIZE-1)/2)

sip_t* tweak_just_do_it(const tan_t* wcs, const starxy_t* imagexy,
                        const double* starxyz,
                        const double* star_ra, const double* star_dec,
                        const double* star_radec,
                        int nstars, double jitter_arcsec,
                        int order, int inverse_order, int iterations,
                        anbool weighted, anbool skip_shift) {
    tweak_t* twee = NULL;
    sip_t* sip = NULL;

    twee = tweak_new();
    twee->jitter = jitter_arcsec;
    twee->sip->a_order  = twee->sip->b_order  = order;
    twee->sip->ap_order = twee->sip->bp_order = inverse_order;
    twee->weighted_fit = weighted;
    if (skip_shift)
        tweak_skip_shift(twee);

    tweak_push_image_xy(twee, imagexy);
    if (starxyz)
        tweak_push_ref_xyz(twee, starxyz, nstars);
    else if (star_ra && star_dec)
        tweak_push_ref_ad(twee, star_ra, star_dec, nstars);
    else if (star_radec)
        tweak_push_ref_ad_array(twee, star_radec, nstars);
    else {
        logerr("Need starxyz, (star_ra and star_dec), or star_radec");
        return NULL;
    }
    tweak_push_wcs_tan(twee, wcs);
    tweak_iterate_to_order(twee, order, iterations);

    // Steal the resulting SIP structure
    sip = twee->sip;
    twee->sip = NULL;

    tweak_free(twee);
    return sip;
}

static void get_dydx_range(double* ximg, double* yimg, int nimg,
                           double* xcat, double* ycat, int ncat,
                           double *mindx, double *mindy,
                           double *maxdx, double *maxdy) {
    int i, j;
    *maxdx = -1e100;
    *mindx = 1e100;
    *maxdy = -1e100;
    *mindy = 1e100;

    for (i = 0; i < nimg; i++) {
        for (j = 0; j < ncat; j++) {
            double dx = ximg[i] - xcat[j];
            double dy = yimg[i] - ycat[j];
            *maxdx = MAX(dx, *maxdx);
            *maxdy = MAX(dy, *maxdy);
            *mindx = MIN(dx, *mindx);
            *mindy = MIN(dy, *mindy);
        }
    }
}

static void get_shift(double* ximg, double* yimg, int nimg,
                      double* xcat, double* ycat, int ncat,
                      double mindx, double mindy, double maxdx, double maxdy,
                      double* xshift, double* yshift) {
    int i, j;
    int themax, themaxind, ys, xs;

    // hough transform
    int hsz = 1000; // hough histogram size (per side)
    int *hough = calloc(hsz * hsz, sizeof(int)); // allocate bins
    int kern[] = {0, 2, 3, 2, 0,  // approximate gaussian smoother
                  2, 7, 12, 7, 2,       // should be KERNEL_SIZE x KERNEL_SIZE
                  3, 12, 20, 12, 3,
                  2, 7, 12, 7, 2,
                  0, 2, 3, 2, 0};

    assert(sizeof(kern) == KERNEL_SIZE * KERNEL_SIZE * sizeof(int));

    for (i = 0; i < nimg; i++) {    // loop over all pairs of source-catalog objs
        for (j = 0; j < ncat; j++) {
            double dx = ximg[i] - xcat[j];
            double dy = yimg[i] - ycat[j];
            int hszi = hsz - 1;
            int iy = hszi * ( (dy - mindy) / (maxdy - mindy) ); // compute deltay using implicit floor
            int ix = hszi * ( (dx - mindx) / (maxdx - mindx) ); // compute deltax using implicit floor

            // check to make sure the point is in the box
            if (KERNEL_MARG <= iy && iy < hsz - KERNEL_MARG &&
                KERNEL_MARG <= ix && ix < hsz - KERNEL_MARG) {
                int kx, ky;
                for (ky = -2; ky <= 2; ky++)
                    for (kx = -2; kx <= 2; kx++)
                        hough[(iy - ky)*hsz + (ix - kx)] += kern[(ky + 2) * 5 + (kx + 2)];
            }
        }
    }

    // find argmax in hough
    themax = 0;
    themaxind = -1;
    for (i = 0; i < hsz*hsz; i++) {
        if (themax < hough[i]) {
            themaxind = i;
            themax = hough[i];
        }
    }
    // which hough bin is the max?
    ys = themaxind / hsz;
    xs = themaxind % hsz;

    *yshift = (ys / (double)hsz) * (maxdy - mindy) + mindy;
    *xshift = (xs / (double)hsz) * (maxdx - mindx) + mindx;
    debug("xs = %d, ys = %d\n", xs, ys);
    debug("get_shift: mindx=%g, maxdx=%g, mindy=%g, maxdy=%g\n", mindx, maxdx, mindy, maxdy);
    debug("get_shift: xs=%g, ys=%g\n", *xshift, *yshift);

    free(hough);
}

static sip_t* do_entire_shift_operation(tweak_t* t, double rho) {
    get_shift(t->x, t->y, t->n,
              t->x_ref, t->y_ref, t->n_ref,
              rho*t->mindx, rho*t->mindy, rho*t->maxdx, rho*t->maxdy,
              &t->xs, &t->ys);
    wcs_shift(&(t->sip->wcstan), t->xs, t->ys);
    return NULL;
}

/* This function is intended only for initializing newly allocated tweak
 * structures, NOT for operating on existing ones.*/
void tweak_init(tweak_t* t) {
    memset(t, 0, sizeof(tweak_t));
    t->sip = sip_create();
}

tweak_t* tweak_new() {
    tweak_t* t = malloc(sizeof(tweak_t));
    tweak_init(t);
    return t;
}

void tweak_iterate_to_order(tweak_t* t, int maxorder, int iterations) {
    int order;
    int k;

    for (order=1; order<=maxorder; order++) {
        logverb("\n");
        logverb("--------------------------------\n");
        logverb("Order %i\n", order);
        logverb("--------------------------------\n");

        t->sip->a_order  = t->sip->b_order  = order;
        //t->sip->ap_order = t->sip->bp_order = order;
        tweak_go_to(t, TWEAK_HAS_CORRESPONDENCES);

        for (k=0; k<iterations; k++) {
            logverb("\n");
            logverb("--------------------------------\n");
            logverb("Iterating tweak: order %i, step %i\n", order, k);
            t->state &= ~TWEAK_HAS_LINEAR_CD;
            tweak_go_to(t, TWEAK_HAS_LINEAR_CD);
            tweak_clear_correspondences(t);
        }
    }
}

#define CHECK_STATE(x) if (state & x) sl_append(s, #x)

static char* state_string(unsigned int state) {
    sl* s = sl_new(4);
    char* str;
    CHECK_STATE(TWEAK_HAS_SIP);
    CHECK_STATE(TWEAK_HAS_IMAGE_XY);
    CHECK_STATE(TWEAK_HAS_IMAGE_XYZ);
    CHECK_STATE(TWEAK_HAS_IMAGE_AD);
    CHECK_STATE(TWEAK_HAS_REF_XY);
    CHECK_STATE(TWEAK_HAS_REF_XYZ);
    CHECK_STATE(TWEAK_HAS_REF_AD);
    CHECK_STATE(TWEAK_HAS_CORRESPONDENCES);
    CHECK_STATE(TWEAK_HAS_COARSLY_SHIFTED);
    CHECK_STATE(TWEAK_HAS_FINELY_SHIFTED);
    CHECK_STATE(TWEAK_HAS_REALLY_FINELY_SHIFTED);
    CHECK_STATE(TWEAK_HAS_LINEAR_CD);
    str = sl_join(s, " ");
    sl_free2(s);
    return str;
}

char* tweak_get_state_string(const tweak_t* t) {
    return state_string(t->state);
}

#undef CHECK_STATE

void tweak_clear_correspondences(tweak_t* t) {
    if (t->state & TWEAK_HAS_CORRESPONDENCES) {
        // our correspondences are also now toast
        assert(t->image);
        assert(t->ref);
        assert(t->dist2);
        il_free(t->image);
        il_free(t->ref);
        dl_free(t->dist2);
        if (t->weight)
            dl_free(t->weight);
        t->image    = NULL;
        t->ref      = NULL;
        t->dist2    = NULL;
        t->weight   = NULL;
        t->state &= ~TWEAK_HAS_CORRESPONDENCES;
    }
    assert(!t->image);
    assert(!t->ref);
    assert(!t->dist2);
    assert(!t->weight);
}

void tweak_clear_on_sip_change(tweak_t* t) {
    // tweak_clear_correspondences(t);
    tweak_clear_image_ad(t);
    tweak_clear_ref_xy(t);
    tweak_clear_image_xyz(t);
}

// ref_xy are the catalog star positions in image coordinates
void tweak_clear_ref_xy(tweak_t* t) {
    if (t->state & TWEAK_HAS_REF_XY) {
        //assert(t->x_ref);
        free(t->x_ref);
        //assert(t->y_ref);
        t->x_ref = NULL;
        free(t->y_ref);
        t->y_ref = NULL;
        t->state &= ~TWEAK_HAS_REF_XY;
    }
    assert(!t->x_ref);
    assert(!t->y_ref);
}

// radec of catalog stars
void tweak_clear_ref_ad(tweak_t* t) {
    if (t->state & TWEAK_HAS_REF_AD) {
        assert(t->a_ref);
        free(t->a_ref);
        t->a_ref = NULL;
        assert(t->d_ref);
        free(t->d_ref);
        t->d_ref = NULL;
        t->n_ref = 0;
        tweak_clear_correspondences(t);
        tweak_clear_ref_xy(t);
        t->state &= ~TWEAK_HAS_REF_AD;
    }
    assert(!t->a_ref);
    assert(!t->d_ref);
}

// image objs in ra,dec according to current tweak
void tweak_clear_image_ad(tweak_t* t) {
    if (t->state & TWEAK_HAS_IMAGE_AD) {
        assert(t->a);
        free(t->a);
        t->a = NULL;
        assert(t->d);
        free(t->d);
        t->d = NULL;
        t->state &= ~TWEAK_HAS_IMAGE_AD;
    }
    assert(!t->a);
    assert(!t->d);
}

void tweak_clear_image_xyz(tweak_t* t) {
    if (t->state & TWEAK_HAS_IMAGE_XYZ) {
        assert(t->xyz);
        free(t->xyz);
        t->xyz = NULL;
        t->state &= ~TWEAK_HAS_IMAGE_XYZ;
    }
    assert(!t->xyz);
}

void tweak_clear_image_xy(tweak_t* t) {
    if (t->state & TWEAK_HAS_IMAGE_XY) {
        assert(t->x);
        free(t->x);
        t->x = NULL;
        assert(t->y);
        free(t->y);
        t->y = NULL;
        t->state &= ~TWEAK_HAS_IMAGE_XY;
    }
    assert(!t->x);
    assert(!t->y);
}

// tell us (from outside tweak) where the catalog stars are
void tweak_push_ref_ad(tweak_t* t, const double* a, const double *d, int n) {
    assert(a);
    assert(d);
    assert(n);
    tweak_clear_ref_ad(t);
    assert(!t->a_ref);
    assert(!t->d_ref);
    t->a_ref = malloc(sizeof(double) * n);
    t->d_ref = malloc(sizeof(double) * n);
    memcpy(t->a_ref, a, n*sizeof(double));
    memcpy(t->d_ref, d, n*sizeof(double));
    t->n_ref = n;
    t->state |= TWEAK_HAS_REF_AD;
}

void tweak_push_ref_ad_array(tweak_t* t, const double* ad, int n) {
    int i;
    assert(ad);
    assert(n);
    tweak_clear_ref_ad(t);
    assert(!t->a_ref);
    assert(!t->d_ref);
    t->a_ref = malloc(sizeof(double) * n);
    t->d_ref = malloc(sizeof(double) * n);
    for (i=0; i<n; i++) {
        t->a_ref[i] = ad[2*i + 0];
        t->d_ref[i] = ad[2*i + 1];
    }
    t->n_ref = n;
    t->state |= TWEAK_HAS_REF_AD;
}

static void ref_xyz_from_ad(tweak_t* t) {
    int i;
    assert(t->state & TWEAK_HAS_REF_AD);
    assert(!t->xyz_ref);
    t->xyz_ref = malloc(sizeof(double) * 3 * t->n_ref);
    assert(t->xyz_ref);
    for (i = 0; i < t->n_ref; i++)
        radecdeg2xyzarr(t->a_ref[i], t->d_ref[i], t->xyz_ref + 3 * i);
    t->state |= TWEAK_HAS_REF_XYZ;
}

static void ref_ad_from_xyz(tweak_t* t) {
    int i, n;
    assert(t->state & TWEAK_HAS_REF_XYZ);
    assert(!t->a_ref);
    assert(!t->d_ref);
    n = t->n_ref;
    t->a_ref = malloc(sizeof(double) * n);
    t->d_ref = malloc(sizeof(double) * n);
    assert(t->a_ref);
    assert(t->d_ref);
    for (i=0; i<n; i++)
        xyzarr2radecdeg(t->xyz_ref + 3*i, t->a_ref + i, t->d_ref + i);
    t->state |= TWEAK_HAS_REF_XYZ;
}

// tell us (from outside tweak) where the catalog stars are
void tweak_push_ref_xyz(tweak_t* t, const double* xyz, int n) {
    assert(xyz);
    assert(n);
    tweak_clear_ref_ad(t);
    assert(!t->xyz_ref);
    t->xyz_ref = malloc(sizeof(double) * 3 * n);
    assert(t->xyz_ref);
    memcpy(t->xyz_ref, xyz, 3*n*sizeof(double));
    t->n_ref = n;
    t->state |= TWEAK_HAS_REF_XYZ;
}

void tweak_push_image_xy(tweak_t* t, const starxy_t* xy) {
    tweak_clear_image_xy(t);
    t->x = starxy_copy_x(xy);
    t->y = starxy_copy_y(xy);
    t->n = starxy_n(xy);
    t->state |= TWEAK_HAS_IMAGE_XY;
}

void tweak_skip_shift(tweak_t* t) {
    t->state |= (TWEAK_HAS_COARSLY_SHIFTED | TWEAK_HAS_FINELY_SHIFTED |
                 TWEAK_HAS_REALLY_FINELY_SHIFTED);
}

// DualTree RangeSearch callback. We want to keep track of correspondences.
// Potentially the matching could be many-to-many; we allow this and hope the
// optimizer can take care of it.
static void dtrs_match_callback(void* extra, int image_ind, int ref_ind, double dist2) {
    tweak_t* t = extra;
    image_ind = kdtree_permute(t->kd_image, image_ind);
    ref_ind   = kdtree_permute(t->kd_ref,   ref_ind);
    il_append(t->image, image_ind);
    il_append(t->ref, ref_ind);
    dl_append(t->dist2, dist2);
    if (t->weight)
        dl_append(t->weight, exp(-dist2 / (2.0 * t->jitterd2)));
}

void tweak_push_correspondence_indices(tweak_t* t, il* image, il* ref, dl* distsq, dl* weight) {
    t->image = image;
    t->ref = ref;
    t->dist2 = distsq;
    t->weight = weight;
    t->state |= TWEAK_HAS_CORRESPONDENCES;
}

// The jitter is in radians
static void find_correspondences(tweak_t* t, double jitter) {
    double dist;
    double* data_image = malloc(sizeof(double) * t->n * 3);
    double* data_ref = malloc(sizeof(double) * t->n_ref * 3);

    assert(t->state & TWEAK_HAS_IMAGE_XYZ);
    assert(t->state & TWEAK_HAS_REF_XYZ);
    tweak_clear_correspondences(t);

    memcpy(data_image, t->xyz, 3*t->n*sizeof(double));
    memcpy(data_ref, t->xyz_ref, 3*t->n_ref*sizeof(double));

    t->kd_image = kdtree_build(NULL, data_image, t->n, 3, 4, KDTT_DOUBLE,
                               KD_BUILD_BBOX);

    t->kd_ref = kdtree_build(NULL, data_ref, t->n_ref, 3, 4, KDTT_DOUBLE,
                             KD_BUILD_BBOX);

    // Storage for correspondences
    t->image = il_new(600);
    t->ref = il_new(600);
    t->dist2 = dl_new(600);
    if (t->weighted_fit)
        t->weight = dl_new(600);

    dist = rad2dist(jitter);

    logverb("search radius = %g arcsec\n", rad2arcsec(jitter));

    // Find closest neighbours
    dualtree_rangesearch(t->kd_image, t->kd_ref,
                         RANGESEARCH_NO_LIMIT, dist, FALSE, NULL,
                         dtrs_match_callback, t,
                         NULL, NULL);

    kdtree_free(t->kd_image);
    kdtree_free(t->kd_ref);
    t->kd_image = NULL;
    t->kd_ref = NULL;
    free(data_image);
    free(data_ref);

    logverb("Number of correspondences: %zu\n", il_size(t->image));
}

static double correspondences_rms_arcsec(tweak_t* t, int weighted) {
    double err2 = 0.0;
    int i;
    double totalweight = 0.0;
    for (i=0; i<il_size(t->image); i++) {
        double imgxyz[3];
        double refxyz[3];
        double weight;
        int refi, imgi;
        if (weighted && t->weight)
            weight = dl_get(t->weight, i);
        else
            weight = 1.0;
        totalweight += weight;

        imgi = il_get(t->image, i);
        sip_pixelxy2xyzarr(t->sip, t->x[imgi], t->y[imgi], imgxyz);

        refi = il_get(t->ref, i);
        radecdeg2xyzarr(t->a_ref[refi], t->d_ref[refi], refxyz);

        err2 += weight * distsq(imgxyz, refxyz, 3);
    }
    return distsq2arcsec( err2 / totalweight );
}

#if 0
// in arcseconds^2 on the sky (chi-sq)
static double figure_of_merit(tweak_t* t, double *rmsX, double *rmsY) {
    double sqerr = 0.0;
    int i;
    for (i = 0; i < il_size(t->image); i++) {
        double a, d;
        double xyzpt[3];
        double xyzpt_ref[3];
        sip_pixelxy2radec(t->sip, t->x[il_get(t->image, i)],
                          t->y[il_get(t->image, i)], &a, &d);

        // xref and yref should be intermediate WC's not image x and y!
        radecdeg2xyzarr(a, d, xyzpt);
        radecdeg2xyzarr(t->a_ref[il_get(t->ref, i)],
                        t->d_ref[il_get(t->ref, i)], xyzpt_ref);
        sqerr += distsq(xyzpt, xyzpt_ref, 3);
    }
    return rad2arcsec(1)*rad2arcsec(1)*sqerr;
}

static double figure_of_merit2(tweak_t* t) {
    // find error in pixels^2
    double sqerr = 0.0;
    int i;
    for (i = 0; i < il_size(t->image); i++) {
        double x, y, dx, dy;
        Unused anbool ok;
        ok = sip_radec2pixelxy(t->sip, t->a_ref[il_get(t->ref, i)], t->d_ref[il_get(t->ref, i)], &x, &y);
        assert(ok);
        dx = t->x[il_get(t->image, i)] - x;
        dy = t->y[il_get(t->image, i)] - y;
        sqerr += dx * dx + dy * dy;
    }
    // convert to arcsec^2
    return sqerr * square(sip_pixel_scale(t->sip));
}
#endif

// FIXME: adapt this function to take as input the correspondences to use VVVVV
//    wic is World Intermediate Coordinates, either along ra or dec
//       i.e. canonical image coordinates
//    wic_corr is the list of corresponding indexes for wip
//    pix_corr is the list of corresponding indexes for pixels
//    pix is raw image pixels (either u or v)
//    siporder is the sip order up to MAXORDER (defined in sip.h)
//    the correspondences are passed so that we can stick RANSAC around the whole
//    thing for better estimation.

// Run a polynomial tweak
static void do_sip_tweak(tweak_t* t) {
    sip_t sipout;
    size_t i, M;

    // a_order and b_order should be the same!
    assert(t->sip->a_order == t->sip->b_order);

    debug("do_sip_tweak starting.\n");
    logverb("RMS error of correspondences: %g arcsec\n",
            correspondences_rms_arcsec(t, 0));
    if (t->weighted_fit)
        logverb("Weighted RMS error of correspondences: %g arcsec\n",
                correspondences_rms_arcsec(t, 1));

    M = il_size(t->image);
    double* starxyz = malloc(M * 3 * sizeof(double));
    double* fieldxy = malloc(M * 2 * sizeof(double));
    double* weights = NULL;
    if (t->weighted_fit)
        weights = malloc(M * sizeof(double));

    int result;
    for (i=0; i<M; i++) {
        int refi;
        int imi;
        imi = il_get(t->image, i);
        fieldxy[2*i + 0] = t->x[imi];
        fieldxy[2*i + 1] = t->y[imi];
        refi = il_get(t->ref, i);
        radecdeg2xyzarr(t->a_ref[refi], t->d_ref[refi], starxyz + i*3);
        if (t->weighted_fit)
            weights[i] = dl_get(t->weight, i);
    }

    int doshift = 1;
    result = fit_sip_wcs(starxyz, fieldxy, weights, M,
                         &(t->sip->wcstan), t->sip->a_order, t->sip->ap_order, doshift,
                         &sipout);
    free(starxyz);
    free(fieldxy);
    free(weights);
    if (result) {
        ERROR("fit_sip_wcs failed\n");
        return;
    }
    memcpy(t->sip, &sipout, sizeof(sip_t));

    // recalc using new SIP
    tweak_clear_on_sip_change(t);
    tweak_go_to(t, TWEAK_HAS_IMAGE_AD);
    tweak_go_to(t, TWEAK_HAS_REF_XY);

    logverb("RMS error of correspondences: %g arcsec\n",
            correspondences_rms_arcsec(t, 0));
    if (t->weighted_fit)
        logverb("Weighted RMS error of correspondences: %g arcsec\n",
                correspondences_rms_arcsec(t, 1));
}

// Really what we want is some sort of fancy dependency system... DTDS!
// Duct-tape dependencey system (DTDS)
#define done(x) t->state |= x; return x;
#define want(x)                                 \
    if (flag == x && t->state & x)              \
        return x;                               \
    else if (flag == x)
#define ensure(x)                               \
    if (!(t->state & x)) {                      \
        return tweak_advance_to(t, x);          \
    }

unsigned int tweak_advance_to(tweak_t* t, unsigned int flag) {
    want(TWEAK_HAS_IMAGE_AD) {
        int jj;
        ensure(TWEAK_HAS_SIP);
        ensure(TWEAK_HAS_IMAGE_XY);
        debug("Satisfying TWEAK_HAS_IMAGE_AD\n");
        // Convert to ra dec
        assert(!t->a);
        assert(!t->d);
        t->a = malloc(sizeof(double) * t->n);
        t->d = malloc(sizeof(double) * t->n);
        for (jj = 0; jj < t->n; jj++)
            sip_pixelxy2radec(t->sip, t->x[jj], t->y[jj], t->a + jj, t->d + jj);
        done(TWEAK_HAS_IMAGE_AD);
    }

    want(TWEAK_HAS_REF_AD) {
        if (!(t->a_ref && t->d_ref)) {
            ensure(TWEAK_HAS_REF_XYZ);
            debug("Satisfying TWEAK_HAS_REF_AD\n");
            ref_ad_from_xyz(t);
        }
        assert(t->a_ref && t->d_ref);
        done(TWEAK_HAS_REF_AD);
    }

    want(TWEAK_HAS_REF_XYZ) {
        if (!t->xyz_ref) {
            ensure(TWEAK_HAS_REF_AD);
            debug("Satisfying TWEAK_HAS_REF_XYZ\n");
            ref_xyz_from_ad(t);
        }
        assert(t->xyz_ref);
        done(TWEAK_HAS_REF_XYZ);
    }

    want(TWEAK_HAS_REF_XY) {
        int jj;
        ensure(TWEAK_HAS_REF_AD);
        debug("Satisfying TWEAK_HAS_REF_XY\n");
        assert(t->state & TWEAK_HAS_REF_AD);
        assert(t->n_ref);
        assert(!t->x_ref);
        assert(!t->y_ref);
        t->x_ref = malloc(sizeof(double) * t->n_ref);
        t->y_ref = malloc(sizeof(double) * t->n_ref);
        for (jj = 0; jj < t->n_ref; jj++) {
            Unused anbool ok;
            ok = sip_radec2pixelxy(t->sip, t->a_ref[jj], t->d_ref[jj],
                                   t->x_ref + jj, t->y_ref + jj);
            assert(ok);
        }
        done(TWEAK_HAS_REF_XY);
    }

    want(TWEAK_HAS_IMAGE_XYZ) {
        int i;
        ensure(TWEAK_HAS_IMAGE_AD);
        debug("Satisfying TWEAK_HAS_IMAGE_XYZ\n");
        assert(!t->xyz);
        t->xyz = malloc(3 * t->n * sizeof(double));
        for (i = 0; i < t->n; i++)
            radecdeg2xyzarr(t->a[i], t->d[i], t->xyz + 3*i);
        done(TWEAK_HAS_IMAGE_XYZ);
    }

    want(TWEAK_HAS_COARSLY_SHIFTED) {
        ensure(TWEAK_HAS_REF_XY);
        ensure(TWEAK_HAS_IMAGE_XY);
        debug("Satisfying TWEAK_HAS_COARSLY_SHIFTED\n");
        get_dydx_range(t->x, t->y, t->n, t->x_ref, t->y_ref, t->n_ref,
                       &t->mindx, &t->mindy, &t->maxdx, &t->maxdy);
        do_entire_shift_operation(t, 1.0);
        tweak_clear_image_ad(t);
        tweak_clear_ref_xy(t);
        done(TWEAK_HAS_COARSLY_SHIFTED);
    }

    want(TWEAK_HAS_FINELY_SHIFTED) {
        ensure(TWEAK_HAS_REF_XY);
        ensure(TWEAK_HAS_IMAGE_XY);
        ensure(TWEAK_HAS_COARSLY_SHIFTED);
        debug("Satisfying TWEAK_HAS_FINELY_SHIFTED\n");
        // Shrink size of hough box
        do_entire_shift_operation(t, 0.3);
        tweak_clear_image_ad(t);
        tweak_clear_ref_xy(t);
        done(TWEAK_HAS_FINELY_SHIFTED);
    }

    want(TWEAK_HAS_REALLY_FINELY_SHIFTED) {
        ensure(TWEAK_HAS_REF_XY);
        ensure(TWEAK_HAS_IMAGE_XY);
        ensure(TWEAK_HAS_FINELY_SHIFTED);
        debug("Satisfying TWEAK_HAS_REALLY_FINELY_SHIFTED\n");
        // Shrink size of hough box
        do_entire_shift_operation(t, 0.03);
        tweak_clear_image_ad(t);
        tweak_clear_ref_xy(t);
        done(TWEAK_HAS_REALLY_FINELY_SHIFTED);
    }

    want(TWEAK_HAS_CORRESPONDENCES) {
        ensure(TWEAK_HAS_REF_XYZ);
        ensure(TWEAK_HAS_IMAGE_XYZ);
        debug("Satisfying TWEAK_HAS_CORRESPONDENCES\n");
        t->jitterd2 = arcsec2distsq(t->jitter);
        find_correspondences(t, 6.0 * arcsec2rad(t->jitter));
        done(TWEAK_HAS_CORRESPONDENCES);
    }

    want(TWEAK_HAS_LINEAR_CD) {
        ensure(TWEAK_HAS_SIP);
        ensure(TWEAK_HAS_REALLY_FINELY_SHIFTED);
        ensure(TWEAK_HAS_REF_XY);
        ensure(TWEAK_HAS_REF_AD);
        ensure(TWEAK_HAS_IMAGE_XY);
        ensure(TWEAK_HAS_CORRESPONDENCES);
        debug("Satisfying TWEAK_HAS_LINEAR_CD\n");
        do_sip_tweak(t);
        tweak_clear_on_sip_change(t);
        done(TWEAK_HAS_LINEAR_CD);
    }

    // small memleak -- but it's a major bug if this happens, so suck it up.
    logerr("die for dependence: %s\n", state_string(flag));
    assert(0);
    return -1;
}

void tweak_push_wcs_tan(tweak_t* t, const tan_t* wcs) {
    memcpy(&(t->sip->wcstan), wcs, sizeof(tan_t));
    t->state |= TWEAK_HAS_SIP;
}

void tweak_go_to(tweak_t* t, unsigned int dest_state) {
    while (!(t->state & dest_state))
        tweak_advance_to(t, dest_state);
}

#define SAFE_FREE(xx) {free((xx)); xx = NULL;}
void tweak_clear(tweak_t* t) {
    if (!t)
        return ;
    SAFE_FREE(t->a);
    SAFE_FREE(t->d);
    SAFE_FREE(t->x);
    SAFE_FREE(t->y);
    SAFE_FREE(t->xyz);
    SAFE_FREE(t->a_ref);
    SAFE_FREE(t->d_ref);
    SAFE_FREE(t->x_ref);
    SAFE_FREE(t->y_ref);
    SAFE_FREE(t->xyz_ref);
    if (t->sip) {
        sip_free(t->sip);
        t->sip = NULL;
    }
    il_free(t->image);
    il_free(t->ref);
    dl_free(t->dist2);
    if (t->weight)
        dl_free(t->weight);
    t->image = NULL;
    t->ref = NULL;
    t->dist2 = NULL;
    t->weight = NULL;
    kdtree_free(t->kd_image);
    kdtree_free(t->kd_ref);
}

void tweak_free(tweak_t* t) {
    tweak_clear(t);
    free(t);
}


#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#ifdef _MSC_VER				/* identifies this as a Microsoft compiler */
#define _USE_MATH_DEFINES	/* added RX2011 */
#endif
#include <math.h>
#include <time.h>
#include <limits.h>
#include <gsl/gsl_multimin.h>
//#include "/Users/tischler/dev/GNU/include/gsl/gsl_multimin.h"
#include "Euler.h"

#ifdef DEBUG_ON
#define DEBUG 1				/* a value of 5 gets everything */
#endif

#define MIN(A,B)	( (A)<(B)?(A):(B) )
#define MAX(A,B)	( (A)>(B)?(A):(B) )

double my_f (const gsl_vector *v, void *params);

/******************************************************************************************
 * return the error between a lattice rotated by Euler angles and the measured spots
 * thus the returned value is a minimum for the optimimal Euler angles
 */
double my_f(
const gsl_vector *v,	/* vector with current value of Euler angles */
void	*pattern)		/* pointer to the parameters, in this case a pointer to the pattern structure */
{
	double  G[3];						/* calculated G vector for each hkl */
	double  recip[3][3];				/* reciprocal space */
	double  M_Euler[3][3];				/* an Euler matrix */
	size_t  N;							/* local copy of number of measured spots to check */
	double a,b,g;						/* local copy of the Euler angles */
	struct patternOfOneGrain *p;		/* local value of pattern */
	double  rms;						/* rms value of error, the returned value */
	size_t  i;							/* loop index */
	double  hkl[3];						/* float value of hkl */
	double  dot;

	/* p = pattern; */
	p = (struct patternOfOneGrain *)pattern;
	N = p->Ni;											/* local value for convenience */
	a = gsl_vector_get(v,0);							/* the Euler angles */
	b = gsl_vector_get(v,1);
	g = gsl_vector_get(v,2);
	MatrixCopy33(recip,p->xtal.recip);					/* copy unrotated reciprocal lattice into recip */

	EulerMatrix(a,b,g,M_Euler);							/* make the rotation matrix M_Euler from Euler angles */
	MatrixMultiply33(M_Euler,recip,recip);				/* rotate recip by Euler angles */
	for (i=0,rms=0.0;i<N;i++) {
		VECTOR_COPY3(hkl,p->hkls[i]);					/* convert integer hkl to float for following MatrixMultiply31 */
if (hkl[0]!=hkl[0]) {
fprintf(stderr,"hkl[0] is NaN\n");
}
		MatrixMultiply31(recip,hkl,G);
if (G[0]!=G[0]) {
fprintf(stderr,"G[0] is NaN\n");
}
		normalize3(G);									/* normalized gvec of hkl[i] rotated by Euler angles */
		dot = dot3(G,p->Ghat[i]);
		dot = MIN(dot,1.0);
		rms += 1. - dot;								/* (1-dot) ~ 0.5*(delta angle)^2*/
	}
	rms = sqrt(rms/(double)N);							/* rms of angle errors */
	return rms;
}

/* When calling optimizeEulerAngles(), be sure to set pattern.xtal:

	e.g.
	pattern.xtal.a = pattern.xtal.b = pattern.xtal.c = 4.05;		// just default to Al for now
	pattern.xtal.alpha = pattern.xtal.beta = pattern.xtal.gamma = M_PI/2.0;

	then call:
		FillRecipInLattice(&(pattern.lattice));
	to set the reciprocal lattice in the lattice structure
*/
/* this routine optimizes approximate Euler angles using the gsl simplex method. */
int optimizeEulerAngles(
double  startStep,					/* starting step size (radians), make this a bit larger than the error */
double  epsAbs,						/* tolerance on the answer (=0.001) */
long	maxIter,					/* maximum number of iterations, stop after this many regardless, perhaps 100 */
struct	patternOfOneGrain *pattern) /* provides G^'s and hkl's for one fitted pattern, and the lattice parameters */
{
	size_t np = 3;					/* number of dimensions of the function, 3 Euler angles */

	/* T tells what kind of minimizer we are using, select the simplex method */
	const gsl_multimin_fminimizer_type *T = gsl_multimin_fminimizer_nmsimplex;
	gsl_multimin_fminimizer *s = NULL;		/* this is the minimizer */
	gsl_vector *ss, *x;						/* pointers to vectors,  ss is stepsize, x is point */
	gsl_multimin_function minex_func;		/* struct defining parts needed for minimizing */
	size_t iter=0, i;
	int status;								/* status of returned gsl function, 0 is OK */
	double size;							/* size of minimizer, sort of maximum error */

	double  G[3];							/* calculated G vector for each hkl */
	double  recip[3][3];					/* reciprocal space */
	double  M_Euler[3][3];					/* an Euler matrix */
	double  hkl[3];							/* float value of hkl */
	double  dot;

	if (!pattern) { fprintf(stderr,"pattern in NULL on entry to optimizeEulerAngles()\n"); return 1; }
	if (pattern->Ni<3) { fprintf(stderr,"less than 3 pairs of Ghat & hkls on entry to optimizeEulerAngles()\n"); return 1; }
	if (!(pattern->Ghat) || !(pattern->hkls)) { fprintf(stderr,"Ghat or hkls are NULL on entry to optimizeEulerAngles()\n"); return 1; }
	if ((pattern->xtal.a)<=0 || (pattern->xtal.b)<=0 || (pattern->xtal.c)<=0) { fprintf(stderr,"invalid lattice in optimizeEulerAngles()\n"); return 1; }

	/* Initial vertex step size vector */
	ss = gsl_vector_alloc(np);				/* step sizes in x and y */

	/* Set all step sizes to startStep, the starting step size */
	gsl_vector_set_all(ss, startStep);		/* set ss[0] = ss[1]  = ss[2] = startStep (radian) */

	/* Starting point */
	x = gsl_vector_alloc(np);				/* allocate space for a vector of length 3 */
	gsl_vector_set(x, 0, pattern->alpha);   /* set to starting guess, input (alpha,beta,gamma) */
	gsl_vector_set(x, 1, pattern->beta);
	gsl_vector_set(x, 2, pattern->gamma);

	/* Initialize the minimizer */
	minex_func.f = &my_f;					/* the function that returns the value */
	minex_func.n = np;						/* number of parameters */
	minex_func.params = (void *)pattern;	/* a generic pointer that points to something with info for my_f() */
	/* Note: my_func.df, and my_func.fdf are not needed since simplex method does not use derivatives */

	s = gsl_multimin_fminimizer_alloc(T, np);	/* allocate for a 2d simplex minimizer named 's' */

	/* Set starting point for the minimizer:
		for the minimizer 's',
		working on a function 'minex_func'
		starting from the initial guess x[2]
		using step sizes of ss[2]  */
	gsl_multimin_fminimizer_set(s, &minex_func, x, ss);

#if (DEBUG)		/******************************************************************************/
	fprintf(fout, "start minimization with Euler angles (%15.10f, %15.10f, %15.10f) (deg.)\n", \
		pattern->alpha*180/M_PI,pattern->beta*180/M_PI,pattern->gamma*180/M_PI);
#endif			/******************************************************************************/

	/* Iterate the minimizer.  This looop iterates the minimization method until it is good enough */
	iter = 0;
	do {
		iter++;
		status = gsl_multimin_fminimizer_iterate(s);	/* do one iteration of the minimizer */
		if (status) break;								/* status=0 is OK, check here for an error */

		size = gsl_multimin_fminimizer_size (s); 		/* get 'size' of minimizer */
		status = gsl_multimin_test_size (size, epsAbs);	/* test size against epsAbs */

#if (DEBUG>1)		/******************************************************************************/
		fprintf(fout,"%5ld  f(", iter);
		for (i=0; i<np; i++) {
			fprintf(fout,"%.8f", gsl_vector_get(s->x, i));
			if (i<np-1) fprintf(fout,", ");
		}
		fprintf(fout,") = %15.8g size = %.5g\n", s->fval, size);
#endif			/******************************************************************************/
	} while (status == GSL_CONTINUE && iter < (size_t)maxIter); /* stop after maxIter regardless of size */

#if (DEBUG)		/******************************************************************************/
	if (status == GSL_SUCCESS) fprintf(fout,"converged to minimum successfully\n");
	else fprintf(fout,"did not converge successfully, status = %d\n",status);
	fprintf(fout, "  optimization changed Euler angles by (%15.10f, %15.10f, %15.10f) (deg.)\n", \
		(gsl_vector_get(s->x,0)-pattern->alpha)*180/M_PI, (gsl_vector_get(s->x,1)-pattern->beta)*180/M_PI, \
		(gsl_vector_get(s->x,2)-pattern->gamma)*180/M_PI);
#endif			/******************************************************************************/

	pattern->alpha = gsl_vector_get(s->x,0);			/* set to final values, resultant (alpha,beta,gamma) returned */
	pattern->beta = gsl_vector_get(s->x,1);
	pattern->gamma = gsl_vector_get(s->x,2);

#if (DEBUG)		/******************************************************************************/
	fprintf(fout, "after minimization, Euler angles are (%15.10f, %15.10f, %15.10f) (deg.)\n", \
		pattern->alpha*180/M_PI,pattern->beta*180/M_PI,pattern->gamma*180/M_PI);
#endif			/******************************************************************************/

	/* recalculate the err[] list in pattern */
	MatrixCopy33(recip,pattern->xtal.recip);		/* copy unrotated reciprocal lattice into recip */
	EulerMatrix(pattern->alpha,pattern->beta,pattern->gamma,M_Euler);/* rotation matrix from Euler angles */
	MatrixMultiply33(M_Euler,recip,recip);			/* rotate recip by Euler angles */
	for (i=0;i<(size_t)(pattern->Ni);i++) {
		VECTOR_COPY3(hkl,pattern->hkls[i]);			/* convert integer hkl to float for following MatrixMultiply31 */
		MatrixMultiply31(recip,hkl,G);
		normalize3(G);								/* normalized gvec of hkl[i] rotated by Euler angles */
		dot = dot3(G,pattern->Ghat[i]);				/* note, Ghat here is the measured, not calculated G */
		(pattern->err)[i] = acos(MIN(dot,1.0));
	}

	gsl_vector_free(x);									/* free vector of x,y positions */
	gsl_vector_free(ss);								/* free vector of step sizes */
	gsl_multimin_fminimizer_free (s);					/* free the minimizer */
	return status;
}


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

#include <gsl/gsl_integration.h>
#include <gsl/gsl_interp.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_errno.h>

#include "ccl.h"

/*----- ROUTINE: dc_NakamuraSuto -----
INPUT: cosmology, scale factor
TASK: Computes the peak threshold: delta_c(z) assuming LCDM.
Cosmology dependence of the critical linear density according to the spherical-collapse model.
Fitting function from Nakamura & Suto (1997; arXiv:astro-ph/9710107).
*/
double dc_NakamuraSuto(ccl_cosmology *cosmo, double a, int *status){

  double Om_mz = ccl_omega_x(cosmo, a, ccl_species_m_label, status);
  double dc0 = (3./20.)*pow(12.*M_PI,2./3.);
  double dc = dc0*(1.+0.012299*log10(Om_mz));

  return dc;

}

/*----- ROUTINE: Dv_BryanNorman -----
INPUT: cosmology, scale factor
TASK: Computes the virial collapse density contrast with respect to the matter density assuming LCDM.
Cosmology dependence of the virial collapse density according to the spherical-collapse model
Fitting function from Bryan & Norman (1998; arXiv:astro-ph/9710107)
*/
double Dv_BryanNorman(ccl_cosmology *cosmo, double a, int *status){

  double Om_mz = ccl_omega_x(cosmo, a, ccl_species_m_label, status);
  double x = Om_mz-1.;
  double Dv0 = 18.*pow(M_PI,2);
  double Dv = (Dv0+82.*x-39.*pow(x,2))/Om_mz;

  return Dv;
}

static double sigmaM_m2r(ccl_cosmology *cosmo, double halomass, int *status)
{
  double rho_m, smooth_radius;

  // Comoving matter density
  rho_m = ccl_rho_x(cosmo, 1., ccl_species_m_label, 1, status);

  smooth_radius = pow((3.0*halomass) / (4*M_PI*rho_m), (1.0/3.0));

  return smooth_radius;
}

void ccl_cosmology_compute_sigma(ccl_cosmology *cosmo, ccl_f2d_t *psp, int *status)
{
  if(cosmo->computed_sigma)
    return;

  int na = cosmo->spline_params.A_SPLINE_NA_SM + cosmo->spline_params.A_SPLINE_NLOG_SM - 1;
  int nm = cosmo->spline_params.LOGM_SPLINE_NM;
  double *m = NULL;
  double *y = NULL;
  double *aa = NULL;

  // create linearly-spaced values of log-mass.
  m = ccl_linear_spacing(cosmo->spline_params.LOGM_SPLINE_MIN,
                         cosmo->spline_params.LOGM_SPLINE_MAX, nm);
  if (m == NULL ||
      (fabs(m[0]-cosmo->spline_params.LOGM_SPLINE_MIN)>1e-5) ||
      (fabs(m[nm-1]-cosmo->spline_params.LOGM_SPLINE_MAX)>1e-5) ||
      (m[nm-1]>10E17)) {
    *status = CCL_ERROR_MEMORY;
    ccl_cosmology_set_status_message(cosmo,
                                     "ccl_massfunc.c: ccl_cosmology_compute_sigma(): "
                                     "Error creating linear spacing in m\n");
  }

  // create scale factor array
  if (*status == 0) {
    aa = ccl_linlog_spacing(cosmo->spline_params.A_SPLINE_MINLOG_SM,
                            cosmo->spline_params.A_SPLINE_MIN_SM,
                            cosmo->spline_params.A_SPLINE_MAX,
                            cosmo->spline_params.A_SPLINE_NLOG_SM,
                            cosmo->spline_params.A_SPLINE_NA_SM);
    if (aa == NULL) {
      *status = CCL_ERROR_MEMORY;
      ccl_cosmology_set_status_message(cosmo,
                                       "ccl_massfunc.c: ccl_cosmology_compute_sigma(): "
                                       "Error creating scale factor array\n");
    }
  }

  // create space for y, to be filled with sigma
  if (*status == 0) {
    y = malloc(sizeof(double)*nm*na);
    if (y == NULL) {
      *status = CCL_ERROR_MEMORY;
      ccl_cosmology_set_status_message(cosmo,
                                       "ccl_massfunc.c: ccl_cosmology_compute_sigma(): "
                                       "memory allocation\n");
    }
  }

  // fill in sigma, if no errors have been triggered at this time.
  if (*status == 0) {
    #pragma omp parallel shared(na, aa, nm, m, y, status, cosmo, psp) \
                         default(none)
    {
      int i, j;
      double a_sf, smooth_radius;
      int local_status = *status;

      #pragma omp for
      for (j=0; j<na; j++) {
        a_sf = aa[j];
        for (i=0; i<nm; i++) {
          smooth_radius = sigmaM_m2r(cosmo, pow(10,m[i]), &local_status);
          y[j*nm + i] = log(ccl_sigmaR(cosmo, smooth_radius, a_sf,
                                       psp, &local_status));
        }
      } //end omp for
      if (local_status) {
        #pragma omp atomic write
        *status = local_status;
      }
    } //end omp parallel
  }

  gsl_spline2d *lsM = NULL;
  if (*status == 0) {
    lsM = gsl_spline2d_alloc(gsl_interp2d_bicubic, nm, na);
    if (lsM == NULL) {
      *status = CCL_ERROR_MEMORY;
      ccl_cosmology_set_status_message(cosmo,
                                       "ccl_massfunc.c: ccl_cosmology_compute_sigma(): "
                                       "error allocating 2D spline\n");
    }
  }

  if(*status == 0) {
    int s2dstatus=gsl_spline2d_init(lsM, m, aa, y, nm, na);
    if (s2dstatus) {
      *status = CCL_ERROR_SPLINE;
      ccl_cosmology_set_status_message(cosmo,
                                       "ccl_massfunc.c: ccl_cosmology_compute_sigma(): "
                                       "error initializing spline\n");
    }
  }

  if (*status == 0) {
    cosmo->computed_sigma = true;
    cosmo->data.logsigma = lsM;
  }
  else
    gsl_spline2d_free(lsM);

  free(aa);
  free(m);
  free(y);
}

/*----- ROUTINE: ccl_sigma_M -----
INPUT: ccl_cosmology * cosmo, double halo mass in units of Msun, double scale factor
TASK: returns sigma from the sigmaM interpolation. Also computes the sigma interpolation if
necessary.
*/
double ccl_sigmaM(ccl_cosmology *cosmo, double log_halomass, double a, int *status)
{
  // Check if sigma has already been calculated
  if (!cosmo->computed_sigma) {
    *status = CCL_ERROR_SIGMA_INIT;
    ccl_cosmology_set_status_message(cosmo,
                                     "ccl_massfunc.c: ccl_sigmaM(): "
                                     "sigma(M) spline has not been computed!");
    return NAN;
  }

  double lgsigmaM;
  int gslstatus = gsl_spline2d_eval_e(cosmo->data.logsigma, log_halomass,
                                      a, NULL, NULL, &lgsigmaM);

  if(gslstatus != GSL_SUCCESS) {
    ccl_raise_gsl_warning(gslstatus, "ccl_massfunc.c: ccl_sigmaM():");
    *status |= gslstatus;
  }

  return exp(lgsigmaM);
}

/*----- ROUTINE: ccl_dlnsigM_dlogM -----
INPUT: ccl_cosmology *cosmo, double halo mass in units of Msun
TASK: returns the value of the derivative of ln(sigma^-1) with respect to log10 in halo mass.
*/
double ccl_dlnsigM_dlogM(ccl_cosmology *cosmo, double log_halomass, double a, int *status)
{
  // Check if sigma has already been calculated
  if (!cosmo->computed_sigma) {
    *status = CCL_ERROR_SIGMA_INIT;
    ccl_cosmology_set_status_message(cosmo,
                                     "ccl_massfunc.c: ccl_dlnsigM_dlogM(): "
                                     "sigma(M) spline has not been computed!");
    return NAN;
  }

  double dlsdlgm;
  int gslstatus = gsl_spline2d_eval_deriv_x_e(cosmo->data.logsigma,
                                              log_halomass, a,
                                              NULL, NULL, &dlsdlgm);
  if(gslstatus) {
    ccl_raise_gsl_warning(gslstatus, "ccl_massfunc.c: ccl_dlnsigM_dlogM():");
    *status |= gslstatus;
  }
  return -dlsdlgm;
}


/**
 * \brief Function Definition Library to implement Sensor Fusion Algorithm
 * \author jaspalsingh@cmail.carleton.ca
 * \version 1.0
 * \date 2019-12-01
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_eigen.h>
#include "sensor_fusion.h"

/**
 * \brief Structure pointer for computed support degree matrix
 * \details Structure pointer for calculating support degree matrix
 * \param[in] sensor_readings array of sensor readings
 * \return Success: pointer to the structure support_degree_matrix
 * \return Failure: NULL
 */
struct support_degree_matrix *compute_support_degree_matrix(double *sensor_readings) {
    struct support_degree_matrix *spd;
    spd = (struct support_degree_matrix *)malloc(sizeof(struct support_degree_matrix));
    // count the number of sensor readings
    int length = sizeof(sensor_readings) / sizeof(sensor_readings[0]);
    spd->sensor_count = length;
    double **arrptr = (double **)malloc(length * sizeof(double *));
    int count = 0;
    for (int i = 0; i < length; i++) {
        arrptr[i] = (double *)malloc(length * sizeof(double));
    }
    // allocate memory space for support degree matrix structure
    spd->sd_matrix = (double *)malloc(sizeof(double) * length * length);
    if (arrptr == NULL || spd == NULL || spd->sd_matrix == NULL) {
        printf("ERROR: Failed to allocate memory at %s\n", __func__);
        return NULL;
    }
    // compute support degree matrix
    for (int i = 0; i < length; i++) {
        for (int j = 0; j < length; j++) {
            arrptr[i][j] = exp(-1 * fabs(sensor_readings[i] - sensor_readings[j]));
            spd->sd_matrix[count] = arrptr[i][j];
            count++;
        }
    }
    printf("INFO : Computed support degree matrix\n");
    for (int i = 0; i < length * length; i++) {
        printf("DEBUG: Value => %f\n", spd->sd_matrix[i]);
    }
    free(arrptr);
    return spd;
}

/**
 * \brief Structure pointer calculate eigen value and eigen vector
 * \details Structure pointer for calculating eigen value and eigen vector
 * \param[in] support_degree_matrix pointer to structure support degree matrix
 * \return Success: pointer to structure eigen_value_vector
 * \return Failure: NULL
 */
struct eigen_value_vector *compute_eigen(struct support_degree_matrix *spd) {
    if ((spd->sd_matrix == NULL) || (spd->sensor_count == 0)) {
        printf("ERROR: Invalid param passed at %s\n", __func__);
        return NULL;
    }
    int length = spd->sensor_count;
    struct eigen_value_vector *eigen;
    // allocation memory space for eigen value vector structure
    eigen = (struct eigen_value_vector *)malloc(sizeof(struct eigen_value_vector));
    eigen->eigen_value = (double *)malloc(sizeof(double) * length);
    eigen->eigen_vector = (double **)malloc(length * sizeof(double *));
    for (int i = 0; i < length; i++) {
        eigen->eigen_vector[i] = (double *)malloc(length * sizeof(double));
    }
    if (eigen == NULL || eigen->eigen_vector == NULL || eigen->eigen_value == NULL) {
        printf("ERROR: Failed to allocate memory at %s\n", __func__);
        return NULL;
    }
    //Refer: https://www.gnu.org/software/gsl/doc/html/eigen.html
    // compute eigen values
    gsl_matrix_view m = gsl_matrix_view_array(spd->sd_matrix, length, length);
    gsl_vector *eval = gsl_vector_alloc(length);
    gsl_matrix *evec = gsl_matrix_alloc(length, length);
    gsl_eigen_symmv_workspace *w = gsl_eigen_symmv_alloc(length);
    gsl_eigen_symmv(&m.matrix, eval, evec, w);
    gsl_eigen_symmv_free(w);
    gsl_eigen_symmv_sort(eval, evec, GSL_EIGEN_SORT_VAL_DESC);
    {
        int i;
        for (i = 0; i < length; i++) {
            eigen->eigen_value[i] = gsl_vector_get(eval, i);
            gsl_vector_view evec_i = gsl_matrix_column(evec, i);
            for (int j = 0; j < length; j++) {
                eigen->eigen_vector[i][j] = *(*(&evec_i.vector.data) + j * length);
            }
        }
    }
    for (int i = 0; i < length; i++) {
        printf("DEBUG: Eigen Value => %g\n", eigen->eigen_value[i]);
        printf("DEBUG: Eigen Vector => \n");
        for (int j = 0; j < length; j++) {
            printf("%g\n", eigen->eigen_vector[i][j]);
        }
    }
    gsl_vector_free(eval);
    gsl_matrix_free(evec);
    return eigen;
}

/**
 * \brief Function to calculate the contribution rate of principal component
 * \details Compute the contribution rate of kth and mth principal components
 * \param[in] eigen_value_vector pointer to structure of eigen_value_vector
 * \param[in] sensor_count total count of sensors
 * \return Success: pointer to contribution_rate of type double
 * \return Failure: NULL
 */
double *compute_contribution_rate(struct eigen_value_vector *eigen,
                                  int sensor_count) {
    if ((sensor_count == 0) || (eigen->eigen_value == NULL) || (eigen->eigen_vector == NULL)) {
        printf("ERROR: Invalid param passed at %s\n", __func__);
        return NULL;
    }
    double *contribution_rate = (double *)malloc(sensor_count * sizeof(double));
    if (contribution_rate == NULL) {
        printf("ERROR: Failed to allocate memory at %s\n", __func__);
        return NULL;
    }
    double sum = 0;
    for (int i = 0; i < sensor_count; i++) {
        sum += eigen->eigen_value[i];
    }
    printf("INFO : Compute contribution rate\n");
    for (int j = 0; j < sensor_count; j++) {
        contribution_rate[j] = eigen->eigen_value[j] / sum;
        printf("DEBUG: Rate => %g\n", contribution_rate[j]);
    }
    return contribution_rate;
}

/**
 * \brief Function to compute contribution rate
 * \details Compute the count of contribution rates to be used for sensors
 * \param[in] contribution_rate
 * \param[in] threshold
 * \param[in] sensor_count number of sensors to be used
 * \return Success: contribution rate count,
 * \return Failure: -1
 */
int compute_contrib_rate_count(double *contribution_rate,
                               float threshold,
                               int sensor_count) {
    if ((contribution_rate == NULL) || (threshold == 0) || (sensor_count == 0)) {
        printf("ERROR: Invalid param passed at %s\n", __func__);
        return -1;
    }
    double sum = 0;
    int m;
    for (int k = 0; k < sensor_count; k++) {
        for (int j = 0; j <= k; j++) {
            sum += contribution_rate[j];
        }
        if (sum <= threshold) {
            m = k - 1;
            printf("DEBUG: Contribution rate to use %d\n", m);
            return m;
        }
    }
    printf("DEBUG: Count of Contribution rates => %d\n", sensor_count);
    return sensor_count;
}

/**
 * \brief Function to compute principal component
 * \details Computed principal component for each of eigen vector
 * \param[in] support_degree_matrix pointer to structure support degree matrix
 * \param[in] eigen_vector pointer to pointer of eigen vector of type double
 * \param[in] contrib_rate_count number of contribution rate to be used
 * \return Success: principal of support degree matrix,
 * \return Failure: NULL
 */
double **compute_principal_component(struct support_degree_matrix *spd,
                                     double **eigen_vector,
                                     int contrib_rate_count) {
    if ((spd->sensor_count == 0) || (contrib_rate_count == 0) || (spd->sd_matrix == NULL) || (eigen_vector == NULL)) {
        printf("ERROR: Invalid param passed at %s\n", __func__);
        return NULL;
    }
    int n = spd->sensor_count;
    int m = contrib_rate_count;
    int count = 0;
    double **arrptr = (double **)malloc(n * sizeof(double *));
    for (int i = 0; i < n; i++) {
        arrptr[i] = (double *)malloc(n * sizeof(double));
    }
    if (arrptr == NULL) {
        printf("ERROR: Failed to allocate memory at %s\n", __func__);
        return NULL;
    }
    printf("INFO : Compute support degree matrix\n");
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            arrptr[i][j] = spd->sd_matrix[count];
            printf("%f\n", arrptr[i][j]);
            count++;
        }
    }
    double **principal_components_matrix = (double **)malloc(m * sizeof(double *));
    for (int i = 0; i < n; i++) {
        principal_components_matrix[i] = (double *)malloc(n * sizeof(double));
    }
    if (principal_components_matrix == NULL) {
        printf("ERROR: Failed to allocate memory at %s\n", __func__);
        return NULL;
    }
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            principal_components_matrix[i][j] = 0;
            for (int k = 0; k < n; k++) {
                principal_components_matrix[i][j] += eigen_vector[i][k] * arrptr[k][j];
            }
        }
    }
    free(arrptr);
    return principal_components_matrix;
}

/**
 * \brief Function to compute integrated support degree
 * \details Computed support degree matrix for each of the sensors
 * \param[in] principle_components pointer to pointer of principle component of type double
 * \param[in] contribution_rate pointer to contribution rate of type double
 * \param[in] contrib_rate_count number of contribution rate to be used
 * \param[in] sensor_count number of sensors
 * \return Success: array of support degree matrix,
 * \return Failure: NULL
 */
double *compute_integrated_support_degree(double **principle_components,
                                          double *contribution_rate,
                                          int contrib_rate_count,
                                          int sensor_count) {
    if ((sensor_count == 0) || (contrib_rate_count == 0) || (principle_components == NULL) || (contribution_rate == NULL)) {
        printf("ERROR: Invalid param passed at %s\n", __func__);
        return NULL;
    }
    double *arr = (double *)malloc(sensor_count * sizeof(double));
    if (arr == NULL) {
        printf("ERROR: Failed to allocate memory at %s\n", __func__);
        return NULL;
    }
    // Calculate integrated support degree matrix
    for (int i = 0; i < sensor_count; i++) {
        for (int j = 0; j < contrib_rate_count; j++) {
            printf("INFO: Principal Component => %d-%f, Contribution Rate => %d-%f\n", i, principle_components[i][j], j, contribution_rate[i]);
            arr[i] += principle_components[i][j] * contribution_rate[i];
        }
    }
    return arr;
}

/**
 * \brief Function to eliminate invalid sensor readings
 * \details Discard the faulty sensor readings from the support degree matrix
 * \param[in,out] integrated_support_degree_matrix pointer to integrated support degree matrix
 * \param[in] fault_tolerance_threshold threshold value to cut off invalid readings
 * \param[in] sensor_count total count of sensors
 * \return Success: 0,
 * \return Failure: -1
 */
int eliminate_incorrect_data(double *integrated_support_degree_matrix,
                             double fault_tolerance_threshold,
                             int sensor_count) {
    int i;
    double mean = 0;
    double sum = 0;
    double *arr = integrated_support_degree_matrix;
    if ((sensor_count == 0) || (arr == NULL)) {
        printf("ERROR: Invalid param passed at %s\n", __func__);
        return -1;
    }
    for (i = 0; i < sensor_count; i++) {
        sum += arr[i];
    }
    mean = sum / (i + 1);
    for (i = 0; i < sensor_count; i++) {
        if (fabs(arr[i]) < fabs(fault_tolerance_threshold * mean)) {
            arr[i] = 0;
        }
    }
    return 0;
}

/**
 * \brief Function to compute weight coefficient for each sensor
 * \details Compute weighted coefficients for each sensors from the support degree matrix
 * \param[in] integrated_support_degree_matrix pointer to integrated support degree matrix
 * \param[in] sensor_count total count of sensors
 * \return Success: array of weighted coefficients,
 * \return Failure: NULL
 */
double *compute_weight_coefficient(double *integrated_support_degree_matrix,
                                   int sensor_count) {
    double sum_weight_coefficient = 0;
    if ((sensor_count == 0) || (integrated_support_degree_matrix == NULL)) {
        printf("ERROR: Invalid param passed at %s\n", __func__);
        return NULL;
    }
    double *arr = (double *)malloc(sensor_count * sizeof(double));
    if (arr == NULL) {
        printf("ERROR: Failed to allocate memory at %s\n", __func__);
        return NULL;
    }
    memset(arr, 0, sensor_count * sizeof(double));
    for (int i = 0; i < sensor_count; i++) {
        sum_weight_coefficient += integrated_support_degree_matrix[i];
    }
    for (int i = 0; i < sensor_count; i++) {
        arr[i] = integrated_support_degree_matrix[i] / sum_weight_coefficient;
    }
    return arr;
}

/**
 * \brief Function to compute fused output
 * \details Compute the aggregated value of sensors from the weighted coefficients
 * \param[in] weight_coefficient pointer to weighted coefficient of sensors
 * \param[in] sensor_data pointer to data of sensors
 * \param[in] sensor_count total count of sensors
 * \return Success: fused value of sensor data
 * \return Failure: -1
 */
double compute_fused_output(double *weight_coefficient,
                            double *sensor_data,
                            int sensor_count) {
    double fused_output = 0;
    if ((sensor_count == 0) || (weight_coefficient == NULL) || (sensor_data == NULL)) {
        printf("ERROR: Invalid param passed at %s\n", __func__);
        return -1;
    }
    for (int i = 0; i < sensor_count; i++) {
        fused_output += weight_coefficient[i] * sensor_data[i];
    }
    return fused_output;
}


/* BalSelSims.c

Simulation of balancing selection simulation with sampling (finite population effects)
Selection; reproduction; mutation based on deterministic recursions, then sampling
Repeats for 'Length' generations

Simulation uses routines found with the GNU Scientific Library (GSL)
(http://www.gnu.org/software/gsl/)
Since GSL is distributed under the GNU General Public License
(http://www.gnu.org/copyleft/gpl.html), you must download it
separately from this file.

This program can be compiled in e.g. GCC using a command like:
gcc BalSelSims -lm -lgsl -lgslcblas -I/usr/local/include -L/usr/local/lib BalSelSims.c

Then run by executing:
./BalSelSims N s rec sex self gc reps
Where:
- N is the population size
- s is the fitness disadvantage of homozygotes
- rec is recombination rate
- sex is rate of sex (a value between 0 = obligate asex, and 1 = obligate sex)
- self is selfing rate
- gc is gene conversion
- reps is how many times to introduce linked neutral allele

Note that haplotypes are defined as:
x1 = ab
x2 = Ab
x3 = aB
x4 = AB

Genotypes defined as:
g11 = g1 = ab/ab
g12 = g2 = Ab/ab
g13 = g3 = aB/ab
g14 = g4 = AB/ab
g22 = g5 = Ab/Ab
g23 = g6 = Ab/aB
g24 = g7 = Ab/AB
g33 = g8 = aB/aB
g34 = g9 = aB/AB
g44 = g10 = AB/AB

*/

/* Preprocessor statements */
#include <stdio.h>
#include <time.h>
#include <math.h>
#include <stddef.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

/* Function prototypes */
void geninit(double *geninit);
void selection(double *geninit);
void reproduction(double *geninit);
void gconv(double *geninit);
void neutinit(double *geninit,const gsl_rng *r);
double ncheck(double *geninit);
double pcheck(double *geninit);

/* Global variable declaration */
unsigned int N = 0;		/* Pop size */
double s = 0;			/* Fitness disadvantage of homozyogtes */
double rec = 0;			/* Recombination rate */
double sex = 0;			/* Rate of sexual reproduction */
double self = 0;		/* Rate of self-fertilisation */
double gc = 0;			/* Rate of gene conversion */

/* Main program */
int main(int argc, char *argv[]){
	unsigned int g, i; 				/* Counters. Reps counter, geno counter */
	unsigned int reps;				/* Length of simulation (no. of introductions of neutral site) */
	double Bcheck = 0;				/* Frequency of B after each reproduction */
	double Acheck = 0;				/* Frequency of polymorphism */
	double Hsum = 0;				/* Summed heterozygosity over transit time of neutral allele */

	/* GSL random number definitions */
	const gsl_rng_type * T;
	gsl_rng * r;

	/* This reads in data from command line. */
	if(argc != 8){
		fprintf(stderr,"Invalid number of input values.\n");
		exit(1);
	}
	N = strtod(argv[1],NULL);
	s = strtod(argv[2],NULL);
	rec = strtod(argv[3],NULL);
	sex = strtod(argv[4],NULL);
	self = strtod(argv[5],NULL);
	gc = strtod(argv[6],NULL);
	reps = strtod(argv[7],NULL);

	/* Arrays definition and memory assignment */
	double *genotype = calloc(10,sizeof(double));				/* Genotype frequencies */
	unsigned int *gensamp = calloc(10,sizeof(unsigned int));	/* New population samples */

	/* create a generator chosen by the
    environment variable GSL_RNG_TYPE */

	gsl_rng_env_setup();
	if (!getenv("GSL_RNG_SEED")) gsl_rng_default_seed = time(0);
	T = gsl_rng_default;
	r = gsl_rng_alloc(T);

	/* Initialising genotypes */
	geninit(genotype);

    /* Run simulation for 2000 generations to create a burn in */
    for(g = 0; g < 2000; g++){

    	/* Selection routine */
    	selection(genotype);

    	/* Reproduction routine */
    	reproduction(genotype);

       	/* Gene conversion routine */
       	gconv(genotype);

       	/* Sampling based on new frequencies */
       	gsl_ran_multinomial(r,10,N,genotype,gensamp);
       	for(i = 0; i < 10; i++){
			*(genotype + i) = (*(gensamp + i))/(1.0*N);
       	}

       	/* Printing out results (for testing) */
		/*
	   	for(i = 0; i < 10; i++){
			printf("%.10lf ", *(genotype + i));
		}
		printf("\n");
		*/
    }

	/* Reintroducing neutral genotype, resetting hap sum */
	neutinit(genotype,r);
    Bcheck = ncheck(genotype);
    Hsum = Bcheck*(1-Bcheck);
    /* printf("%.10lf %.10lf\n",Bcheck,Hsum); */

    /* Introduce and track neutral mutations 'reps' times */
    g = 0;
    while(g < reps){

    	/* Selection routine */
    	selection(genotype);

    	/* Reproduction routine */
    	reproduction(genotype);

    	/* Gene conversion routine */
       	gconv(genotype);

    	/* Sampling based on new frequencies */
       	gsl_ran_multinomial(r,10,N,genotype,gensamp);
       	for(i = 0; i < 10; i++){
       		*(genotype + i) = (*(gensamp + i))/(1.0*N);
       	}

       	/* Checking state of haplotypes: if B fixed reset so can start fresh next time */
		Bcheck = ncheck(genotype);
		Hsum += Bcheck*(1-Bcheck);
		/* printf("%.10lf %.10lf\n",Bcheck,Hsum); */

		/* If polymorphism fixed then abandon simulation */
       	Acheck = pcheck(genotype);
       	if(Acheck == 0){
       		g = reps;
       	}

       	if(Bcheck == 0 || Bcheck == 1){
       		printf("%.10lf\n",Hsum);
       		g++;
       		/* printf("Rep Number %d\n",g); */

       		if(Bcheck == 1){
       			/* Reset genotypes so B becomes ancestral allele */
       			*(genotype + 0) = *(genotype + 7);
       			*(genotype + 1) = *(genotype + 8);
       			*(genotype + 4) = *(genotype + 9);
       			*(genotype + 7) = 0;
       			*(genotype + 8) = 0;
       			*(genotype + 9) = 0;
       		}

     		/* Reintroducing neutral genotype, resetting hap sum */
	    	neutinit(genotype,r);
			Bcheck = ncheck(genotype);
       		Hsum = Bcheck*(1-Bcheck);
       	}

	}	/* End of simulation */

	/* Freeing memory and wrapping up */
 	gsl_rng_free(r);
 	free(gensamp);
	free(genotype);
	/* printf("The End!\n"); */
	return 0;
}

/* Initialising genotypes */
void geninit(double *geninit){

	/* Basic idea: before neutral mutation introduced, g0 = 1/4; g1 = 1/2; g4 = 1/4. Deterministic frequencies. */

	*(geninit + 0) = 0.25;
	*(geninit + 1) = 0.50;
	*(geninit + 2) = 0;
	*(geninit + 3) = 0;
	*(geninit + 4) = 0.25;
	*(geninit + 5) = 0;
	*(geninit + 6) = 0;
	*(geninit + 7) = 0;
	*(geninit + 8) = 0;
	*(geninit + 9) = 0;

}	/* End of gen initiation routine */

/* Initialising NEUTRAL allele */
void neutinit(double *geninit, const gsl_rng *r){

	double *probin = calloc(4,sizeof(double));					/* Probability inputs, determine location of neutral allele */
	unsigned int *probout = calloc(4,sizeof(unsigned int));		/* Output from multinomial sampling */

	/* Basic idea: g11 at freq p^2; g12 at freq 2pq; g22 at freq q^2.
	So weighted sampling to determine which genotype the neutral allele arises on
	*/

	/* Prob definitions */
	*(probin + 0) = *(geninit + 0);
	*(probin + 1) = (*(geninit + 1))/2.0;
	*(probin + 2) = (*(geninit + 1))/2.0;
	*(probin + 3) = *(geninit + 4);

	gsl_ran_multinomial(r,4,1,probin,probout);

	/* Redefining genotypes depending on outcome */
	if(*(probout + 0) == 1){
		*(geninit + 0) = (*(geninit + 0)) - 1/(1.0*N);
		*(geninit + 2) = (*(geninit + 2)) + 1/(1.0*N);
	}
	else if(*(probout + 1) == 1){
		*(geninit + 1) = (*(geninit + 1)) - 1/(1.0*N);
		*(geninit + 3) = (*(geninit + 3)) + 1/(1.0*N);
	}
	else if(*(probout + 2) == 1){
		*(geninit + 1) = (*(geninit + 1)) - 1/(1.0*N);
		*(geninit + 5) = (*(geninit + 5)) + 1/(1.0*N);
	}
	else if(*(probout + 3) == 1){
		*(geninit + 4) = (*(geninit + 4)) - 1/(1.0*N);
		*(geninit + 6) = (*(geninit + 6)) + 1/(1.0*N);
	}

 	free(probout);
	free(probin);

}	/* End of gen initiation routine */

/* Selection routine */
void selection(double *geninit){
	double Waa, WAa, WAA;		/* Fitness of locus A (bal sel locus) */
	double Wmean;				/* Mean fitness */

	Waa = 1-s;
	WAa = 1;
	WAA = 1-s;

	/* Mean fitness calculation */
	Wmean = ((*(geninit + 0))*Waa) + ((*(geninit + 1))*WAa) + ((*(geninit + 2))*Waa) + ((*(geninit + 3))*WAa) + ((*(geninit + 4))*WAA) + ((*(geninit + 5))*WAa) + ((*(geninit + 6))*WAA) + ((*(geninit + 7))*Waa) + ((*(geninit + 8))*WAa) + ((*(geninit + 9))*WAA);

	/* Changing frequencies by selection */
	*(geninit + 0) = ((*(geninit + 0))*Waa)/Wmean;
	*(geninit + 1) = ((*(geninit + 1))*WAa)/Wmean;
	*(geninit + 2) = ((*(geninit + 2))*Waa)/Wmean;
	*(geninit + 3) = ((*(geninit + 3))*WAa)/Wmean;
	*(geninit + 4) = ((*(geninit + 4))*WAA)/Wmean;
	*(geninit + 5) = ((*(geninit + 5))*WAa)/Wmean;
	*(geninit + 6) = ((*(geninit + 6))*WAA)/Wmean;
	*(geninit + 7) = ((*(geninit + 7))*Waa)/Wmean;
	*(geninit + 8) = ((*(geninit + 8))*WAa)/Wmean;
	*(geninit + 9) = ((*(geninit + 9))*WAA)/Wmean;

}	/* End of selection routine */

/* Reproduction routine */
void reproduction(double *geninit){
	/* Fed-in genotype frequencies (for ease of programming) */
	double g11s, g12s, g13s, g14s, g22s, g23s, g24s, g33s, g34s, g44s;
	/* Genotype frequencies after sex (outcross and selfing) */
	double g11SX, g12SX, g13SX, g14SX, g22SX, g23SX, g24SX, g33SX, g34SX, g44SX;
	/* Genotype frequencies after ASEX */
	double g11AS, g12AS, g13AS, g14AS, g22AS, g23AS, g24AS, g33AS, g34AS, g44AS;
	/* Haplotypes */
	double x1, x2, x3, x4;

	/* Initial definition of genotypes */
	g11s = *(geninit + 0);
	g12s = *(geninit + 1);
	g13s = *(geninit + 2);
	g14s = *(geninit + 3);
	g22s = *(geninit + 4);
	g23s = *(geninit + 5);
	g24s = *(geninit + 6);
	g33s = *(geninit + 7);
	g34s = *(geninit + 8);
	g44s = *(geninit + 9);

	/* Baseline change in haplotype frequencies */
	x1 = g11s + (g12s + g13s + g14s)/2.0 - ((g14s - g23s)*rec)/2.0;
	x2 = g22s + (g12s + g23s + g24s)/2.0 + ((g14s - g23s)*rec)/2.0;
	x3 = g33s + (g13s + g23s + g34s)/2.0 + ((g14s - g23s)*rec)/2.0;
	x4 = g44s + (g14s + g24s + g34s)/2.0 - ((g14s - g23s)*rec)/2.0;

	/* Change in SEXUAL frequencies (both outcrossing and selfing) */
	g11SX = (g11s + (g12s + g13s + g14s*pow((1 - rec),2) + g23s*pow(rec,2))/4.0)*self*sex + (1 - self)*pow(x1,2)*sex;
	g22SX = (g22s + (g12s + g24s + g23s*pow((1 - rec),2) + g14s*pow(rec,2))/4.0)*self*sex + (1 - self)*pow(x2,2)*sex;
	g33SX = (g33s + (g13s + g34s + g23s*pow((1 - rec),2) + g14s*pow(rec,2))/4.0)*self*sex + (1 - self)*pow(x3,2)*sex;
	g44SX = (g44s + (g24s + g34s + g14s*pow((1 - rec),2) + g23s*pow(rec,2))/4.0)*self*sex + (1 - self)*pow(x4,2)*sex;
	g12SX = ((g12s + (g14s + g23s)*(1 - rec)*rec)*self*sex)/2.0 + 2.0*(1 - self)*x1*x2*sex;
	g13SX = ((g13s + (g14s + g23s)*(1 - rec)*rec)*self*sex)/2.0 + 2.0*(1 - self)*x1*x3*sex;
	g14SX = ((g14s*pow((1 - rec),2) + g23s*pow(rec,2))*self*sex)/2.0 + 2.0*(1 - self)*x1*x4*sex;
	g23SX = ((g23s*pow((1 - rec),2) + g14s*pow(rec,2))*self*sex)/2.0 + 2.0*(1 - self)*x2*x3*sex;
	g24SX = ((g24s + (g14s + g23s)*(1 - rec)*rec)*self*sex)/2.0 + 2.0*(1 - self)*x2*x4*sex;
	g34SX = ((g34s + (g14s + g23s)*(1 - rec)*rec)*self*sex)/2.0 + 2.0*(1 - self)*x3*x4*sex;

	/* Change in ASEXUAL frequencies */
	g11AS = g11s*(1 - sex);
	g12AS = g12s*(1 - sex);
	g13AS = g13s*(1 - sex);
	g14AS = g14s*(1 - sex);
	g22AS = g22s*(1 - sex);
	g23AS = g23s*(1 - sex);
	g24AS = g24s*(1 - sex);
	g33AS = g33s*(1 - sex);
	g34AS = g34s*(1 - sex);
	g44AS = g44s*(1 - sex);

	/* Combining to give overall frequency change following reproduction */
	*(geninit + 0) = g11AS + g11SX;
	*(geninit + 1) = g12AS + g12SX;
	*(geninit + 2) = g13AS + g13SX;
	*(geninit + 3) = g14AS + g14SX;
	*(geninit + 4) = g22AS + g22SX;
	*(geninit + 5) = g23AS + g23SX;
	*(geninit + 6) = g24AS + g24SX;
	*(geninit + 7) = g33AS + g33SX;
	*(geninit + 8) = g34AS + g34SX;
	*(geninit + 9) = g44AS + g44SX;

}	/* End of reproduction routine */

/* Gene conversion routine */
void gconv(double *geninit){

	/* Fed-in genotype frequencies (for ease of programming) */
	double g11r, g12r, g13r, g14r, g22r, g23r, g24r, g33r, g34r, g44r;
	/* Frequencies after gene conversion */
	double g11gc, g12gc, g13gc, g14gc, g22gc, g23gc, g24gc, g33gc, g34gc, g44gc;

	/* Initial definition of genotypes */
	g11r = *(geninit + 0);
	g12r = *(geninit + 1);
	g13r = *(geninit + 2);
	g14r = *(geninit + 3);
	g22r = *(geninit + 4);
	g23r = *(geninit + 5);
	g24r = *(geninit + 6);
	g33r = *(geninit + 7);
	g34r = *(geninit + 8);
	g44r = *(geninit + 9);

	/* Gene conversion equations */
	g11gc = g11r + (gc*g12r)/4.0 + (gc*g13r)/4.0;
	g12gc = g12r*(1 - gc/2.0) + ((g14r + g23r)*gc)/4.0;
	g13gc = g13r*(1 - gc/2.0) + ((g14r + g23r)*gc)/4.0;
	g14gc = (1-gc)*g14r;
	g22gc = g22r + (g12r*gc)/4.0 + (g24r*gc)/4.0;
	g23gc = (1-gc)*g23r;
	g24gc = g24r*(1 - gc/2.0) + ((g14r + g23r)*gc)/4.0;
	g33gc = g33r + (g13r*gc)/4.0 + (g34r*gc)/4.0;
	g34gc = g34r*(1 - gc/2.0) + ((g14r + g23r)*gc)/4.0;
	g44gc = g44r + (g24r*gc)/4.0 + (g34r*gc)/4.0;

	/* Output */
	*(geninit + 0) = g11gc;
	*(geninit + 1) = g12gc;
	*(geninit + 2) = g13gc;
	*(geninit + 3) = g14gc;
	*(geninit + 4) = g22gc;
	*(geninit + 5) = g23gc;
	*(geninit + 6) = g24gc;
	*(geninit + 7) = g33gc;
	*(geninit + 8) = g34gc;
	*(geninit + 9) = g44gc;
}		/* End of gene conversion routine */

/* Has neutral allele fixed or not?  Measuring freq of B */
double ncheck(double *geninit){
	/* Fed-in genotype frequencies (for ease of programming) */
	double g11s, g12s, g13s, g14s, g22s, g23s, g24s, g33s, g34s, g44s;
	/* Haplotypes ONLY CONTAINING B */
	double x3, x4;
	double Btot = 0;        /* Total frequency of B */

	/* Initial definition of genotypes */
	g11s = *(geninit + 0);
	g12s = *(geninit + 1);
	g13s = *(geninit + 2);
	g14s = *(geninit + 3);
	g22s = *(geninit + 4);
	g23s = *(geninit + 5);
	g24s = *(geninit + 6);
	g33s = *(geninit + 7);
	g34s = *(geninit + 8);
	g44s = *(geninit + 9);

	/* Calculation of haplotypes containing B */

	x3 = g33s + (g13s + g23s + g34s)/2.0;
	x4 = g44s + (g14s + g24s + g34s)/2.0;

	/* Checking */
	Btot = x3 + x4;
	return Btot;

}	/* End of B check routine */

/* Checking if balancing polymorphism lost or not */
double pcheck(double *geninit){
	/* Fed-in genotype frequencies (for ease of programming) */
	double g11s, g12s, g13s, g14s, g22s, g23s, g24s, g33s, g34s, g44s;
	double x2, x4;
	double Atot = 0;        /* Total frequency of A */

	/* Initial definition of genotypes */
	g11s = *(geninit + 0);
	g12s = *(geninit + 1);
	g13s = *(geninit + 2);
	g14s = *(geninit + 3);
	g22s = *(geninit + 4);
	g23s = *(geninit + 5);
	g24s = *(geninit + 6);
	g33s = *(geninit + 7);
	g34s = *(geninit + 8);
	g44s = *(geninit + 9);

	/* Calculation of haplotypes containing B */

	x2 = g22s + (g12s + g23s + g24s)/2.0;
	x4 = g44s + (g14s + g24s + g34s)/2.0;

	/* Checking */
	Atot = x2 + x4;
	return Atot;

}	/* End of A check routine */

/* End of program */


#include <math.h>
#define FORCE_OPENBLAS_COMPLEX_STRUCT
#include <cblas.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef int MKL_INT;
typedef openblas_complex_double MKL_Complex16;

#define AVX512_ALIGN 64

#include "gene_bruit_rayleigh_scalaire.c"

int main(void)
{

  int     val;
  int 	  M, N, K;
  int 	  lda, ldb, ldc;
  int     kboucle;
  float   charge;
  char    *p;

  FILE *fichier1;
  fichier1 = fopen("../results/results_openblas_gemm_fp64.dat", "w");

  int has_param_m = 0, has_param_n = 0, has_param_k = 0;
  char transa = CblasNoTrans, transb = CblasNoTrans;

  if ((p = getenv("MATMUL_M"))) {
    M = atoi(p);
    has_param_m=1;
  }
  if ((p = getenv("MATMUL_N"))) {
    N = atoi(p);
    has_param_n=1;
  }
  if ((p = getenv("MATMUL_K"))) {
    K = atoi(p);
    has_param_k=1;
  }

  if ((p = getenv("TRANSA"))) {
    if(*p == 'T')
		transa = CblasTrans;
  }
  if ((p = getenv("TRANSB"))) {
    if(*p == 'T')
		transb = CblasTrans;
  }

  /* Begin huge loop */
  for (kboucle = 1; kboucle < 31; kboucle++) {

    val = 100 * kboucle;
    M = has_param_m?M:val;
    N = has_param_n?N:val;
	K = has_param_k?K:val;

	lda = (transa == CblasNoTrans)?M:K;
    ldb = (transb == CblasNoTrans)?K:N;
    ldc = M;

    printf(">>>>> Matrix size %dx%d (K=%d, TransA=%d TransB=%d) <<<<<<\n", M, N, K,\
		transa - CblasNoTrans, transb - CblasNoTrans);

    int    iboucle, jboucle;
    int    param = 20;
    float *mat_real, *mat_imag;

    // Timers
    struct timespec tpdeb, tpfin, tpcour;
    clockid_t clock_id = CLOCK_REALTIME;
    int status2;

    double dureeloc, dureetot;
    dureetot = 0.0;

    // BLAS
    MKL_Complex16 alpha, beta;
    MKL_Complex16 *A, *B, *C;

    alpha.real = 1.0;
    alpha.imag = 0.0;
    beta.real  = 0.0;
    beta.imag  = 0.0;

    mat_real = (float *)aligned_alloc(AVX512_ALIGN, val * val * sizeof(float));
    mat_imag = (float *)aligned_alloc(AVX512_ALIGN, val * val * sizeof(float));

    A = (MKL_Complex16 *)aligned_alloc(AVX512_ALIGN, M * K * sizeof(MKL_Complex16));
    B = (MKL_Complex16 *)aligned_alloc(AVX512_ALIGN, K * N * sizeof(MKL_Complex16));
    C = (MKL_Complex16 *)aligned_alloc(AVX512_ALIGN, M * N * sizeof(MKL_Complex16));

    for (iboucle = 0; iboucle < val * val; iboucle++) {
      gene_bruit_rayleigh_scalaire(param, mat_real + iboucle,
                                   mat_imag + iboucle);
    }

    // Force hermitian matrix
    for (iboucle = 0; iboucle < K; iboucle++) {
      for (jboucle = 0; jboucle < N; jboucle++) {
        B[(iboucle * N) + jboucle].real = (double)(*(mat_real + (jboucle * K) + iboucle));

        B[(iboucle * N) + jboucle].imag =
            -1.0 * (double)(*(mat_imag + (jboucle * K) + iboucle));
      }
    }

    for (iboucle = 0; iboucle < M * K; iboucle++) {
      A[iboucle].real = (double)mat_real[iboucle];
      A[iboucle].imag = (double)mat_imag[iboucle];
    }

    status2 = clock_gettime(clock_id, &tpdeb); // start timer;

    cblas_zgemm(CblasColMajor, transa, transb, \
        M, N, K, \
        &alpha, A, lda, \
        B, ldb, &beta, C, ldc);

    status2 = clock_gettime(clock_id, &tpfin); // stop timer

    dureeloc = (float)(tpfin.tv_sec - tpdeb.tv_sec) +
               (float)(tpfin.tv_nsec - tpdeb.tv_nsec) * 1.e-9;
    dureetot = dureetot + dureeloc;

    printf("execution time = %f ms\n", dureeloc * 1000);
    charge = (float)M;
    charge = 8.0f * charge * charge * charge;
    charge = charge / dureeloc;
    printf("%f GFLOPS\n", charge / 1e9);
    fprintf(fichier1, "%5.2f\n", charge / 1e9);

    free(mat_real);
    free(mat_imag);

    free(A);
    free(B);
    free(C);
  }
  fclose(fichier1);

  return 0;
}


/* multifit_nlinear/mcholesky.c
 *
 * Copyright (C) 2015, 2016 Patrick Alken
 *
 * 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 3 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 Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/*
 * This module calculates the solution of the normal equations least squares
 * system:
 *
 * [ J~^T J~ + mu D~^T D~ ] p~ = -J~^T f
 *
 * using the modified Cholesky decomposition. Quantities are scaled
 * according to:
 *
 * J~ = J S
 * D~ = D S
 * p~ = S^{-1} p
 *
 * where S is a diagonal matrix and S_jj = || J_j || and J_j is column
 * j of the Jacobian. This balancing transformation seems to be more
 * numerically stable for some Jacobians.
 */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_multifit_nlinear.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_permutation.h>

#include "common.c"

typedef struct
{
  gsl_matrix *JTJ;           /* J^T J */
  gsl_matrix *work_JTJ;      /* copy of J^T J */
  gsl_vector *rhs;           /* -J^T f, size p */
  gsl_permutation *perm;     /* permutation matrix for modified Cholesky */
  gsl_vector *work3p;        /* workspace, size 3*p */
  double mu;                 /* current regularization parameter */
} mcholesky_state_t;

static void *mcholesky_alloc (const size_t n, const size_t p);
static int mcholesky_init(const void * vtrust_state, void * vstate);
static int mcholesky_presolve(const double mu, const void * vtrust_state, void * vstate);
static int mcholesky_solve(const gsl_vector * f, gsl_vector *x,
                           const  void * vtrust_state, void *vstate);
static int mcholesky_solve_rhs(const gsl_vector * b, gsl_vector *x, mcholesky_state_t *state);
static int mcholesky_regularize(const double mu, const gsl_vector * diag, gsl_matrix * A,
                                mcholesky_state_t * state);

static void *
mcholesky_alloc (const size_t n, const size_t p)
{
  mcholesky_state_t *state;

  (void)n;

  state = calloc(1, sizeof(mcholesky_state_t));
  if (state == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate mcholesky state", GSL_ENOMEM);
    }

  state->JTJ = gsl_matrix_alloc(p, p);
  if (state->JTJ == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for JTJ", GSL_ENOMEM);
    }

  state->work_JTJ = gsl_matrix_alloc(p, p);
  if (state->work_JTJ == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for JTJ workspace",
                      GSL_ENOMEM);
    }

  state->rhs = gsl_vector_alloc(p);
  if (state->rhs == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for rhs", GSL_ENOMEM);
    }

  state->perm = gsl_permutation_alloc(p);
  if (state->perm == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for perm", GSL_ENOMEM);
    }

  state->work3p = gsl_vector_alloc(3 * p);
  if (state->work3p == NULL)
    {
      GSL_ERROR_NULL ("failed to allocate space for work3p", GSL_ENOMEM);
    }

  state->mu = -1.0;

  return state;
}

static void
mcholesky_free(void *vstate)
{
  mcholesky_state_t *state = (mcholesky_state_t *) vstate;

  if (state->JTJ)
    gsl_matrix_free(state->JTJ);

  if (state->work_JTJ)
    gsl_matrix_free(state->work_JTJ);

  if (state->rhs)
    gsl_vector_free(state->rhs);

  if (state->perm)
    gsl_permutation_free(state->perm);

  if (state->work3p)
    gsl_vector_free(state->work3p);

  free(state);
}

static int
mcholesky_init(const void * vtrust_state, void * vstate)
{
  const gsl_multifit_nlinear_trust_state *trust_state =
    (const gsl_multifit_nlinear_trust_state *) vtrust_state;
  mcholesky_state_t *state = (mcholesky_state_t *) vstate;

  /* compute J^T J */
  gsl_blas_dsyrk(CblasLower, CblasTrans, 1.0, trust_state->J, 0.0, state->JTJ);

  return GSL_SUCCESS;
}

/*
mcholesky_presolve()
  Compute the modified Cholesky decomposition of J^T J + mu D^T D.
Modified Cholesky is used in case mu = 0 and there are rounding
errors in forming J^T J which could lead to an indefinite matrix.

Inputs: mu     - LM parameter
        vstate - workspace

Notes:
1) On output, state->work_JTJ contains the Cholesky decomposition of
J^T J + mu D^T D
*/

static int
mcholesky_presolve(const double mu, const void * vtrust_state, void * vstate)
{
  const gsl_multifit_nlinear_trust_state *trust_state =
    (const gsl_multifit_nlinear_trust_state *) vtrust_state;
  mcholesky_state_t *state = (mcholesky_state_t *) vstate;
  gsl_matrix *JTJ = state->work_JTJ;
  const gsl_vector *diag = trust_state->diag;
  int status;

  /* copy lower triangle of A to workspace */
  gsl_matrix_tricpy(CblasLower, CblasNonUnit, JTJ, state->JTJ);

  /* augment normal equations: A -> A + mu D^T D */
  status = mcholesky_regularize(mu, diag, JTJ, state);
  if (status)
    return status;

  /* compute modified Cholesky decomposition */
  status = gsl_linalg_mcholesky_decomp(JTJ, state->perm, NULL);
  if (status)
    return status;

  state->mu = mu;

  return GSL_SUCCESS;
}

/*
mcholesky_solve()
  Compute (J^T J + mu D^T D) x = -J^T f

Inputs: f      - right hand side vector f
        x      - (output) solution vector
        vstate - mcholesky workspace
*/

static int
mcholesky_solve(const gsl_vector * f, gsl_vector *x,
               const  void * vtrust_state, void *vstate)
{
  const gsl_multifit_nlinear_trust_state *trust_state =
    (const gsl_multifit_nlinear_trust_state *) vtrust_state;
  mcholesky_state_t *state = (mcholesky_state_t *) vstate;
  int status;

  /* compute rhs = -J^T f */
  gsl_blas_dgemv(CblasTrans, -1.0, trust_state->J, f, 0.0, state->rhs);

  status = mcholesky_solve_rhs(state->rhs, x, state);
  if (status)
    return status;

  return GSL_SUCCESS;
}

static int
mcholesky_rcond(double * rcond, void * vstate)
{
  int status;
  mcholesky_state_t *state = (mcholesky_state_t *) vstate;
  double rcond_JTJ;

  if (state->mu != 0)
    {
      /*
       * Cholesky decomposition hasn't been computed yet, or was computed
       * with mu > 0 - recompute Cholesky decomposition of J^T J
       */

      /* copy lower triangle of JTJ to workspace */
      gsl_matrix_tricpy(CblasLower, CblasNonUnit, state->work_JTJ, state->JTJ);

      /* compute modified Cholesky decomposition */
      status = gsl_linalg_mcholesky_decomp(state->work_JTJ, state->perm, NULL);
      if (status)
        return status;
    }

  status = gsl_linalg_mcholesky_rcond(state->work_JTJ, state->perm, &rcond_JTJ, state->work3p);
  if (status == GSL_SUCCESS)
    *rcond = sqrt(rcond_JTJ);

  return status;
}

/* solve: (J^T J + mu D^T D) x = b */
static int
mcholesky_solve_rhs(const gsl_vector * b, gsl_vector *x, mcholesky_state_t *state)
{
  int status;
  gsl_matrix *JTJ = state->work_JTJ;

  status = gsl_linalg_mcholesky_solve(JTJ, state->perm, b, x);
  if (status)
    return status;

  return GSL_SUCCESS;
}

/* A <- A + mu D^T D */
static int
mcholesky_regularize(const double mu, const gsl_vector * diag, gsl_matrix * A,
                     mcholesky_state_t * state)
{
  (void) state;

  if (mu != 0.0)
    {
      size_t i;

      for (i = 0; i < diag->size; ++i)
        {
          double di = gsl_vector_get(diag, i);
          double *Aii = gsl_matrix_ptr(A, i, i);
          *Aii += mu * di * di;
        }
    }

  return GSL_SUCCESS;
}

static const gsl_multifit_nlinear_solver mcholesky_type =
{
  "mcholesky",
  mcholesky_alloc,
  mcholesky_init,
  mcholesky_presolve,
  mcholesky_solve,
  mcholesky_rcond,
  mcholesky_free
};

const gsl_multifit_nlinear_solver *gsl_multifit_nlinear_solver_mcholesky = &mcholesky_type;


/*
  NAME:
     splitnmergegauss
  PURPOSE:
     split one gaussian and merge two other gaussians
  CALLING SEQUENCE:
     splitnmergegauss(struct gaussian * gaussians,int K, gsl_matrix * qij,
     int j, int k, int l)
  INPUT:
     gaussians   - model gaussians
     K           - number of gaussians
     qij         - matrix of log(posterior likelihoods)
     j,k         - gaussians that need to be merged
     l           - gaussian that needs to be split
  OUTPUT:
     updated gaussians
  REVISION HISTORY:
     2008-09-21 - Written Bovy
*/
#include <math.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_linalg.h>
#include "proj_gauss_mixtures.h"

void splitnmergegauss(struct gaussian * gaussians,int K,
		      gsl_matrix * qij, int j, int k, int l){
  //get the gaussians to be split 'n' merged
  int d = (gaussians->VV)->size1;//dim of mm
  //int partial_indx[]= {-1,-1,-1};/* dummy argument for logsum */
  //bool * dummy_allfixed = (bool *) calloc(K,sizeof(bool));
  //j,k,l gaussians
  struct gaussian gaussianj, gaussiank, gaussianl;
  gaussianj.mm = gsl_vector_alloc(d);
  gaussianj.VV = gsl_matrix_alloc(d,d);
  gaussiank.mm = gsl_vector_alloc(d);
  gaussiank.VV = gsl_matrix_alloc(d,d);
  gaussianl.mm = gsl_vector_alloc(d);
  gaussianl.VV = gsl_matrix_alloc(d,d);

  gsl_matrix * unitm = gsl_matrix_alloc(d,d);
  gsl_matrix_set_identity(unitm);
  gsl_vector * eps = gsl_vector_alloc(d);
  double qjj,qjk,detVVjl;
  int kk;
  for (kk = 0; kk != K; ++kk){
    if (kk == j){
      gaussianj.alpha = gaussians->alpha;
      gsl_vector_memcpy(gaussianj.mm,gaussians->mm);
      gsl_matrix_memcpy(gaussianj.VV,gaussians->VV);
      qjj = exp(logsum(qij,j,false));//,dummy_allfixed));
    }
    if (kk == k){
      gaussiank.alpha = gaussians->alpha;
      gsl_vector_memcpy(gaussiank.mm,gaussians->mm);
      gsl_matrix_memcpy(gaussiank.VV,gaussians->VV);
      qjk = exp(logsum(qij,k,false));//,dummy_allfixed));
    }
    if (kk == l){
      gaussianl.alpha = gaussians->alpha;
      gsl_vector_memcpy(gaussianl.mm,gaussians->mm);
      gsl_matrix_memcpy(gaussianl.VV,gaussians->VV);
    }
    ++gaussians;
  }
  gaussians -= K;

  //merge j & k
  gaussianj.alpha += gaussiank.alpha;
  if (qjk == 0. && qjj == 0){
    gsl_vector_add(gaussianj.mm,gaussiank.mm);
    gsl_vector_scale(gaussianj.mm,0.5);
    gsl_matrix_add(gaussianj.VV,gaussiank.VV);
    gsl_matrix_scale(gaussianj.VV,0.5);
  }
  else{
    gsl_vector_scale(gaussianj.mm,qjj/(qjj+qjk));
    gsl_vector_scale(gaussiank.mm,qjk/(qjj+qjk));
    gsl_vector_add(gaussianj.mm,gaussiank.mm);
    gsl_matrix_scale(gaussianj.VV,qjj/(qjj+qjk));
    gsl_matrix_scale(gaussiank.VV,qjk/(qjj+qjk));
    gsl_matrix_add(gaussianj.VV,gaussiank.VV);
  }

  //split l
  gaussianl.alpha /= 2.;
  gaussiank.alpha = gaussianl.alpha;
  detVVjl = bovy_det(gaussianl.VV);
  detVVjl= pow(detVVjl,1./d);
  gsl_matrix_scale(unitm,detVVjl);
  gsl_matrix_memcpy(gaussiank.VV,unitm);
  gsl_matrix_memcpy(gaussianl.VV,unitm);
  gsl_vector_memcpy(gaussiank.mm,gaussianl.mm);
  bovy_randvec(eps,d,sqrt(detVVjl));
  gsl_vector_add(gaussiank.mm,eps);
  bovy_randvec(eps,d,sqrt(detVVjl));
  gsl_vector_add(gaussianl.mm,eps);

  //copy everything back into the right gaussians
  for (kk = 0; kk != K; ++kk){
    if (kk == j){
      gaussians->alpha = gaussianj.alpha;
      gsl_vector_memcpy(gaussians->mm,gaussianj.mm);
      gsl_matrix_memcpy(gaussians->VV,gaussianj.VV);
    }
    if (kk == k){
      gaussians->alpha = gaussiank.alpha;
      gsl_vector_memcpy(gaussians->mm,gaussiank.mm);
      gsl_matrix_memcpy(gaussians->VV,gaussiank.VV);
    }
    if (kk == l){
      gaussians->alpha = gaussianl.alpha;
      gsl_vector_memcpy(gaussians->mm,gaussianl.mm);
      gsl_matrix_memcpy(gaussians->VV,gaussianl.VV);
    }
    ++gaussians;
  }
  gaussians -= K;

  //cleanup
  gsl_matrix_free(unitm);
  gsl_vector_free(eps);
  //free(dummy_allfixed);

  return ;
}


/* randist/gauss.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 James Theiler, Brian Gough
 *
 * 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <config.h>
#include <math.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

/* Of the two methods provided below, I think the Polar method is more
 * efficient, but only when you are actually producing two random
 * deviates.  We don't produce two, because then we'd have to save one
 * in a static variable for the next call, and that would screws up
 * re-entrant or threaded code, so we only produce one.  This makes
 * the Ratio method suddenly more appealing.  There are further tests
 * one can make if the log() is slow.  See Knuth for details */

/* Both methods pass the statistical tests; but the polar method
 * seems to be a touch faster on my home Pentium, EVEN though we
 * are only using half of the available random deviates!
 */

/* Polar (Box-Mueller) method; See Knuth v2, 3rd ed, p122 */

double
gsl_ran_gaussian (const gsl_rng * r, const double sigma)
{
  double x, y, r2;

  do
    {
      /* choose x,y in uniform square (-1,-1) to (+1,+1) */

      x = -1 + 2 * gsl_rng_uniform (r);
      y = -1 + 2 * gsl_rng_uniform (r);

      /* see if it is in the unit circle */
      r2 = x * x + y * y;
    }
  while (r2 > 1.0 || r2 == 0);

  /* Box-Muller transform */
  return sigma * y * sqrt (-2.0 * log (r2) / r2);
}

/* Ratio method (Kinderman-Monahan); see Knuth v2, 3rd ed, p130 */
/* K+M, ACM Trans Math Software 3 (1977) 257-260. */

double
gsl_ran_gaussian_ratio_method (const gsl_rng * r, const double sigma)
{
  double u, v, x;

  do
    {
      v = gsl_rng_uniform (r);
      do
	{
	  u = gsl_rng_uniform (r);
	}
      while (u == 0);
      /* Const 1.715... = sqrt(8/e) */
      x = 1.71552776992141359295 * (v - 0.5) / u;
    }
  while (x * x > -4.0 * log (u));

  return sigma * x;
}

double
gsl_ran_gaussian_pdf (const double x, const double sigma)
{
  double u = x / fabs (sigma);
  double p = (1 / (sqrt (2 * M_PI) * fabs (sigma))) * exp (-u * u / 2);
  return p;
}

double
gsl_ran_ugaussian (const gsl_rng * r)
{
  return gsl_ran_gaussian (r, 1.0);
}

double
gsl_ran_ugaussian_ratio_method (const gsl_rng * r)
{
  return gsl_ran_gaussian_ratio_method (r, 1.0);
}

double
gsl_ran_ugaussian_pdf (const double x)
{
  return gsl_ran_gaussian_pdf (x, 1.0);
}