/* mcmc.h
 * David W. Pearson
 * 17 July 2018
 *
 * This header file will be responsible for running Markov Chain Monte Carlo fitting.
 */

#ifndef _MCMC_H_
#define _MCMC_H_

#include "bispectrum_model.h"
#include <iostream>
#include <fstream>
#include <random>
#include <vector>
#include <string>
#include <cmath>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>

std::random_device seeder;
std::mt19937_64 gen(seeder());
std::uniform_real_distribution<double> dist(-1.0, 1.0);

class bkmcmc{
    int num_data, num_pars;
    std::vector<double> data; // These should have size of num_data
    std::vector<std::vector<double>> Psi; // num_data vectors of size num_data
    std::vector<double> theta_0, theta_i, param_vars, min, max; // These should all have size of num_pars
    std::vector<float4> k; // This should have size of num_data
    std::vector<bool> limit_pars; // This should have size of num_pars
    double chisq_0, chisq_i;

    // Sets the values of theta_i.
    void get_param_real(); // done

    // Calculates the chi^2 for the current proposal, theta_i
    double calc_chi_squared(); // done

    // Performs one MCMC trial. Returns true if proposal accepted, false otherwise
    bool trial(float4 *ks, double *Bk, double &L, double &R); // done

    // Writes the current accepted parameters to the screen
    void write_theta_screen(); // done

    // Burns the requested number of parameter realizations to move to a higher likelihood region
    void burn_in(int num_burn, float4 *ks, double *Bk); // done

    // Changes the initial guesses for the search range around parameters until acceptance = 0.234
    void tune_vars(float4 *ks, double *Bk); // done

    public:
        std::vector<double> model; // These should have size num_data

        // Initializes most of the data members and gets an initial chisq_0
        bkmcmc(std::string data_file, std::string cov_file, std::vector<double> &pars,
               std::vector<double> &vars, float4 *ks, double *Bk); // done

        // Displays information to the screen to check that the vectors are all the correct size
        void check_init(); // done

        // Sets which parameters should be limited and what the limits are
        void set_param_limits(std::vector<bool> &lim_pars, std::vector<double> &min_in,
                              std::vector<double> &max_in); // done

        // Runs the MCMC chain for num_draws realizations, writing to reals_file
        void run_chain(int num_draws, int num_burn, std::string reals_file, float4 *ks, double *Bk, bool new_chain);

};

void bkmcmc::get_param_real() {
    for (int i = 0; i < bkmcmc::num_pars; ++i) {
        if (bkmcmc::limit_pars[i]) {
            if (bkmcmc::theta_0[i] + bkmcmc::param_vars[i] > bkmcmc::max[i]) {
                double center = bkmcmc::max[i] - bkmcmc::param_vars[i];
                bkmcmc::theta_i[i] = center + dist(gen)*bkmcmc::param_vars[i];
            } else if (bkmcmc::theta_0[i] - bkmcmc::param_vars[i] < bkmcmc::min[i]) {
                double center = bkmcmc::min[i] + bkmcmc::param_vars[i];
                bkmcmc::theta_i[i] = center + dist(gen)*bkmcmc::param_vars[i];
            } else {
                bkmcmc::theta_i[i] = bkmcmc::theta_0[i] + dist(gen)*bkmcmc::param_vars[i];
            }
        } else {
            bkmcmc::theta_i[i] = bkmcmc::theta_0[i] + dist(gen)*bkmcmc::param_vars[i];
        }
    }
}

double bkmcmc::calc_chi_squared() {
    double chisq = 0.0;
    for (int i = 0; i < bkmcmc::num_data; ++i) {
        for (int j = i; j < bkmcmc::num_data; ++j) {
            if (bkmcmc::data[i] > 0 && bkmcmc::data[j] > 0) {
                chisq += (bkmcmc::data[i] - bkmcmc::model[i])*Psi[i][j]*(bkmcmc::data[j] - bkmcmc::model[j]);
            }
        }
    }
    return chisq;
}

bool bkmcmc::trial(float4 *ks, double *d_Bk, double &L, double &R) {
    bkmcmc::get_param_real();
    model_calc(bkmcmc::theta_i, ks, d_Bk, bkmcmc::model);
    bkmcmc::chisq_i = bkmcmc::calc_chi_squared();

    L = exp(0.5*(bkmcmc::chisq_0 - bkmcmc::chisq_i));
    R = (dist(gen) + 1.0)/2.0;

    if (L > R) {
        for (int i = 0; i < bkmcmc::num_pars; ++i)
            bkmcmc::theta_0[i] = bkmcmc::theta_i[i];
        bkmcmc::chisq_0 = bkmcmc::chisq_i;
        return true;
    } else {
        return false;
    }
}

void bkmcmc::write_theta_screen() {
    std::cout.precision(6);
    for (int i = 0; i < bkmcmc::num_pars; ++i) {
        std::cout.width(15);
        std::cout << bkmcmc::theta_0[i];
    }
    std::cout.width(15);
    std::cout << pow(bkmcmc::theta_0[3]*bkmcmc::theta_0[4]*bkmcmc::theta_0[4],1.0/3.0);
    std::cout.width(15);
    std::cout << bkmcmc::chisq_0;
    std::cout.flush();
}

void bkmcmc::burn_in(int num_burn, float4 *ks, double *d_Bk) {
    std::cout << "Burning the first " << num_burn << " trials to move to higher likelihood..." << std::endl;
    double L, R;
    for (int i = 0; i < num_burn; ++i) {
        bool move = bkmcmc::trial(ks, d_Bk, L, R);
        if (true) {
            std::cout << "\r";
            std::cout.width(5);
            std::cout << i;
            bkmcmc::write_theta_screen();
            std::cout.width(15);
            std::cout << L;
            std::cout.width(15);
            std::cout << R;
            std::cout.flush();
        }
    }
    std::cout << std::endl;
}

void bkmcmc::tune_vars(float4 *ks, double *d_Bk) {
    std::cout << "Tuning acceptance ratio..." << std::endl;
    double acceptance = 0.0;
    while (acceptance <= 0.233 || acceptance >= 0.235) {
        int accept = 0;
        double L, R;
        for (int i = 0; i < 10000; ++i) {
            bool move = bkmcmc::trial(ks, d_Bk, L, R);
            if (move) {
                std::cout << "\r";
                bkmcmc::write_theta_screen();
                accept++;
            }
        }
        std::cout << std::endl;
        acceptance = double(accept)/10000.0;

        if (acceptance <= 0.233) {
            for (int i = 0; i < bkmcmc::num_pars; ++i)
                bkmcmc::param_vars[i] *= 0.99;
        }
        if (acceptance >= 0.235) {
            for (int i = 0; i < bkmcmc::num_pars; ++i)
                bkmcmc::param_vars[i] *= 1.01;
        }
        std::cout << "acceptance = " << acceptance << std::endl;
    }
    std::ofstream fout;
    fout.open("variances.dat", std::ios::out);
    for (int i = 0; i < bkmcmc::num_pars; ++i)
        fout << bkmcmc::param_vars[i] << " ";
    fout << "\n";
    fout.close();
}

bkmcmc::bkmcmc(std::string data_file, std::string cov_file, std::vector<double> &pars,
               std::vector<double> &vars, float4 *ks, double *d_Bk) {
    std::ifstream fin;
    std::ofstream fout;

    std::cout << "Reading in and storing data file..." << std::endl;
    if (std::ifstream(data_file)) {
        fin.open(data_file.c_str(), std::ios::in);
        while (!fin.eof()) {
            float4 kt;
            double B;
            fin >> kt.w >> kt.x >> kt.y >> kt.z >> B;
            if (!fin.eof()) {
                bkmcmc::k.push_back(kt);
                bkmcmc::data.push_back(B);
                bkmcmc::model.push_back(0.0);
            }
        }
        fin.close();
    } else {
        std::stringstream message;
        message << "Could not open " << data_file << std::endl;
        throw std::runtime_error(message.str());
    }

    bkmcmc::num_data = bkmcmc::data.size();
    std::cout << "num_data = " << bkmcmc::num_data << std::endl;

    gsl_matrix *cov = gsl_matrix_alloc(bkmcmc::num_data, bkmcmc::num_data);
    gsl_matrix *psi = gsl_matrix_alloc(bkmcmc::num_data, bkmcmc::num_data);
    gsl_permutation *perm = gsl_permutation_alloc(bkmcmc::num_data);

    std::cout << "Reading in covariance and computing its inverse..." << std::endl;
    if (std::ifstream(cov_file)) {
        fin.open(cov_file.c_str(), std::ios::in);
        for (int i = 0; i < bkmcmc::num_data; ++i) {
            for (int j = 0; j < bkmcmc::num_data; ++j) {
                double element;
                fin >> element;
                gsl_matrix_set(cov, i, j, element);
            }
        }
        fin.close();
    } else {
        std::stringstream message;
        message << "Could not open " << cov_file << std::endl;
        throw std::runtime_error(message.str());
    }

    int s;
    gsl_linalg_LU_decomp(cov, perm, &s);
    gsl_linalg_LU_invert(cov, perm, psi);

    for (int i = 0; i < bkmcmc::num_data; ++i) {
        std::vector<double> row;
        row.reserve(bkmcmc::num_data);
        for (int j = 0; j < bkmcmc::num_data; ++j) {
            row.push_back((1.0 - double(bkmcmc::num_data + 1.0)/2048.0)*gsl_matrix_get(psi, i, j));
        }
        bkmcmc::Psi.push_back(row);
    }

    gsl_matrix_free(cov);
    gsl_matrix_free(psi);
    gsl_permutation_free(perm);

    gpuErrchk(cudaMemcpy(ks, bkmcmc::k.data(), bkmcmc::num_data*sizeof(float4), cudaMemcpyHostToDevice));

    gpuErrchk(cudaMemcpy(d_Bk, bkmcmc::model.data(), bkmcmc::num_data*sizeof(double),
                         cudaMemcpyHostToDevice));

    bkmcmc::num_pars = pars.size();
    std::cout << "num_pars = " << bkmcmc::num_pars << std::endl;

    for (int i = 0; i < bkmcmc::num_pars; ++i) {
        bkmcmc::theta_0.push_back(pars[i]);
        bkmcmc::theta_i.push_back(0.0);
        bkmcmc::limit_pars.push_back(false);
        bkmcmc::max.push_back(0.0);
        bkmcmc::min.push_back(0.0);
        bkmcmc::param_vars.push_back(vars[i]);
    }

    std::cout << "Calculating initial model and chi^2..." << std::endl;
    model_calc(bkmcmc::theta_0, ks, d_Bk, bkmcmc::model);
    bkmcmc::chisq_0 = bkmcmc::calc_chi_squared();

    fout.open("Bk_mod_check.dat", std::ios::out);
    for (int i =0; i < bkmcmc::num_data; ++i) {
        fout.precision(3);
        fout << bkmcmc::k[i].x << " " << bkmcmc::k[i].y << " " << bkmcmc::k[i].z << " ";
        fout.precision(15);
        fout << bkmcmc::data[i] << " " << bkmcmc::model[i] << "\n";
    }
    fout.close();
}

void bkmcmc::check_init() {
    std::cout << "Number of data points: " << bkmcmc::num_data << std::endl;
    std::cout << "    data.size()      = " << bkmcmc::data.size() << std::endl;
    std::cout << "    model.size()     = " << bkmcmc::model.size() << std::endl;
    std::cout << "    Psi.size()       = " << bkmcmc::Psi.size() << std::endl;
    std::cout << "Number of parameters:  " << bkmcmc::num_pars << std::endl;
    std::cout << "    theta_0.size()   = " << bkmcmc::theta_0.size() << std::endl;
    std::cout << "    theta_i.size()   = " << bkmcmc::theta_i.size() << std::endl;
    std::cout << "    limit_pars.size()= " << bkmcmc::limit_pars.size() << std::endl;
    std::cout << "    min.size()       = " << bkmcmc::min.size() << std::endl;
    std::cout << "    max.size()       = " << bkmcmc::max.size() << std::endl;
    std::cout << "    param_vars.size()= " << bkmcmc::param_vars.size() << std::endl;
}

void bkmcmc::set_param_limits(std::vector<bool> &lim_pars, std::vector<double> &min_in,
                              std::vector<double> &max_in) {
    for (int i = 0; i < bkmcmc::num_pars; ++i) {
        bkmcmc::limit_pars[i] = lim_pars[i];
        bkmcmc::max[i] = max_in[i];
        bkmcmc::min[i] = min_in[i];
    }
}

void bkmcmc::run_chain(int num_draws, int num_burn, std::string reals_file, float4 *ks, double *d_Bk, bool new_chain) {
    int num_old_rels = 0;
    if (new_chain) {
        std::cout << "Starting new chain..." << std::endl;
        bkmcmc::burn_in(num_burn, ks, d_Bk);
        bkmcmc::tune_vars(ks, d_Bk);
    } else {
        std::cout << "Resuming previous chain..." << std::endl;
        std::ifstream fin;
        fin.open("variances.dat", std::ios::in);
        for (int i = 0; i < bkmcmc::num_pars; ++i) {
            double var;
            fin >> var;
            bkmcmc::param_vars[i] = var;
        }
        fin.close();
        fin.open(reals_file.c_str(), std::ios::in);
        while (!fin.eof()) {
            double alpha;
            num_old_rels++;
            std::cout << "\r";
            for (int i = 0; i < bkmcmc::num_pars; ++i) {
                fin >> bkmcmc::theta_0[i];
                std::cout.width(10);
                std::cout << bkmcmc::theta_0[i];
            }
            fin >> alpha;
            fin >> bkmcmc::chisq_0;
            std::cout.width(10);
            std::cout << alpha;
            std::cout.width(10);
            std::cout << bkmcmc::chisq_0;
        }
        fin.close();
        num_old_rels--;
    }

    std::ofstream fout;
    double L, R;
    fout.open(reals_file.c_str(), std::ios::app);
    fout.precision(15);
    for (int i = 0; i < num_draws; ++i) {
        bool move = bkmcmc::trial(ks, d_Bk, L, R);
        for (int par = 0; par < bkmcmc::num_pars; ++par) {
            fout << bkmcmc::theta_0[par] << " ";
        }
        double alpha = pow(bkmcmc::theta_0[3]*bkmcmc::theta_0[4]*bkmcmc::theta_0[4], 1.0/3.0);
        fout << alpha << " " << bkmcmc::chisq_0 << "\n";
        if (move) {
            std::cout << "\r";
            std::cout.width(15);
            std::cout << i + num_old_rels;
            bkmcmc::write_theta_screen();
        }
    }
    std::cout << std::endl;
    fout.close();
}

