/**
    Exosphere simulation code as described in Kegerreis et al. (2016) and
    the README.pdf documentation.

    Publically available at icc.dur.ac.uk/index.php?content=Research/Topics/O13

    Functions header.
**/


/*
    0. Includes etc.
    1. Functions
        1.1 RNG
        1.2 Migration
        1.3 Cold Traps
        1.4 Misc
        1.5 Admin etc.
    2. Template Functions
    3. Macros
*/


#ifndef FUNCTIONS_H_INCLUDED
#define FUNCTIONS_H_INCLUDED
////////////////////////////////////////////////////////////////////////////////


////////////////////////////////////////////////////////////////////////////////
//                              0. Includes etc.                              //
////////////////////////////////////////////////////////////////////////////////


#include <iostream>
#include <fstream>
#include <iterator>
#include <vector>
#include <string>
#include <cstdarg>
#include <numeric>
#include <cmath>
#include <ctime>
#include <algorithm>

using namespace std;


////////////////////////////////////////////////////////////////////////////////
//                              1. Functions                                  //
////////////////////////////////////////////////////////////////////////////////


//============================================================================//
//                              1.1 RNG                                       //
//============================================================================//


//=========================
// Sources
//=========================
double  rand_phi_iso(double xi);

double  rand_theta_iso(double xi);

double  rand_phi_solar(double xi_1, double xi_2, double xi_3);

double  rand_theta_solar(double xi_1, double xi_2);


//=========================
// Velocity
//=========================
double  Maxwell_cdf(int T, double v, const double m);

double  Maxwell_flux_cdf(int T, double v, const double m);

vector< vector<double> >    make_A2_Maxwell_cdf(double T_0, double T_1,
                                                double T_min, const double m);

vector< vector<double> >    make_A2_Maxwell_flux_cdf(double T_0, double T_1,
                                                     double T_min,
                                                     const double m);

double  rand_vel(double T, double xi_1, double xi_2, double T_min,
                 vector< vector<double> > *Maxwell_cdf_array);

double  rand_alpha(double xi);

double  rand_alpha_flux(double xi);

double  rand_eta(double xi);


//============================================================================//
//                              1.2 Migration                                 //
//============================================================================//


double  phi_prime(double phi, double time);

double  residence_time(double T, double Q);


//=========================
// Temperature
//=========================
vector< vector<double> >    make_A2_longitude_offset(int seed,
                                                     double sig_offset,
                                                     int num_phi_,
                                                     int num_theta);

vector< vector<double> >    make_A2_H_temperature(double T_0, double T_1,
                                                  int num_phi_, int num_theta);


//=========================
// Hops
//=========================
double  get_a(double v_0);

double  get_e(double a, double alpha);

double  get_max_height(double a, double e);

double  get_delta(double a, double e);

double  get_end_theta(double theta_0, double delta, double eta);

double  get_epsilon(double theta, double theta_0, double delta);

double  get_end_phi(double phi_0, double epsilon, double eta);

double  get_flight_time(double a, double e, double delta);


//=========================
// At altitude
//=========================
double  get_delta_p(double a, double e, double alt);

double  get_v(double a, double alt);

double  get_alpha(double a, double e, double alt);


//============================================================================//
//                              1.3 Cold Traps                                //
//============================================================================//


bool    test_capt_at_theta(double theta, double xi, double f_N, double f_S);

bool    test_capt_local_traps(double theta, double phi, double theta_trap_N_1,
                              double theta_trap_N_2, double phi_trap_N_1,
                              double phi_trap_N_2, double theta_trap_S_1,
                              double theta_trap_S_2, double phi_trap_S_1,
                              double phi_trap_S_2);

bool    test_capt_seasonal(double theta, double xi, double time,
                           double f_season_N, double f_season_S);

double  get_time_seasonal_rel(double theta, double time, double xi);


//============================================================================//
//                              1.4 Misc                                      //
//============================================================================//


int     round_to_int(double x);

int     get_index(double val, double max_val, int num_val);

