/*
Copyright (C) 2019-2020 JingWeiZhangHuai <jingweizhanghuai@163.com>
Licensed under the Apache License, Version 2.0; you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
*/

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

#include <cblas.h>
#include "morn_tensor.h"

struct TensorResizePara
{
    MLayer *prev;

    int height;
    int width;

    int res_valid;
};

void *mTensorResizePara(MList *ini,char *name)
{
    struct TensorResizePara *para = (struct TensorResizePara *)mMalloc(sizeof(struct TensorResizePara));

    para->prev = mNetworkLayer(ini,mINIRead(ini,name,"prev"));
    mException((para->prev == NULL),EXIT,"invalid prev");
    para->res_valid = (strcmp("Input",mLayerType(para->prev))!=0);

    para->height= DFLT; mINIRead(ini,name,"height","%d",&(para->height));
    para->width = DFLT; mINIRead(ini,name,"width" ,"%d",&(para->width ));

    return para;
}

struct HandleTensorResize
{
    int *lx;
    int *ly;
    float *wx;
    float *wy;
};
void endTensorResize(void *info)
{
    struct HandleTensorResize *handle = (struct HandleTensorResize *)info;
    if(handle->lx !=NULL) mFree(handle->lx);
    if(handle->ly !=NULL) mFree(handle->ly);
    if(handle->wx !=NULL) mFree(handle->wx);
    if(handle->wy !=NULL) mFree(handle->wy);
}
#define HASH_TensorResize 0xb6f42d3a