#endif


/*
   Compile as

   gcc compGraph2graph6.c -O2 -lgsl -lgslcblas -o compGraph2graph6
*/
#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#include <gsl/gsl_linalg.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_eigen.h>
#include <gsl/gsl_blas.h>

/* Order of input graphs - note this code will only work for orders smaller than 62 */
#define N 19

/* Maximum order we expect for the the comparability graph */
#define MAX_ORDER 75000

/* Forbidden eigenvalue of the star complement */
#define SC_EIG 2

/* Specified precision used while doing floating point operations we treat EPS as zero */
#define EPS 0.000001

#define EXPECTED_CLIQUE 57

/* graph6 related things */

#define SIZELEN(n) ((n)<=SMALLN?1:((n)<=SMALLISHN?4:8))
#define G6LEN(n) (SIZELEN(n) \
         + ((size_t)(n)/12)*((size_t)(n)-1) + (((size_t)(n)%12)*((size_t)(n)-1)+11)/12)

#define SMALLN 62
#define BIAS6 63
#define TOPBIT6 32
#define SMALLISHN 258047
#define MAXBYTE 126
#define C6MASK 63

static gsl_matrix *mat, *mat_inv;
static gsl_permutation *perm;

static gsl_vector *vecs[1<<N]; //vecs and vecs_prod are arrays of vectors with size 2^N
static gsl_vector *vecs_prod[1<<N];

static gsl_vector *vec_j;

/* storing the current graph6 string of our graph */
static char line[G6LEN(N)+2];

static char gcode[G6LEN(MAX_ORDER)+3]; //gcode is string with enough room to store graph of order MAX_ORDER

/* number of currently processed graphs. */
static unsigned nproc = 0;
static unsigned skipped = 0;

static FILE *outFile;

static void stringtomat(char *s) {

	char *p;
	int i,j,k,x = 0;

    /* Clear the adjacency matrix */
    gsl_matrix_set_zero(mat);

    p = s + 1;

    k = 1;

    gsl_matrix_set(mat, 0, 0, SC_EIG);

    for (j = 1; j < N; ++j) {
        gsl_matrix_set(mat, j, j, SC_EIG);
        for (i = 0; i < j; ++i) {
            if (--k == 0) {
        		k = 6;
		        x = *(p++) - BIAS6;
            }

            if (x & TOPBIT6) {
                gsl_matrix_set(mat, i, j, -1);
                gsl_matrix_set(mat, j, i, -1);
            }
            x <<= 1;
        }
    }

    int signum;
    gsl_linalg_LU_decomp(mat, perm , &signum);
    gsl_linalg_LU_invert(mat, perm, mat_inv);

}

/* Some graph6 string thingie */
static void encodegraphsize(const int n, char **pp) {
    char *p;

    p = *pp;
    if (n <= SMALLN)
        *p++ = BIAS6 + n;
    else {
        *p++ = MAXBYTE;
        *p++ = BIAS6 + (n >> 12);
        *p++ = BIAS6 + ((n >> 6) & C6MASK);
        *p++ = BIAS6 + (n & C6MASK);
    }
    *pp = p;
}

/* Given that mat_inv is the M = SC_EIG*I - A we
   compute all binary vectors x such that x M x^t == SC_EIG
   and x M j^t == -1.

   Finally for any pair x,y of such vectors we add an edge to our
   graph if x M y^t is either 0 or -1.
*/

/* This is a global thingie since declaring it localy as
 * an array of size 1<<N breaks the stack limit */
gsl_vector *verts[1<<N];

static void constructGraph(void) {

    double res; //allocate memory for double-sized result
    unsigned i, j; //declare variables used for looping

    /* After the first iteration this holds the number of
       vertices of the obtained graph. The respective vertices
       are stored in vecs_prod[0]...vecs_prod[cache_size-1].
    */
    unsigned cache_size = 0; //initialize cache_size to 0

    for (i = 0; i < 1<<N; i++) { //loop through all 2^N vectors in vecs, looking to fill in the vertices
        gsl_blas_dsymv(CblasUpper, 1, mat_inv, vecs[i], 0, vecs_prod[cache_size]); //set vecs_prod[cache_size] to mat_inv*vecs[i], assuming mat_inv is symmetrical and using upper triangle ((rI-A_H)^-1*u), this is the bilinear form <u,_>

        gsl_blas_ddot(vecs_prod[cache_size], vec_j, &res); //calculate <u,j> and store it in space allocated for result

        if (fabs(res+1) < EPS) { //if <u,j> = -1 (within EPS precision)
            gsl_blas_ddot(vecs_prod[cache_size], vecs[i], &res); //calculate <u,u> and store it in space allocated for result

            if (fabs(res-SC_EIG) < EPS) { //if <u,u> = r (within EPS precision)
                verts[cache_size] = vecs[i]; //it's a vertex, store vecs[i] in verts
                cache_size+=1; //search for the next vector
            }
        }
    }
    //we now have matched arrays vert[] and vec_prod[], with valid vertices in the former and <u,_> in the latter
    if (cache_size < EXPECTED_CLIQUE) { //if there aren't enough elements to form target clique, no point in continuing, skip to next star complement candidate
        skipped++;
        return;
    }

    assert(cache_size < MAX_ORDER); //check if order is less than maximum allowed order
	//prepare to output compatibility graph
    char *p = gcode; //point *p at start of gcode
    encodegraphsize(cache_size, &p);

    int k = 6, x = 0; //initialize bit written counter and bit to be written

    for (i = 1; i < cache_size; i++) { //loop through all vertices u_i
        for (j = 0; j < i; j++) { //loop through <u_i,u_1> up to <u_i,u_i>, since <u,v> = <v,u> we don't have to check the rest as they will be checked in a later iteration
            x <<= 1; //ready next bit
            gsl_blas_ddot(vecs_prod[i], verts[j], &res); //calculate <u_i,u_j> and store it in space allocated for result

            /* We have an edge */
            if (fabs(res) <= EPS || fabs(res+1) <= EPS) { //if <u_i,u_j> = 0 or -1 (it's an edge)
                x |= 1; //set last bit of x to 1
            }
            if (--k == 0) { //decrease bit written counter
                *p++ = BIAS6 + x; //if we wrote 6 bits, add padding for ASCII character and write byte to where *p is pointing and advance pointer *p forward by 1 character
                k = 6; //reset counters
                x = 0;
            }

        }
    }
	//all full bytes written to gcode, prepare to write to file
    if (k != 6) { //if last byte is incomplete (still waiting to be written)
        *p++ = BIAS6 + (x << k); //push last bits to front, add padding, write to *p, advance *p
    }
    *p++ = '\n'; //stitch on newline, advance *p
    *p = '\0'; //stitch on null (to end string)

    fputs(gcode, outFile); //write to output file
}

static void init_vectors(void) {

    unsigned i,j;

    vec_j = gsl_vector_alloc(N); //allocate memory for vec_j with size N

    assert(vec_j); //check if vec_j was created successfully

    gsl_vector_set_all(vec_j, 1); //initialize vec_j as all 1's

    for (i = 0; i < 1<<N; i++) { //loop through all 2^N vectors in vecs[] and vecs_prod[]
        vecs[i] = gsl_vector_calloc(N); //initialize each vecs[i] to all 0
        vecs_prod[i] = gsl_vector_alloc(N); //allocate memory for vectors of size N

        assert(vecs[i] && vecs_prod[i]); //check if vectors were initialized successfully

        /* We fill the i'th vector of vecs */
        for (j = 0; j < N; j++) { //loop through length of vecs[i]
            if ( i & (1<<j) ) { //set vecs[i] to binary representation of i
                gsl_vector_set(vecs[i], j, 1);
            }
        }
    }
}

int main(int argc, char **argv) { //start here

    static FILE *infile; //declare pointer to input file

    assert (argc > 1); //check if arguments were provided (you can provide multiple, only the first will be used though)

    infile = fopen(argv[1], "r"); //open first argument in read mode and point *infile at it
    mat = gsl_matrix_alloc(N, N); //allocate memory for mat as NxN matrix (not initialized yet, contents are garbage)
    mat_inv = gsl_matrix_alloc(N, N); //allocate memory for mat_inv as NxN matrix
    perm = gsl_permutation_calloc(N); //initialize perm as identity permutation of length N

    char buf[512]; //declare string of 511+1 characters

    snprintf(buf, sizeof(buf), "%s.out", argv[1]); //fill buf with name of input file, append .out to the name
    outFile = fopen(buf, "w"); //create file with name contained in buf in write mode, point *outFile at it

    assert (outFile && infile && mat && mat_inv && perm); //check if everything exists

    init_vectors();

    while (1) {
        if (!fgets(line, sizeof(line), infile))
            break;
        stringtomat(line);
        constructGraph();
        nproc++;
	}

    printf("Successfuly processed: %u graphs. Skipped: %u\n" , nproc, skipped);

    return 0;
}


/*

Excited States software: KGS
Contributors: See CONTRIBUTORS.txt
Contact: kgs-contact@simtk.org

Copyright (C) 2009-2017 Stanford University

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

This entire text, including the above copyright notice and this permission notice
shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
IN THE SOFTWARE.

*/

#ifndef JACOBIANRELATED_H
#define JACOBIANRELATED_H

#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_multimin.h>

///These values have to be chosen according to the numerical analysis
static const double SINGVAL_TOL = 0.000000000001; // only generic 10^-12
//static const double SINGVAL_TOL = 0.000001; // include non-generic 10^-10
static const double RIGID_TOL = 0.0000000001; //depends on molecule, but 10^-10 seems a good fit!
//static const double RIGID_TOL = 0.000001; //depends on molecule ToDo: Adapt to obtain on-the-fly parameter selection

class NullSpaceRet
{
public:
//	static gsl_matrix* Ut;//The Ut from the SVD
	static gsl_matrix* V;//The V from the SVD
	static gsl_vector* singularValues;
	int nullspaceSize;
	int m, n;

//	gsl_matrix* P;//Null-space projection matrix
	gsl_matrix* m_nullspaceBasis; //basis of the nullspace
//	gsl_matrix* Jinv;//Jacobian pseudo-inverse

    gsl_vector* m_rigidAngles;
    gsl_vector* m_rigidHBonds;

    int m_numCoordinated;
    int m_numRigid;
    int m_numRigidHBonds;

	NullSpaceRet (int input_m, int input_n);
	~NullSpaceRet ();
	void print();

	// A function to make the projection operation more efficient;
	void NullSpacePostCompute();

	/// A function that identifies rigidified dihedral angles
	void RigidityAnalysis(gsl_matrix* HBondJacobian);

	//A function to project a vector on the nullspace
	void ProjectOnNullSpace (gsl_vector *to_project, gsl_vector *after_project);

	// To manipulate owner configations of NullSpaceRet instances
	void numOwners(unsigned int numOwners);
	unsigned int numOwners() const;
private:
	unsigned int numOwners_;
};

void nr_svd(gsl_matrix* a_g, int m, int n, double w[], gsl_matrix* v_out);
void ComputeNullSpace(gsl_matrix* Jac, NullSpaceRet* Ret);
//void ProjectOnNullSpace (NullSpaceRet* nullspace, gsl_vector *e, gsl_vector *to_project, gsl_vector *after_project);
void TorsionUpdate(NullSpaceRet* nullspace, gsl_vector *e, double lambda, gsl_vector *dTheta);

#endif