double  get_flux(double v, double alpha);


//=========================
// Check location
//=========================
bool    is_north_pole(double theta);

bool    is_south_pole(double theta);

bool    is_in_day(double phi_);

bool    is_recording_time(double time, int N_bounce, const double start_time,
                          const int start_N_bounce);

bool    is_in_mare(double theta, double phi, double theta_m_1, double theta_m_2,
                   double phi_m_1, double phi_m_2);

bool    is_LADEE_lat(double theta);


//=========================
// Bin Vector Values
//=========================
vector<int> bin_values(vector<double> values, int num_bins, double max_value);


//============================================================================//
//                              1.5 Admin etc.                                //
//============================================================================//


//=========================
// Check Inputs
//=========================
void    check_arguments(string source, string emerge_dist, string losses,
                        string T_dist, string cold_traps, string record_dist,
                        string surf_intn, string misc, string rates,
                        string backup_results);


//=========================
// Printing Information
//=========================
string  get_date_time(char date_y_n, char time_y_n);

void    print_date_time(char date, char time);

void    print_progress_MPI(int i, int num_loops, int percent, clock_t t_start,
                           int world_rank, int world_size);


////////////////////////////////////////////////////////////////////////////////
//                              2. Template Functions                         //
////////////////////////////////////////////////////////////////////////////////


//=========================
// Statistics
//=========================
template    <typename T>
double      get_mean(vector<T> const & vec)
{
    /// Returns the mean of the elements in a vector.
    double  mean=0;

    mean    = (double) accumulate(vec.begin(), vec.end(), 0.0) / vec.size();

    return  mean;
}
//========


template    <typename T>
double      get_stdev(vector<T> const & vec)
{
    /// Returns the standard deviation of the elements in a vector.
    double  mean, square_sum, stdev;

    mean        = accumulate(vec.begin(), vec.end(), 0.0) / vec.size();
    square_sum  = inner_product(vec.begin(), vec.end(), vec.begin(), 0.0);
    stdev       = sqrt(square_sum/vec.size() - mean*mean);

    return  stdev;
}
//========


//=========================
// Overloaded Vector Operators
//=========================
template    <typename T>
vector<T>   operator+(vector<T> const & a, vector<T> const & b)
{
    /// Overloaded operator + for vectors.
    vector<T>   result;
    result.reserve(a.size());

    transform(a.begin(), a.end(), b.begin(), back_inserter(result), plus<T>());

    return  result;
}
//========


template    <typename T>
vector<T>   operator-(vector<T> const & a, vector<T> const & b)
{
    /// Overloaded operator - for vectors.
    vector<T>   result;
    result.reserve(a.size());

    transform(a.begin(), a.end(), b.begin(), back_inserter(result), minus<T>());

    return  result;
}
//========


template    <typename T>
void        operator+=(vector<T> & a, vector<T> const & b)
{
    /// Overloaded operator += for vectors.
    a   = a + b;
}
//========


template    <typename T>
void        operator-=(vector<T> & a, vector<T> const & b)
{
    /// Overloaded operator -= for vectors.
    a   = a - b;
}
//========


//=========================
// Saving to Files
//=========================
template    <typename T>
void        save(string filename, char print_y_n, T const & var)
{
    /// Saves a variable to a new file, adds ".txt" if not included. \
        Doesn't print saved message if print_y_n == 'n'.
    // If filename doesn't include ".txt", then append it
    string  extension   = ".txt";
    // Add ".txt" on end if ".txt" not found
    if (filename.find(extension) == string::npos) {
        filename.append(extension);
    }
    ofstream    file;
    file.open(filename.c_str());
    file << var << endl;
    if (print_y_n != 'n') {
        cout << "data saved to \"" << filename.c_str() << "\"" << endl;
    }
}
//========