void TensorResizeSet(MLayer *layer)
{
    if(layer->state != DFLT) return;
    struct TensorResizePara *para = (struct TensorResizePara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *res= para->prev->res;
    MTensor *out= layer->tns;

    MHandle *hdl=mHandle(out,TensorResize);
    struct HandleTensorResize *handle = (struct HandleTensorResize *)(hdl->handle);

    if(para->height<=0) para->height= in->height;
    if(para->width <=0) para->width = in->width;

    mTensorRedefine(out,in->batch,in->channel,para->height,para->width,NULL);
    if(morn_network_flag==MORN_TRAIN)
    {
        if(INVALID_TENSOR(res)) mTensorRedefine(res,in->batch,in->channel,in->height,in->width,in->data);
        else                    mTensorRedefine(res,in->batch,in->channel,in->height,in->width,NULL);
    }

    if(handle->lx !=NULL) {mFree(handle->lx);} handle->lx=(int   *)mMalloc(para->width *sizeof(int));
    if(handle->ly !=NULL) {mFree(handle->ly);} handle->ly=(int   *)mMalloc(para->height*sizeof(int));
    if(handle->wx !=NULL) {mFree(handle->wx);} handle->wx=(float *)mMalloc(para->width *sizeof(float));
    if(handle->wy !=NULL) {mFree(handle->wy);} handle->wy=(float *)mMalloc(para->height*sizeof(float));

    float kx = (float)(in->width)/(float)(out->width);
    for(int i=0;i<out->width;i++)
    {
        float l=kx*((float)i-0.5);
        if(l<0) l=0.0f;else if(l>=in->width-1) l=(float)(in->width)-1.00001;
        handle->lx[i] = (int)l;
        handle->wx[i] = (float)(handle->lx[i]+1)-l;
    }

    float ky = (float)(in->height)/(float)(out->height);
    for(int i=0;i<out->height;i++)
    {
        float l=ky*((float)i-0.5);
        if(l<0) l=0.0f;else if(l>=in->height-1) l=(float)(in->height)-1.00001;
        handle->ly[i] = (int)l;
        handle->wy[i] = (float)(handle->ly[i]+1)-l;
    }

    hdl->valid = 1;
}

void mTensorResizeForward(MLayer *layer)
{
    mException(INVALID_POINTER(layer),EXIT,"invalid input");
    mException(strcmp("Resize",mLayerType(layer)),EXIT,"invalid layer type");
    struct TensorResizePara *para = (struct TensorResizePara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *out=layer->tns;

    TensorResizeSet(layer);

    MHandle *hdl=mHandle(out,TensorResize);
    struct HandleTensorResize *handle = (struct HandleTensorResize *)(hdl->handle);

    int  in_size = in->width* in->height;
    int out_size =out->width*out->height;

    for(int h=0;h<out->height;h++)for(int w=0;w<out->width;w++)
    {
        int out_idx = h*out->width+w;
        int in_idx1 =(handle->ly[h]  )*in->width + (handle->lx[w]  );
        int in_idx2 =(handle->ly[h]  )*in->width + (handle->lx[w]+1);
        int in_idx3 =(handle->ly[h]+1)*in->width + (handle->lx[w]  );
        int in_idx4 =(handle->ly[h]+1)*in->width + (handle->lx[w]+1);
        float w1 = (    handle->wy[h])*(    handle->wx[w]);
        float w2 = (    handle->wy[h])*(1.0-handle->wx[w]);
        float w3 = (1.0-handle->wy[h])*(    handle->wx[w]);
        float w4 = (1.0-handle->wy[h])*(1.0-handle->wx[w]);

        for(int b=0;b<in->batch;b++)for(int c=0;c<in->channel;c++)
        {
            out->data[b][c*out_size+out_idx] = in->data[b][c*in_size+in_idx1]*w1
                                              +in->data[b][c*in_size+in_idx2]*w2
                                              +in->data[b][c*in_size+in_idx3]*w3
                                              +in->data[b][c*in_size+in_idx4]*w4;
        }
    }

    layer->state = MORN_FORWARD;
}

void mTensorResizeBackward(MLayer *layer)
{
    mException(INVALID_POINTER(layer),EXIT,"invalid input");
    mException(strcmp("Resize",mLayerType(layer)),EXIT,"invalid layer type");
    struct TensorResizePara *para = (struct TensorResizePara *)(layer->para);
    if(para->res_valid==0) return;
    MTensor *res= para->prev->res;
    MTensor *out= layer->res;

    MHandle *hdl=mHandle(layer->tns,TensorResize);
    struct HandleTensorResize *handle = (struct HandleTensorResize *)(hdl->handle);
    mException((hdl->valid == 0),EXIT,"no forward operate");

    int  in_size =res->width*res->height;
    int out_size =out->width*out->height;

    if(para->prev->state == MORN_FORWARD)
    {
        for(int b=0;b<res->batch;b++) memset(res->data[b],0,in_size*res->channel*sizeof(float));
        para->prev->state = MORN_BACKWARD;
    }

    for(int h=0;h<out->height;h++)for(int w=0;w<out->width;w++)
    {
        int out_idx = h*out->width+w;
        int in_idx1 =(handle->ly[h]  )*res->width + (handle->lx[w]  );
        int in_idx2 =(handle->ly[h]  )*res->width + (handle->lx[w]+1);
        int in_idx3 =(handle->ly[h]+1)*res->width + (handle->lx[w]  );
        int in_idx4 =(handle->ly[h]+1)*res->width + (handle->lx[w]+1);
        float w1 = (    handle->wy[h])*(    handle->wx[w]);
        float w2 = (    handle->wy[h])*(1.0-handle->wx[w]);
        float w3 = (1.0-handle->wy[h])*(    handle->wx[w]);
        float w4 = (1.0-handle->wy[h])*(1.0-handle->wx[w]);

        for(int b=0;b<res->batch;b++)for(int c=0;c<res->channel;c++)
        {
            res->data[b][c*in_size+in_idx1] += out->data[b][c*out_size+out_idx]*w1;
            res->data[b][c*in_size+in_idx2] += out->data[b][c*out_size+out_idx]*w2;
            res->data[b][c*in_size+in_idx3] += out->data[b][c*out_size+out_idx]*w3;
            res->data[b][c*in_size+in_idx4] += out->data[b][c*out_size+out_idx]*w4;
        }
    }
}


/*********************
 * Time intergral KDK scheme.
 * kick and drifts.
 *
 * This code was initially modified by Jun Koda,
 * from the original serial COLA code
 * by Svetlin Tassev.
 *
 * The kick and drift still supports a COLA compat-mode.
 * Most of the nasty factors are for COLA compat-mode
 * (not needed in PM)
 * We also added a 2LPT mode that does just 2LPT.
 *
 *  Yu Feng <rainwoodman@gmail.com>
 *
 */

#include <math.h>
#include <string.h>
#include <assert.h>
#include <alloca.h>
#include <mpi.h>

#include <gsl/gsl_integration.h>
#include <gsl/gsl_roots.h>
#include <gsl/gsl_sf_hyperg.h>
#include <gsl/gsl_errno.h>

#include <fastpm/libfastpm.h>
#include <fastpm/logging.h>

#include "pmpfft.h"
#include "vpm.h"

static double
Sq(double ai, double af, double aRef, double nLPT, FastPMCosmology * c, int USE_NONSTDDA);

static double
Sphi(double ai, double af, double aRef, double nLPT, FastPMCosmology * c, int USE_NONSTDDA);

static inline void
fastpm_drift_lookup(FastPMDriftFactor * drift, double af, double * dyyy, double * da1, double * da2)
{
    double ind;

    if(af == drift->af) {
        *dyyy = drift->dyyy[drift->nsamples - 1];
        *da1  = drift->da1[drift->nsamples - 1];
        *da2  = drift->da2[drift->nsamples - 1];
        return;
    }
    if(af == drift->ai) {
        *dyyy = drift->dyyy[0];
        *da1  = drift->da1[0];
        *da2  = drift->da2[0];
        return;
    }
    {
        ind = (af - drift->ai) / (drift->af - drift->ai) * (drift->nsamples - 1);
        int l = floor(ind);
        double u = l + 1 - ind;
        double v = ind - l;
        if(l + 1 >= drift->nsamples) {
            fastpm_raise(-1, "drift beyond factor's available range. ");
        }
        *dyyy = drift->dyyy[l] * u + drift->dyyy[l + 1] * v;
        *da1  = drift->da1[l] * u + drift->da1[l + 1] * v;
        *da2  = drift->da2[l] * u + drift->da2[l + 1] * v;
    }
}

inline void
fastpm_drift_one(FastPMDriftFactor * drift, FastPMStore * p, ptrdiff_t i, double xo[3], double af)
{

    double dyyy_f, da1_f, da2_f;
    double dyyy_i, da1_i, da2_i;
    double dyyy, da1, da2;

    fastpm_drift_lookup(drift, af, &dyyy_f, &da1_f, &da2_f);
    fastpm_drift_lookup(drift, p->meta.a_x, &dyyy_i, &da1_i, &da2_i);

    dyyy = dyyy_f - dyyy_i;
    da1 = da1_f - da1_i;
    da2 = da2_f - da2_i;

    int d;
    for(d = 0; d < 3; d ++) {
        double v;
        switch(drift->forcemode) {
            case FASTPM_FORCE_2LPT:
                xo[d] = p->x[i][d] + p->dx1[i][d] * da1 + p->dx2[i][d] * da2;
            break;
            case FASTPM_FORCE_ZA:
                xo[d] = p->x[i][d] + p->dx1[i][d] * da1;
            break;
            case FASTPM_FORCE_FASTPM:
            case FASTPM_FORCE_PM:
                xo[d] = p->x[i][d] + p->v[i][d] * dyyy;
            break;
            case FASTPM_FORCE_COLA:
                /* For cola, remove the lpt velocity to find the residual velocity v*/
                v = p->v[i][d] - (p->dx1[i][d]*drift->Dv1 + p->dx2[i][d]*drift->Dv2);
                xo[d] = p->x[i][d] + v * dyyy;
                xo[d] += p->dx1[i][d] * da1 + p->dx2[i][d] * da2;
            break;
        }
        /* if PGDCorrection is enabled, add it */
        if(p->pgdc) {
            /* no drift; to protect the pgdc line */
            if (drift->ai == drift->af) continue;
            xo[d] += 0.5 * p->pgdc[i][d] * dyyy / drift->dyyy[drift->nsamples-1];
        }
    }
}
static inline void
fastpm_kick_lookup(FastPMKickFactor * kick, double af, double * dda, double * Dv1, double * Dv2)
{
    double ind;

    if(af == kick->af) {
        *dda = kick->dda[kick->nsamples - 1];
        *Dv1 = kick->Dv1[kick->nsamples - 1];
        *Dv2 = kick->Dv2[kick->nsamples - 1];
        return;
    }
    if(af == kick->ai) {
        *dda = kick->dda[0];
        *Dv1 = kick->Dv1[0];
        *Dv2 = kick->Dv2[0];
        return;
    }
    {
        ind = (af - kick->ai) / (kick->af - kick->ai) * (kick->nsamples - 1);
        int l = floor(ind);
        double u = l + 1 - ind;
        double v = ind - l;
        if(l + 1 >= kick->nsamples) {
            fastpm_raise(-1, "kick beyond factor's available range. ");
        }
        *dda = kick->dda[l] * u + kick->dda[l + 1] * v;
        *Dv1 = kick->Dv1[l] * u + kick->Dv1[l + 1] * v;
        *Dv2 = kick->Dv2[l] * u + kick->Dv2[l + 1] * v;
    }
}

inline void
fastpm_kick_one(FastPMKickFactor * kick, FastPMStore * p, ptrdiff_t i, float vo[3], double af)
{
    double dda_i, Dv1_i, Dv2_i;
    double dda_f, Dv1_f, Dv2_f;
    double dda, Dv1, Dv2;

    fastpm_kick_lookup(kick, af, &dda_f, &Dv1_f, &Dv2_f);
    fastpm_kick_lookup(kick, p->meta.a_v, &dda_i, &Dv1_i, &Dv2_i);
    dda = dda_f - dda_i;
    Dv1 = Dv1_f - Dv1_i;
    Dv2 = Dv2_f - Dv2_i;

    int d;
    for(d = 0; d < 3; d++) {
        float ax = p->acc[i][d];       // unlike a_x, which means a at which x is calcd
        if(kick->forcemode == FASTPM_FORCE_COLA) {
            ax += (p->dx1[i][d]*kick->q1 + p->dx2[i][d]*kick->q2);
        }
        vo[d] = p->v[i][d] + ax * dda;
        if(kick->forcemode == FASTPM_FORCE_COLA) {
            vo[d] += (p->dx1[i][d] * Dv1 + p->dx2[i][d] * Dv2);
        }
    }
}

// Leap frog time integration

void
fastpm_kick_store(FastPMKickFactor * kick,
    FastPMStore * pi, FastPMStore * po, double af)
{
    int np = pi->np;

    // Kick using acceleration at a= ac
    // Assume forces at a=ac is in particles->force

    int i;
#pragma omp parallel for
    for(i=0; i<np; i++) {
        int d;
        float vo[3];
        fastpm_kick_one(kick, pi, i, vo, af);
        for(d = 0; d < 3; d++) {
            po->v[i][d] = vo[d];
        }
    }

    //velocity is now at a= avel1
    po->meta.a_v = af;
}

static double G_p(FastPMGrowthInfo * growth_info)
{
    /* integral of G_p */
    return growth_info->D1;
}
static double g_p(FastPMGrowthInfo * growth_info)
{
    return DGrowthFactorDa(growth_info);
}

static double G_f(FastPMGrowthInfo * growth_info)
{
    /* integral of g_f */
    double a = growth_info->a;
    return a * a * a * HubbleEa(a, growth_info->c) * g_p(growth_info);
}

static double g_f(FastPMGrowthInfo * growth_info)
{
    double a = growth_info->a;
    FastPMCosmology * c = growth_info->c;

    double E = HubbleEa(a, c);
    double dEda = DHubbleEaDa(a, c);

    double dDda = g_p(growth_info);
    double d2Dda2 = D2GrowthFactorDa2(growth_info);

    double g_f = 3 * a * a * E * dDda
                   + a * a * a * dEda * dDda
                   + a * a * a * E * d2Dda2;
    return g_f;
}

void fastpm_kick_init(FastPMKickFactor * kick, FastPMSolver * fastpm, double ai, double ac, double af)
{
    FastPMCosmology * c = fastpm->cosmology;
    kick->forcemode = fastpm->config->FORCE_TYPE;

    FastPMGrowthInfo gi_i;
    FastPMGrowthInfo gi_c;
    FastPMGrowthInfo gi_e;

    fastpm_growth_info_init(&gi_i, ai, c);
    fastpm_growth_info_init(&gi_c, ac, c);

    double E_i = HubbleEa(ai, c);
    double E_c = HubbleEa(ac, c);

    double D1_i = gi_i.D1;
    double D2_i = gi_i.D2;
    double f1_i = gi_i.f1;
    double f2_i = gi_i.f2;

    double D1_c = gi_c.D1;
    double D2_c = gi_c.D2;

    double Omega_m0 = Omega_source(1, c);
    double Omega_mc = Omega_source(ac, c);

    // kick->q1,2 are used for the COLA force implementation.
    // growth_mode = ODE and LCDM should match for an LCDM background,
    // but neither is guaranteed accurate for a background with radiaiton.
    // We advise using LCDM mode for forcemode = FASTPM_FORCE_COLA, as in the
    // original implementation of FastPM.
    kick->q1 = D1_c;
    switch (c->growth_mode){
        case FASTPM_GROWTH_MODE_LCDM:
            kick->q2 = D1_c*D1_c * (1.0 + 7.0/3.0 * pow(Omega_mc, 1.0/143.0));
        break;
        case FASTPM_GROWTH_MODE_ODE:
            kick->q2 = D1_c*D1_c * (1 - D1_c*D1_c/D2_c);
        break;
        default:
            fastpm_raise(-1, "Please enter a valid growth mode.\n");
    }

    kick->nsamples = 32;
    int i;

    double Dv1i = D1_i * ai * ai * E_i * f1_i;
    double Dv2i = D2_i * ai * ai * E_i * f2_i;
    for(i = 0; i < kick->nsamples; i ++) {
        double ae = ai * (1.0 * (kick->nsamples - 1 - i) / (kick->nsamples - 1))
                  + af * (1.0 * i / (kick->nsamples - 1));

        fastpm_growth_info_init(&gi_e, ae, c);
        double D1_e = gi_e.D1;
        double f1_e = gi_e.f1;
        double D2_e = gi_e.D2;
        double f2_e = gi_e.f2;
        double E_e = HubbleEa(ae, c);

        if(kick->forcemode == FASTPM_FORCE_FASTPM) {
            kick->dda[i] = -1.5 * Omega_mc * ac
               * E_c
               * (G_f(&gi_e) - G_f(&gi_i)) / g_f(&gi_c);
        } else {
            kick->dda[i] = -1.5 * Omega_m0
                * Sphi(ai, ae, ac, fastpm->config->nLPT, c, kick->forcemode == FASTPM_FORCE_COLA);
        }
        kick->Dv1[i] = D1_e * ae * ae * E_e * f1_e - Dv1i;
        kick->Dv2[i] = D2_e * ae * ae * E_e * f2_e - Dv2i;
    }

    kick->ai = ai;
    kick->ac = ac;
    kick->af = af;

    /* Output growth and FastPM factor at af for reference.
    This is a weird place to put this, but it's convenient because G and g are static */
    fastpm_info("Growth/FastPM factors at a = %6.4f: D1=%g, D2=%g, f1=%g, f2=%g, G_p=%g, G_f=%g, g_p=%g, g_f=%g\n",
               ai,
               gi_i.D1,
               gi_i.D2,
               gi_i.f1,
               gi_i.f2,
               G_p(&gi_i),
               G_f(&gi_i),
               g_p(&gi_i),
               g_f(&gi_i));
}

void
fastpm_drift_init(FastPMDriftFactor * drift, FastPMSolver * fastpm,
                double ai, double ac, double af)
{
    FastPMCosmology * c = fastpm->cosmology;
    drift->forcemode = fastpm->config->FORCE_TYPE;

    FastPMGrowthInfo gi_i;
    FastPMGrowthInfo gi_c;
    FastPMGrowthInfo gi_e;

    fastpm_growth_info_init(&gi_i, ai, c);
    fastpm_growth_info_init(&gi_c, ac, c);

    double E_c = HubbleEa(ac, c);

    double D1_i = gi_i.D1;
    double D2_i = gi_i.D2;

    double D1_c = gi_c.D1;
    double D2_c = gi_c.D2;
    double f1_c = gi_c.f1;
    double f2_c = gi_c.f2;

    drift->nsamples = 32;
    int i;

    for(i = 0; i < drift->nsamples; i ++ ) {
        double ae = ai * (1.0 * (drift->nsamples - 1 - i) / (drift->nsamples - 1))
                  + af * (1.0 * i / (drift->nsamples - 1));

        fastpm_growth_info_init(&gi_e, ae, c);     // overwrite each iteration
        double D1_e = gi_e.D1;
        double D2_e = gi_e.D2;

        if (drift->forcemode == FASTPM_FORCE_FASTPM) {
            drift->dyyy[i] = 1 / (ac * ac * ac * E_c)
                        * (G_p(&gi_e) - G_p(&gi_i)) / g_p(&gi_c);
        } else {
            drift->dyyy[i] = Sq(ai, ae, ac, fastpm->config->nLPT, c, drift->forcemode == FASTPM_FORCE_COLA);
        }
        drift->da1[i] = D1_e - D1_i;    // change in D_1lpt
        drift->da2[i] = D2_e - D2_i;  // change in D_2lpt
    }
    drift->af = af;
    drift->ai = ai;
    drift->ac = ac;
    drift->Dv1 = D1_c * ac * ac * E_c * f1_c;
    drift->Dv2 = D2_c * ac * ac * E_c * f2_c;
}

void
fastpm_drift_store(FastPMDriftFactor * drift,
               FastPMStore * pi, FastPMStore * po,
               double af)
{
    int np = pi->np;

    int i;
    // Drift
#pragma omp parallel for
    for(i=0; i<np; i++) {
        double xo[3] = {0};
        fastpm_drift_one(drift, pi, i, xo, af);
        int d;
        for(d = 0; d < 3; d ++) {
            po->x[i][d] = xo[d];
        }
    }
    po->meta.a_x = af;
}

//
// Functions for our modified time-stepping (used when StdDA=0):
//
struct iparam {
    FastPMCosmology * cosmology;
    double nLPT;
};

double gpQ(double a, double nLPT) {
    return pow(a, nLPT);
}

static double stddriftfunc (double a, struct iparam * iparam) {
    return 1 / (pow(a, 3) * HubbleEa(a, iparam->cosmology));
}

static double nonstddriftfunc (double a, struct iparam * iparam) {
    return gpQ(a, iparam->nLPT)/(pow(a, 3) * HubbleEa(a, iparam->cosmology));
}

static double stdkickfunc (double a, struct iparam * iparam) {
    return 1/ (pow(a, 2) * HubbleEa(a, iparam->cosmology));
}

static double integrand(double a, void * params) {
    void ** p = (void**) params;
    double (*func)(double a, struct iparam * s) = p[0];
    struct iparam * s = p[1];
    return func(a, s);
}

double integrate(double ai, double af,
        struct iparam * iparam,
        double (*func)(double , struct iparam * )) {

    gsl_integration_workspace * w
        = gsl_integration_workspace_alloc (5000);

    gsl_function F;
    double error;
    double result;

    F.params = (void*[]){func, iparam};
    F.function = integrand;

    gsl_integration_qag (&F, ai, af, 0, 1e-8, 5000, 6,
            w, &result, &error);

    gsl_integration_workspace_free (w);
    return result;
}

/*
       When StdDA=0, one needs to set nLPT.
       assumes time dep. for velocity = B a^nLPT
       nLPT is a real number. Sane values lie in the range (-4,3.5). Cannot be 0, but of course can be -> 0 (say 0.001).
       See Section A.3 of TZE.
       */

static double
Sq(double ai, double af, double aRef, double nLPT, FastPMCosmology * c, int USE_NONSTDDA)
{
    double resultstd, result;
    struct iparam iparam[1];
    iparam->cosmology = c;
    iparam->nLPT = nLPT;

    resultstd = integrate(ai, af, iparam, stddriftfunc);

    result = integrate(ai, af, iparam, nonstddriftfunc);
    result /= gpQ(aRef, nLPT);

    /*
    fastpm_info("ref time = %6.4f, std drift =%g, non std drift = %g \n",
        aRef, resultstd, result); */

    if (USE_NONSTDDA)
        return result;
    else
        return resultstd;
}

double DERgpQ(double a, double nLPT) {
    /* This must return d(gpQ)/da */
    return nLPT*pow(a, nLPT-1);
}

static double
Sphi(double ai, double af, double aRef, double nLPT, FastPMCosmology * c, int USE_NONSTDDA)
{
    double result;
    double resultstd;

    struct iparam iparam[1];
    iparam->cosmology = c;
    iparam->nLPT = nLPT;

    result = (gpQ(af, nLPT) - gpQ(ai, nLPT)) * aRef
        / (pow(aRef, 3) * HubbleEa(aRef, c) * DERgpQ(aRef, nLPT));

    resultstd = integrate(ai, af, iparam, stdkickfunc);

    /*
    fastpm_info("ref time = %6.4f, std kick = %g, non std kick = %g\n",
            aRef, resultstd, result); */

    if (USE_NONSTDDA) {
        return result;
    } else {
        return resultstd;
    }
}


/* #################################################
 * # stochastic simulation of adapting populations #
 * # with variing constraint type                  #
 * # and exponential mutation kernel               #
 * #################################################
 *
 * 2012-2014, Lukas Geyrhofer
 *
 *
 * #################################################
 *
 * simplest usage (population size 1e7):
 * ./travelingwavepeak_exp -N 1e7 1> out.txt 2> conf.txt
 *
 * ################################################# */

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

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_vector.h>

int space = 300;
int space0 = 100;
double dx = 1e-2;
int maxSteps = 1000;

int outputstep = 100;
int quiet = 0;

int noise = 0;
double populationsize = 1.;
double populationvariance;
int correctformeanfitness =0;
int printhistotype = 1;

int allshifts = 0;
int shiftthreshold = 1;
double current_mean_fitness = 0.;

int read_from_file = 0;
int write_to_file = 0;
char c_infile[128],c_outfile[128];

int have_u_infile = 0;
char u_infile[128];
double *u_read,*u;
int dens_ustar_latticeratio = 1;
int u_space,u_space0;
double u_dx;
double wavespeed = 0.;
double speedprefactor;

double epsilon = 1e-2;
double twoepssqrt;
double *nn;
double *tmp;
double *x;

double mutationrate = 1e-5;
double *mutation_inflow;
double mutation_outflow;
double mutation_sigma = 1e-2;

const gsl_rng* rg; /* gsl, global generator */
const gsl_rng_type* T;
unsigned long int randseed = 0;

double popdens_0thmom;
double popdens_1stmom;
double popdens_2ndmom;

int averagepopdens = 0;
int averagepopdens_center = 1;
int averagepopdens_resolution = 1;
int averagepopdens_havefile = 0;
double *averagepopdens_dens;
char averagepopdens_outputfile[128];
double averagepopdens_count = 0.;
double averagepopdens_dx;
int averagepopdens_space;
int averagepopdens_space0;
int averagepopdens_lower;
int averagepopdens_higher;

int fixationextinction_events = 0;
double fe_expected_fixationprobability = .5;
double fe_predicted_fixationprobability;
double *ww,*vv;
double *fe_start_ww, *fe_start_vv;
double *tmpw,*tmpv;
double fe_final_threshold = 1e-4; // if (current_fixation_prob > 1 - 1e-4) or (current_fixation_prob < 1e-4), then restart...
double fe_subpopsize;
int fe_count_extinctions = 0;
int fe_count_fixations = 0;
double current_fixation_prob;

int print_error(char *msg) {
  fprintf(stderr,"ERROR: %s\n",msg);
  exit(1);
}

// ************************************************************
// **   parameters
// ************************************************************

void parsecomamndline(int argn, char *argv[]) {
  char c;
  while((c = getopt(argn,argv,"s:z:d:S:e:D:O:qQi:o:R:N:CM:u:U:T:H:h:PE:")) != -1) {
    switch(c) {
      case 's':	space = atoi(optarg);
		break;
      case 'z':	space0 = atoi(optarg);
		break;
      case 'd':	dx = atof(optarg);
		break;
      case 'i':	strcpy(c_infile,optarg);
		read_from_file = 1;
		break;
      case 'o':	strcpy(c_outfile,optarg);
		write_to_file = 1;
		break;
      case 'u':	strcpy(u_infile,optarg);
		if(noise > 0)
		  print_error("Only a single constraint-type can be used (either option -N POPSIZE or -u FILENAME)");
		noise = 2;
		break;
      case 'U':	if(noise == 2) {
		  dens_ustar_latticeratio = atoi(optarg);
		}else{
		  print_error("option -u FILENAME needed before option -U RATIO");
		}
		break;
      case 'S':	maxSteps = atoi(optarg);
		break;
      case 'e': epsilon = atof(optarg);
		break;
      case 'D':	mutationrate = atof(optarg);
		break;
      case 'O':	outputstep = atoi(optarg);
		break;
      case 'q':	quiet = 2;
		break;
      case 'Q':	quiet = 1;
		break;
      case 'R':	randseed =atoi(optarg);
		break;
      case 'N': populationsize = atof(optarg);
		if(noise > 0)
		  print_error("Only a single constraint-type can be used (either option -N POPSIZE or -u FILENAME)");
		noise = 1;
		break;
      case 'C':	correctformeanfitness = 1;
		break;
      case 'M':	mutation_sigma = atof(optarg);
		break;
      case 'T':	shiftthreshold = atoi(optarg);
		break;
      case 'h':	strcpy(averagepopdens_outputfile,optarg);
		averagepopdens_havefile = 1;
		averagepopdens = 1;
		break;
      case 'H':	averagepopdens = 1;
		averagepopdens_resolution = atoi(optarg);
		if(averagepopdens_resolution < 0) {
		  averagepopdens_resolution *= -1;
		  averagepopdens_center = 0;
		}
		break;
      case 'P':	printhistotype = 0;
		break;
      case 'E':	fe_expected_fixationprobability = atof(optarg);
		break;
    }
  }
  if(randseed==0)randseed=time(NULL);
}

// ************************************************************
// **   input and output for configuration files
// ************************************************************

void read_popdens(int importparameters) {
  int i;
  int icount,dcount;
  int *ival;
  double *dval;
  int c_space,c_space0;
  double c_dx;
  FILE *fp;

  fp = fopen(c_infile,"rb");
  if(fp != NULL) {
    fread(&icount,sizeof(int),1,fp);
    fread(&dcount,sizeof(int),1,fp);
    fread(&c_dx,sizeof(double),1,fp);
    fread(&c_space,sizeof(int),1,fp);
    fread(&c_space0,sizeof(int),1,fp);

    if(importparameters) {
      space = c_space;
      space0 = c_space0;
      dx = c_dx;
    }

    if(icount>0) {
      ival = (int*)malloc(icount*sizeof(int));
      fread(ival,sizeof(int),icount,fp);
      free(ival);
    }

    if(dcount>0) {
      dval = (double*)malloc(dcount*sizeof(int));
      fread(dval,sizeof(double),dcount,fp);
      free(dval);
    }

    if(space != c_space)print_error("lattice does not match!");
    nn = (double*)malloc(space*sizeof(double));
    fread(nn,sizeof(double),space,fp);
    for(i=0;i<space;i++) {
      nn[i] *= dx;
    }
    fclose(fp);
  }else{
    print_error("could not open c-infile");
  }

}

void write_popdens() {
  int icount=0,dcount=0;
  int i;
  FILE *fpc;

  fpc=fopen(c_outfile,"wb");
  if(fpc != NULL) {
    fwrite(&icount,sizeof(int),1,fpc);
    fwrite(&dcount,sizeof(int),1,fpc);
    fwrite(&dx,sizeof(double),1,fpc);
    fwrite(&space,sizeof(int),1,fpc);
    fwrite(&space0,sizeof(int),1,fpc);
    for(i=0;i<space;i++)nn[i] /= dx;
    fwrite(nn,sizeof(double),space,fpc);
    fclose(fpc);
  }else{
    print_error("could not open c-outfile");
  }
}

void read_constraint(int importparameters) {
  int icount,dcount;
  int *ival;
  double *dval;
  FILE *fp;

  fp = fopen(u_infile,"rb");
  if(fp!=NULL) {
    fread(&icount,sizeof(int),1,fp);
    fread(&dcount,sizeof(int),1,fp);
    fread(&u_dx,sizeof(double),1,fp);
    fread(&u_space,sizeof(int),1,fp);
    fread(&u_space0,sizeof(int),1,fp);

    if(icount>0) {
      ival = (int*)malloc(icount*sizeof(int));
      fread(ival,sizeof(int),icount,fp);
      free(ival);
    }

    if(dcount>=3) {
      dval = (double*)malloc(dcount*sizeof(int));
      fread(dval,sizeof(double),dcount,fp);
      mutationrate = dval[0];
      wavespeed = dval[1];
      mutation_sigma = dval[2];
      free(dval);
    }else{
      print_error("not enough values in constraint file! need at least 3 double parameters: mutationrate, wavespeed, mutationsigma!");
    }

    if(importparameters) {
      space = u_space/dens_ustar_latticeratio;
      space0 = u_space0/dens_ustar_latticeratio;
      dx = u_dx*dens_ustar_latticeratio;
    }else{
      if((space*dens_ustar_latticeratio != u_space)||(space0*dens_ustar_latticeratio != u_space0))print_error("lattice does not match! u");
    }

    u_read = (double*)malloc(u_space*sizeof(double));
    fread(u_read,sizeof(double),u_space,fp);

    fclose(fp);
  }

  u = (double*)malloc(space*sizeof(double));

}

void flat_constraint(double size) {
  int i;
  u=(double*)malloc(space*sizeof(double));
  for(i=0;i<space;i++)u[i] = 1./size;
}

void initialize_with_gaussian_popdens() {
  int i;
  double startvariance;
  if(noise == 2) {
    startvariance = wavespeed;
  }else{
    // Good et al., PNAS (2012)
    startvariance = mutation_sigma*mutation_sigma*2*log(mutation_sigma*populationsize)/(log(mutation_sigma/mutationrate)*log(mutation_sigma/mutationrate));
  }
  for(i=0;i<space;i++) {
    nn[i] = exp(-(i-space0)*(i-space0)*dx*dx/(2.*startvariance));
  }
}

// ************************************************************
// **   screen output
// ************************************************************

void print_populationdensity(int timestep) {
  int i;
  double corr = allshifts*dx;
  if(correctformeanfitness)corr += current_mean_fitness;
  for(i=0;i<space;i++) {
    fprintf(stderr,"%lf %14.10lf %20.10e %20.10e\n",timestep*epsilon,(i-space0)*dx+corr,ww[i]+vv[i],ww[i]);
    if(printhistotype)fprintf(stderr,"%lf %14.10lf %20.10e %20.10e\n",timestep*epsilon,(i-space0+1)*dx+corr,ww[i]+vv[i],ww[i]);
  }
  fprintf(stderr,"\n");
}

// ************************************************************
// **   initialization
// ************************************************************

int initialize() {
  int i;

  if(read_from_file) {
    read_popdens(1);
    if(noise == 0)flat_constraint(1.);
    if(noise == 1)flat_constraint(populationsize);
    if(noise == 2)read_constraint(0);
  }else{
    if(noise == 0)flat_constraint(1.);
    if(noise == 1)flat_constraint(populationsize);
    if(noise == 2)read_constraint(1);
    nn = (double*)calloc(space,sizeof(double));
    initialize_with_gaussian_popdens();
  }
  tmp = (double*)malloc(space*sizeof(double));

  gsl_rng_env_setup();
  T = gsl_rng_default;
  rg = gsl_rng_alloc(T);
  gsl_rng_set(rg, randseed);

  twoepssqrt = sqrt(2.*epsilon);
  speedprefactor = wavespeed/dx*epsilon;

  x = (double*)malloc(space*sizeof(double));
  mutation_inflow = (double*)malloc(space*sizeof(double));
  mutation_outflow = 0;
  for(i=1;i<space;i++) {
    x[i] = (i-space0)*dx;
    mutation_inflow[i] = epsilon*exp(-i*dx/mutation_sigma)/mutation_sigma*mutationrate*dx;
    mutation_outflow += mutation_inflow[i];
  }

  if(quiet<2) {
    printf("#################################################################################\n");
    printf("# stochastic simulation of adapting population with exponential mutation kernel #\n");
    printf("#################################################################################\n");
    printf("#    mutationrate     = %e\n",mutationrate);
    printf("#    mutation_sigma   = %e\n",mutation_sigma);
    if(noise==0) {
      printf("#    constraint       = deterministic\n");
    }
    if(noise==1) {
      printf("#    constraint       = fixedN\n");
      printf("#    populationsize   = %e\n",populationsize);
    }
    if(noise==2) {
      printf("#    constraint       = ustar\n");
      printf("#    wavespeed        = %e\n",wavespeed);
    }
    printf("#    (lattice) space  = %d\n",space);
    printf("#    (lattice) space0 = %d\n",space0);
    printf("#    (lattice) dx     = %e\n",dx);
    printf("#    randseed         = %d\n",randseed);
  }

}

// ************************************************************
// **   average popdens
// ************************************************************

void init_averagepopdens() {

  averagepopdens_dx = dx/(1.*averagepopdens_resolution);
  averagepopdens_space = space*averagepopdens_resolution;
  averagepopdens_space0 = space0*averagepopdens_resolution;

  averagepopdens_dens = (double*)calloc(averagepopdens_space,sizeof(double));

  averagepopdens_lower = (int)(-averagepopdens_resolution/2);
  averagepopdens_higher = (int)(averagepopdens_resolution/2);
  if(averagepopdens_higher - averagepopdens_lower < averagepopdens_resolution)averagepopdens_higher++;

}

void update_averagepopdens() {
  int i,j,idx;
  int offset=0;

  if(averagepopdens_center) {
    offset = (int)(current_mean_fitness/dx*averagepopdens_resolution);
  }

  for(i=0;i<space;i++) {
    for(j=averagepopdens_lower;j<averagepopdens_higher;j++) {
      idx = i*averagepopdens_resolution+j-offset;
      if((averagepopdens_space > idx) &&( idx >= 0))
	averagepopdens_dens[idx] += nn[i];
    }
  }

  averagepopdens_count += 1.;

}

void write_averagepopdens() {
  int i;
  FILE *fp;

  if(averagepopdens_havefile == 1) {
    fp = fopen(averagepopdens_outputfile,"w");
  }else{
    fp = stdout;
  }

  for(i=0;i<averagepopdens_space;i++) {
    fprintf(fp,"%lf %e\n",(i-averagepopdens_space0)*averagepopdens_dx,averagepopdens_dens[i]/averagepopdens_count);
  }

  if(averagepopdens_havefile == 1) {
    fclose(fp);
  }

  free(averagepopdens_dens);
}

// ************************************************************
// **   main algorithm
// ************************************************************

void update_u(int timestep) {
   // bug if shiftthreshold -T >1, use shiftthreshold only with fixedN
  int i,j;
  int baseshift;
  double fracshift;

  if(dens_ustar_latticeratio > 1) {
    baseshift = (int)(timestep*epsilon*wavespeed/u_dx) - allshifts*dens_ustar_latticeratio;
    fracshift = timestep*epsilon*wavespeed/u_dx - 1.*(baseshift + allshifts*dens_ustar_latticeratio);
    i = 0;
    j = 0;
    while( (i<space) && (j<u_space) ) {
      if((j+baseshift)%dens_ustar_latticeratio == 0) {
	u[i] = (fracshift) * u_read[j];
      }else if((j+baseshift)%dens_ustar_latticeratio == dens_ustar_latticeratio - 1) {
	u[i] += (1.-fracshift) * u_read[j];
	u[i] /= (1.*dens_ustar_latticeratio);
	i++;
      }else{
	u[i] += u_read[j];
      }
      j++;
    }
  }else{
    fracshift = timestep*epsilon*wavespeed/u_dx - 1.*allshifts;
    // ignore u[0], as it is unlikely that the population has a significant subpop there...
    for(i=1;i<space;i++)u[i] = fracshift * u_read[i-1] + (1.-fracshift)*u_read[i];
  }
}

void update_mean_fit() {
  int i;
  current_mean_fitness = 0.;
  populationsize = 0.;
  for(i=0;i<space;i++) {
    current_mean_fitness += (i-space0)*(ww[i]+vv[i]);
    populationsize += ww[i] + vv[i];
  }
  current_mean_fitness *= dx/populationsize;
}

void shift_population_backward(int step) {
  int i;
  for(i=0;i<space-step;i++){
    ww[i] = ww[i+step];
    vv[i] = vv[i+step];
  }
  for(i=space-step;i<space;i++) {
    ww[i] = 0.;
    vv[i] = 0.;
  }
}

void shift_population_forward(int step) {
  int i;
  for(i=space-1;i>step;i--) {
    ww[i] = ww[i-step];
    vv[i] = vv[i-step];
  }
  for(i=step;i>=0;i--) {
    ww[i] = 0.;
    vv[i] = 0.;
  }
}

void shift_population(int timestep) {
  int shift = (int)floor(current_mean_fitness/dx);
  if(shift >= shiftthreshold) {shift_population_backward(shift);}
  if(shift <= -shiftthreshold) {shift_population_forward(shift);}
  allshifts += shift;
  current_mean_fitness -= shift*dx;
}

void reproduce(int timestep) {
  int i,j;
  double tmpn;

  if(noise < 2) {
    update_mean_fit();
  }else if(noise == 2) {
    current_mean_fitness = timestep*wavespeed*epsilon - allshifts*dx;
  }

  if((current_mean_fitness > shiftthreshold*dx) || (current_mean_fitness < -shiftthreshold*dx)) {
    shift_population(timestep);
  }

  tmpw[0] = 0;
  tmpw[space-1] = 0;

  tmpv[0] = 0;
  tmpv[space-1] = 0;

  for(i=1;i<space-1;i++) {
    tmpw[i] = ww[i]*(1.+(x[i]-mutationrate-current_mean_fitness)*epsilon);
    tmpv[i] = vv[i]*(1.+(x[i]-mutationrate-current_mean_fitness)*epsilon);
    for(j = 1;j<=i;j++) {
      tmpw[i] += mutation_inflow[j]*ww[i-j];
      tmpv[i] += mutation_inflow[j]*vv[i-j];
    }
    if(tmpw[i]<0) {
      tmpw[i] = 0;
    }else if(noise > 0) {
      if(tmpw[i] < 1e9) { // Poisson-RNG breaks down for parameters > 1e9. see GSL doc.
			 // use smaller bins if this occurs too often or population size too large
	tmpw[i] += twoepssqrt*(gsl_ran_poisson(rg,tmpw[i])-tmpw[i]);
      }else{
	tmpw[i] = 1e9;
      }
    }

    if(tmpv[i]<0) {
      tmpv[i] = 0;
    }else if(noise > 0) {
      if(tmpv[i] < 1e9) { // Poisson-RNG breaks down for parameters > 1e9. see GSL doc.
			 // use smaller bins if this occurs too often or population size too large
	tmpv[i] += twoepssqrt*(gsl_ran_poisson(rg,tmpv[i])-tmpv[i]);
      }else{
	tmpv[i] = 1e9;
      }
    }

  }

  memcpy(ww,tmpw,space*sizeof(double));
  memcpy(vv,tmpv,space*sizeof(double));

}

void populationconstraint(int timestep) {
  int i;
  double sum = 0., inv;

  if(noise == 2) {
    update_u(timestep);
  }
  current_fixation_prob = 0;
  popdens_0thmom = 0;
  popdens_1stmom = 0;
  popdens_2ndmom = 0;
  fe_subpopsize = 0;
  for(i=0;i<space;i++) {
    sum += (ww[i] + vv[i])*u[i];
    current_fixation_prob += ww[i]*u[i];
    fe_subpopsize += ww[i];
    popdens_0thmom += ww[i] + vv[i];
    popdens_1stmom += (i-space0)*(ww[i] + vv[i]);
    popdens_2ndmom += (i-space0)*(i-space0)*(ww[i] + vv[i]);
  }
  popdens_1stmom *= dx;
  popdens_2ndmom *= dx*dx;

  populationsize = popdens_0thmom;
  populationvariance = popdens_2ndmom/popdens_0thmom - popdens_1stmom*popdens_1stmom/(popdens_0thmom*popdens_0thmom);

  inv = 1./(sum);
  current_fixation_prob *= inv;
  for(i=0;i<space;i++) {
    ww[i] *= inv;
    vv[i] *= inv;
  }
}

void initial_constraint() {
  int i;
  double sum=0;
  double inv;
  for(i=0;i<space;i++)sum += nn[i] * u[i];
  inv = 1./sum;
  for(i=0;i<space;i++)nn[i] *= inv;
}

// ************************************************************
// **   extinction and fixation events
// ************************************************************

int init_fe() {
  int i,j;
  int startindex = space;
  double fixprob = 0.,frac_lastbin;

  ww = (double*)malloc(space*sizeof(double));
  vv = (double*)malloc(space*sizeof(double));
  fe_start_ww = (double*)malloc(space*sizeof(double));
  fe_start_vv = (double*)malloc(space*sizeof(double));
  tmpw = (double*)malloc(space*sizeof(double));
  tmpv = (double*)malloc(space*sizeof(double));

  update_u(0);
  initial_constraint();
  while (fixprob < fe_expected_fixationprobability) {
    startindex--;
    fixprob += nn[startindex]*u[startindex];
  }

  frac_lastbin = 1. + (fe_expected_fixationprobability - fixprob)/(nn[startindex]*u[startindex]);
  fe_predicted_fixationprobability = fixprob;
  for(j=0;j<space;j++) {
    if(j<startindex) {
      ww[j] = 0;
      vv[j] = nn[j];
    }else if(j==startindex){
      ww[j] = frac_lastbin*nn[j];
      vv[j] = (1-frac_lastbin)*nn[j];
    }else{
      ww[j] = nn[j];
      vv[j] = 0;
    }
  }

  memcpy(fe_start_ww,ww,space*sizeof(double));
  memcpy(fe_start_vv,vv,space*sizeof(double));
}

double reset_fe() {
  allshifts = 0;
  memcpy(ww,fe_start_ww,space*sizeof(double));
  memcpy(vv,fe_start_vv,space*sizeof(double));
}

// ************************************************************
// **   cleanup
// ************************************************************

void cleanup() {
  free(nn);
  free(tmp);
  free(x);
  free(mutation_inflow);
  free(u);
  if(noise == 2) {
    free(u_read);
  }
  free(ww);
  free(vv);
  free(tmpw);
  free(tmpv);
  free(fe_start_ww);
  free(fe_start_vv);
}

// ************************************************************
// **   main
// ************************************************************

int main(int argn, char *argv[]) {
  int i=1;
  double v;

  parsecomamndline(argn,argv);
  initialize();
  init_fe();
  populationconstraint(0);
  if(averagepopdens) {
    init_averagepopdens();
  }
  if (quiet<2) fprintf(stdout,"%10.3lf %20.10e %.10e %.10e %.10e %.10e\n",i*epsilon,0.,populationvariance,populationsize,fe_subpopsize,current_fixation_prob);

//   print_populationdensity(0);
//   exit(1);

  while(fe_count_extinctions + fe_count_fixations < maxSteps) {
    reproduce(i);
    populationconstraint(i);

    if(i%outputstep == 0) {
      if (quiet<2)fprintf(stdout,"%10.3lf %20.10e %.10e %.10e %.10e %.10e\n",i*epsilon,allshifts*dx + popdens_1stmom/popdens_0thmom,populationvariance,populationsize,fe_subpopsize,current_fixation_prob);
      if (quiet==0) print_populationdensity(i);
      if (averagepopdens) update_averagepopdens();
    }
    if(current_fixation_prob < fe_final_threshold) {
      reset_fe();
      fe_count_extinctions++;
      printf("# event ( %d of %d ): extinction %lf\n",fe_count_extinctions+fe_count_fixations,maxSteps,i*epsilon);
      i=0;
    }
    if(current_fixation_prob > 1-fe_final_threshold) {
      reset_fe();
      fe_count_fixations++;
      printf("# event ( %d of %d ): fixation %lf\n",fe_count_extinctions+fe_count_fixations,maxSteps,i*epsilon);
      i=0;
    }
    i++;
  }
  if(write_to_file)write_popdens();
  if(averagepopdens)write_averagepopdens();
  cleanup();
  return 0;
}


#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;
}

/*----- ROUTINE: r_delta -----
INPUT: cosmology, halo mass, scale factor, halo overdensity
TASK: Computes comoving halo radius assuming the overdensity criteria
*/
double r_delta(ccl_cosmology *cosmo, double halomass, double a, double odelta, int *status){

  double rho_matter = ccl_rho_x(cosmo, 1., 1, 1, status);

  return pow(halomass*3.0/(4.0*M_PI*rho_matter*odelta),1.0/3.0);

}

// This checks to make sure all necessary halo mass function parameters have been set-up,
// as well as associated splines.
void ccl_cosmology_compute_hmfparams(ccl_cosmology *cosmo, int *status) {
  if(cosmo->computed_hmfparams)
    return;

  gsl_spline* alphahmf = NULL;
  gsl_spline* betahmf = NULL;
  gsl_spline* gammahmf = NULL;
  gsl_spline* phihmf = NULL;
  gsl_spline* etahmf = NULL;

  // declare parameter splines on case-by-case basis
  switch(cosmo->config.mass_function_method) {
  case ccl_tinker10:{
    double delta[9] = {200.0, 300.0, 400.0, 600.0, 800.0, 1200.0, 1600.0, 2400.0, 3200.0};
    double lgdelta[9];
    double alpha[9] = {0.368, 0.363, 0.385, 0.389, 0.393, 0.365, 0.379, 0.355, 0.327};
    double beta[9] = {0.589, 0.585, 0.544, 0.543, 0.564, 0.623, 0.637, 0.673, 0.702};
    double gamma[9] ={0.864, 0.922, 0.987, 1.09, 1.20, 1.34, 1.50, 1.68, 1.81};
    double phi[9] = {-0.729, -0.789, -0.910, -1.05, -1.20, -1.26, -1.45, -1.50, -1.49};
    double eta[9] = {-0.243, -0.261, -0.261, -0.273, -0.278, -0.301, -0.301, -0.319, -0.336};
    int nd = 9;
    int i;

    for(i=0; i<nd; i++) {
      lgdelta[i] = log10(delta[i]);
    }

    if (*status == 0) {
      alphahmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (alphahmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating alpha(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(alphahmf, lgdelta, alpha, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE ;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating alpha(D) spline\n");
      }
    }

    if (*status == 0) {
      betahmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (betahmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating beta(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(betahmf, lgdelta, beta, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating beta(D) spline\n");
      }
    }

    if (*status == 0) {
      gammahmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (gammahmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating gamma(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(gammahmf, lgdelta, gamma, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating gamma(D) spline\n");
      }
    }

    if (*status == 0) {
      phihmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (phihmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating phi(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(phihmf, lgdelta, phi, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating phi(D) spline\n");
      }
    }

    if (*status == 0) {
      etahmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (etahmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating eta(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(etahmf, lgdelta, eta, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating eta(D) spline\n");
      }
    }

    if (*status == 0) {
      cosmo->data.alphahmf = alphahmf;
      cosmo->data.betahmf = betahmf;
      cosmo->data.gammahmf = gammahmf;
      cosmo->data.phihmf = phihmf;
      cosmo->data.etahmf = etahmf;
      cosmo->computed_hmfparams = true;
    }
    break;
  }
  case ccl_tinker:{
    double delta[9] = {200.0, 300.0, 400.0, 600.0, 800.0, 1200.0, 1600.0, 2400.0, 3200.0};
    double lgdelta[9];
    double alpha[9] = {0.186, 0.200, 0.212, 0.218, 0.248, 0.255, 0.260, 0.260, 0.260};
    double beta[9] = {1.47, 1.52, 1.56, 1.61, 1.87, 2.13, 2.30, 2.53, 2.66};
    double gamma[9] ={2.57, 2.25, 2.05, 1.87, 1.59, 1.51, 1.46, 1.44, 1.41};
    double phi[9] = {1.19, 1.27, 1.34, 1.45, 1.58, 1.80, 1.97, 2.24, 2.44};
    int nd = 9;
    int i;

    for(i=0; i<nd; i++) {
      lgdelta[i] = log10(delta[i]);
    }

    if (*status == 0) {
      alphahmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (alphahmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating alpha(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(alphahmf, lgdelta, alpha, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating alpha(D) spline\n");
      }
    }

    if (*status == 0) {
      betahmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (betahmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating beta(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(betahmf, lgdelta, beta, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating beta(D) spline\n");
      }
    }

    if (*status == 0) {
      gammahmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (gammahmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating gamma(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(gammahmf, lgdelta, gamma, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating gamma(D) spline\n");
      }
    }

    if (*status == 0) {
      phihmf = gsl_spline_alloc(cosmo->spline_params.D_SPLINE_TYPE, nd);
      if (phihmf == NULL) {
        *status = CCL_ERROR_MEMORY;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error allocating phi(D) spline\n");
      }
    }
    if (*status == 0) {
      *status = gsl_spline_init(phihmf, lgdelta, phi, nd);
      if (*status) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_massfunc.c: ccl_cosmology_compute_hmfparams(): Error creating phi(D) spline\n");
      }
    }

    if (*status == 0) {
      cosmo->data.alphahmf = alphahmf;
      cosmo->data.betahmf = betahmf;
      cosmo->data.gammahmf = gammahmf;
      cosmo->data.phihmf = phihmf;
      cosmo->computed_hmfparams = true;
    }

    break;
  }
  default:
    // Error message could go here if we decide to make this public facing.
    // Currently not accessible from the API though.
    break;
  }

  if (*status) {
    gsl_spline_free(alphahmf);
    gsl_spline_free(betahmf);
    gsl_spline_free(gammahmf);
    gsl_spline_free(phihmf);
    gsl_spline_free(etahmf);
  }
}

//TODO: some of these are unused, many are included in ccl.h

/*----- ROUTINE: ccl_massfunc_f -----
INPUT: cosmology+parameters, a halo mass, and scale factor
TASK: Outputs fitting function for use in halo mass function calculation;
  currently only supports:
    ccl_tinker (arxiv 0803.2706 )
    ccl_tinker10 (arxiv 1001.3162 )
    ccl_angulo (arxiv 1203.3216 )
    ccl_watson (arxiv 1212.0095 )
    ccl_shethtormen (arxiv 9901122)
-*/
static double massfunc_f(ccl_cosmology *cosmo, double halomass, double a, double odelta, int *status)
{
  double fit_A, fit_a, fit_b, fit_c, fit_d, fit_p, overdensity_delta;
  double Omega_m_a;
  double delta_c_Tinker, nu;

  double sigma=ccl_sigmaM(cosmo, log10(halomass), a, status);
  int gslstatus;

  switch(cosmo->config.mass_function_method) {

  // Equation (10) in arxiv: 9901122
  // Note that Sheth & Tormen (1999) use nu=(dc/sigma)^2 whereas we use nu=dc/sigma
  case ccl_shethtormen:

    // Check if odelta is outside the interpolated range
    if (odelta != Dv_BryanNorman(cosmo, a, status)) {
      *status = CCL_ERROR_HMF_DV;
      ccl_cosmology_set_status_message(
        cosmo, "ccl_massfunc.c: massfunc_f(): Sheth-Tormen called with not virial Delta_v\n");
      return NAN;
    }

    // ST mass function fitting parameters
    fit_A = 0.21616;
    fit_p = 0.3;
    fit_a = 0.707;

    // nu = delta_c(z) / sigma(M)
    nu = dc_NakamuraSuto(cosmo, a, status)/ccl_sigmaM(cosmo, log10(halomass), a, status);

    return nu*fit_A*(1.+pow(fit_a*pow(nu,2),-fit_p))*exp(-fit_a*pow(nu,2)/2.);

  case ccl_tinker:

    // Check if odelta is outside the interpolated range
    if ((odelta < 200) || (odelta > 3200)) {
      *status = CCL_ERROR_HMF_INTERP;
      ccl_cosmology_set_status_message(
        cosmo,
        "ccl_massfunc.c: massfunc_f(): Tinker 2008 only "
        "supported in range of Delta = 200 to Delta = 3200.\n");
      return NAN;
    }

    if (!cosmo->computed_hmfparams) {
      *status = CCL_ERROR_HMF_INIT;
      ccl_cosmology_set_status_message(
        cosmo,
        "ccl_massfunc.c: massfunc_f(): mass function parameters splines have not been computed!");
      return NAN;
    }

    // Compute HMF parameter (alpha, beta, gamma, phi) splines if they haven't
    // been computed already
    gslstatus = gsl_spline_eval_e(cosmo->data.alphahmf, log10(odelta), NULL, &fit_A);
    gslstatus |= gsl_spline_eval_e(cosmo->data.betahmf, log10(odelta), NULL, &fit_a);
    gslstatus |= gsl_spline_eval_e(cosmo->data.gammahmf, log10(odelta), NULL, &fit_b);
    gslstatus |= gsl_spline_eval_e(cosmo->data.phihmf, log10(odelta), NULL, &fit_c);
    fit_d = pow(10, -1.0*pow(0.75 / log10(odelta / 75.0), 1.2));

    fit_A = fit_A*pow(a, 0.14);
    fit_a = fit_a*pow(a, 0.06);
    fit_b = fit_b*pow(a, fit_d);
    if(gslstatus != GSL_SUCCESS) {
      ccl_raise_gsl_warning(gslstatus, "ccl_massfunc.c: massfunc_f():");
      *status |= gslstatus;
      ccl_cosmology_set_status_message(
        cosmo, "ccl_massfunc.c: massfunc_f(): interpolation error for Tinker MF\n");
      return NAN;
    }
    return fit_A*(pow(sigma/fit_b,-fit_a)+1.0)*exp(-fit_c/sigma/sigma);
    break;
  case ccl_tinker10:
    // this version uses f(nu) parameterization from Eq. 8 in Tinker et al. 2010
    // use this for consistency with Tinker et al. 2010 fitting function for halo bias

    // Check if odelta is outside the interpolated range
    if ((odelta < 200) || (odelta > 3200)) {
      *status = CCL_ERROR_HMF_INTERP;
      ccl_cosmology_set_status_message(
        cosmo,
        "ccl_massfunc.c: massfunc_f(): Tinker 2010 only "
        "supported in range of Delta = 200 to Delta = 3200.\n");
      return 0;
    }

    //critical collapse overdensity assumed in this model
    delta_c_Tinker = 1.686;
    nu = delta_c_Tinker/(sigma);

    if (!cosmo->computed_hmfparams) {
      *status = CCL_ERROR_HMF_INIT;
      ccl_cosmology_set_status_message(
        cosmo,
        "ccl_massfunc.c: massfunc_f(): mass function parameters splines have not been computed!");
      return NAN;
    }

    gslstatus = gsl_spline_eval_e(cosmo->data.alphahmf, log10(odelta), NULL, &fit_A); //alpha in Eq. 8
    gslstatus |= gsl_spline_eval_e(cosmo->data.etahmf, log10(odelta), NULL, &fit_a); //eta in Eq. 8
    gslstatus |= gsl_spline_eval_e(cosmo->data.betahmf, log10(odelta), NULL, &fit_b); //beta in Eq. 8
    gslstatus |= gsl_spline_eval_e(cosmo->data.gammahmf, log10(odelta), NULL, &fit_c); //gamma in Eq. 8
    gslstatus |= gsl_spline_eval_e(cosmo->data.phihmf, log10(odelta), NULL, &fit_d); //phi in Eq. 8;

    fit_a *=pow(a, -0.27);
    fit_b *=pow(a, -0.20);
    fit_c *=pow(a, 0.01);
    fit_d *=pow(a, 0.08);
    if(gslstatus != GSL_SUCCESS) {
      ccl_raise_gsl_warning(gslstatus, "ccl_massfunc.c: massfunc_f():");
      *status |= gslstatus;
      ccl_cosmology_set_status_message(
        cosmo, "ccl_massfunc.c: massfunc_f(): interpolation error for Tinker 2010 MF\n");
      return NAN;
    }
    return nu*fit_A*(1.+pow(fit_b*nu,-2.*fit_d))*pow(nu, 2.*fit_a)*exp(-0.5*fit_c*nu*nu);
    break;

  case ccl_watson:
    if(odelta!=200.) {
      *status = CCL_ERROR_HMF_INTERP;
      ccl_cosmology_set_status_message(
        cosmo, "ccl_massfunc.c: massfunc_f(): Watson HMF only supported for Delta = 200.\n");
      return NAN;
    }
    // these parameters from: Angulo et al 2012 (arxiv 1203.3216 )
    Omega_m_a = ccl_omega_x(cosmo, a, ccl_species_m_label,status);
    fit_A = Omega_m_a*(0.990*pow(a,3.216)+0.074);
    fit_a = Omega_m_a*(5.907*pow(a,3.599)+2.344);
    fit_b = Omega_m_a*(3.136*pow(a,3.058)+2.349);
    fit_c = 1.318;

    return fit_A*(pow(sigma/fit_b,-fit_a)+1.0)*exp(-fit_c/sigma/sigma);

  case ccl_angulo:
    if(odelta!=200.) {
      *status = CCL_ERROR_HMF_INTERP;
      ccl_cosmology_set_status_message(
        cosmo, "ccl_massfunc.c: massfunc_f(): Angulo HMF only supported for Delta = 200.\n");
      return NAN;
    }
    // these parameters from: Watson et al 2012 (arxiv 1212.0095 )
    fit_A = 0.201;
    fit_a = 2.08;
    fit_b = 1.7;
    fit_c = 1.172;

    return fit_A*pow( (fit_a/sigma)+1.0, fit_b)*exp(-fit_c/sigma/sigma);

  default:
    *status = CCL_ERROR_MF;
    ccl_cosmology_set_status_message(cosmo ,
	    "ccl_massfunc.c: massfunc_f(): Unknown or non-implemented mass function method: %d \n",
	    cosmo->config.mass_function_method);
    return NAN;
  }
}

static double ccl_halo_b1(ccl_cosmology *cosmo, double halomass, double a, double odelta, int *status)
{
  double fit_A, fit_B, fit_C, fit_a, fit_b, fit_c, fit_p, overdensity_delta, y;
  double delta_c_Tinker, nu;
  double sigma=ccl_sigmaM(cosmo,log10(halomass),a, status);
  switch(cosmo->config.mass_function_method) {

  // Equation (12) in  arXiv: 9901122
  // Derived using the peak-background split applied to the mass function in the same paper
  // Note that Sheth & Tormen (1999) use nu=(dc/sigma)^2 whereas we use nu=dc/sigma
  case ccl_shethtormen:

    // Check if Delta_v is the virial Delta_v
    if (odelta != Dv_BryanNorman(cosmo, a, status)) {
      *status = CCL_ERROR_HMF_DV;
      ccl_cosmology_set_status_message(
        cosmo, "ccl_massfunc.c: halo_b1(): Sheth-Tormen called with not virial Delta_v\n");
      return NAN;
    }

    // ST bias fitting parameters (which are the same as for the mass function)
    fit_p = 0.3;
    fit_a = 0.707;

    // Cosmology dependent delta_c and nu
    double delta_c = dc_NakamuraSuto(cosmo, a, status);
    nu = delta_c/ccl_sigmaM(cosmo, log10(halomass), a, status);

    return 1.+(fit_a*pow(nu,2)-1.+2.*fit_p/(1.+pow(fit_a*pow(nu,2),fit_p)))/delta_c;

    //this version uses b(nu) parameterization, Eq. 6 in Tinker et al. 2010
    // use this for consistency with Tinker et al. 2010 fitting function for halo bias
  case ccl_tinker10:
    y = log10(odelta);
    //critical collapse overdensity assumed in this model
    delta_c_Tinker = 1.686;
    //peak height - note that this factorization is incorrect for e.g. massive neutrino cosmologies
    nu = delta_c_Tinker/(sigma);
    // Table 2 in https://arxiv.org/pdf/1001.3162.pdf
    fit_A = 1.0 + 0.24*y*exp(-pow(4./y,4.));
    fit_a = 0.44*y-0.88;
    fit_B = 0.183;
    fit_b = 1.5;
    fit_C = 0.019+0.107*y+0.19*exp(-pow(4./y,4.));
    fit_c = 2.4;

    return 1.-fit_A*pow(nu,fit_a)/(pow(nu,fit_a)+pow(delta_c_Tinker,fit_a))+fit_B*pow(nu,fit_b)+fit_C*pow(nu,fit_c);
    break;

  default:
    *status = CCL_ERROR_MF;
    ccl_cosmology_set_status_message(cosmo ,
	    "ccl_massfunc.c: ccl_halo_b1(): No b(M) fitting function implemented for mass_function_method: %d \n",
      cosmo->config.mass_function_method);
    return 0;
  }
}

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

  // create linearly-spaced values of the mass.
  int nm = cosmo->spline_params.LOGM_SPLINE_NM;
  double *m = NULL;
  double *y = NULL;
  double smooth_radius;
  double na, nb;

  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_cosmology_compute_sigmas(): Error creating linear spacing in m\n");
  }

  if (*status == 0) {
    // create space for y, to be filled with sigma and dlnsigma_dlogm
    y = malloc(sizeof(double)*nm);
    if (y == NULL) {
      *status = CCL_ERROR_MEMORY;
    }
  }

  // start up of GSL pointers
  int gslstatus = 0;
  gsl_spline *logsigma = NULL;
  gsl_spline *dlnsigma_dlogm = NULL;

  // fill in sigma, if no errors have been triggered at this time.
  if (*status == 0) {
    for (int i=0; i<nm; i++) {
      smooth_radius = ccl_massfunc_m2r(cosmo, pow(10,m[i]), status);
      y[i] = log(ccl_sigmaR(cosmo, smooth_radius, 1., status));
    }
    logsigma = gsl_spline_alloc(cosmo->spline_params.M_SPLINE_TYPE, nm);
    if (logsigma == NULL) {
      *status = CCL_ERROR_MEMORY;
    }
  }

  if (*status == 0) {
    gslstatus = gsl_spline_init(logsigma, m, y, nm);
    if (gslstatus != GSL_SUCCESS) {
      *status = CCL_ERROR_SPLINE ;
      ccl_cosmology_set_status_message(
        cosmo,
        "ccl_massfunc.c: ccl_cosmology_compute_sigma(): Error creating sigma(M) spline\n");
    }
  }

  // again, making splines assuming nothing bad has happened to this point
  if (*status == 0 ) {
    for (int i=0; i<nm; i++) {
      if(i==0) {
        gslstatus |= gsl_spline_eval_e(logsigma, m[i], NULL,&na);
        gslstatus |= gsl_spline_eval_e(logsigma, m[i]+cosmo->spline_params.LOGM_SPLINE_DELTA/2., NULL,&nb);
        y[i] = 2.*(na-nb)*y[i] / cosmo->spline_params.LOGM_SPLINE_DELTA;
      }
      else if (i==nm-1) {
        gslstatus |= gsl_spline_eval_e(logsigma, m[i]-cosmo->spline_params.LOGM_SPLINE_DELTA/2., NULL,&na);
        gslstatus |= gsl_spline_eval_e(logsigma, m[i], NULL,&nb);
        y[i] = 2.*(na-nb)*y[i] / cosmo->spline_params.LOGM_SPLINE_DELTA;
      }
      else {
        gslstatus |= gsl_spline_eval_e(logsigma, m[i]-cosmo->spline_params.LOGM_SPLINE_DELTA/2., NULL,&na);
        gslstatus |= gsl_spline_eval_e(logsigma, m[i]+cosmo->spline_params.LOGM_SPLINE_DELTA/2., NULL,&nb);
        y[i] = (na-nb) / cosmo->spline_params.LOGM_SPLINE_DELTA;
      }
    }

    if(gslstatus != GSL_SUCCESS ) {
      ccl_raise_gsl_warning(
        gslstatus, "ccl_massfunc.c: ccl_cosmology_compute_sigma():");
      *status = CCL_ERROR_SPLINE;
      ccl_cosmology_set_status_message(
        cosmo,
        "ccl_massfunc.c: ccl_cosmology_compute_sigma(): "
        "Error evaluating grid points for dlnsigma/dlogM spline\n");
    }
  }

  if (*status == 0) {
    dlnsigma_dlogm = gsl_spline_alloc(cosmo->spline_params.M_SPLINE_TYPE, nm);
    if (dlnsigma_dlogm == NULL) {
      *status = CCL_ERROR_MEMORY;
    }
  }

  if (*status == 0) {
    gslstatus = gsl_spline_init(dlnsigma_dlogm, m, y, nm);
    if (gslstatus != GSL_SUCCESS) {
      *status = CCL_ERROR_SPLINE ;
      ccl_cosmology_set_status_message(
        cosmo,
        "ccl_massfunc.c: ccl_cosmology_compute_sigma(): Error creating dlnsigma/dlogM spline\n");
    }
  }

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

  free(m);
  free(y);
}

/*----- ROUTINE: ccl_massfunc -----
INPUT: ccl_cosmology * cosmo, double halo mass in units of Msun, double scale factor
TASK: returns halo mass function as dn/dlog10(m) in comoving Msun^-1 Mpc^-3 (haloes per mass interval per volume)
*/
double ccl_massfunc(ccl_cosmology *cosmo, double halomass, double a, double odelta, int *status)
{
  if (cosmo->params.N_nu_mass>0){
	  *status = CCL_ERROR_NOT_IMPLEMENTED;
	  ccl_cosmology_set_status_message(cosmo, "ccl_background.c: ccl_cosmology_compute_growth(): Support for the halo mass function in cosmologies with massive neutrinos is not yet implemented.\n");
	  return NAN;
  }

  double logmass;
  logmass = log10(halomass);
  if (logmass > cosmo->spline_params.LOGM_SPLINE_MAX || logmass < cosmo->spline_params.LOGM_SPLINE_MIN){
    *status = CCL_ERROR_HMF_INTERP;
    ccl_cosmology_set_status_message(cosmo, "ccl_massfunc(): The specified halo mass is outside of the range.");
    return NAN;
  }

  if (fabs(cosmo->params.mu_0)>1e-14 || fabs(cosmo->params.sigma_0)>1e-14){
	  *status = CCL_ERROR_NOT_IMPLEMENTED;
	  strcpy(cosmo->status_message,"ccl_massfunc.c: ccl_massfunc(): The halo mass funcion is not implemented the mu / Sigma modified gravity parameterisation.\n");
	  return NAN;
  }

  double f, rho_m;

  rho_m = ccl_constants.RHO_CRITICAL*cosmo->params.Omega_m*cosmo->params.h*cosmo->params.h;
  f=massfunc_f(cosmo,halomass,a,odelta,status);

  return f*rho_m*ccl_dlnsigM_dlogM(cosmo,logmass,status)/halomass;
}

/*----- ROUTINE: ccl_halob1 -----
INPUT: ccl_cosmology * cosmo, double halo mass in units of Msun, double scale factor
TASK: returns dimensionless linear halo bias
*/
double ccl_halo_bias(ccl_cosmology *cosmo, double halomass, double a, double odelta, int *status)
{
  if (cosmo->params.N_nu_mass>0){
	  *status = CCL_ERROR_NOT_IMPLEMENTED;
	  ccl_cosmology_set_status_message(cosmo, "ccl_background.c: ccl_cosmology_compute_growth(): Support for the halo bias in cosmologies with massive neutrinos is not yet implemented.\n");
	  return NAN;
  }

  if (fabs(cosmo->params.mu_0)>1e-14 || fabs(cosmo->params.sigma_0)>1e-14){
	  *status = CCL_ERROR_NOT_IMPLEMENTED;
	  strcpy(cosmo->status_message,"ccl_massfunc.c: ccl_halobias(): The halo bias is not implemented the mu / Sigma modified gravity parameterisation.\n");
	  return NAN;
  }

  double f;
  f = ccl_halo_b1(cosmo,halomass,a,odelta, status);
  return f;
}
/*---- ROUTINE: ccl_massfunc_m2r -----
INPUT: ccl_cosmology * cosmo, halomass in units of Msun
TASK: takes halo mass and converts to halo radius
  in units of Mpc.
*/
double ccl_massfunc_m2r(ccl_cosmology *cosmo, double halomass, int *status)
{
  double rho_m, smooth_radius;

  // Comoving matter density
  //rho_m = ccl_constants.RHO_CRITICAL*cosmo->params.Omega_m*cosmo->params.h*cosmo->params.h;
  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;
}

/*----- 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)
{
  double sigmaM;
  // 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(): linear power spctrum has not been computed!");
    return NAN;
  }

  double lgsigmaM;

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

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

  // Interpolate to get sigma
  sigmaM = exp(lgsigmaM)*ccl_growth_factor(cosmo, a, status);
  return sigmaM;
}

/*----- 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, 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(): linear power spctrum has not been computed!");
    return NAN;
  }

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


#ifndef SVGP_H
#define SVGP_H

#define M_PI 3.14159265358979323846

#include "Vector.h"
#include "Matrix.h"
#include "GaussianProcess.h"
#include <gsl/gsl_math.h>

class SVGP //: public GaussianProcess
	{
	protected:
		Matrix X; ///< Samples (N x M)
		Vector Y; ///< Output (N x 1)
		int N; ///< Number of samples
		int d;

		// Kernel parameters
		real noise_var;
		real sig_var;
		Vector scale_length;

		/// SVI parameters
		int num_inducing; ///< inducing inputs (J)
		Matrix Z; ///< Hidden variables (N x J)
		Vector u; ///< Global variables (N x 1) (targets)

		// variational distribution parametrization
		Matrix S;
		Vector m;

		real l; //step length
		int samples; //not currently used.. but should be

		Vector p_mean;
		Matrix p_var;

		Matrix Kmm;
		Matrix Knm;
		Matrix Kmn;
		Matrix Knn;
		Matrix invKmm;
		Matrix K_tilde;

		real Beta;
		real KL;
		real L;

        Matrix currentSample;
        Vector currentObservation;

        //preferably these could be done in minibatches but only with single examples for now
        //the example is repeated <subsamples> times as in Hoffman et al [2013]
        virtual void local_update(); //updating Z (local/latent variables)
        virtual void global_update(const Matrix& X_samples, const Vector& Y_samples); //updating m, S (which parametrizes q(u) and in turn gives global param u)
		//virtual Vector optimize_Z(int max_iters);
        virtual void init();
        //virtual real LogLikelihood(double* data); //computes the likelihood given a Z*, to use for optimizing the position of Z
        virtual void getSamples(Matrix& X_, Vector& Y_);
	public:
		SVGP(Matrix& X, Vector& Y, Matrix& Z, real noise_var, real sig_var, Vector scale_length, int samples);
        //initializes Z through k-means
		SVGP(Matrix& X, Vector& Y, int num_inducing, real noise_var, real sig_var, Vector scale_length, int samples);
		virtual Matrix Kernel(const Matrix& A, const Matrix& B);
		virtual void Prediction(const Vector& x, real& mean, real& var);
		virtual void UpdateGaussianProcess(); //update
		virtual void FullUpdateGaussianProcess();
		virtual real LogLikelihood();
        //virtual real LogLikelihood(const gsl_vector *v);
		virtual void AddObservation(const Vector& x, const real& y);
		virtual void AddObservation(const std::vector<Vector>& x, const std::vector<real>& y);
		virtual void Clear();

};
#endif


/*
Copyright (C) 2019-2020 JingWeiZhangHuai <jingweizhanghuai@163.com>
Licensed under the Apache License, Version 2.0; you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
*/

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

#include <cblas.h>
#include "morn_tensor.h"

void ConvTensorToMatData(MTensor *tns,int bc,float *mdata,int knl_height,int knl_width,int y_stride,int x_stride);
void ConvMatDataToTensor(float *mdata,MTensor *tns,int bc,int knl_height,int knl_width,int y_stride,int x_stride);

struct TensorSampleConvPara
{
    MLayer *prev;

    int knl_num;
    int knl_height;
    int knl_width;

    int x_stride;
    int y_stride;

    float sample_ratio;

    int res_valid;

    float rate;
    float decay;
    float momentum;
};
void *mTensorSampleConvPara(MSheet *ini,char *name)
{
    struct TensorSampleConvPara *para = (struct TensorSampleConvPara *)mMalloc(sizeof(struct TensorSampleConvPara));

    char *value = mINIRead(ini,name,"prev");
    para->prev = mNetworkLayer(ini,value);
    mException((para->prev == NULL),EXIT,"invalid prev");

    para->res_valid = (strcmp("Input",mLayerType(para->prev))!=0);

    value = mINIRead(ini,name,"knl_num");
    if(value != NULL) para->knl_num= atoi(value);else para->knl_num= 1;

    value = mINIRead(ini,name,"knl_height");
    if(value != NULL) para->knl_height= atoi(value);else para->knl_height= 1;

    value = mINIRead(ini,name,"knl_width");
    if(value != NULL) para->knl_width= atoi(value);else para->knl_width= 1;

    value = mINIRead(ini,name,"x_stride");
    if(value != NULL) para->x_stride= atoi(value);else para->x_stride= 1;

    value = mINIRead(ini,name,"y_stride");
    if(value != NULL) para->y_stride= atoi(value);else para->y_stride= 1;

    value = mINIRead(ini,name,"sample_ratio");
    if(value != NULL) para->sample_ratio = atof(value);else para->sample_ratio = 0.5;

    value = mINIRead(ini,name,"rate");
    if(value != NULL) para->rate = atof(value);
    else
    {
        value = mINIRead(ini,"para","rate");
        if(value != NULL) para->rate = atof(value);
        else              para->rate = 0.001;
    }

    value = mINIRead(ini,name,"decay");
    if(value != NULL) para->decay = atof(value);
    else
    {
        value = mINIRead(ini,"para","decay");
        if(value != NULL) para->decay = atof(value);
        else              para->decay = 0.01;
    }
    mException((para->decay<0.0f)||(para->decay>=1.0f),EXIT,"invalid para decay");

    value = mINIRead(ini,name,"momentum");
    if(value != NULL) para->momentum = atof(value);
    else
    {
        value = mINIRead(ini,"para","momentum");
        if(value != NULL) para->momentum = atof(value);
        else              para->momentum = 0.9;
    }

    return para;
}

struct HandleTensorSampleConv
{
    float *mat;
    float *data;
    float *kernel;
    float *update;
    int   *locate;
};
void endTensorSampleConv(void *info)
{
    struct HandleTensorSampleConv *handle = (struct HandleTensorSampleConv *)info;
    if(handle->mat   != NULL) mFree(handle->mat   );
    if(handle->data  != NULL) mFree(handle->data  );
    if(handle->kernel!= NULL) mFree(handle->kernel);
    if(handle->update!= NULL) mFree(handle->update);
    if(handle->locate!= NULL) mFree(handle->locate);
}
#define HASH_TensorSampleConv 0x8a31af2
void TensorSampleConvSet(MLayer *layer)
{
    if(layer->state != DFLT) return;

    struct TensorSampleConvPara *para = (struct TensorSampleConvPara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *res= para->prev->res;
    MTensor *out= layer->tns;

    MHandle *hdl=mHandle(out,TensorSampleConv);
    struct HandleTensorSampleConv *handle = (struct HandleTensorSampleConv *)(hdl->handle);

    int out_height= in->height/para->y_stride;
    int out_width = in->width /para->x_stride;
    int mheight = (out_height*out_width);
    int mwidth = (para->knl_height*para->knl_width*in->channel)*para->sample_ratio;

    int data_size = para->knl_num*(mwidth+1);

    mTensorRedefine(out,in->batch,para->knl_num,out_height,out_width,NULL);
    if(morn_network_flag == MORN_TRAIN)
    {
        if(INVALID_TENSOR(res)) mTensorRedefine(res,in->batch,in->channel,in->height,in->width,in->data);
        else                    mTensorRedefine(res,in->batch,in->channel,in->height,in->width,NULL);

        if(handle->update != NULL) mFree(handle->update);
        handle->update =(float *)mMalloc(data_size*sizeof(float));
        memset(handle->update,0,data_size*sizeof(float));
    }

    if(handle->kernel !=NULL) mFree(handle->kernel);
    handle->kernel = (float *)mMalloc(data_size*sizeof(float));

    if(morn_network_parafile==NULL)
    {
        float scale = sqrt(2.0f/mwidth);
        for(int i=0;i<data_size;i++)
            handle->kernel[i] = mNormalRand(0.0f,1.0f)*scale;
    }
    else
    {
        mNetworkParaRead(layer,"kernel",handle->kernel,data_size*sizeof(float));
    }

    int matwidth = para->knl_height*para->knl_width*in->channel+1;
    if(handle->mat!=NULL) mFree(handle->mat);
    handle->mat = (float *)mMalloc(mheight*matwidth*sizeof(float));

    if(handle->data!=NULL) mFree(handle->data);
    handle->data= (float *)mMalloc(mheight*(mwidth+1)*sizeof(float));

    if(handle->locate!=NULL) mFree(handle->locate);
    handle->locate = (int *)mMalloc(mwidth*para->knl_num*sizeof(int));
    if(morn_network_parafile==NULL)
    {
        char *flag = (char *)mMalloc(para->knl_height*para->knl_width*in->channel*sizeof(char));
        for(int j=0;j<para->knl_num;j++)
        {
            memset(flag,1,para->knl_height*para->knl_width*in->channel*sizeof(char));
            int *locate = handle->locate+j*mwidth;

            for(int i=0;i<mwidth;)
            {
                locate[i]=mRand(0,para->knl_height*para->knl_width*in->channel);
                if(flag[i]==0) continue;
                flag[i]=0; i++;
            }

            mAscSortS32(locate,NULL,locate,NULL,mwidth);
        }
        mFree(flag);
    }
    else
    {
        mNetworkParaRead(layer,"locate",handle->locate,mwidth*para->knl_num*sizeof(int));
    }

    hdl->valid = 1;
}

void mTensorSampleConvForward(MLayer *layer)
{
    mException(INVALID_POINTER(layer),EXIT,"invalid input");
    mException(strcmp("SampleConv",mLayerType(layer)),EXIT,"invalid layer type");
    struct TensorSampleConvPara *para = (struct TensorSampleConvPara *)(layer->para);

    MTensor *in = para->prev->tns;
    MTensor *out=layer->tns;

    TensorSampleConvSet(layer);

    MHandle *hdl=mHandle(out,TensorSampleConv);
    struct HandleTensorSampleConv *handle = (struct HandleTensorSampleConv *)(hdl->handle);

    int mheight = (out->height*out->width);
    int mwidth = (para->knl_height*para->knl_width*in->channel)*para->sample_ratio;
    int matwidth = para->knl_height*para->knl_width*in->channel+1;

    float *mat_data = handle->mat;
    float *in_data = handle->data;

    for(int b=0;b<in->batch;b++)
    {
        ConvTensorToMatData(in,b,mat_data,para->knl_height,para->knl_width,para->y_stride,para->x_stride);

        for(int k=0;k<para->knl_num;k++)
        {
            int   *locate = handle->locate+k* mwidth;
            float *kernel = handle->kernel+k*(mwidth+1);
            float *out_data = out->data[b]+k* mheight;
            for(int j=0;j<mheight;j++)
            {
                for(int i=0;i<mwidth;i++) in_data[j*mwidth+i]=mat_data[j*matwidth+locate[i]];
                in_data[j*mwidth+mwidth]=1.0f;
            }

            cblas_sgemm(CblasRowMajor,CblasNoTrans,CblasTrans,
                        1,mheight,mwidth+1,
                        1.0f,
                          kernel,mwidth+1,
                         in_data,mwidth+1,
                   0.0f,out_data,mheight);
        }
    }

    layer->state = MORN_FORWARD;
}

void mTensorSampleConvBackward(MLayer *layer)
{
    mException(INVALID_POINTER(layer),EXIT,"invalid input");
    mException(strcmp("SampleConv",mLayerType(layer)),EXIT,"invalid layer type");
    struct TensorSampleConvPara *para = (struct TensorSampleConvPara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *res= para->prev->res;
    MTensor *out=layer->res;

    MHandle *hdl=mHandle(layer->tns,TensorSampleConv);
    struct HandleTensorSampleConv *handle = (struct HandleTensorSampleConv *)(hdl->handle);
    mException((hdl->valid == 0),EXIT,"no forward operate");

    int mheight = (out->height*out->width);
    int mwidth = (para->knl_height*para->knl_width*in->channel)*para->sample_ratio;
    int matwidth = para->knl_height*para->knl_width*in->channel+1;

    float *mat_data = handle->mat;
    float *in_data  = handle->data;
    float *res_data = handle->data;

    mNetworkParaWrite(layer,"locate",handle->locate,mwidth*para->knl_num*sizeof(int));
    mNetworkParaWrite(layer,"kernel",handle->kernel,para->knl_num*mwidth*sizeof(float));

    for(int b=0;b<out->batch;b++)
    {
        ConvTensorToMatData(in,b,mat_data,para->knl_height,para->knl_width,para->y_stride,para->x_stride);

        for(int k=0;k<para->knl_num;k++)
        {
            int   *locate = handle->locate+k* mwidth;
            float *update = handle->update+k*(mwidth+1);
            float *out_data = out->data[b]+k* mheight;
            for(int j=0;j<mheight;j++)
            {
                for(int i=0;i<mwidth;i++) in_data[j*mwidth+i]=mat_data[j*matwidth+locate[i]];
                in_data[j*mwidth+mwidth]=1.0f;
            }

            cblas_sgemm(CblasRowMajor,CblasNoTrans,CblasNoTrans,
                        1,mwidth,mheight,
                        1.0f/mheight,
                        out_data,mheight,
                        in_data,mwidth,
                        (b==0)?para->momentum:1.0f,
                        update,mwidth);

        }
    }

    cblas_saxpby(para->knl_num*mwidth,
                 (0.0f-(para->rate/(float)(in->batch))),handle->update,1,
                 (1.0f-(para->decay*para->rate))       ,handle->kernel,1);

    if(para->res_valid==0) return;

    if(para->prev->state == MORN_FORWARD)
    {
        for(int b=0;b<res->batch;b++)
            memset(res->data[b],0,in->height*in->width*in->channel*sizeof(float));
        para->prev->state = MORN_BACKWARD;
    }

    for(int b=0;b<in->batch;b++)
    {
        memset(mat_data,0,matwidth*mheight*sizeof(float));

        for(int k=0;k<para->knl_num;k++)
        {
            int   *locate = handle->locate+k* mwidth;
            float *kernel = handle->kernel+k*(mwidth+1);
            float *out_data = out->data[b]+k* mheight;

            cblas_sgemm(CblasRowMajor,CblasTrans,CblasNoTrans,
                        mheight,mwidth,1,
                        1.0f,
                        out_data,mheight,
                          kernel,mwidth,
                    0.0,res_data,mwidth);

            for(int j=0;j<mheight;j++)
            {
                for(int i=0;i<mwidth;i++) mat_data[j*matwidth+locate[i]]+=res_data[j*mwidth+i];
            }
        }

        ConvMatDataToTensor(mat_data,res,b,para->knl_height,para->knl_width,para->y_stride,para->x_stride);
    }
}

void DirConvTensorToMatData(MTensor *tns,int bc,float *mdata,int knl_r,int knl_dir,int y_stride,int x_stride)
{
    int height = tns->height;
    int width  = tns->width;
    int channel= tns->channel;

    int out_width = width/x_stride;
    int out_height=height/y_stride;
    int mwidth = (knl_r+knl_r+1)*channel+1;
    int mheight= out_height*out_width;

    int sx=0,sy=0;
         if(knl_dir==0) {sx=0;sy= 1;}
    else if(knl_dir==1) {sx=1;sy= 1;}
    else if(knl_dir==2) {sx=1;sy= 0;}
    else if(knl_dir==3) {sx=1;sy=-1;}

    int x0=(width -(out_width -1)*x_stride)/2;
    int y0=(height-(out_height-1)*y_stride)/2;

    float *tdata = tns->data[bc];
    int tsize = tns->height*tns->width;

    int i,j,c;
    for(j=out_width;j<mheight;j++)
    {
        int n=y0+j/out_width*y_stride-knl_r*sy;
        int m=x0+j%out_width*x_stride-knl_r*sx;
        for(i=0;i<mwidth-1;i+=channel)
        {
            int h=n;if(h<0)h=0;else if(h>=height)h=height-1;
            int w=m;if(w<0)w=0;else if(w>= width)w= width-1;
            for(c=0;c<channel;c++) mdata[(j*mwidth)+i+c]=tdata[c*tsize+h*width+w];
            n=n+sy;m=m+sx;
        }
        mdata[(j*mwidth)+mwidth-1]=1.0f;
    }
}

void DirConvMatDataToTensor(float *mdata,MTensor *tns,int bc,int knl_r,int knl_dir,int y_stride,int x_stride)
{
    int height = tns->height;
    int width  = tns->width;
    int channel= tns->channel;

    int out_width = width/x_stride;
    int out_height=height/y_stride;
    int mwidth = (knl_r+knl_r+1)*channel+1;
    int mheight= out_height*out_width;

    int sx=0,sy=0;
         if(knl_dir==0) {sx=0;sy= 1;}
    else if(knl_dir==1) {sx=1;sy= 1;}
    else if(knl_dir==2) {sx=1;sy= 0;}
    else if(knl_dir==3) {sx=1;sy=-1;}

    int x0=(width -(out_width -1)*x_stride)/2;
    int y0=(height-(out_height-1)*y_stride)/2;

    float *tdata = tns->data[bc];
    int tsize = tns->height*tns->width;

    int i,j,c;
    for(j=0;j<mheight;j++)
    {
        int n=y0+j/out_width*y_stride-knl_r*sy;
        int m=x0+j%out_width*x_stride-knl_r*sx;
        for(i=0;i<mwidth-1;i+=channel)
        {
            int h=n;if(h<0)h=0;else if(h>=height)h=height-1;
            int w=m;if(w<0)w=0;else if(w>= width)w= width-1;
            for(c=0;c<channel;c++) tdata[c*tsize+h*width+w]+=mdata[(j*mwidth)+i+c];
            n=n+sy;m=m+sx;
        }
    }
}

struct TensorDirConvPara
{
    MLayer *prev;

    int knl_num;
    int knl_r;

    int dir;

    int x_stride;
    int y_stride;

    int res_valid;

    float rate;
    float decay;
    float momentum;
};
void *mTensorDirConvPara(MSheet *ini,char *name)
{
    struct TensorDirConvPara *para = (struct TensorDirConvPara *)mMalloc(sizeof(struct TensorDirConvPara));

    char *value = mINIRead(ini,name,"prev");
    para->prev = mNetworkLayer(ini,value);
    mException((para->prev == NULL),EXIT,"invalid prev");

    para->res_valid = (strcmp("Input",mLayerType(para->prev))!=0);

    value = mINIRead(ini,name,"knl_num");
    if(value != NULL) para->knl_num= atoi(value);else para->knl_num= 1;

    value = mINIRead(ini,name,"knl_r");
    if(value != NULL) para->knl_r= atoi(value);else para->knl_r= 1;

    value = mINIRead(ini,name,"dir");
    mException((value==NULL),EXIT,"invalid convolution direction");
         if((strcmp(value, "u-d" )==0)||(strcmp(value, "d-u" )==0)) para->dir=0;
    else if((strcmp(value,"lu-rd")==0)||(strcmp(value,"rd-lu")==0)) para->dir=1;
    else if((strcmp(value, "l-r" )==0)||(strcmp(value, "r-l" )==0)) para->dir=2;
    else if((strcmp(value,"ru-ld")==0)||(strcmp(value,"ld-ru")==0)) para->dir=3;
    else mException(1,EXIT,"invalid convolution direction");

    value = mINIRead(ini,name,"x_stride");
    if(value != NULL) para->x_stride= atoi(value);else para->x_stride= 1;

    value = mINIRead(ini,name,"y_stride");
    if(value != NULL) para->y_stride= atoi(value);else para->y_stride= 1;

    value = mINIRead(ini,name,"rate");
    if(value != NULL) para->rate = atof(value);
    else
    {
        value = mINIRead(ini,"para","rate");
        if(value != NULL) para->rate = atof(value);
        else              para->rate = 0.001;
    }

    value = mINIRead(ini,name,"decay");
    if(value != NULL) para->decay = atof(value);
    else
    {
        value = mINIRead(ini,"para","decay");
        if(value != NULL) para->decay = atof(value);
        else              para->decay = 0.01;
    }
    mException((para->decay<0.0f)||(para->decay>=1.0f),EXIT,"invalid para decay");

    value = mINIRead(ini,name,"momentum");
    if(value != NULL) para->momentum = atof(value);
    else
    {
        value = mINIRead(ini,"para","momentum");
        if(value != NULL) para->momentum = atof(value);
        else              para->momentum = 0.9;
    }

    return para;
}

struct HandleTensorDirConv
{
    float *mat;
    float *kernel;
    float *update;
};
void endTensorDirConv(void *info)
{
    struct HandleTensorDirConv *handle = (struct HandleTensorDirConv *)info;
    if(handle->mat   != NULL) mFree(handle->mat);
    if(handle->kernel!= NULL) mFree(handle->kernel);
    if(handle->update!= NULL) mFree(handle->update);
}
#define HASH_TensorDirConv 0x9e2674b3
void TensorDirConvSet(MLayer *layer)
{
    if(layer->state != DFLT) return;

    struct TensorDirConvPara *para = (struct TensorDirConvPara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *res= para->prev->res;
    MTensor *out= layer->tns;

    MHandle *hdl=mHandle(out,TensorDirConv);
    struct HandleTensorDirConv *handle = (struct HandleTensorDirConv *)(hdl->handle);

    int out_height= in->height/para->y_stride;
    int out_width = in->width /para->x_stride;
    int mheight = (out_height*out_width);
    int mwidth = (para->knl_r+para->knl_r+1)*in->channel+1;
    int data_size = para->knl_num*mwidth;

    mTensorRedefine(out,in->batch,para->knl_num,out_height,out_width,NULL);
    if(morn_network_flag == MORN_TRAIN)
    {
        if(INVALID_TENSOR(res)) mTensorRedefine(res,in->batch,in->channel,in->height,in->width,in->data);
        else                    mTensorRedefine(res,in->batch,in->channel,in->height,in->width,NULL);

        if(handle->update != NULL) mFree(handle->update);
        handle->update =(float *)mMalloc(data_size*sizeof(float));
        memset(handle->update,0,data_size*sizeof(float));
    }

    if(handle->kernel !=NULL) mFree(handle->kernel);
    handle->kernel = (float *)mMalloc(data_size*sizeof(float));

    if(morn_network_parafile==NULL)
    {
        float scale = sqrt(2.0f/mwidth);
        for(int i=0;i<data_size;i++)
            handle->kernel[i] = mNormalRand(0.0f,1.0f)*scale;
    }
    else
    {
        mNetworkParaRead(layer,"kernel",handle->kernel,data_size*sizeof(float));
    }

    if(handle->mat!=NULL) mFree(handle->mat);
    handle->mat = (float *)mMalloc(mheight*mwidth*sizeof(float));

    hdl->valid = 1;
}

void mTensorDirConvForward(MLayer *layer)
{
    mException(INVALID_POINTER(layer),EXIT,"invalid input");
    mException(strcmp("DirConv",mLayerType(layer)),EXIT,"invalid layer type");
    struct TensorDirConvPara *para = (struct TensorDirConvPara *)(layer->para);

    MTensor *in = para->prev->tns;
    MTensor *out=layer->tns;

    TensorDirConvSet(layer);

    MHandle *hdl=mHandle(out,TensorDirConv);
    struct HandleTensorDirConv *handle = (struct HandleTensorDirConv *)(hdl->handle);

    int mheight = (out->height*out->width);
    int mwidth = (para->knl_r+para->knl_r+1)*in->channel+1;

    float *kernel_data= handle->kernel;
    float *in_data = handle->mat;

    for(int b=0;b<in->batch;b++)
    {
        DirConvTensorToMatData(in,b,in_data,para->knl_r,para->dir,para->y_stride,para->x_stride);
        float *out_data = out->data[b];

        cblas_sgemm(CblasRowMajor,CblasNoTrans,CblasTrans,
                    para->knl_num,mheight,mwidth,
                    1.0f,
                    kernel_data,mwidth,
                        in_data,mwidth,
               0.0f,   out_data,mheight);
    }

    layer->state = MORN_FORWARD;
}

void mTensorDirConvBackward(MLayer *layer)
{
    mException(INVALID_POINTER(layer),EXIT,"invalid input");
    mException(strcmp("DirConv",mLayerType(layer)),EXIT,"invalid layer type");
    struct TensorDirConvPara *para = (struct TensorDirConvPara *)(layer->para);
    MTensor *in = para->prev->tns;
    MTensor *res= para->prev->res;
    MTensor *out=layer->res;

    MHandle *hdl=mHandle(layer->tns,TensorDirConv);
    struct HandleTensorDirConv *handle = (struct HandleTensorDirConv *)(hdl->handle);
    mException((hdl->valid == 0),EXIT,"no forward operate");

    int mheight = (out->height*out->width);
    int mwidth = (para->knl_r+para->knl_r+1)*in->channel+1;

    float *kernel_data= handle->kernel;
    float *update_data= handle->update;
    float *    in_data= handle->mat;
    float *   res_data= handle->mat;

    mNetworkParaWrite(layer,"kernel",kernel_data,para->knl_num*mwidth*sizeof(float));

    for(int b=0;b<out->batch;b++)
    {
        DirConvTensorToMatData(in,b,in_data,para->knl_r,para->dir,para->y_stride,para->x_stride);
        float *out_data = out->data[b];

        cblas_sgemm(CblasRowMajor,CblasNoTrans,CblasNoTrans,
                    para->knl_num,mwidth,mheight,
                    1.0f/mheight,
                       out_data,mheight,
                        in_data,mwidth,
                    (b==0)?para->momentum:1.0f,
                    update_data,mwidth);
    }

    cblas_saxpby(para->knl_num*mwidth,
                 (0.0f-(para->rate/(float)(in->batch))),update_data,1,
                 (1.0f-(para->decay*para->rate))       ,kernel_data,1);

    if(para->res_valid==0) return;

    if(para->prev->state == MORN_FORWARD)
    {
        for(int b=0;b<res->batch;b++)
            memset(res->data[b],0,in->height*in->width*in->channel*sizeof(float));
        para->prev->state = MORN_BACKWARD;
    }

    for(int b=0;b<in->batch;b++)
    {
        float *out_data = out->data[b];

        cblas_sgemm(CblasRowMajor,CblasTrans,CblasNoTrans,
                    mheight,mwidth,para->knl_num,
                    1.0f,
                       out_data,mheight,
                    kernel_data,mwidth,
                0.0,   res_data,mwidth);

        DirConvMatDataToTensor(res_data,res,b,para->knl_r,para->dir,para->y_stride,para->x_stride);
    }
}


/**
 *
 * @file core_ctrtri.c
 *
 *  PLASMA core_blas kernel
 *  PLASMA is a software package provided by Univ. of Tennessee,
 *  Univ. of California Berkeley and Univ. of Colorado Denver
 *
 * @version 2.6.0
 * @author Julien Langou
 * @author Henricus Bouwmeester
 * @author Mathieu Faverge
 * @date 2010-11-15
 * @generated c Tue Jan  7 11:44:46 2014
 *
 **/
#include <lapacke.h>
#include "common.h"

/***************************************************************************//**
 *
 * @ingroup CORE_PLASMA_Complex32_t
 *
 *  CORE_ctrtri - Computes the inverse of a complex upper or lower
 *  triangular matrix A.
 *
 *******************************************************************************
 *
 * @param[in] uplo
 *          = PlasmaUpper: Upper triangle of A is stored;
 *          = PlasmaLower: Lower triangle of A is stored.
 *
 * @param[in] diag
 *          = PlasmaNonUnit: A is non-unit triangular;
 *          = PlasmaUnit:    A is unit triangular.
 *
 * @param[in] N
 *          The order of the matrix A. N >= 0.
 *
 * @param[in,out] A
 *          On entry, the triangular matrix A.  If UPLO = 'U', the
 *          leading N-by-N upper triangular part of the array A
 *          contains the upper triangular matrix, and the strictly
 *          lower triangular part of A is not referenced.  If UPLO =
 *          'L', the leading N-by-N lower triangular part of the array
 *          A contains the lower triangular matrix, and the strictly
 *          upper triangular part of A is not referenced.  If DIAG =
 *          'U', the diagonal elements of A are also not referenced and
 *          are assumed to be 1.  On exit, the (triangular) inverse of
 *          the original matrix.
 *
 * @param[in] LDA
 *          The leading dimension of the array A. LDA >= max(1,N).
 *
 * @param[out] info
 *          - 0 on successful exit
 *          - <0 if -i, the i-th argument had an illegal value
 *          - >0 if i, A(i,i) is exactly zero.  The triangular
 *          matrix is singular and its inverse can not be computed.
 *
 ******************************************************************************/
#if defined(PLASMA_HAVE_WEAK)
#pragma weak CORE_ctrtri = PCORE_ctrtri
#define CORE_ctrtri PCORE_ctrtri
#endif
void CORE_ctrtri(PLASMA_enum uplo, PLASMA_enum diag, int N, PLASMA_Complex32_t *A, int LDA, int *info)
{
    *info = LAPACKE_ctrtri_work(
        LAPACK_COL_MAJOR,
        lapack_const(uplo), lapack_const(diag),
        N, A, LDA);
}