/*
Project 2: A Simple Ocean Model
This code models a set of equations for a 2-layer model of Ekman heat transport defined by F. Codron 2012: Ekman Heat Transport for Slab Oceans
The model is designed to take input from the Met Office Unified Model (UM) in .dat file form. As a result, this model uses the same spatial
horizontal resolution as the UM (90 latitude grid points and 144 longitude points, using spacing of 2.0 and 2.5 degrees respectively). The code
solves the heat transport using a finite difference method, and a sparse linear algebra GSL solver.
Model can be ran on an Ubuntu terminal using the following set of commands:
	gcc -Wall -I/GSL_PATH/gsl/include -c pr2_ocean_transport_v7.c
	gcc -L/GSL_PATH/gsl/lib pr2_ocean_transport_v7.o -lgsl -lgslcblas -lm
	./a.out [-v {NO_TRANSPORT, DIFUSION_ONLY, DIFUSION_EKMAN}]
where the options for -v are integers defined in enum Transport.
Created by Jake K. Eager
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <limits.h>
#include <unistd.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_spmatrix.h>
#include <gsl/gsl_splinalg.h>

#define C_D 0.0013 // drag coefficient, dimensionless
#define C_V 4000.0 // specific heat capacity of water J/kg/K
#define EMISSIVITY 0.985 // effective blackbody emissivity of ocean water
#define HEAT_TRANSFER_COEFFICIENT 3000. /* heat transfer coefficient for water W/m2/K 50-3000*/
#define RHO_AIR 1.22 // air density kg/m3
#define RHO_WATER 1027 // km/m3
#define R_PLANET (1.12384*6.371e6) // radius of planet in m (as fraction of Earth radius)
#define EPSILON (0.00001) // 0.00001 pretty good need to find appropriate value
#define OMEGA (6.501e-6) // angular frequency planet, about its central axis
// #define OMEGA (7.272e-5) // Earth value for omega
#define SIGMA (5.67e-8) // stefan boltzman constant
#define H_S 50.0 // thickness of surface layer in m
#define H_D 150.0 // thickness of deep layer below
#define D (25000.0) // horizontal diffusion coefficient m2/s
// #define D (25000.0*50.0/(H_S+H_D)) // horizontal diffusion coefficient m2/s MIGHT NEED THIS FIX, UNSURE...
#define N_LATS 90 /* number of latitude points */
#define N_LONS 144 /* number of longitude points */
#define LAT_MIN -89.0
#define DELTA_LAT 2.0
#define LON_MIN 1.25
#define DELTA_LON 2.5
#define HOURS_PER_DAY 24.
#define MINUTES_PER_HOUR 60.
#define SECONDS_PER_MINUTE 60.
#define DELTA_T (12.0*MINUTES_PER_HOUR*SECONDS_PER_MINUTE) // time step of 1 hours
#define TIME_STEPS (20000*2) // for time step of 12 hours, 20,000 time steps needed for 10,000 day run
#define TIME_OUTPUT_FREQ (25.0*HOURS_PER_DAY*MINUTES_PER_HOUR*SECONDS_PER_MINUTE) // print update frequency in seconds, but first term is number of days
#define DATA_OUTPUT_FREQ (100.0*HOURS_PER_DAY*MINUTES_PER_HOUR*SECONDS_PER_MINUTE) // output frequency in seconds, but first term is number of days
#define TIME_OUTPUT_TOL (1e-6) // days
#define T_OFFSET 0.0 // IF you want to restart run, specify start time here, so it doesn't save over stuff, and change input temp files to output ones
#define TOL (1e-6)
#define MAX_ITER 100
// input data files
// #define LAND_MASK_DATA "input_data/land_mask_day_cont_100_180.dat"
#define LAND_MASK_DATA "input_data/land_mask_no_land.dat"
#define MAX_FILE_LINE_SIZE 4000
#define SW_FLUX_NET_DOWN_DATA "input_data/ProCb/surface_net_downward_shortwave_flux.dat"
#define LW_FLUX_DOWN_DATA "input_data/ProCb/surface_downwelling_longwave_flux_in_air.dat"
// #define LW_FLUX_DOWN_DATA "input_data/ProCb/surface_net_downward_longwave_flux.dat"
#define LATENT_UP_FLUX_DATA "input_data/ProCb/surface_upward_latent_heat_flux.dat"
#define SENSIBLE_UP_FLUX_DATA "input_data/ProCb/surface_upward_sensible_heat_flux.dat"
// #define INITIAL_SURFACE_TEMP_DATA "output_data/ProCb/T_surf_1000_days.dat"
// #define INITIAL_DEEP_TEMP_DATA "output_data/ProCb/T_deep_1000_days.dat"
#define INITIAL_SURFACE_TEMP_DATA "input_data/ProCb/surface_temperature.dat"
#define INITIAL_DEEP_TEMP_DATA "input_data/ProCb/surface_temperature.dat"
#define U_WIND_DATA "input_data/ProCb/x_wind.dat"
#define V_WIND_DATA "input_data/ProCb/y_wind.dat"
#define X_STRESS_DATA "input_data/ProCb/surface_downward_eastward_stress.dat"
#define Y_STRESS_DATA "input_data/ProCb/surface_downward_northward_stress.dat"
#define LATS_FILE "input_data/lats.dat"
#define LONS_FILE "input_data/lons.dat"
// output data files
#define MAX_FNAME_CHAR 100
// #define PYTHON_EXE "python"
// #define PYTHON_SCRIPT 		"plotting_scripts/temperature_contour_plot.py"
#define OUTPUT_UPWARD_Q_FLUX "output_data/ProCb/upward_surface_Q_flux_"
#define OUTPUT_SURFACE_TEMP_DATA "output_data/ProCb/T_surf_"
#define OUTPUT_DEEP_TEMP_DATA "output_data/ProCb/T_deep_"
#define OUTPUT_X_MASS_FLUX_DATA "output_data/ProCb/eastward_flow_"
#define OUTPUT_Y_MASS_FLUX_DATA "output_data/ProCb/northward_flow_"
#define OUTPUT_VERITCAL_FLUX_DATA "output_data/ProCb/vertical_mass_flux_"
#define OUTPUT_DATA_EXT "_days.dat"

typedef enum coords { THETA, PHI, N_COORDS } Coords;
typedef enum depths { SURFACE, DEEP, N_DEPTHS } Depths;
typedef enum transport {NO_TRANSPORT, DIFUSION_ONLY, DIFUSION_EKMAN} Transport;
typedef enum errors { MALLOC_ERR, FREE_ERR, FOPEN_ERR, DEPTH_ERR, DERIVATIVE_ERR, READ_DATA_ERR, SOLVER_ERR, VERSION_ERR } Errors;

typedef struct grid_vals { // values of latitude and longitude
	double lat[N_LATS];
	double lon[N_LONS];
} Grid_vals;

static int select_version(int argc, char **argv);
static void time_stepper(int n_times, int n_lats, int n_lons, int version);

int main(int argc, char **argv)
{
	int version = select_version(argc, argv);
	time_stepper(TIME_STEPS,N_LATS,N_LONS,version);

	return 0;
}

/* Selects the whether to solve using no horizontal transport, diffusion only or diffusion and Ekman transport. */
static int select_version(int argc, char **argv)
{
	extern char *optarg;
	int c, err = 0;
	// char *case_number=NULL;
	int version=DIFUSION_EKMAN; // default case is the step function test case
	static char usage[] = "usage: %s [-v version_number] [name ...]\n";

	while ((c = getopt(argc, argv, "v:")) != -1)
		switch (c) {
		case 'v':
			version=atoi(optarg);
			if(version!=NO_TRANSPORT && version!=DIFUSION_EKMAN && version!=DIFUSION_ONLY){
				fprintf(stderr, "Version selected not valid, please enter [-v {%i, %i, %i}]\n", NO_TRANSPORT, DIFUSION_ONLY, DIFUSION_EKMAN);
				exit(VERSION_ERR);
			}
			break;
		case '?':
			err = 1;
			break;
		}
	if (err) {
		fprintf(stderr, usage, argv[0]);
		exit(1);
	}
	return version;
}

/* Checks memory allocation has been successful */
static void *xmalloc(size_t n)
{
    void *p = malloc(n);
    if(p == NULL)   {
        fprintf(stderr,"Out of memory.\n");
        exit(MALLOC_ERR);
    }
    return p;
}

/* checks pointer being freed is not NULL */
static void xfree(void *p)
{
	if(p == NULL)   {
        fprintf(stderr,"Pointer is NULL, so cannot free.\n");
        exit(FREE_ERR);
    }
	free(p);
}

/* allocates a 2d array of pointers */
static int **create_2d_int_pointer(int n, int m)
{
	int **p = (int **)xmalloc(n*sizeof(int*));
	for(int i=0;i<n;i++){
		p[i]=(int *)xmalloc(m*sizeof(int));
	}
	return p;
}

/* frees a 2d array of pointers */
static void destroy_2d_int_pointer(int **p, int n)
{
	for(int i=0;i<n;i++){
		xfree(p[i]);
	}
	xfree(p);
}

/* allocates a 2d array of pointers */
static double **create_2d_pointer(int n, int m)
{
	double **p = (double **)xmalloc(n*sizeof(double*));
	for(int i=0;i<n;i++){
		p[i]=(double *)xmalloc(m*sizeof(double));
	}
	return p;
}

/* frees a 2d array of pointers */
static void destroy_2d_pointer(double **p, int n)
{
	for(int i=0;i<n;i++){
		xfree(p[i]);
	}
	xfree(p);
}

/* allocates a 3d array of pointers */
static double ***create_3d_pointer(int n, int m, int l)
{
	double ***p = (double ***)xmalloc(n*sizeof(double**));
	for(int i=0;i<n;i++){
		p[i]=(double **)xmalloc(m*sizeof(double*));
		for(int j=0;j<m;j++){
			p[i][j]=(double *)xmalloc(l*sizeof(double));
		}
	}
	return p;
}

/* frees a 3d array of pointers */
static void destroy_3d_pointer(double ***p, int n, int m)
{
	for(int i=0;i<n;i++){
		for(int j=0;j<m;j++){
			xfree(p[i][j]);
		}
		xfree(p[i]);
	}
	xfree(p);
}

/* checks file opened correctly */
static void *xfopen(char *filename, char *command)
{
	FILE *outfile;
	outfile = fopen(filename,command);
	if(outfile==NULL){
		fprintf(stderr, "%s could not be opened.\n", filename);
		exit(FOPEN_ERR);
	}
	return outfile;
}

static void construct_fname(char* fname, size_t fname_length, const char* base, const char* time, const char* ext)
{ // construct filename from base, time of run and appropriate extention
    strcpy(fname, base);
    strcat(fname,time);
    strcat(fname,ext);
}

/* reads in a .dat files called filename, that contains data with n_lats rows and n_lons collumns, ie UM output data */
static void read_input_data(double **var, char *filename, int n_lats, int n_lons)
{
	char line[MAX_FILE_LINE_SIZE];
	char *data = line;
	FILE *input;
	input = xfopen(filename, "r" );
	int offset;
	for(int i=0;i<n_lats;i++){
		fgets(line, MAX_FILE_LINE_SIZE, input);
		for(int j=0;j<n_lons;j++){
			int found = sscanf(data, " %le%n", &var[i][j], &offset);
			if(found!=1){
				fprintf(stderr, "%s does not contain enough longitudinal points.\n", filename);
				exit(READ_DATA_ERR);
			}
			data+=offset;
		}
		data=line;
	}
  	fclose(input);
}

/* reads in a .dat files called filename, that contains integer data with n_lats rows and n_lons collumns, ie UM output data */
static void read_input_int_data(int **var, char *filename, int n_lats, int n_lons)
{
	char line[MAX_FILE_LINE_SIZE];
	char *data = line;
	FILE *input;
	input = xfopen(filename, "r" );
	int offset;
	for(int i=0;i<n_lats;i++){
		fgets(line, MAX_FILE_LINE_SIZE, input);
		for(int j=0;j<n_lons;j++){
			int found = sscanf(data, " %d%n", &var[i][j], &offset);
			if(found!=1){
				fprintf(stderr, "%s does not contain enough longitudinal points.\n", filename);
				exit(READ_DATA_ERR);
			}
			data+=offset;
		}
		data=line;
	}
  	fclose(input);
}

/* writes out put data to filename on a grid of n_lats by n_lons. */
static void write_data(double **var, char *fbasename, char *time_str, Grid_vals *grid, int n_lats, int n_lons)
{
	char *fname = (char*) xmalloc((MAX_FNAME_CHAR)*sizeof(char));
	construct_fname(fname, MAX_FNAME_CHAR, fbasename, time_str, OUTPUT_DATA_EXT);
	FILE *output;
	output = xfopen(fname, "w" );
	for(int i=0;i<n_lats;i++){
		for(int j=0;j<n_lons;j++){
			fprintf(output, "%lg\t", var[i][j]);
		}
		fprintf(output, "\n");
	}
	fclose(output);
	xfree(fname);
}

static void calc_Q_flux(double **upward_Q_flux, double **T_surf, Grid_vals *grid, int n_lats, int n_lons)
{
	double **F_net_sw_down = create_2d_pointer(n_lats,n_lons);
	double **F_lw_down = create_2d_pointer(n_lats,n_lons);
	double **F_latent_up = create_2d_pointer(n_lats,n_lons);
	double **F_sensible_up = create_2d_pointer(n_lats,n_lons);
	read_input_data(F_net_sw_down,SW_FLUX_NET_DOWN_DATA,n_lats,n_lons);
	read_input_data(F_lw_down,LW_FLUX_DOWN_DATA,n_lats,n_lons);
	read_input_data(F_latent_up,LATENT_UP_FLUX_DATA,n_lats,n_lons);
	read_input_data(F_sensible_up,SENSIBLE_UP_FLUX_DATA,n_lats,n_lons);
	for(int i=0;i<n_lats;i++){
		for(int j=0;j<n_lons;j++){
			upward_Q_flux[i][j] = F_latent_up[i][j] + F_sensible_up[i][j] + EMISSIVITY*SIGMA*pow(T_surf[i][j],4) - (F_net_sw_down[i][j] + F_lw_down[i][j]);
		}
	}
	destroy_2d_pointer(F_net_sw_down,n_lats);
	destroy_2d_pointer(F_lw_down,n_lats);
	destroy_2d_pointer(F_latent_up,n_lats);
	destroy_2d_pointer(F_sensible_up,n_lats);
}

/* convert angle in degrees to radians for evaluating trigonometric functions in calculus */
static double deg_to_rad(double theta)
{
	return theta*M_PI/180.0;
}

/* calculates the first spatial derivative of A with respect to a (which may be theta or phi), on a spherical surface. Used mainly to
calculate derivative of mass transport components M_theta and M_phi. Boundary conditions are used for a spherical geometry, using A[i][0]
to be adjacent to A[i][n_lons-1], and also A[0][j] is adjacent to the opposite point in longitude (j+-n_lats/2) at the same latitude. */
static double dA_da(double **A, int i, int j, int a, int n_lats, int n_lons)
{
	double d_theta=deg_to_rad(DELTA_LAT);
	double d_phi=deg_to_rad(DELTA_LON);
	if(a==THETA){
		if(i==n_lats-1){
			int j_temp = j+n_lons/2;
			if(j_temp>=n_lons){
				j_temp -=n_lons;
			}
			return (A[i][j_temp]-A[i-1][j])/(2*d_theta);
		}
		else if(i==0){
			int j_temp = j+n_lons/2;
			if(j_temp>=n_lons){
				j_temp -=n_lons;
			}
			return (A[i+1][j]-A[i][j_temp])/(2*d_theta);
		}
		else{
			return (A[i+1][j]-A[i-1][j])/(2*d_theta);
		}
	}
	else if(a==PHI){
		if(j==n_lons-1){
			return (A[i][0]-A[i][j-1])/(2*d_phi);
		}
		else if(j==0){
			return (A[i][j+1]-A[i][n_lons-1])/(2*d_phi);
		}
		else{
			return (A[i][j+1]-A[i][j-1])/(2*d_phi);
		}
	}
	else{
		fprintf(stderr, "Specified derivative is not respect to either theta or phi\n");
		exit(DERIVATIVE_ERR);
	}
}

/* Returns the horizontal divergence of quanitity M. determines whether there is an up or downward flow between the ocean layers */
static double calculate_div_M(double ***M, Grid_vals *grid, int i, int j, int n_lats, int n_lons)
{
	double theta=deg_to_rad(grid->lat[i]);
	double dM_theta_dtheta = dA_da(M[THETA],i,j,THETA,n_lats,n_lons);
	double dM_phi_dphi = dA_da(M[PHI],i,j,PHI,n_lats,n_lons);
	double div_M=1./R_PLANET*(dM_theta_dtheta-tan(theta)*M[THETA][i][j]+1./cos(theta)*dM_phi_dphi);
	return div_M;
}