template    <typename T>
void        save_append(string filename, char print_y_n, T const & var)
{
    /// Saves a variable to an existing file, adds ".txt" if not included. \
        (Does work with strings, unlike save_mult)
    // If filename doesn't include ".txt", then append it
    string  extension   = ".txt";
    // Add ".txt" on end if ".txt" not found
    if (filename.find(extension) == string::npos) {
        filename.append(extension);
    }
    ofstream    file;
    // Open in append mode
    file.open(filename.c_str(), ios_base::app);
    file << var << endl;
    if (print_y_n != 'n') {
        cout << "data appended to \"" << filename.c_str() << "\"" << endl;
    }
}
//========


template    <typename T>
void        save_mult(string filename, char print_y_n, int num_vars,
                      T const & var, ...)
{
    /// Saves var to a new file, then appends other variables of the same type \
        to the same file. (Doesn't work with strings!)
    va_list elems;
    // Initialise elems to store args after var
    va_start(elems, var);
    // if filename doesn't include ".txt", then append it
    string  extension   = ".txt";
    if (filename.find(extension) == string::npos) {
        filename.append(extension);
    }

    // Save the first variable
    ofstream    file;
    file.open(filename.c_str());
    file << var << endl;
    file.close();
    // Save the other variables
    // Open in append mode
    file.open(filename.c_str(), ios_base::app);
    for (int i=0; i<(num_vars-1); i++) {
        file << va_arg(elems, T) << endl;
    }
    if (print_y_n != 'n') {
        cout << "data saved to \"" << filename.c_str() << "\"" << endl;
    }
}
//========


template    <typename T>
void        save_mult_append(string filename, char print_y_n, int num_vars,
                             T const & var, ...)
{
    /// As save_mult() but all appended to an existing file.
    va_list elems;
    // Initialise elems to store args after var
    va_start(elems, var);

    // if filename doesn't include ".txt", then append it
    string  extension   = ".txt";
    if (filename.find(extension) == string::npos) {
        filename.append(extension);
    }

    // Save the first variable
    ofstream    file;
    // Open in append mode
    file.open(filename.c_str(), ios_base::app);
    file << var << endl;
    file.close();
    // Save the other variables
    // Open in append mode
    file.open(filename.c_str(), ios_base::app);
    for (int i=0; i<(num_vars-1); i++) {
        file << va_arg(elems, T) << endl;
    }
    if (print_y_n != 'n') {
        cout << "data appended to \"" << filename.c_str() << "\"" << endl;
    }
}
//========


template    <typename T>
void        save_to_dir(string filename, string dir, char print_y_n,
                        T const & var)
{
    /// Saves a variable to a new file in a specified directory, \
        adds ".txt" if not included in filename. \
        Doesn't print saved message if print_y_n == 'n'. \
        Can't get to work with mult or append functions.
    // If filename doesn't include ".txt", then append it
    string  extension   = ".txt";
    if (filename.find(extension) == string::npos) {
        filename.append(extension);
    }
    string      file_in_dir_name    = dir.append(filename);
    ofstream    file_in_dir(file_in_dir_name.c_str());
    file_in_dir << var << endl;
    if (print_y_n != 'n') {
        cout << "data saved to \"" << file_in_dir_name.c_str() << "\"" << endl;
    }
}
//========


//========
// Save Vectors
//========
template    <typename T>
void        save_to_dir_v(string filename, string dir, char print_y_n,
                          vector<T> const & vec)
{
    /// Saves a vector to a new file in a specified directory, \
        adds ".txt" if not included.
    // If filename doesn't include ".txt", then append it
    string  extension   = ".txt";
    if (filename.find(extension) == string::npos) {
        filename.append(extension);
    }
    string      file_in_dir_name    = dir.append(filename);
    ofstream    file_in_dir(file_in_dir_name.c_str());
    // Write each vector element to file
    for (int i=0; i<vec.size(); i++) {
        file_in_dir << vec[i] << ' ';
    }
    if (print_y_n != 'n') {
        cout << "data saved to \"" << file_in_dir_name.c_str() << "\"" << endl;
    }
}
//========


template    <typename T>
void        save_to_dir_vv(string filename, string dir, char print_y_n,
                           vector< vector<T> > const & vecvec)
{
    /// Saves a vector of vectors to a new file in a specified directory, \
        one vector per line, adds ".txt" if not included.
    // If filename doesn't include ".txt", then append it:
    string  extension   = ".txt";
    if (filename.find(extension) == string::npos) {
        filename.append(extension);
    }
    string      file_in_dir_name    = dir.append(filename);
    ofstream    file_in_dir(file_in_dir_name.c_str());

    // Loop through 1st dimension
    for (int i=0; i<vecvec.size(); i++) {
        // Write each vector element to file
        for (int j=0; j<vecvec[i].size(); j++) {
            file_in_dir << vecvec[i][j] << ' ';
        }
        file_in_dir << endl;
    }
    if (print_y_n != 'n') {
        cout << "data saved to \"" << file_in_dir_name.c_str() << "\"" << endl;
    }
}
//========


template    <typename T>
void        save_to_dir_vvv(string filename, string dir, char print_y_n,
                            vector< vector<T> > const & vecvecvec)
{
    /// Saves a vector of vectors of vectors to a new file in a specified \
        directory, one vector per line, adds ".txt" if not included.
    // If filename doesn't include ".txt", then append it:
    string  extension   = ".txt";
    if (filename.find(extension) == string::npos) {
        filename.append(extension);
    }
    string      file_in_dir_name    = dir.append(filename);
    ofstream    file_in_dir(file_in_dir_name.c_str());

    // Loop through 1st and 2nd dimensions
    for (int i=0; i<vecvecvec.size(); i++) {
        for (int j=0; j<vecvecvec[i].size(); j++) {
            // Write each vector element to file
            for (int k=0; k<vecvecvec[i][j].size(); k++) {
                file_in_dir << vecvecvec[i][j][k] << ' ';
            }
        }
        file_in_dir << endl;
    }
    if (print_y_n != 'n') {
        cout << "data saved to \"" << file_in_dir_name.c_str() << "\"" << endl;
    }
}
//========


////////////////////////////////////////////////////////////////////////////////
//                              3. Macros                                     //
////////////////////////////////////////////////////////////////////////////////


//=========================
// Printing Variable Names and Values
//=========================
// Use scientific full precision
#define print_prec(var) \
    cout << #var << " = " << scientific << (var) << endl

#define print_debug(var) \
    cout << #var << " = " << (var) << endl

#define print_debug2(var_1, var_2) \
    cout << #var_1 << " = " << (var_1) << endl \
         << #var_2 << " = " << (var_2) << endl

#define print_debug3(var_1, var_2, var_3) \
    cout << #var_1 << " = " << (var_1) << endl \
         << #var_2 << " = " << (var_2) << endl \
         << #var_3 << " = " << (var_3) << endl

#define print_debug4(var_1, var_2, var_3, var_4) \
    cout << #var_1 << " = " << (var_1) << endl \
         << #var_2 << " = " << (var_2) << endl \
         << #var_3 << " = " << (var_3) << endl \
         << #var_4 << " = " << (var_4) << endl

#define print_debug5(var_1, var_2, var_3, var_4, var_5) \
    cout << #var_1 << " = " << (var_1) << endl \
         << #var_2 << " = " << (var_2) << endl \
         << #var_3 << " = " << (var_3) << endl \
         << #var_4 << " = " << (var_4) << endl \
         << #var_5 << " = " << (var_5) << endl

#define print_debug6(var_1, var_2, var_3, var_4, var_5, var_6) \
    cout << #var_1 << " = " << (var_1) << endl \
         << #var_2 << " = " << (var_2) << endl \
         << #var_3 << " = " << (var_3) << endl \
         << #var_4 << " = " << (var_4) << endl \
         << #var_5 << " = " << (var_5) << endl \
         << #var_6 << " = " << (var_6) << endl
//========


////////////////////////////////////////////////////////////////////////////////
#endif // FUNCTIONS_H_INCLUDED