static void calc_flow_Sv(double ***Sv_flow, double ***M, Grid_vals *grid, int n_lats, int n_lons)
{
	for(int i=0;i<n_lats;i++){
		double theta=deg_to_rad(grid->lat[i]);
		for(int j=0;j<n_lons;j++){
			Sv_flow[PHI][i][j] = M[PHI][i][j] * R_PLANET * DELTA_LAT * M_PI / 180. / RHO_WATER / 1.0e6;
			Sv_flow[THETA][i][j] = M[THETA][i][j] * R_PLANET * cos(theta) * deg_to_rad(DELTA_LON) / RHO_WATER / 1.0e6;
		}
	}
}

static void process_output_data(double ***T, double ***M, Grid_vals *grid, int n_lats, int n_lons, double time)
{ // processes output data
    double days = T_OFFSET+time/(HOURS_PER_DAY*MINUTES_PER_HOUR*SECONDS_PER_MINUTE);
    days=round(days);
    printf("Data outputted at %lg\n", days);
    int day_int = days+TIME_OUTPUT_TOL;
    char time_str[10];
    sprintf(time_str, "%d", day_int);
    // char *fname = (char*) xmalloc((MAX_FNAME_CHAR)*sizeof(char));
    double **upward_Q_flux= create_2d_pointer(n_lats,n_lons);
    calc_Q_flux(upward_Q_flux,T[SURFACE],grid,n_lats,n_lons);
    // construct_fname(fname, MAX_FNAME_CHAR, OUTPUT_UPWARD_LW_SURFACE_FLUX, time_str, OUTPUT_DATA_EXT);
    write_data(upward_Q_flux, OUTPUT_UPWARD_Q_FLUX, time_str, grid, n_lats, n_lons);
    destroy_2d_pointer(upward_Q_flux,n_lats);
    double ***Sv_flow = create_3d_pointer(N_COORDS,N_LATS,N_LONS);
    calc_flow_Sv(Sv_flow,M,grid,n_lats,n_lons);
    write_data(Sv_flow[PHI], OUTPUT_X_MASS_FLUX_DATA, time_str, grid, n_lats, n_lons);
    write_data(Sv_flow[THETA], OUTPUT_Y_MASS_FLUX_DATA, time_str, grid, n_lats, n_lons);
 	destroy_3d_pointer(Sv_flow,N_COORDS,N_LATS);
    double **div_M= create_2d_pointer(n_lats,n_lons);
    for(int i=0;i<n_lats;i++){
		for(int j=0;j<n_lons;j++){
			div_M[i][j] = calculate_div_M(M,grid,i,j,n_lats,n_lons);
		}
	}
	write_data(div_M, OUTPUT_VERITCAL_FLUX_DATA, time_str, grid, n_lats, n_lons);
	destroy_2d_pointer(div_M,n_lats);
    // construct_fname(fname, MAX_FNAME_CHAR, OUTPUT_SURFACE_TEMP_DATA, time_str, OUTPUT_DATA_EXT);
    write_data(T[SURFACE], OUTPUT_SURFACE_TEMP_DATA, time_str, grid, n_lats, n_lons);
    // construct_fname(fname, MAX_FNAME_CHAR, OUTPUT_DEEP_TEMP_DATA, time_str, OUTPUT_DATA_EXT);
    write_data(T[DEEP], OUTPUT_DEEP_TEMP_DATA, time_str, grid, n_lats, n_lons);
    // xfree(fname);
}

/* returns height of layer given index h, from enum heights. */
static double height_of_slab(int h)
{
	if(h==SURFACE){
		return H_S;
	}
	else if(h==DEEP){
		return H_D;
	}
	else{
		fprintf(stderr, "Index for depth is out of range.\n");
		exit(DEPTH_ERR);
	}
}

/* Constructs a grid with latitude and longitude points on the sphere, starting at LAT_MIN and LON_MIN, increasing in increments of DELTA_LAT
and DELTA_LON */
static void construct_grid(Grid_vals *grid, int n_lats, int n_lons, char *lats_file, char *lons_file)
{
	FILE *lats_out;
	lats_out = xfopen(lats_file, "w" );
	for(int i=0;i<n_lats;i++){
		grid->lat[i]=LAT_MIN+(i*DELTA_LAT);
		fprintf(lats_out, "%i\t%lg\n", i, grid->lat[i]);
	}
	fclose(lats_out);

	FILE *lons_out;
	lons_out = xfopen(lons_file, "w" );
	for(int j=0;j<n_lons;j++){
		grid->lon[j]=LON_MIN+(j*DELTA_LON);
		fprintf(lons_out, "%i\t%lg\n", j, grid->lon[j]);
	}
	fclose(lons_out);
}

/* Calculates the surface stress on the ocean surface as a result of the atmospheric winds near the surface.
Reads in velocity data from .dat files */
static void calculate_surface_stress(double ***tau, Grid_vals *grid, int n_lats, int n_lons)
{
	double ***v = create_3d_pointer(N_COORDS,N_LATS,N_LONS);
	read_input_data(v[THETA],V_WIND_DATA,n_lats,n_lons);
	read_input_data(v[PHI],U_WIND_DATA,n_lats,n_lons);
	for(int coord=THETA;coord<N_COORDS;coord++){
		for(int i=0;i<n_lats;i++){
			for(int j=0;j<n_lons;j++){
				tau[coord][i][j]=C_D*RHO_AIR*v[coord][i][j]*fabs(v[coord][i][j]); // retains sign of velocity
			}
		}
	}
	destroy_3d_pointer(v,N_COORDS,N_LATS);
}

/* Returns the coriolis parameter at an latitude theta */
static double calculate_coriolis_parameter(double theta)
{ // f=2*Omega*sin(theta)
	return 2*OMEGA*sin(deg_to_rad(theta));
}

/* calculates the mass flux due to Ekman transport */
static void calculate_mass_flux(double ***M, int **land_mask, Grid_vals *grid, int n_lats, int n_lons)
{
	double ***tau = create_3d_pointer(N_COORDS,N_LATS,N_LONS);
	// read_input_data(tau[THETA],Y_STRESS_DATA,n_lats,n_lons);
	// read_input_data(tau[PHI],X_STRESS_DATA,n_lats,n_lons);
	calculate_surface_stress(tau, grid, n_lats, n_lons);
	for(int i=0;i<n_lats;i++){
		double f = calculate_coriolis_parameter(grid->lat[i]);
		for(int j=0;j<n_lons;j++){
			if(land_mask[i][j]==1){
				M[THETA][i][j]=0.;
				M[PHI][i][j]=0.;
			}
			else{
				M[THETA][i][j]=(EPSILON*tau[THETA][i][j]+f*tau[PHI][i][j])/(pow(EPSILON,2)+pow(f,2));
				M[PHI][i][j]=(EPSILON*tau[PHI][i][j]-f*tau[THETA][i][j])/(pow(EPSILON,2)+pow(f,2));
			}
		}
	}
	destroy_3d_pointer(tau,N_COORDS,N_LATS);
}

/* Plots the script called filename in python */
// static void plot_data(char *filename)
// {
// 	char command[PATH_MAX];
// 	snprintf(command, sizeof(command), "%s %s", PYTHON_EXE, filename);
// 	system( command );
// }

/* Calculates the index for the matrix used in the solver, before being compressed */
static int calculate_matrix_index(int height, int lat, int lon, int n_lats, int n_lons)
{
	return height*n_lats*n_lons+lat*n_lons+lon;
}

/* At maximum or minimum latitutde, the adjacent latitude point above or bellow, respectively, is located at the same latitude,
but at the new longitudinal point (from lon to new_lon) */
static int calculate_new_lon(int lon, int n_lons)
{
	int new_lon = lon+n_lons/2;
	if(new_lon>=n_lons){
		new_lon -= n_lons;
	}
	return new_lon;
}

/* calculates and sets the matrix elements for a 2 layer ocean with 144 longitude points and 90 latitude points.
version determines whether the matrix is calculated for a no horizontal transport model, diffusion only model or the full Ekman model. */
static void calculate_matrix(gsl_spmatrix *A, double ***M, int **land_mask, Grid_vals *grid, int n_lats, int n_lons, int version)
{
	double d_theta=deg_to_rad(DELTA_LAT);
	double d_phi=deg_to_rad(DELTA_LON);
	double s_dfsn = DELTA_T*D/pow(R_PLANET,2);

	for(int height=SURFACE; height<N_DEPTHS; height++){
		double thickness = height_of_slab(height);
		double s_ekmn = DELTA_T/(RHO_WATER*thickness*R_PLANET);
		for(int lat=0; lat<n_lats; lat++){
			double theta = deg_to_rad(grid->lat[lat]);
			for(int lon=0; lon<n_lons; lon++){
				double Aii=0.0;
				int i;
				i = calculate_matrix_index(height,lat,lon,n_lats,n_lons);
				/* central point - check if land */
				if(land_mask[lat][lon]==1){
					Aii=1.;
					gsl_spmatrix_set(A, i, i, Aii);
					continue;
				}
				double dM_theta_dtheta = dA_da(M[THETA],lat,lon,THETA,n_lats,n_lons);
				double dM_phi_dphi = dA_da(M[PHI],lat,lon,PHI,n_lats,n_lons);
				double M_theta = M[THETA][lat][lon];
				double M_phi = M[PHI][lat][lon];
				double div_M = calculate_div_M(M,grid,lat,lon,n_lats,n_lons);
				double Aij=0.0;
				int j;
				/* central point */
				Aii=1.;
				if(version != NO_TRANSPORT){
					// if(height==SURFACE){
					Aii+= 2.*s_dfsn*(1./pow(d_theta,2)+1./pow(cos(theta)*d_phi,2));
					// }
					if(version==DIFUSION_EKMAN){
						if(height==SURFACE){
							Aii+= s_ekmn*(dM_theta_dtheta-M_theta*tan(theta)+1./cos(theta)*dM_phi_dphi);
							if(div_M<0.0){
								Aii+=(-DELTA_T/(RHO_WATER*thickness)*div_M);
							}
						}
						else if(height==DEEP){
							Aii+= -s_ekmn*(dM_theta_dtheta-M_theta*tan(theta)+1./cos(theta)*dM_phi_dphi);
							if(div_M>0.0){
								Aii+=(+DELTA_T/(RHO_WATER*thickness)*div_M);
							}
						}
					}
				}
				gsl_spmatrix_set(A, i, i, Aii);
				/* Neigbouring points */
				int lon_j, lat_j;
				if(version != NO_TRANSPORT){
					/* neighbouring latitude - below */
					if(lat==0){
						lon_j = calculate_new_lon(lon,n_lons);
						lat_j = lat;
						j = calculate_matrix_index(height,lat_j,lon_j,n_lats,n_lons);
					}
					else{
						lon_j = lon;
						lat_j = lat-1;
						j = calculate_matrix_index(height,lat_j,lon_j,n_lats,n_lons);
					}
					Aij=-s_dfsn*(1./pow(d_theta,2)+tan(theta)/(2*d_theta));
					if(version==DIFUSION_EKMAN){
						if(height==SURFACE){
							Aij+= s_ekmn*(-M_theta/(2.*d_theta));
						}
						else{
							Aij+= -s_ekmn*(-M_theta/(2.*d_theta));
						}
					}
					if(land_mask[lat_j][lon_j]==1){
						Aii+=Aij;
						Aij=0.;
					}
					gsl_spmatrix_set(A, i, j, Aij);
					/* neighbouring latitude - above */
					if(lat==n_lats-1){
						lon_j = calculate_new_lon(lon,n_lons);
						lat_j = lat;
						j = calculate_matrix_index(height,lat_j,lon_j,n_lats,n_lons);
					}
					else{
						lon_j = lon;
						lat_j = lat+1;
						j = calculate_matrix_index(height,lat_j,lon_j,n_lats,n_lons);
					}
					Aij=-s_dfsn*(1./pow(d_theta,2)-tan(theta)/(2*d_theta));
					if(version==DIFUSION_EKMAN){
						if(height==SURFACE){
							Aij+= s_ekmn*(M_theta/(2.*d_theta));
						}
						else{
							Aij+= -s_ekmn*(M_theta/(2.*d_theta));
						}
					}
					if(land_mask[lat_j][lon_j]==1){
						Aii+=Aij;
						Aij=0.;
					}
					gsl_spmatrix_set(A, i, j, Aij);
					/* neighbouring longitude - left */
					if(lon==0){
						lon_j = n_lons-1;
						lat_j = lat;
						j = calculate_matrix_index(height,lat_j,lon_j,n_lats,n_lons);
					}
					else{
						lon_j = lon-1;
						lat_j = lat;
						j = calculate_matrix_index(height,lat_j,lon_j,n_lats,n_lons);
					}
					Aij=-s_dfsn*1./pow(cos(theta)*d_phi,2);
					if(version==DIFUSION_EKMAN){
						if(height==SURFACE){
							Aij+= s_ekmn*(-M_phi/(2.*d_phi*cos(theta)));
						}
						else{
							Aij+= -s_ekmn*(-M_phi/(2.*d_phi*cos(theta)));
						}
					}
					if(land_mask[lat_j][lon_j]==1){
						Aii+=Aij;
						Aij=0.;
					}
					gsl_spmatrix_set(A, i, j, Aij);
					/* neighbouring longitude - right */
					if(lon==n_lons-1){
						lon_j = 0;
						lat_j = lat;
						j = calculate_matrix_index(height,lat_j,lon_j,n_lats,n_lons);
					}
					else{
						lon_j = lon+1;
						lat_j = lat;
						j = calculate_matrix_index(height,lat_j,lon_j,n_lats,n_lons);
					}
					Aij=-s_dfsn*1./pow(cos(theta)*d_phi,2);
					if(version==DIFUSION_EKMAN){
						if(height==SURFACE){
							Aij+= s_ekmn*(M_phi/(2.*d_phi*cos(theta)));
						}
						else{
							Aij+= -s_ekmn*(M_phi/(2.*d_phi*cos(theta)));
						}
					}
					if(land_mask[lat_j][lon_j]==1){
						Aii+=Aij;
						Aij=0.;
					}
					gsl_spmatrix_set(A, i, j, Aij);
					gsl_spmatrix_set(A, i, i, Aii);
					/* interaction between deep and surface layer */
					if(version==DIFUSION_EKMAN){
						if(height==SURFACE && div_M>0.0){
							Aij=(-DELTA_T/(RHO_WATER*thickness)*div_M);
							j = calculate_matrix_index(DEEP,lat,lon,n_lats,n_lons);
							gsl_spmatrix_set(A, i, j, Aij);
						}
						else if(height==DEEP && div_M<0.0){
							Aij=(+DELTA_T/(RHO_WATER*thickness)*div_M);
							j = calculate_matrix_index(SURFACE,lat,lon,n_lats,n_lons);
							gsl_spmatrix_set(A, i, j, Aij);
						}
					}
				}
			}
		}
	}
}

/* Reads in atmospheric fluxes, net stellar sw flux (subrtracts reflected sw radiation from the incident)
and incoming lw flux from atmospheric emission */
static void calculate_F_a(double **F_a, int n_lats, int n_lons)
{
	double **F_net_sw_down = create_2d_pointer(n_lats,n_lons);
	double **F_lw_down = create_2d_pointer(n_lats,n_lons);
	double **F_latent_up = create_2d_pointer(n_lats,n_lons);
	double **F_sensible_up = create_2d_pointer(n_lats,n_lons);
	read_input_data(F_net_sw_down,SW_FLUX_NET_DOWN_DATA,n_lats,n_lons);
	read_input_data(F_lw_down,LW_FLUX_DOWN_DATA,n_lats,n_lons);
	read_input_data(F_latent_up,LATENT_UP_FLUX_DATA,n_lats,n_lons);
	read_input_data(F_sensible_up,SENSIBLE_UP_FLUX_DATA,n_lats,n_lons);
	for(int i=0;i<n_lats;i++){
		for(int j=0;j<n_lons;j++){
			F_a[i][j] = (F_net_sw_down[i][j] + F_lw_down[i][j]) - (F_latent_up[i][j] + F_sensible_up[i][j]);
		}
	}
	destroy_2d_pointer(F_net_sw_down,n_lats);
	destroy_2d_pointer(F_lw_down,n_lats);
	destroy_2d_pointer(F_latent_up,n_lats);
	destroy_2d_pointer(F_sensible_up,n_lats);
}

/* Checks if convetion condition is met, then calculates the associated flux. */
static void calculate_F_c(double **F_c, double ***T, int n_lats, int n_lons)
{
	for(int i=0;i<n_lats;i++){
		for(int j=0;j<n_lons;j++){
			if(T[DEEP][i][j]>T[SURFACE][i][j]){ /* convection condition */
				F_c[i][j] = HEAT_TRANSFER_COEFFICIENT*(T[DEEP][i][j]-T[SURFACE][i][j]);
			}
			else{
				F_c[i][j] = 0.0;
			}
		}
	}
}

/* calculated and sets the vector b in matrix Ax=b. EMISSIVITY*SIGMA*T^4 is a blackbody emission from one of the layers,
to ensure an energy exchange between the layers. 1.0/(RHO_WATER*C_V*H_S) prefactor converts a flux in W/m2 to K/s */
static void calculate_vector_b(double ***T, int **land_mask, gsl_vector *b, int n_lats, int n_lons)
{ // calculates b from matrix equation Ax=b
	double **F_a = create_2d_pointer(n_lats,n_lons);
	calculate_F_a(F_a,n_lats,n_lons);
	double **F_c = create_2d_pointer(n_lats,n_lons);
	calculate_F_c(F_c,T,n_lats,n_lons);
	double b_hij = 0.0;
	for(int h=0;h<N_DEPTHS;h++){
		for(int i=0;i<n_lats;i++){
			for(int j=0;j<n_lons;j++){
				if(land_mask[i][j]==1){
					b_hij = T[h][i][j];
				}
				else{
					double depth = height_of_slab(h);
					if(h==SURFACE){
						b_hij = DELTA_T/(RHO_WATER*C_V*depth)*(F_c[i][j]+F_a[i][j]-EMISSIVITY*SIGMA*pow(T[SURFACE][i][j],4))+ T[h][i][j];
						// b_hij = DELTA_T/(RHO_WATER*C_V*depth)*(F_c[i][j]+F_a[i][j]) + T[h][i][j];
					}
					else{
						b_hij = DELTA_T/(RHO_WATER*C_V*depth)*(-F_c[i][j]) + T[h][i][j];
					}
				}
				gsl_vector_set(b, (h*n_lats*n_lons+(i*n_lons+j)), b_hij);
			}
		}
	}
	destroy_2d_pointer(F_a,n_lats);
	destroy_2d_pointer(F_c,n_lats);
}

/* Uses GSL sparse linear algebra to solve Ax=b */
static void calculate_new_T(double ***T, double ***M, int **land_mask, Grid_vals *grid, int n_lats, int n_lons, int version)
{
	int n = N_DEPTHS*n_lats*n_lons;
	gsl_spmatrix *A = gsl_spmatrix_alloc(n ,n);
	gsl_spmatrix *C;
	gsl_vector *b = gsl_vector_alloc(n);
	gsl_vector *x = gsl_vector_alloc(n);

	calculate_matrix(A,M,land_mask,grid,n_lats,n_lons,version);
	calculate_vector_b(T,land_mask,b,n_lats,n_lons);
	C = gsl_spmatrix_ccs(A);

	/* now solve the system with the GMRES iterative solver */
    const double tol = TOL;  /* solution relative tolerance */
    const size_t max_iter = MAX_ITER; /* maximum iterations */
    const gsl_splinalg_itersolve_type *S = gsl_splinalg_itersolve_gmres;
    gsl_splinalg_itersolve *work = gsl_splinalg_itersolve_alloc(S, n, 0);
    size_t iter = 0;
    double residual;
    int status;

    /* initial guess x = 0 */
    gsl_vector_set_zero(x);
 //    for(int h=0;h<N_DEPTHS;h++){
	// 	for(int i=0;i<n_lats;i++){
	// 		for(int j=0;j<n_lons;j++){
	// 		gsl_vector_set(x, (h*n_lats*n_lons+(i*n_lons+j)), T[h][i][j]);
	// 		}
	// 	}
	// }

    /* solve the system A x = b */
    do{
		status = gsl_splinalg_itersolve_iterate(C, b, tol, x, work);

		if (status != GSL_SUCCESS && iter==max_iter-1){ /* Solver exits if solution diverges */
			residual = gsl_splinalg_itersolve_normr(work);
			fprintf(stderr, "iter %zu residual = %.12e\n", iter, residual);
		  	fprintf(stderr, "Diverged\n");
		  	exit(SOLVER_ERR);
		}
	}
	while (status == GSL_CONTINUE && ++iter < max_iter);

    /* output solution */
    for(int h=0;h<N_DEPTHS;h++){ /* sets the new temperature from the gsl vector */
	    for(int i=0;i<n_lats;i++){
			for(int j=0;j<n_lons;j++){
				T[h][i][j]=gsl_vector_get(x,h*n_lats*n_lons+(i*n_lons+j));
			}
		}
	}

    gsl_splinalg_itersolve_free(work);

	gsl_spmatrix_free(A);
	gsl_spmatrix_free(C);
	gsl_vector_free(b);
	gsl_vector_free(x);
}

/* Function that steps through selected time. First creates the horizontal grid, then read in temperatures,
and steps through n_times time steps. Writes the final Temperature data to files, and plots a selected plotting script. */
static void time_stepper(int n_times, int n_lats, int n_lons, int version)
{
	Grid_vals *grid = xmalloc(sizeof(Grid_vals));
	construct_grid(grid,n_lats,n_lons,LATS_FILE,LONS_FILE);
	double ***T=create_3d_pointer(N_DEPTHS,N_LATS,N_LONS);
	read_input_data(T[SURFACE],INITIAL_SURFACE_TEMP_DATA,n_lats,n_lons);
	read_input_data(T[DEEP],INITIAL_DEEP_TEMP_DATA,n_lats,n_lons);
	int **land_mask = create_2d_int_pointer(n_lats,n_lons);
	read_input_int_data(land_mask,LAND_MASK_DATA,n_lats,n_lons);
	double ***M=create_3d_pointer(N_COORDS,N_LATS,N_LONS);
	double time = 0.;
	double days = 0.;
	for(int t=0;t<n_times;t++){
		time = ((double)t+1.)*DELTA_T;
		calculate_mass_flux(M,land_mask,grid,n_lats,n_lons);
		calculate_new_T(T,M,land_mask,grid,n_lats,n_lons,version);
		if (fmod(time,TIME_OUTPUT_FREQ)<TIME_OUTPUT_TOL){
			days = T_OFFSET+time/(HOURS_PER_DAY*MINUTES_PER_HOUR*SECONDS_PER_MINUTE);
			printf("Days passed = %lg\n", days);
		}
		if (fmod(time,DATA_OUTPUT_FREQ)<TIME_OUTPUT_TOL){
			process_output_data(T, M, grid, n_lats, n_lons, time);
		}
	}
	destroy_2d_int_pointer(land_mask,n_lats);
	process_output_data(T, M, grid, n_lats, n_lons, time);
	// plot_data(PYTHON_SCRIPT);
	destroy_3d_pointer(M,N_COORDS,N_LATS);
	destroy_3d_pointer(T,N_DEPTHS,N_LATS);
	xfree(grid);
}


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

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

#include "nrsrc/nrutil.h"

#include "prototypes.h"
#include "globvars.h"

double comp_DF_halo_exact(double E, double L);
double comp_DF_bulge_exact(double E, double L);

double comp_DF_H_iso(double E);
double comp_DF_H_ani(double Q);

double comp_DF_halo_Eddington(double Q);
double comp_DF_bulge_Eddington(double Q);

void comp_DF_init(void);
void compute_d2rhodpsi2_halo(void);
void compute_d2rhodpsi2_bulge(void);

double eddington_integrand(double t, void *params);

#define WORKSIZE 100000
gsl_integration_workspace *Workspace;

double *xi, *yi, *eddint;
double *list_radius;
double *list_E;                            /* tabulated energies */
double *DistFunc_halo, *DistFunc_bulge;    /* distribution function(s) */
double *psi_R;
double *rho_R_halo,*rho_R_bulge;
double *drhodpsi_halo,*d2rhodpsi2_halo;
double *drhodpsi_bulge,*d2rhodpsi2_bulge;

/* size of DF look-up table */
int DFSIZE= RSIZE * 2;

void compute_DF_lookuptable(void)
{
  FILE *fd;
  char dffile[100]="";
  int i;

  printf("Start computing distribution function.\n"); fflush(stdout);

#ifdef MAXWELLIAN
  AnisotropyRadius= -1;
  printf("MAXWELLIAN is turned on\n");
#else
  printf("MAXWELLIAN is turned off\n");
#endif

#ifdef DF_H_MODEL
  printf("DF_H_MODEL is turned on\n");
#else
  printf("DF_H_MODEL is turned off\n");
#endif

#ifdef DF_EDDINGTON
  printf("DF_EDDINGTON is turned on\n");
#else
  printf("DF_EDDINGTON is turned off\n");
#endif

  printf("AnisotropyRadius= %g\n", AnisotropyRadius);
  printf("ra= %g\n", AnisotropyRadius * RH);

  comp_DF_init();

  for(i = 0; i <= DFSIZE; i++)
    {
	printf("Computing DF table bin %d/%d (e= %g)\n", i, DFSIZE, -1.0*psi_R[i]); fflush(stdout);

	list_E[i]= -1.0 * psi_R[i];

	DistFunc_halo[i] = comp_DF_halo_exact(list_E[i], 0.0);
	DistFunc_bulge[i] = comp_DF_bulge_exact(list_E[i], 0.0);
    }
  printf("\n");

  /* write this to a file so we have a record of it */
  if(strstr(OutputFile,".hdf5"))
        strncpy(dffile, OutputFile, strlen(OutputFile)-5);
  if(strstr(OutputFile,".dat"))
        strncpy(dffile, OutputFile, strlen(OutputFile)-4);
  strcat(dffile, ".df");
  printf("DF look-up table saved in file: %s\n",dffile);
  if((fd = fopen(dffile,"w")))
    {
	fprintf(fd,"   Energy       f_halo(E)    f_bulge(E)\n");
	for(i= 0; i<= DFSIZE; i++)
	  fprintf(fd," %10.5e   %10.5e   %10.5e\n",list_E[i],DistFunc_halo[i],DistFunc_bulge[i]);
	fclose(fd);
    }

}

double comp_DF_halo_exact(double E, double L)
{
  double Q, f_e= 0.0;

  if((N_HALO == 0) || (M_HALO == 0))
    return 0;

  if (AnisotropyRadius > 0.0)
	Q= - E - (L*L/2./(AnisotropyRadius * RH)/(AnisotropyRadius * RH));
  else
	Q= - E;

#ifdef DF_H_MODEL
	/* calculate f_e exactly from the H90 analytic formula */
      if (AnisotropyRadius > 0.0)
        f_e= comp_DF_H_ani(Q);
      else
        f_e= comp_DF_H_iso(E);
#endif

#ifdef DF_EDDINGTON
	/* calculate f_e using Eddington's formula */
	f_e= comp_DF_halo_Eddington(Q);
	printf("f_e_edd(Q)/f_H_iso(E)= %g\n",f_e/comp_DF_H_iso(E));
	printf("f_e_edd(Q)/f_H_ani(Q)= %g\n",f_e/comp_DF_H_ani(Q)); fflush(stdout);
	printf("------\n");
#endif

  return f_e;
}

double comp_DF_bulge_exact(double E, double L)
{
  double Q, f_e= 0.0;

  if((N_BULGE == 0) || (M_BULGE == 0))
    return 0;

/* for the moment, we only assume isotropic DF and use Eddington's formula */

/*
  if (AnisotropyRadius > 0.0)
        Q= - E - (L*L/2./(AnisotropyRadius * RH)/(AnisotropyRadius * RH));
  else
*/
        Q= - E;

  f_e= comp_DF_bulge_Eddington(Q);

  return f_e;
}

/*
   interpolate the halo f_e from the pre-computed table
*/
double comp_DF_halo(double E, double L)
{
  double f_e= 0.0, Q, ee, ue;
  int ie;

  ee= E;

  if(AnisotropyRadius > 0.0)
    {
      Q= - E - (L*L/2./(AnisotropyRadius * RH)/(AnisotropyRadius * RH));

      /* in this case, table is stored in terms of -Q */
      ee= -1.0 * Q;
    }

  /* if e>0 or q<0 then dist. func. is zero */
  if(ee>0.0)
    return 0.0;

  ie= find_idx(ee, list_E, DFSIZE);

/*
  printf("\nE= %g  Q= %g ee= %g\n",E, Q, ee);
  printf("ie= %d  list_E[0]= %g   list_E[DFSIZE]= %g \n",ie, list_E[0], list_E[DFSIZE]);
  printf("list_E[%d]= %g\n", ie-1, list_E[ie-1]);
  printf("list_E[%d]= %g\n", ie, list_E[ie]);
  printf("list_E[%d]= %g\n", ie+1, list_E[ie+1]);
  fflush(stdout);
*/

  if(ie < 0 || ie > DFSIZE)
    ie= 0;

  ue = (ee - list_E[ie]) / (list_E[ie + 1] - list_E[ie]);

  if(ee>list_E[0])
    {
      /*
      printf("\nparticle has energy very close to zero\n");
      printf("ee=  %g, table min list_E[0]= %g\n",ee,list_E[0]);
      printf("return 0.\n");
      */
      ue= 0.0;
    }

  if(ie < 0 || ie >= DFSIZE || ue < 0 || ue > 1)
    {
      printf("\nee= %g   (in bounds? list_E[0]= %g  list_E[DFSIZE]= %g\n",ee, list_E[0], list_E[DFSIZE]);
      printf("in: comp_DF_fromtable\nerror: ie=%d out of range ue=%g\nstopping\n", ie, ue);
      exit(0);
    }

  f_e= DistFunc_halo[ie] * (1-ue) + DistFunc_halo[ie+1] * ue;

  return f_e;
}

/*
   interpolate the bulge f_e from the pre-computed table
*/
double comp_DF_bulge(double E, double L)
{
  double f_e= 0.0, Q, ee, ue;
  int ie;

  ee= E;

  /* for the moment, we'll assume the bulge is isotropic */
/*
  if(AnisotropyRadius > 0.0)
    {
      Q= - E - (L*L/2./(AnisotropyRadius * RH)/(AnisotropyRadius * RH));
      ee= -1.0 * Q;
    }
*/

  /* if e>0 or q<0 then dist. func. is zero */
  if(ee>0.0)
    return 0.0;

  ie= find_idx(ee, list_E, DFSIZE);

  if(ie < 0 || ie > DFSIZE)
    ie= 0;

  ue = (ee - list_E[ie]) / (list_E[ie + 1] - list_E[ie]);

  if(ee>list_E[0])
    {
      /*
      printf("\nparticle has energy very close to zero\n");
      printf("ee=  %g, table min list_E[0]= %g\n",ee,list_E[0]);
      printf("return 0.\n");
      */
      ue= 0.0;
    }

  if(ie < 0 || ie >= DFSIZE || ue < 0 || ue > 1)
    {
      printf("\nee= %g   (in bounds? list_E[0]= %g  list_E[DFSIZE]= %g\n",ee, list_E[0], list_E[DFSIZE]);
      printf("in: comp_DF_fromtable\nerror: ie=%d out of range ue=%g\nstopping\n", ie, ue);
      exit(0);
    }

  f_e= DistFunc_bulge[ie] * (1-ue) + DistFunc_bulge[ie+1] * ue;

  return f_e;
}

/* ----------------------------------------------------------------- */

/* the following are analytic forms of the distribution
   function, currently, it's only Hernquist 1990 */

/* Isotropic Hernquist Model */

double comp_DF_H_iso(double E)
{
  double vg, q, prefac, f_e;

  if (E >= 0.0)
	return 0.0;

  vg= sqrt(G * M_HALO / RH);
  q= sqrt(-1.0 * E * RH / G / M_HALO);

  if (q >= 1.0)
	return (1.0e+30);

  prefac= M_HALO / (8.0 * sqrt(2.0) * PI * PI * PI * RH * RH * RH * vg * vg * vg);

  f_e= pow((1-q*q), -5./2.) * ( 3.*asin(q) + q*pow((1.-q*q),0.5)*(1.-2.*q*q) * (8.*q*q*q*q - 8.*q*q - 3.));
/*
  printf("               E= %g  vg= %g  q= %g  prefac= %g f_e %g\n", E, vg, q, prefac, f_e); fflush(stdout);
*/

  return prefac * f_e;
}

/* Anisotropic Hernquist Model */

double comp_DF_H_ani(double Q)
{

  if (AnisotropyRadius <= 0.0)
	return 0;

  double vg, qbar, prefac, f_e, f_iso;
  double ra;

  ra= AnisotropyRadius * RH;

  if (Q <= 0.0)
	return 0;

  vg= sqrt(G * M_HALO / RH);
  qbar= sqrt(RH * Q / G / M_HALO);

  if (qbar >= 1.0)
	return 1.0e+30;

  prefac= M_HALO / (sqrt(2.0) * PI * PI * PI * RH * RH * RH * vg * vg * vg);

  f_e= (RH * RH / ra / ra) * qbar * (1. - 2.*qbar*qbar);

  f_iso= comp_DF_H_iso(-1.0*Q);

/*
printf("DF_ani  Q= %g   qbar= %g  f_iso= %g  prefac= %g  f_e= %g\n",Q,qbar,f_iso,prefac,f_e); fflush(stdout);
*/

  return f_iso + prefac * f_e;
}

/* -----------------------------------------------------------------

     Now, calculate DF using the Eddington formula.

   ----------------------------------------------------------------- */

gsl_spline * d2rhodpsi2_spline_halo;
gsl_interp_accel * d2rhodpsi2_spline_acc_halo;

gsl_spline * d2rhodpsi2_spline_bulge;
gsl_interp_accel * d2rhodpsi2_spline_acc_bulge;

// Parameter struct for the Eddington integration function.
typedef struct {
  const gsl_spline* s;
  gsl_interp_accel* a;
  double Q;
} spline_params;

/* The integrand for the Eddington integration. Now uses
   the GSL spline functionality rather than the NR, which is evil. */
double eddington_integrand(double t, void *params)
{
  spline_params* p = params;

//printf("t= %g\t", t);
//printf("Q= %g\t",p->Q);
//printf("sqrt(p->Q-t)= %g   \t", sqrt(p->Q-t));
//printf("d2rhodpsi2= %g  \t", gsl_spline_eval(p->s, t, p->a));
  double v2 = gsl_spline_eval(p->s, t, p->a) / sqrt(p->Q - t);
//printf("v2= %g  \n",v2); fflush(stdout);

  return v2;
}

double comp_DF_halo_Eddington(double Q)
{
  double f_e;
  int i;
  double result, abserr;

  /* first, we'll generate the spline for d2rho/dpsi2 term in integrand */

  if(Q<=0)
    return 0.0;

  /* load spline for d2rho/dpsi2 term into function parameters */
  spline_params p;
  p.s = d2rhodpsi2_spline_halo;
  p.a = d2rhodpsi2_spline_acc_halo;
  p.Q = Q;

  /* error checking: write entire DF integrand and the spline */
/*
  FILE *fd;
  char tempfile[100]="";
  sprintf(tempfile, "df_integrands/%g.txt", Q);
  if((fd = fopen(tempfile,"w")))
    {
        fprintf(fd,"Q= %g \n",Q);
        fprintf(fd,"       xi            yi          d2rho_dpsi2            Q           Q-psiR     sqrt(Q-psiR) \n");
        for(i = 0; i <= DFSIZE; i++)
          fprintf(fd," %8.5e     %8.5e     %8.5e    %8.5e    %8.5e      %8.5e     \n",xi[i],yi[i],d2rho_dpsi2[i], Q, Q-psi_R[i], sqrt(Q-psi_R[i]));
        fclose(fd);
    }
*/

  gsl_function F;
  F.function = &eddington_integrand;
  F.params = &p;

  double epsabs=0;
  double epsrel=1e-4;
  int intstatus;

  /* do our own error handling, so that we can adjust the accuracy if needed (i.e., at small r) */
  gsl_error_handler_t * old_handler = gsl_set_error_handler_off(); /* turn off the error-handler */

  do
  {
    intstatus= gsl_integration_qags(&F, 0.0, Q, epsabs, epsrel, WORKSIZE, Workspace, &result, &abserr);
    printf("epsrel = %g     result = %g +/- %g   (%g %)\n", epsrel, result, abserr, 100.0*abserr/result);
    epsrel *= 2.0;
  } while(intstatus == GSL_EROUND);

  //printf("Workspace.intervals = %d\n", Workspace->size);
  //printf("Workspace.maximum_level = %d\n", Workspace->maximum_level);

  gsl_set_error_handler(old_handler);          /*  turn it on again */

  f_e = result / sqrt(8.0) / PI / PI;

  if(f_e<0)
    f_e= 0.0;

  return f_e;
}

double comp_DF_bulge_Eddington(double Q)
{
  double f_e;
  int i;
  double result, abserr;

  if(Q<=0)
    return 0.0;

  spline_params p;
  p.s = d2rhodpsi2_spline_bulge;
  p.a = d2rhodpsi2_spline_acc_bulge;
  p.Q = Q;

  gsl_function F;
  F.function = &eddington_integrand;
  F.params = &p;

  double epsabs=0;
  double epsrel=1e-4;
  int intstatus;

  gsl_error_handler_t * old_handler = gsl_set_error_handler_off(); /* turn off the error-handler */

  do
  {
    intstatus= gsl_integration_qags(&F, 0.0, Q, epsabs, epsrel, WORKSIZE, Workspace, &result, &abserr);
    //printf("epsrel = %g     result = %g +/- %g   (%g %)\n", epsrel, result, abserr, 100.0*abserr/result);
    //printf("error = %g\n", abserr);
    epsrel *= 2.0;
  } while(intstatus == GSL_EROUND);

  gsl_set_error_handler(old_handler);          /*  turn it on again */

  f_e = result / sqrt(8.0) / PI / PI;

  if(f_e<0)
    f_e= 0.0;

  return f_e;
}

/*
   Computes d2_rho / dpsi2, which is the primary term
   in the integral when using the Eddington formulation
   to calculate the distribution function.

   We do this for the halo and bulge components separately.

*/
void compute_d2rhodpsi2_halo()
{
  FILE *fd;
  char drhodpsifile[100]="";
  int i;

  double R, RplusdR, RminusdR, dR;
  double rho_RplusdR, rho_RminusdR;
  double psi_RplusdR, psi_RminusdR;
  double slope1, slope2, ra;

  // normalize just in case we get really small densities
  double rho_norm, psi_norm;

  for(i = 0; i <= DFSIZE; i++)
    {
	if(i < RSIZE)
	    /* continues same scaling as list_R
            R= list_R[RSIZE] * exp((RSIZE-i) * (log(LL) - log(Baselen)) / (RSIZE - 1));   */
	    /* more modest scaling to zero energy (ie larger radii) */
            R= list_R[RSIZE] * pow(10., 5.0 * (RSIZE-i) / RSIZE);
            //R= list_R[RSIZE] * pow(10., 7.0 * (RSIZE-i) / RSIZE);
	else
            R= list_R[DFSIZE-i];

        if(i == DFSIZE)
            R= list_R[1] * 0.5;

	/* interval of integration */
	dR= R * exp(1.0e-2 * (log(LL) - log(Baselen)) / (RSIZE - 1)) - R;
	if(dR>1.5)
	  dR= 1.5;
	if((dR/R)<1.0e-4)
	  dR=1.0e-4*R;
        RplusdR= R + dR;
        RminusdR= R - dR;

	list_radius[i]= R;

        /* compute psi (=-phi, i.e., negative the potential) along the z axis, although
           the full DF assumes spherical symmetry */
        psi_R[i]= -1.0 * comp_phi(0, R);;
        psi_RplusdR= -1.0 * comp_phi(0, RplusdR);
        psi_RminusdR= -1.0 * comp_phi(0, RminusdR);

        /* -------------------------------
	    DF for halo component  */
        rho_R_halo[i]= comp_rho_halo(R,0);
        rho_RplusdR= comp_rho_halo(RplusdR,0);
        rho_RminusdR= comp_rho_halo(RminusdR,0);

	if(AnisotropyRadius>0)
	  {
	    ra= AnisotropyRadius * RH;
            rho_R_halo[i] *= (1.0 + (R * R / ra / ra ));
            rho_RplusdR *= (1.0 + (RplusdR * RplusdR / ra / ra));
            rho_RminusdR *= (1.0 + (RminusdR * RminusdR / ra / ra));
	  }

	// normalize
	rho_norm= rho_R_halo[i];
	rho_R_halo[i] /= rho_norm;
	rho_RplusdR /= rho_norm;
	rho_RminusdR /= rho_norm;
	psi_norm= psi_R[i];
	psi_R[i] /= psi_norm;
	psi_RplusdR /= psi_norm;
	psi_RminusdR /= psi_norm;

        slope1= (rho_RplusdR - rho_R_halo[i]) / (psi_RplusdR - psi_R[i]);
        slope2= (rho_R_halo[i] - rho_RminusdR) / (psi_R[i] - psi_RminusdR);

        drhodpsi_halo[i]= 0.5*(slope1+slope2);

        d2rhodpsi2_halo[i]= (rho_RplusdR+rho_RminusdR-2.*rho_R_halo[i])/(psi_RplusdR-psi_R[i])/(psi_R[i]-psi_RminusdR);

	// now put units back in
	psi_R[i] *= psi_norm;
	rho_R_halo[i] *= rho_norm;
	drhodpsi_halo[i] *= rho_norm / psi_norm;
	d2rhodpsi2_halo[i] *= rho_norm / psi_norm / psi_norm;

//printf("i= %d, R= %g,  second derivative denominators: %g, %g    diff= %g\n",i,R,(psi_RplusdR-psi_R[i]),(psi_R[i]-psi_RminusdR),(psi_RplusdR-psi_R[i])/(psi_R[i]-psi_RminusdR));
//printf("i= %d, R= %g,  rho= %g   drhodpsi_halo= %g    d2rhodpsi2_halo= %g     second derivative denominators: %g, %g    diff= %g\n",i,R,rho_R_halo[i],drhodpsi_halo[i],d2rhodpsi2_halo[i],(psi_RplusdR-psi_R[i]),(psi_R[i]-psi_RminusdR),(psi_RplusdR-psi_R[i])/(psi_R[i]-psi_RminusdR));
//printf("i= %d, R= %g,    rho_0= %g   ani_factor= %g  rho_R_halo= %g\n",i,R,comp_rho_halo(R,0),(1.0 + (R * R / ra / ra )),rho_R_halo[i]);

    }

  xi[0] = 0;
  yi[0] = 0;
  /* generate spline for the d2rho / dpsi2 term */
  for(i = 0; i <= DFSIZE; i++)
    {
	xi[i + 1] = psi_R[i];
	yi[i + 1] = d2rhodpsi2_halo[i];
	if( (xi[i+1] - xi[i]) < 0)  xi[i+1] =  xi[i] + 1e-5;

	printf("xi[i] = %f  yi[i] = %f delta_x = %f\n",xi[i + 1], yi[i + 1], xi[i+1] - xi[i]);
    }

  printf("check_abcdefg\n");

  gsl_spline_init(d2rhodpsi2_spline_halo, xi, yi, DFSIZE+1);

  printf("check_abcdefg done\n");

  /* write this to a file so we have a record of it */
  strcpy(drhodpsifile,"");
  if(strstr(OutputFile,".hdf5"))
        strncpy(drhodpsifile, OutputFile, strlen(OutputFile)-5);
  if(strstr(OutputFile,".dat"))
        strncpy(drhodpsifile, OutputFile, strlen(OutputFile)-4);
  strcat(drhodpsifile, ".drhodpsi_halo");
  if((fd = fopen(drhodpsifile,"w")))
    {
        fprintf(fd,"# drhodpsi file, halo component\n");
        fprintf(fd,"# n= %d \n",DFSIZE);
        fprintf(fd,"# \n");
        fprintf(fd,"#   R (kpc)         psi               rho         drho/dpsi      d2rho/dpsi2   spline(psi)   \n");
        fprintf(fd,"# \n");
        for(i = 0; i <= DFSIZE; i++)
          fprintf(fd," %8.5e     %8.5e      %8.5e     %8.5e     %8.5e    %8.5e  \n",list_radius[i],psi_R[i],rho_R_halo[i],drhodpsi_halo[i],d2rhodpsi2_halo[i],gsl_spline_eval(d2rhodpsi2_spline_halo, 0.999*psi_R[i], d2rhodpsi2_spline_acc_halo));

        fclose(fd);
    }

}

void compute_d2rhodpsi2_bulge()
{
  FILE *fd;
  char drhodpsifile[100]="";
  int i;

  double R, RplusdR, RminusdR, dR;
  double rho_RplusdR, rho_RminusdR;
  double psi_RplusdR, psi_RminusdR;
  double slope1, slope2, ra;

  // normalize just in case we get really small densities
  double rho_norm, psi_norm;

  if(M_BULGE == 0)
    return;

  for(i = 0; i <= DFSIZE; i++)
    {
	if(i < RSIZE)
	    /* continues same scaling as list_R
            R= list_R[RSIZE] * exp((RSIZE-i) * (log(LL) - log(Baselen)) / (RSIZE - 1));   */
	    /* more modest scaling to zero energy (ie larger radii) */
            R= list_R[RSIZE] * pow(10., 5.0 * (RSIZE-i) / RSIZE);
            //R= list_R[RSIZE] * pow(10., 7.0 * (RSIZE-i) / RSIZE);
	else
            R= list_R[DFSIZE-i];

        if(i == DFSIZE)
            R= list_R[1] * 0.5;

	/* interval of integration */
	dR= R * exp(1.0e-2 * (log(LL) - log(Baselen)) / (RSIZE - 1)) - R;
	if(dR>1.5)
	  dR= 1.5;
	if((dR/R)<1.0e-4)
	  dR=1.0e-4*R;
        RplusdR= R + dR;
        RminusdR= R - dR;

	list_radius[i]= R;

        /* compute psi (=-phi, i.e., negative the potential) along the z axis, although
           the full DF assumes spherical symmetry */
        psi_R[i]= -1.0 * comp_phi(0, R);;
        psi_RplusdR= -1.0 * comp_phi(0, RplusdR);
        psi_RminusdR= -1.0 * comp_phi(0, RminusdR);

        /* -------------------------------
	    DF for bulge component  */
        rho_R_bulge[i]= comp_rho_bulge(R,0);
        rho_RplusdR= comp_rho_bulge(RplusdR,0);
        rho_RminusdR= comp_rho_bulge(RminusdR,0);

	if(AnisotropyRadius>0)
	  {
	    ra= AnisotropyRadius * RH;
            rho_R_bulge[i] *= (1.0 + (R * R / ra / ra ));
            rho_RplusdR *= (1.0 + (RplusdR * RplusdR / ra / ra));
            rho_RminusdR *= (1.0 + (RminusdR * RminusdR / ra / ra));
	  }

	// normalize
	rho_norm= rho_R_bulge[i];
	rho_R_bulge[i] /= rho_norm;
	rho_RplusdR /= rho_norm;
	rho_RminusdR /= rho_norm;
	psi_norm= psi_R[i];
	psi_R[i] /= psi_norm;
	psi_RplusdR /= psi_norm;
	psi_RminusdR /= psi_norm;

        slope1= (rho_RplusdR - rho_R_bulge[i]) / (psi_RplusdR - psi_R[i]);
        slope2= (rho_R_bulge[i] - rho_RminusdR) / (psi_R[i] - psi_RminusdR);

        drhodpsi_bulge[i]= 0.5*(slope1+slope2);

        d2rhodpsi2_bulge[i]= (rho_RplusdR+rho_RminusdR-2.*rho_R_bulge[i])/(psi_RplusdR-psi_R[i])/(psi_R[i]-psi_RminusdR);

	// now put units back in
	psi_R[i] *= psi_norm;
	rho_R_bulge[i] *= rho_norm;
	drhodpsi_bulge[i] *= rho_norm / psi_norm;
	d2rhodpsi2_bulge[i] *= rho_norm / psi_norm / psi_norm;

//printf("i= %d, R= %g,  second derivative denominators: %g, %g    diff= %g\n",i,R,(psi_RplusdR-psi_R[i]),(psi_R[i]-psi_RminusdR),(psi_RplusdR-psi_R[i])/(psi_R[i]-psi_RminusdR));
//printf("i= %d, R= %g,  rho= %g   drhodpsi_bulge= %g    d2rhodpsi2_bulge= %g     second derivative denominators: %g, %g    diff= %g\n",i,R,rho_R_bulge[i],drhodpsi_bulge[i],d2rhodpsi2_bulge[i],(psi_RplusdR-psi_R[i]),(psi_R[i]-psi_RminusdR),(psi_RplusdR-psi_R[i])/(psi_R[i]-psi_RminusdR));
//printf("i= %d, R= %g,    rho_0= %g   ani_factor= %g  rho_R_bulge= %g\n",i,R,comp_rho_bulge(R,0),(1.0 + (R * R / ra / ra )),rho_R_bulge[i]);

    }

  /* generate spline for the d2rho / dpsi2 term */
  for(i = 0; i <= DFSIZE; i++)
    {
	xi[i + 1] = psi_R[i];
	yi[i + 1] = d2rhodpsi2_bulge[i];
        if( (xi[i+1] - xi[i]) < 0)  xi[i+1] =  xi[i] + 1e-5;

    }

  gsl_spline_init(d2rhodpsi2_spline_bulge, xi, yi, DFSIZE+1);

  /* write this to a file so we have a record of it */
  strcpy(drhodpsifile,"");
  if(strstr(OutputFile,".hdf5"))
        strncpy(drhodpsifile, OutputFile, strlen(OutputFile)-5);
  if(strstr(OutputFile,".dat"))
        strncpy(drhodpsifile, OutputFile, strlen(OutputFile)-4);
  strcat(drhodpsifile, ".drhodpsi_bulge");
  if((fd = fopen(drhodpsifile,"w")))
    {
        fprintf(fd,"# drhodpsi file, bulge component\n");
        fprintf(fd,"# n= %d \n",DFSIZE);
        fprintf(fd,"# \n");
        fprintf(fd,"#   R (kpc)         psi               rho         drho/dpsi      d2rho/dpsi2   spline(psi)   \n");
        fprintf(fd,"# \n");
        for(i = 0; i <= DFSIZE; i++)
          fprintf(fd," %8.5e     %8.5e      %8.5e     %8.5e     %8.5e    %8.5e  \n",list_radius[i],psi_R[i],rho_R_bulge[i],drhodpsi_bulge[i],d2rhodpsi2_bulge[i],gsl_spline_eval(d2rhodpsi2_spline_bulge, 0.999*psi_R[i], d2rhodpsi2_spline_acc_bulge));

        fclose(fd);
    }

}

void comp_DF_init(void)
{
  /* stores spline on the Eddington integrand */
  xi= vector(1, DFSIZE+1);
  yi= vector(1, DFSIZE+1);
  eddint= vector(1, DFSIZE+1);

  /* arrays for energies, radii, d2rho/dpsi2 term in Eddington integrand, and pre-computed DF  */
  list_radius= vector(0, DFSIZE);
  list_E= vector(0, DFSIZE);
  psi_R= vector(0, DFSIZE);

  /* halo specific */
  rho_R_halo= vector(0, DFSIZE);
  drhodpsi_halo= vector(0, DFSIZE);
  d2rhodpsi2_halo= vector(0, DFSIZE);
  DistFunc_halo= vector(0, DFSIZE);

  /* bulge specific */
  rho_R_bulge= vector(0, DFSIZE);
  drhodpsi_bulge= vector(0, DFSIZE);
  d2rhodpsi2_bulge= vector(0, DFSIZE);
  DistFunc_bulge= vector(0, DFSIZE);

  /* setup the d2rho_dpsi2 array */

  /* for the halo */
  printf("allocating for halo\n");
  d2rhodpsi2_spline_halo = gsl_spline_alloc(gsl_interp_cspline, DFSIZE+1);
  d2rhodpsi2_spline_acc_halo = gsl_interp_accel_alloc();
  printf("evaluating for halo\n");
  compute_d2rhodpsi2_halo();

  /* for the bulge */
  printf("allocating for bulge\n");
  d2rhodpsi2_spline_bulge = gsl_spline_alloc(gsl_interp_cspline, DFSIZE+1);
  d2rhodpsi2_spline_acc_bulge = gsl_interp_accel_alloc();
  compute_d2rhodpsi2_bulge();

  printf("some workplace alloc\n");
  Workspace = gsl_integration_workspace_alloc(WORKSIZE);
}

void comp_DF_Eddington_close(void)
{
  free_vector(xi, 1, DFSIZE+1);
  free_vector(yi, 1, DFSIZE+1);
  free_vector(eddint, 1, DFSIZE+1);
  free_vector(psi_R, 1, DFSIZE);
  free_vector(list_radius, 1, DFSIZE);
  free_vector(list_E, 1, DFSIZE);
  free_vector(rho_R_halo, 1, DFSIZE);
  free_vector(rho_R_bulge, 1, DFSIZE);
  free_vector(drhodpsi_halo, 1, DFSIZE);
  free_vector(drhodpsi_bulge, 1, DFSIZE);
  free_vector(d2rhodpsi2_halo, 1, DFSIZE);
  free_vector(d2rhodpsi2_bulge, 1, DFSIZE);
  free_vector(DistFunc_halo, 1, DFSIZE);
  free_vector(DistFunc_bulge, 1, DFSIZE);

  gsl_integration_workspace_free(Workspace);

  gsl_spline_free(d2rhodpsi2_spline_halo);
  printf("check_1234\n");
  gsl_interp_accel_free(d2rhodpsi2_spline_acc_halo);

  gsl_spline_free(d2rhodpsi2_spline_bulge);
  printf("check_abcd\n");
  gsl_interp_accel_free(d2rhodpsi2_spline_acc_bulge);
}

/* ----------------------------------------------------------------- */

double compute_ani_beta(double r)
{
  double beta;

  /* isotropic model */
  beta = 0.0;

  if (AnisotropyRadius > 0)
    {
        double ra;

        ra= AnisotropyRadius * RH;

        /* Osipkov & Merritt anisotropy dependency */
        beta = r * r / (r * r + ra * ra);
    }

  return beta;
}


#include <stdio.h>
#include <gsl/gsl_sf_legendre.h>
#include <gsl/gsl_sf_trig.h>
#include <gsl/gsl_pow_int.h>
#include <math.h>

double sin_term(int l,int m,double phi,double theta){
 double out = gsl_pow_int(-1.0,m)*gsl_sf_legendre_sphPlm(l,m,gsl_sf_cos(theta))*gsl_sf_sin(m*phi);
 return out;
}

double cos_term(int l,int m, double phi, double theta){
  double out= gsl_pow_int(-1.0,m)*gsl_sf_legendre_sphPlm(l,m,gsl_sf_cos(theta))*gsl_sf_cos(m*phi);
  return out;
}

float acos_fast(float x) {
  float negate = (float)x < 0;
  x =  x < 0 ? -x : x;
  float ret = -0.0187293;
  ret = ret * x;
  ret = ret + 0.0742610;
  ret = ret * x;
  ret = ret - 0.2121144;
  ret = ret * x;
  ret = ret + 1.5707288;
  ret = ret * sqrt(1.0-x);
  ret = ret - 2 * negate * ret;
  return negate * 3.14159265358979 + ret;
}

float fast_grf(int L, float x, float y,float z,float rands[]){
  float result = 0.0;
  float phi =atan2(y,x);
  float theta = acos_fast(z);
  float temp ;
  register int l = 0;
  register int m = 1;
  for (l=0; l<= L;l++){
    result += gsl_sf_legendre_sphPlm(l,0,gsl_sf_cos(theta))*rands[2*L*l];
    temp = 0.0;
    for(m=1;m<=l;m++){
      temp += (1 - ((m & 1) << 1))*gsl_sf_legendre_sphPlm(l,m,gsl_sf_cos(theta))*(rands[2*m+2*L*l+l]*gsl_sf_sin(m*phi)+rands[1+2*m+2*L*l]*gsl_sf_cos(m*phi));
    }
    result += 1.41421*temp;
  }
  return result;
}


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

#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>

#include "utils.h"
#include "norm.h"

/*
 * * Program that reads reference spectra and measure spectra data
 * and linearly unmixes the data using method of least squares fit.
 * Assumes readings is set to 400 and file path is preset to:
 *
 * syntax: ./linearUnmixer MUfilename ref1filename ref2filename ref3filename ...
 *  returns: completion status, prints composition of ref1, ref2, ... , refn
 *
 *  * ****************************************************************
 *  * Author       Dept.           Date            Notes
 *  * ****************************************************************
 *  * Jake Z       Bio. Eng.       Mar 12 2020     Initial version
 *  * Jake Z            "          Apr 03 2020     It works now!
 *  TODO: investigate the discrepancy between what you set $readings and the amount
 *  of space you allocate to the matrices as you forgot to account for the initial
 *  data clearing of the title headers in the data.
 * */

const int readings = 1000;	//num of data points in each reading
const int skip = 50;		//num of data points to discard from beginning

int main(int argc, char *argv[])
{
	if ( argc <= 2 )
	{
		fprintf( stderr, "error, incorrect usage.\n correct usage: ./unmix unknown ref1 ref2 ... refn\n");
		return 3;
	}

	//allocate memory for vars
	int numRefs = argc - 2;
	char *filePath = "./assets/";	// create prompt/arg for desired assets

	gsl_vector *muVector = gsl_vector_alloc(readings);
	gsl_vector *fVector = gsl_vector_alloc(readings);
	gsl_vector *x = gsl_vector_alloc(numRefs);

	gsl_matrix *refMatrix = gsl_matrix_alloc(readings, numRefs);
	gsl_matrix *C = gsl_matrix_alloc( numRefs, numRefs );
	gsl_matrix *Cinverse = gsl_matrix_alloc( numRefs, numRefs );

	int s;
	gsl_permutation *p = gsl_permutation_alloc(numRefs);

	//initialize vars to 0

	gsl_vector_set_zero(muVector);
	gsl_vector_set_zero(fVector);
	gsl_vector_set_zero(x);
	gsl_matrix_set_zero(refMatrix);
	gsl_matrix_set_zero(C);
	gsl_permutation_init(p);	//INITIALISED THE PERMUTATION

	readFileVector( muVector, filePath, argv[1] );	//read vals of unknown arg
	readFileMatrix( refMatrix, filePath, argc, argv );	//read the refs into matrix

	normalizeVector( muVector );

	//create fn that reads the files

	FILE *csvfile = fopen( "./muVector.csv", "wt" );
	gsl_vector_fprintf( csvfile, muVector, "%f" );
	fclose(csvfile);

	FILE *csvfile2 = fopen( "./refVectors.txt", "wt" );
	gsl_matrix_fprintf( csvfile2, refMatrix, "%f" );
	fclose(csvfile2);

	//now we have variable muVector with vals and refMatrix with vals.

	gsl_blas_dgemm( CblasTrans, CblasNoTrans, 1, refMatrix, refMatrix, 0, C);       //matrix C is now A^T A
	gsl_blas_dgemv( CblasTrans, 1, refMatrix, muVector, 0, fVector);        //fVector is 'b'
	gsl_linalg_LU_decomp( C, p, &s);
	gsl_linalg_LU_invert( C, p, Cinverse );

	//TODO maybe implement the other way you can solve for basis using inverse on both sides, (A^T A )^-1 A^T A x = (a^t a)^-1 (a^t b) = x

	gsl_blas_dgemv( CblasNoTrans, 1, Cinverse, fVector, 0, x);

	gsl_vector_fprintf( stdout, x, "%f" );

	gsl_matrix_free(Cinverse);
	gsl_matrix_free(C);
	gsl_matrix_free(refMatrix);
	gsl_vector_free(muVector);
	gsl_vector_free(fVector);
	gsl_vector_free(x);

	return 0;
}


/*! \file
  Compute 3D shape features and moments.

The following functions compute several shape features, including
central moments, center of gravity, and volume size.

\par Author:
Gabriele Lohmann, MPI-CBS
*/

/* From the Vista library: */
#include <viaio/Vlib.h>
#include <viaio/mu.h>
#include <via/via.h>

#include <gsl/gsl_cblas.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_eigen.h>

/* From the standard C libaray: */
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define SQR(x) ((x)*(x))
#define ABS(x) ((x) > 0 ? (x) : -(x))

static double power(double a, int k)
{
  register double res;
  register int i;

  res = 1.0;
  for (i=0; i<k; i++)
    res *= a;

  return res;
}

/*!
\fn  void VolumeCentroid(Volume v,double mean[3])
\brief Compute center of gravity of a volume.
\param v  input volume.
\param *mean output array.
The center of gravity is returned in mean[0],mean[1],mean[2].
Sufficient memory must have been allocated for <mean> prior to the call.
*/
void VolumeCentroid(Volume v, double mean[3])
{
  double npixels;
  long c,i,ce;
  VTrack t;

  mean[0] = 0;
  mean[1] = 0;
  mean[2] = 0;
  npixels = 0;

  for (i=0; i<v->nbuckets; i++) {
    for (t = v->bucket[i].first; t != NULL; t = t->next)  {
      mean[0] += (double) (t->band * t->length);
      mean[1] += (double) (t->row * t->length);

      ce = t->col + t->length;
      for (c = t->col; c<ce; c++)
	mean[2] += (double) c;

      npixels += t->length;
    }
  }
  if (npixels > 0) {
    mean[0] /= npixels;
    mean[1] /= npixels;
    mean[2] /= npixels;
  }
}

/*!
\fn double VolumeMoment(Volume v,double mean[3],int m0,int m1,int m2)
\brief Compute central moments of a volume.
\param v  input volume
\param *mean input array containing
the center of gravity as mean[0],mean[1],mean[2]  (slice,row,column).
If mean is NULL, the center of gravity is taken to be (0,0,0).
\param m0  first index of moment
\param m1  second index of moment
\param m2  third index of moment
*/
double VolumeMoment(Volume v,double mean[3],long m0,long m1,long m2)
{
  double res=0;
  double a = 1.0,sum = 0;
  long i,ca,ce,ci;
  double b,r,c;
  VTrack t;
  double g0,g1,g2;

  g0 = mean[0];
  g1 = mean[1];
  g2 = mean[2];

  res = 0;
  for (i=0; i<v->nbuckets; i++) {
    for (t = v->bucket[i].first; t != NULL; t = t->next)  {
      b  = (double) t->band - g0;
      r  = (double) t->row  - g1;
      a  = power(b,m0) * power(r,m1);
      ca = t->col;
      ce = ca + t->length;
      sum = 0;
      for (ci=ca; ci<ce; ci++) {
	c = (double)(ci - g2);
	sum += power(c,m2);
      }
      res += sum * a;
    }
  }
  return res;
}

/*!
\fn  void VBinCentroid(VImage src,double mean[3])
\brief Compute central moments of a binary raster image.
\param src  input image (bit repn)
\param *mean output array.
The center of gravity is returned in mean[0],mean[1],mean[2].
Sufficient memory must have been allocated for <mean> prior to the call.
*/
void VBinCentroid(VImage src,double mean[3])
{
  double npixels;
  long b,r,c;

  mean[0] = 0;
  mean[1] = 0;
  mean[2] = 0;
  npixels = 0;
  for (b=0; b<VImageNBands(src); b++) {
    for (r=0; r<VImageNRows(src); r++) {
      for (c=0; c<VImageNColumns(src); c++) {
	if (VPixel(src,b,r,c,VBit) > 0) {
	  mean[0] += b;
	  mean[1] += r;
	  mean[2] += c;
	  npixels++;
	}
      }
    }
  }
  if (npixels > 0) {
    mean[0] /= npixels;
    mean[1] /= npixels;
    mean[2] /= npixels;
  }
}

/*!
\fn  double VBinMoment(VImage src,double mean[3],int m0, int m1, int m2)
\brief Compute central moments of a binary raster image.
\param src  input image (bit repn)
\param *mean input array containing
the center of gravity as mean[0],mean[1],mean[2]  (slice,row,column).
If mean is NULL, the center of gravity is taken to be (0,0,0);
\param m0  first index of moment
\param m1  second index of moment
\param m2  third index of moment
*/
double VBinMoment(VImage src, double mean[3], long m0, long m1, long m2)
{
  double res;
  long b,r,c;
  double g0,g1,g2;

  g0 = mean[0];
  g1 = mean[1];
  g2 = mean[2];

  res = 0;
  for (b=0; b<VImageNBands(src); b++) {
    for (r=0; r<VImageNRows(src); r++) {
      for (c=0; c<VImageNColumns(src); c++) {
	if (VPixel(src,b,r,c,VBit) > 0)
	  res += power((double)(b - g0),m0)
	    * power((double)(r - g1),m1)
	    * power((double)(c - g2),m2);
      }
    }
  }
  return res;
}

/*!
  \fn long VolumeSize(Volume v)
  \brief compute volume size, output the number of voxels of a single volume
  \param v   a single volume
*/
long VolumeSize(Volume v)
{
  long i,isize;
  VTrack t;

  isize = 0;
  for (i=0; i<VolumeNBuckets(v); i++) {
    for (t = VFirstTrack(v,i); VTrackExists(t); t = VNextTrack(t))
      isize += VTrackLength(t);
  }
  return isize;
}

/*!
  \fn long VBinSize(VImage src)
  \brief count number of foreground voxels.
  \param src  input image (bit repn)
*/
long VBinSize(VImage src)
{
  long n;
  long i;
  VBit *bin_pp;

  if (VPixelRepn(src) != VBitRepn) VError(" input image must be bit repn");

  n = 0;
  bin_pp = (VBit *) VPixelPtr(src,0,0,0);
  for (i=0; i<VImageNPixels(src); i++) {
    if (*bin_pp > 0) n++;
    bin_pp++;
  }
  return n;
}

/*!
  \fn float VolumeDir(Volume vol,float *e,float x[3])
  \brief computer principal direction of a volume from its interia matrix
  \param vol   input volume
  \param e     output largest eigenvalue
  \param x     output first eigenvector
*/
float VolumeDir(Volume vol,float *e,float x[3])
{
  gsl_matrix *a=NULL;
  static gsl_matrix *evec=NULL;
  static gsl_vector *eval=NULL;
  static gsl_eigen_symmv_workspace *workspace=NULL;
  double m020,m002,m200,m110,m101,m011;
  double norm,angle;
  double mean[3];
  float tiny=1.0e-5;

  mean[0] = mean[1] = mean[2] = 0;
  VolumeCentroid(vol,mean);

  m020 = VolumeMoment(vol,mean,0,2,0);
  m002 = VolumeMoment(vol,mean,0,0,2);
  m200 = VolumeMoment(vol,mean,2,0,0);

  m110 = VolumeMoment(vol,mean,1,1,0);
  m101 = VolumeMoment(vol,mean,1,0,1);
  m011 = VolumeMoment(vol,mean,0,1,1);

  /* inertia matrix */
  if (a == NULL) {
    a = gsl_matrix_calloc(3,3);
    workspace = gsl_eigen_symmv_alloc(3);
  }
  gsl_matrix_set(a,0,0,m020 + m002);
  gsl_matrix_set(a,1,1,m200 + m002);
  gsl_matrix_set(a,2,2,m200 + m020);

  gsl_matrix_set(a,0,1,-m110);
  gsl_matrix_set(a,1,0,-m110);

  gsl_matrix_set(a,0,2,-m101);
  gsl_matrix_set(a,2,0,-m101);

  gsl_matrix_set(a,1,2,-m011);
  gsl_matrix_set(a,2,1,-m011);

  gsl_eigen_symmv(a,eval,evec,workspace);
  gsl_eigen_symmv_sort(eval,evec,GSL_EIGEN_SORT_VAL_DESC);

  x[0] = gsl_matrix_get(evec,0,0);
  x[1] = gsl_matrix_get(evec,1,0);
  x[2] = gsl_matrix_get(evec,2,0);

  *e = gsl_vector_get(eval,0);

  angle = 0;
  norm = sqrt((double)(SQR(x[0]) + SQR(x[1]) + SQR(x[2])));
  if (norm > tiny) {
    angle = ABS(x[1])/ norm;
  }

  return angle;
}

/*!
  \fn void VolumeEigen(Volume vol,gsl_vector *eval,gsl_matrix *evec)
  \brief computer principal directions of a volume from its interia matrix
  \param vol    input volume
  \param eval   output eigenvalues
  \param evec   output matrix of eigenvectors (columns)
*/
void VolumeEigen(Volume vol,gsl_vector *eval,gsl_matrix *evec)
{
  static gsl_matrix *a=NULL;
  static gsl_eigen_symmv_workspace *workspace=NULL;
  double m020,m002,m200,m110,m101,m011;
  double mean[3];

  mean[0] = mean[1] = mean[2] = 0;
  VolumeCentroid(vol,mean);

  m020 = VolumeMoment(vol,mean,0,2,0);
  m002 = VolumeMoment(vol,mean,0,0,2);
  m200 = VolumeMoment(vol,mean,2,0,0);

  m110 = VolumeMoment(vol,mean,1,1,0);
  m101 = VolumeMoment(vol,mean,1,0,1);
  m011 = VolumeMoment(vol,mean,0,1,1);

  /* inertia matrix */
  if (a == NULL) {
    a = gsl_matrix_calloc(3,3);
    workspace = gsl_eigen_symmv_alloc(3);
  }

  gsl_matrix_set(a,0,0,m020 + m002);
  gsl_matrix_set(a,1,1,m200 + m002);
  gsl_matrix_set(a,2,2,m200 + m020);

  gsl_matrix_set(a,0,1,-m110);
  gsl_matrix_set(a,1,0,-m110);

  gsl_matrix_set(a,0,2,-m101);
  gsl_matrix_set(a,2,0,-m101);

  gsl_matrix_set(a,1,2,-m011);
  gsl_matrix_set(a,2,1,-m011);

  gsl_eigen_symmv(a,eval,evec,workspace);
  gsl_eigen_symmv_sort(eval,evec,GSL_EIGEN_SORT_VAL_DESC);

  return;
}