/*
 * Copyright (C) 2018-2020
 * HERE Europe B.V.
 * Utrecht University (The Netherlands).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 * License-Filename: LICENSE
 */

//
// Created by Wouter Jongeling (wouter.jongeling@gmail.com)
// Modified by Mitra, Aniket on 2019-06-18.
//

/*!
 * @file BrownianBridge.h
 * @brief
 * @authors Aniket Mitra (aniket.mitra@here.com)
 *  Wouter Jongeling (wouter.jongeling@gmail.com)
 */

#ifndef MOVETK_BROWNIANBRIDGE__H
#define MOVETK_BROWNIANBRIDGE__H

#include <iterator>
#include <cmath>
#include "movetk/utils/Requirements.h"
#include "movetk/utils/Iterators.h"
#include "movetk/geom/trajectory_to_interface.h"
#include "movetk/utils/Iterators.h"
#include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h>
#include "gsl/gsl_min.h"

/*!
 * @file AlgorithmTraits.h
 * @brief The Brownian Bridge Movement Model (BBMM)
 * @authors Aniket Mitra (aniket.mitra@here.com)
 */

/*!
 *  @brief a collection of algorithms in movetk
 */
namespace movetk_algorithms {

    // based on
    // https://www.semanticscholar.org/paper/Computational-Movement-Analysis-Using-Brownian-Sijben/c6104979e609e12e1514b8b605a7d14fa6b5d81a

    /*!
    *  @brief a collection of classes for BBMM
    */
    namespace brownian_bridge {

        /*!
         * @struct ParameterTraits
         * @brief This traits class serves as a collection of types
         * for parameterization of BBMM
         * @tparam _GeometryTraits -  This class is a collection of movetk
         *  geometry types. For example @refitem movetk_core::MovetkGeometryKernel
         * @tparam TrajectoryIterator - An iterator type.
         * For example @refitem TabularTrajectory::TrajectoryIterator
         */
        template<class _GeometryTraits, class TrajectoryIterator>
        struct ParameterTraits {
            typedef _GeometryTraits GeometryTraits;
            /*!*
             * @typedef ::NT
             * @brief typedef  of the number type defined in GeometryTraits.
             * For example @refitem  movetk_core::MovetkGeometryKernel::NT
             * */
            typedef typename GeometryTraits::NT NT;
            /*!*
             * @typedef ::Point
             * @brief typedef of the point type defined in GeometryTraits.
             * For example @refitem movetk_core::MovetkGeometryKernel::MovetkPoint
             *  */
            typedef typename GeometryTraits::MovetkPoint Point;
            /*!* @typedef ::Vector
             * @brief typedef of the vector type defined in GeometryTraits.
             * For example @refitem  movetk_core::MovetkGeometryKernel::MovetkVector
             * */
            typedef typename GeometryTraits::MovetkVector Vector;
            /*!* @typedef ::Parameters
             */
            typedef std::tuple<Point, Point, NT, TrajectoryIterator, TrajectoryIterator> Parameters;
            /*!
             * @enum ParameterColumns
             *
             */
            enum ParameterColumns {
                POINT,  /*!< global enum value 1 */
                MU,     /*!< global enum value 1 */
                SIGMA_SQUARED,  /*!< global enum value 1 */
                BEGIN,  /*!< global enum value 1 */
                END    /*!< global enum value 1 */
            };

        };

        /*!
         * @class template<class GeometryTraits, class ParameterTraits,
         *  class Norm, class InputIterator, std::size_t MaxIter> MLE
         * @brief The Maximum Likelihood Estimator
         * @tparam GeometryTraits -   This class is a collection of movetk
         *  geometry types. For example @refitem movetk_core::MovetkGeometryKernel
         * @tparam ParameterTraits -  This traits class serves as a collection of types
         * for parameterization of BBMM
         * @tparam Norm - The type that models Euclidean distance
         * For example @refitem movetk_support::FiniteNorm
         * @tparam InputIterator - A random access iterator type over a collection of
         * tuples where each tuple is a model of type ParameterTraits::Parameters
         * @tparam MaxIter - A positive integer that bounds the number of iterations in the MLE
         * algorithm
         */
        template<class GeometryTraits, class ParameterTraits,
                class Norm, class InputIterator, std::size_t MaxIter,
                typename = movetk_core::requires_random_access_iterator<InputIterator>,
                typename = movetk_core::requires_tuple<
                        typename InputIterator::value_type>,
                typename = movetk_core::requires_tuple_element_as_movetk_point<GeometryTraits,
                        ParameterTraits::ParameterColumns::POINT,
                        typename InputIterator::value_type>,
                typename = movetk_core::requires_tuple_element_as_movetk_point<GeometryTraits,
                        ParameterTraits::ParameterColumns::MU,
                        typename InputIterator::value_type>,
                typename = movetk_core::requires_tuple_element_as_NT<GeometryTraits,
                        ParameterTraits::ParameterColumns::SIGMA_SQUARED,
                        typename InputIterator::value_type>>
        class MLE {
        private:
            typedef typename GeometryTraits::NT NT;
            const gsl_min_fminimizer_type *T = gsl_min_fminimizer_goldensection;
            gsl_min_fminimizer *s = gsl_min_fminimizer_alloc(T);
            gsl_function F;
            NT estimated_parameter;
            std::size_t iter = 0, status;

            struct fn_params {
                InputIterator first;
                InputIterator beyond;
            };

            static double fn(double sigma_squared, void *p) {
                Norm norm;
                struct fn_params *params = (fn_params *) p;
                InputIterator it = params->first;
                NT log_likelihood = 0;
                while (it != params->beyond) {
                    typename GeometryTraits::MovetkVector v = std::get<ParameterTraits::ParameterColumns::POINT>(*it)
                                                              - std::get<ParameterTraits::ParameterColumns::MU>(*it);
                    NT operand1 = -LOG_TWO_PI - log(static_cast<NT>(sigma_squared));
                    NT operand2 = -norm(v) / (2 * sigma_squared);
                    log_likelihood += operand1 + operand2;
                    it++;
                }
                return -1 * log_likelihood;
            }

            void operator()(InputIterator first, InputIterator beyond,
                            NT result, NT x_lower, NT x_upper, NT eps) {
                struct fn_params params;
                params.first = first;
                params.beyond = beyond;
                F.function = &MLE::fn;
                F.params = &params;
                status = gsl_min_fminimizer_set(s, &F, result, x_lower, x_upper);
                do {
                    iter++;
                    status = gsl_min_fminimizer_iterate(s);
                    result = gsl_min_fminimizer_x_minimum(s);
                    x_lower = gsl_min_fminimizer_x_lower(s);
                    x_upper = gsl_min_fminimizer_x_upper(s);

                    status = gsl_min_test_interval(x_lower, x_upper, eps, 0.0);

                    if (status == GSL_SUCCESS) {
                        estimated_parameter = result;
                        break;
                    }
                } while (status == GSL_CONTINUE && iter < MaxIter);

                if (iter >= MaxIter)
                    estimated_parameter = result;
            }

        public:
            /*!
             *
             * @param first
             * @param beyond
             * @param result
             * @param x_lower
             * @param x_upper
             * @param eps
             */
            MLE(InputIterator first, InputIterator beyond,
                NT result, NT x_lower, NT x_upper, NT eps) {
                Norm norm;
                std::size_t num_elements = std::distance(first, beyond);
                if (num_elements == 1) {
                    typename GeometryTraits::MovetkVector v = std::get<ParameterTraits::ParameterColumns::POINT>(*first)
                                                              - std::get<ParameterTraits::ParameterColumns::MU>(*first);
                    NT l = norm(v);
                    if (l < MOVETK_EPS) {
                        estimated_parameter = MOVETK_EPS;
                        return;
                    }
                }
                (*this)(first, beyond, result, x_lower, x_upper, eps);
            }

            /*!
             *
             * @param first
             * @param beyond
             */
            MLE(InputIterator first, InputIterator beyond) {
                Norm norm;
                InputIterator it = first;
                NT upper_bound = 0, squared_length = 0;
                std::size_t num_elements = std::distance(first, beyond);
                if (num_elements == 1) {
                    typename GeometryTraits::MovetkVector v = std::get<ParameterTraits::ParameterColumns::POINT>(*first)
                                                              - std::get<ParameterTraits::ParameterColumns::MU>(*first);
                    NT l = norm(v);
                    if (l < MOVETK_EPS) {
                        estimated_parameter = MOVETK_EPS;
                        return;
                    }
                }
                while (it != beyond) {
                    typename GeometryTraits::MovetkVector v = std::get<ParameterTraits::ParameterColumns::POINT>(*it)
                                                              - std::get<ParameterTraits::ParameterColumns::MU>(*it);
                    NT l = norm(v);
                    squared_length += l;
                    if (l > upper_bound)
                        upper_bound = l;
                    it++;
                }
                NT initial_estimate = squared_length / (2 * num_elements);
                (*this)(first, beyond, initial_estimate, MOVETK_EPS, upper_bound, MOVETK_EPS);
            }

            /*!
             *
             * @return
             */
            NT operator()() {
                return estimated_parameter;
            }

            /*!
             *
             */
            ~MLE() {
                gsl_min_fminimizer_free(s);
            }
        };

        /*!
         *
         * @tparam GeometryTraits
         * @tparam ProbeTraits
         * @tparam ParameterTraits
         * @tparam Norm
         * @tparam GeoProjection
         */
        template<class GeometryTraits, class ProbeTraits, class ParameterTraits, class Norm,
                class GeoProjection>
        class Model {
        private:

            typedef typename GeometryTraits::NT NT;
            typedef typename GeometryTraits::MovetkPoint Point;
            typedef typename GeometryTraits::MovetkVector Vector;
            movetk_core::MakePoint<GeometryTraits> make_point;

            NT joint_log_likelihood = 0;

            Point mean(Point &p1, Point &p2, NT alpha) {
                Vector v = p2 - p1;
                v *= alpha;
                return (p1 + v);
            }

        public:
            /*!
             *
             * @tparam TrajectoryIterator
             * @tparam OutputIterator
             * @param first
             * @param beyond
             * @param result
             */
            template<class TrajectoryIterator, class OutputIterator,
                    typename = movetk_core::requires_random_access_iterator<TrajectoryIterator>,
                    typename = movetk_core::requires_tuple<
                            typename TrajectoryIterator::value_type>,
                    typename = movetk_core::requires_tuple_element_as_arithmetic<
                            ProbeTraits::ProbeColumns::LAT, typename TrajectoryIterator::value_type>,
                    typename = movetk_core::requires_tuple_element_as_arithmetic<
                            ProbeTraits::ProbeColumns::LON, typename TrajectoryIterator::value_type>,
                    typename = movetk_core::requires_tuple_element_as_size_t<
                            ProbeTraits::ProbeColumns::SAMPLE_DATE, typename TrajectoryIterator::value_type>,
                    typename = movetk_core::requires_output_iterator<OutputIterator>,
                    typename = movetk_core::requires_tuple<
                            typename OutputIterator::value_type>,
                    typename = movetk_core::requires_tuple_element_as_movetk_point<GeometryTraits,
                            ParameterTraits::ParameterColumns::POINT,
                            typename OutputIterator::value_type>,
                    typename = movetk_core::requires_tuple_element_as_movetk_point<GeometryTraits,
                            ParameterTraits::ParameterColumns::MU,
                            typename OutputIterator::value_type>,
                    typename = movetk_core::requires_tuple_element_as_NT<GeometryTraits,
                            ParameterTraits::ParameterColumns::SIGMA_SQUARED,
                            typename OutputIterator::value_type>>
            Model(TrajectoryIterator first, TrajectoryIterator beyond, OutputIterator result) {

                auto reflat = std::get<ProbeTraits::ProbeColumns::LAT>(*first);
                auto reflon = std::get<ProbeTraits::ProbeColumns::LON>(*first);
                GeoProjection ref(reflat, reflon);

                TrajectoryIterator tit = first;
                std::size_t NumPoints = std::distance(first, beyond);

                if (NumPoints == 1) {
                    auto lat = std::get<ProbeTraits::ProbeColumns::LAT>(*first);
                    auto lon = std::get<ProbeTraits::ProbeColumns::LON>(*first);
                    auto projected_point = ref.project(lat, lon);
                    Point p1 = make_point(std::cbegin(projected_point), std::cend(projected_point));
                    *result = std::make_tuple(p1, p1, 0, first, first);
                    return;
                }

                std::size_t multiple = (NumPoints - 1) / 2;
                std::size_t remainder = (NumPoints - 1) % 2;
                TrajectoryIterator last = first + 2 * multiple;

                while (tit != last) {
                    auto lat = std::get<ProbeTraits::ProbeColumns::LAT>(*tit);
                    auto lon = std::get<ProbeTraits::ProbeColumns::LON>(*tit);
                    auto projected_point = ref.project(lat, lon);
                    Point p1 = make_point(std::cbegin(projected_point), std::cend(projected_point));

                    lat = std::get<ProbeTraits::ProbeColumns::LAT>(*(tit + 1));
                    lon = std::get<ProbeTraits::ProbeColumns::LON>(*(tit + 1));
                    projected_point = ref.project(lat, lon);
                    Point p2 = make_point(std::cbegin(projected_point), std::cend(projected_point));

                    lat = std::get<ProbeTraits::ProbeColumns::LAT>(*(tit + 2));
                    lon = std::get<ProbeTraits::ProbeColumns::LON>(*(tit + 2));
                    projected_point = ref.project(lat, lon);
                    Point p3 = make_point(std::cbegin(projected_point), std::cend(projected_point));

                    auto ts1 = std::get<ProbeTraits::ProbeColumns::SAMPLE_DATE>(*tit);
                    auto ts2 = std::get<ProbeTraits::ProbeColumns::SAMPLE_DATE>(*(tit + 1));
                    auto ts3 = std::get<ProbeTraits::ProbeColumns::SAMPLE_DATE>(*(tit + 2));
                    NT alpha = static_cast<NT>(  ts3 - ts2 ) / static_cast<NT>( ts3 - ts1 );

                    Point mu = mean(p1, p3, alpha);
                    *result = std::make_tuple(p2, mu, 0, tit, tit + 2);
                    tit += 2;
                }

                if (remainder == 1) {
                    auto lat = std::get<ProbeTraits::ProbeColumns::LAT>(*last);
                    auto lon = std::get<ProbeTraits::ProbeColumns::LON>(*last);
                    auto projected_point = ref.project(lat, lon);
                    Point p1 = make_point(std::cbegin(projected_point), std::cend(projected_point));

                    lat = std::get<ProbeTraits::ProbeColumns::LAT>(*(last + 1));
                    lon = std::get<ProbeTraits::ProbeColumns::LON>(*(last + 1));
                    projected_point = ref.project(lat, lon);
                    Point p2 = make_point(std::cbegin(projected_point), std::cend(projected_point));
                    Point mu = mean(p1, p2, 0.5);
                    *result = std::make_tuple(mu, mu, 0, last, last + 1);
                }

            }

        };

        /*!
         *
         * @tparam GeometryTraits
         * @tparam ParameterTraits
         */
        template<class GeometryTraits, class ParameterTraits>
        class ParameterSelector {
            std::size_t SIZE;
        public:
            /*!
             *
             * @param size
             */
            ParameterSelector(std::size_t size) : SIZE(size) {}

            /*!
             *
             * @tparam InputIterator
             * @tparam OutputIterator
             * @param first
             * @param beyond
             * @param result
             */
            template<class InputIterator, class OutputIterator,
                    typename = movetk_core::requires_random_access_iterator<InputIterator>,
                    typename = movetk_core::requires_output_iterator<OutputIterator>,
                    typename = movetk_core::requires_tuple<typename InputIterator::value_type>,
                    typename = movetk_core::requires_tuple_element_as_NT<GeometryTraits,
                            ParameterTraits::ParameterColumns::SIGMA_SQUARED,
                            typename InputIterator::value_type>,
                    typename = movetk_core::requires_NT<GeometryTraits,
                            typename OutputIterator::value_type >>
            void operator()(InputIterator first, InputIterator beyond, OutputIterator result) {
                assert(std::distance(first, beyond) >= SIZE);
                std::vector<typename ParameterTraits::NT> coeffs;
                std::transform(first, beyond, std::back_inserter(coeffs), [](auto a) {
                    return std::get<ParameterTraits::ParameterColumns::SIGMA_SQUARED>(a);
                });
                std::sort(std::begin(coeffs), std::end(coeffs));
                std::size_t coeff_interval = std::distance(first, beyond) / SIZE;
                for (std::size_t i = 0; i < SIZE; i++) {
                    std::size_t index = coeff_interval / 2 + coeff_interval * i;
                    *result = coeffs[index];
                }
            }
        };

        template<class GeometryTraits, class ParameterTraits, class Norm>
        class LogLikelihood {
            typedef typename GeometryTraits::NT NT;
            typedef typename ParameterTraits::Parameters Parameters;
        public:
            LogLikelihood() = default;

            template<class InputIterator, class OutputIterator,
                    typename = movetk_core::requires_random_access_iterator<InputIterator>,
                    typename = movetk_core::requires_output_iterator<OutputIterator>,
                    typename = movetk_core::requires_NT<GeometryTraits,
                            typename InputIterator::value_type>,
                    typename = movetk_core::requires_NT<GeometryTraits,
                            typename OutputIterator::value_type >>
            void operator()(const Parameters &params, InputIterator first,
                            InputIterator beyond, OutputIterator result) {
                Norm norm;
                typename GeometryTraits::MovetkVector v = std::get<ParameterTraits::ParameterColumns::POINT>(params)
                                                          - std::get<ParameterTraits::ParameterColumns::MU>(params);
                InputIterator pit = first;
                NT squared_length = norm(v);
                while (pit != beyond) {
                    NT operand1 = -LOG_TWO_PI - log(*pit);
                    NT operand2 = -squared_length / (2 * (*pit));
                    NT log_likelihood = operand1 + operand2;
                    *result = log_likelihood;
                    pit++;
                }
            }

            NT operator()(const Parameters &params, NT sigma_squared) {
                Norm norm;
                typename GeometryTraits::MovetkVector v = std::get<ParameterTraits::ParameterColumns::POINT>(params)
                                                          - std::get<ParameterTraits::ParameterColumns::MU>(params);
                NT squared_length = norm(v);
                NT operand1 = -LOG_TWO_PI - log(sigma_squared);
                NT operand2 = -squared_length / (2 * sigma_squared);
                return (operand1 + operand2);
            }
        };

    }

}

#endif //MOVETK_BROWNIANBRIDGE__H


/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2014 Boone "Bea" Adkins
 *
 * 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:
 *
 * 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 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.
 */
/**
 * @class pomdp::MultinomialPolicy
 *
 * CompoundPolicy whose action is determined by sampling a <s>multinomial</s> categorical distribution across
 * sub-policies.
 *
 * @date Aug 10, 2013
 * @author Bea Adkins
 */

#ifndef MULTINOMIAL_POLICY_H
#define MULTINOMIAL_POLICY_H

#include <boost/static_assert.hpp>
#include <boost/type_traits.hpp>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

#include "assert_rng.h"
#include "compound_policy.h"

namespace pomdp
{
template <class Base_Policy>
class MultinomialPolicy : public CompoundPolicySpec<Base_Policy>
{
public:
  MultinomialPolicy() : val_table_(NULL)
  {
    rng_ = gsl_rng_alloc(gsl_rng_taus2);

    // Randomly seed the RNG.
    srand(time(NULL));
    gsl_rng_set (rng_, rand());
  }

  virtual ~MultinomialPolicy()
  {
    gsl_rng_free(rng_);
    if(CompoundPolicySpec<Base_Policy>::isInit())
    {
      gsl_ran_discrete_free(val_table_);
    }
  }

  virtual bool loadFromFiles(const std::vector<std::string>& paths)
  {
    if(!CompoundPolicySpec<Base_Policy>::loadFromFiles(paths))
      return false;
    CompoundPolicySpec<Base_Policy>::is_init_ = false; // Not yet initialized.

    // Now that model number is known, initialize random number variables.
    freshenDistribution();

    CompoundPolicySpec<Base_Policy>::is_init_ = true;
    return true;
  }

  virtual int applyPolicy(const FeatureValue& current) const
  {
    // Sample a policy index.
    int i_sampled_policy = sample();
    if(i_sampled_policy == -1)
      return -1;
    ROS_ASSERT(0 <= i_sampled_policy && i_sampled_policy < CompoundPolicySpec<Base_Policy>::getNumModels());

    // Apply that policy.
    ROS_DEBUG_STREAM("Appying policy #" << i_sampled_policy << ":");
    return CompoundPolicySpec<Base_Policy>::models_[i_sampled_policy]->applyPolicy(current);
  }

  virtual bool setModelLikelihoods(const std::vector<double>& model_likelihoods)
  {
    if(!CompoundPolicySpec<Base_Policy>::setModelLikelihoods(model_likelihoods))
      return false;

    freshenDistribution(); // With new model likelihoods
    return true;
  }
protected:
  /**
   * Random number generator.
   */
  gsl_rng* rng_;

  /**
   * Table for quickly sampling values.
   */
  gsl_ran_discrete_t* val_table_;

  /**
   * @return int within [0, numModels] or -1 on error.
   */
  int sample() const
  {
    if(!CompoundPolicySpec<Base_Policy>::isInit() || CompoundPolicySpec<Base_Policy>::getNumModels() == 0)
      return -1;

    // Sample a value.
    return gsl_ran_discrete(rng_, val_table_);
  }

  /**
   * Call this whenever model_likelihoods_ change, in value or length.
   */
  void freshenDistribution()
  {
    ROS_ASSERT(CompoundPolicySpec<Base_Policy>::model_likelihoods_.size() > 0);

    if(CompoundPolicySpec<Base_Policy>::isInit())
      gsl_ran_discrete_free(val_table_);
    val_table_ = gsl_ran_discrete_preproc(CompoundPolicySpec<Base_Policy>::model_likelihoods_.size(),
                                          &CompoundPolicySpec<Base_Policy>::model_likelihoods_[0]);
  }
};

typedef MultinomialPolicy<MDPPolicy> MultinomialMDPPolicy;
typedef MultinomialPolicy<POMDPPolicy> MultinomialPOMDPPolicy;

} /* namespace pomdp */
#endif /* MULTINOMIAL_POLICY_H */


#pragma once

#include <string>
#include <vector>

#include <gsl/gsl_complex.h>
#include <MathEngine.h>

extern MathEngine engine;

bool printDifference(const std::string& input, const expression expr, const double output, const double expected);
void requireIsEqual(const std::string& input, const double expected);

bool printDifference(const std::string& input, const expression expr, const std::string& output, const std::string& expected);
void requireIsEqual(const std::string& input, const std::string& expected, bool evaluate = false);

bool printDifference(const std::string& input, const expression expr, const expression output, const std::string& expected);
void requireExprIsEqual(const std::string& input, const std::string& expected);

void requireErrorIsEqual(const std::string& input, const std::string& expected);


/*****************************************************************************/
#define this_is "FDBinary v.3 beta (test version of 21 Feb 2010)"
/*****************************************************************************/

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

#include <gsl/gsl_matrix.h>
#include <gsl/gsl_sf.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_multimin.h>

#include "mxfuns.h"
#include "fd3sep.h"
#include "triorb.h"

/*****************************************************************************/

/* function and macro to kill the program with a short message */
#define FDBErrorString "\nError in FDBianry"
#define DIE(s) {fprintf(stderr,"%s: %s\n",FDBErrorString,s);fdbfailure();};
void fdbfailure(void) { exit ( EXIT_FAILURE ); }

/*****************************************************************************/

/* macros for reading values from keyboard */
#define GETDBL(x) {if(1!=scanf("%lg",x)) DIE("failed reading double");}
#define GETINT(x) {if(1!=scanf("%d", x)) DIE("failed reading int");}
#define GETLNG(x) {if(1!=scanf("%ld",x)) DIE("failed reading long");}
#define GETSTR(x) {if(1!=scanf("%s", x)) DIE("failed reading string");}

/*****************************************************************************/

#define SPEEDOFLIGHT 2.998E5 /* speed of light in km/s */

static char *triorb_strings[] = {
  /*  0 */  "period of AB--C",                    "day", "%.9lg",
  /*  1 */  "time of periast pass in AB--C",      "day", "%.9lg",
  /*  2 */  "eccentricity of AB--C",                "1", "%.9lg",
  /*  3 */  "periast long of AB in AB--C",        "deg", "%.9lg",
  /*  4 */  "a sin(incl) of AB in AB--C",   "light-day", "%.9lg",
  /*  5 */  "a sin(incl) of C in AB--C",    "light-day", "%.9lg",
  /*  6 */  "period A--B",                        "day", "%.9lg",
  /*  7 */  "time of periast pass in A--B",       "day", "%.9lg",
  /*  8 */  "eccentricity of A--B",                 "1", "%.9lg",
  /*  9 */  "periast long of A in A--B",          "deg", "%.9lg",
  /* 10 */  "rv semiamp of A in A--B",           "km/s", "%.9lg",
  /* 11 */  "rv semiamp of B in A--B",           "km/s", "%.9lg",
  /* 12 */  "periast long adv per cycle in A--B", "deg", "%.9lg"
};

static char *triorb_pname   (long j) { return triorb_strings[3*j];   }
static char *triorb_punit   (long j) { return triorb_strings[3*j+1]; }
static char *triorb_pformat (long j) { return triorb_strings[3*j+2]; }

/*****************************************************************************/

static long   K, M, N, Ndft, ksw[3], nfp, opsw[TRIORB_NP];
static double **dftobs, **dftmod, **dftres;
static double rvstep, *otimes, *rvcorr, *sig, **lfm, **rvm;
static double op0[TRIORB_NP], dop0[TRIORB_NP];
static double meritfngsl ( const gsl_vector *v, void *params );
static double meritfn ( double *op );

#define MX_FDBINARY_FORMAT "%15.8E   "
static char *mxfd3fmts=MX_FDBINARY_FORMAT;

/* FDBinary */

int main ( void ) {

  long i, i0, i1, j, k, vc, vlen, nruns, niter;
  double **masterobs, **mod, **res, **obs, z0, z1, stoprat;
  char obsfn[1024], resfn[1024], modfn[1024], rvsfn[1024], logfn[1024];
  char *starcode[] = {"A","B","C"};
  FILE *logfp;

  setbuf ( stdout, NULL );
  MxError( FDBErrorString, stdout, fdbfailure );
  MxFormat( mxfd3fmts );

  printf ( "\n  %s\n", this_is );

  printf ( "\n  SECTION 1: LOADING OBSERVED SPECTRA\n\n" );
  printf ( "  observed spectra master file = " );
  GETSTR ( obsfn ); printf ( "%s\n", obsfn );
  vc=0; vlen=0; masterobs = MxLoad ( obsfn, &vc, &vlen ); M = vc - 1;
  printf ( "  loaded %ld spectra with %ld data pts per spectrum\n", M, vlen );
  z0 = **masterobs; z1 = *(*masterobs+vlen-1);
  rvstep = SPEEDOFLIGHT * ( - 1 + exp ((z1-z0)/(vlen-1)) );
  printf ( "  lnlambda range from %lf to %lf\n", z0, z1 );
  printf ( "  rv step %lg km/s per data point\n", rvstep );
  printf ( "  use data from lnlambda = " );
  GETDBL ( &z0 ); printf ( "%lg\n", z0 );
  printf ( "             to lnlambda = " );
  GETDBL ( &z1 ); printf ( "%lg\n", z1 );
  i0 = 0; while ( *(*masterobs+i0) < z0 ) i0++;
  i1 = vlen-1; while ( z1 < *(*masterobs+i1) ) i1--;
  N = i1 - i0 + 1;
  printf ( "  selected %ld data points per spectrum\n", N );
  obs = MxAlloc ( M+1, N );
  for ( i = 0 ; i < N ; i++ ) for ( j = 0 ; j <= M ; j++ )
    *(*(obs+j)+i) = *(*(masterobs+j)+i0+i);
  MxFree ( masterobs, vc, vlen );
  printf ( "  write selected data to file = " );
  GETSTR ( obsfn ); printf ( "%s\n", obsfn ); MxWrite ( obs, M+1, N, obsfn );

  printf ( "\n  SECTION 2: COMPONENT SPECTRA MODELING SWITCHES (0/1)\n\n" );
  for ( K = i = 0 ; i < 3 ; i++ ) {
    int sw;
    GETINT(&sw);
    sw = sw ? 1 : 0;
    printf ( "  %s = %d\n", starcode[i], sw );
    if ( sw ) ksw[K++] = i;
  }
  printf ( "\n  number of components to be resolved is %ld\n", K );

  /* allocating memory */
  Ndft = 2*(N/2 + 1);
  dftobs = MxAlloc ( M, Ndft ); dft_fwd ( M, N, obs+1, dftobs );
  otimes = *MxAlloc ( 1, M );
  rvcorr = *MxAlloc ( 1, M );
  sig = *MxAlloc ( 1, M );
  res = MxAlloc ( M+1, N ); dftres = MxAlloc ( M, Ndft );
  mod = MxAlloc ( K+1, N ); dftmod = MxAlloc ( K, Ndft );
  rvm = MxAlloc ( K, M );
  lfm = MxAlloc ( K, M );
  for ( i = 0 ; i < N ; i++ ) { *(*res+i) = *(*obs+i); *(*mod+i) = *(*obs+i); }

  printf ( "\n  SECTION 3: DESCRIPTORS TO OBSERVED SPECTRA\n" );
  printf ( "             AND LIGHT-FACTORS ASSIGNED TO COMPONENTS\n\n" );
  printf ( "  %14s%12s%12s", "t_obs", "rv_corr", "noise_rms" );
  for ( k = 0 ; k < K ; k++ ) printf ( "      lf_%s", starcode[ksw[k]] );
  printf ( "\n" );
  printf ( "  %14s%12s%12s", "[day]", "[km/s]", "[1]" );
  for ( k = 0 ; k < K ; k++ ) printf ( "       [1]" );
  printf ( "\n\n" );
  for ( j = 0 ; j < M ; j++ )
  {
    GETDBL(otimes+j); printf ( "  %14.5lf", *(otimes+j) );
    GETDBL(rvcorr+j); printf ( "%12.4lf", *(rvcorr+j) );
    GETDBL(sig+j); printf ( "%12.4lf", *(sig+j) );
    for ( k = 0; k < K ; k++ )
      { GETDBL(*(lfm+k)+j); printf ( "%10.4lf", *(*(lfm+k)+j) ); }
    putchar ( '\n' );
  }

  printf ( "\n  SECTION 4: PARAMETERS OF ORBITS\n" );
  for ( nfp = i = 0 ; i < TRIORB_NP ; i++ ) {
    if ( 0 == i ) printf ( "\n  wide (AB--C) orbit\n\n" );
    if ( 6 == i ) printf ( "\n  tight (A--B) orbit\n\n" );
    printf ( "  %s [%s] = ", triorb_pname(i), triorb_punit(i) );
    GETDBL(op0+i); printf ( triorb_pformat(i), *(op0+i) );
    printf ( " +/- " );
    GETDBL(dop0+i); printf ( triorb_pformat(i), *(dop0+i) );
    if ( *(dop0+i) ) { opsw[nfp++] = i; printf ( " *** free ***" ); }
    if ( 2 == i || 8 == i ) *(op0+i) = *(op0+i) / (1-*(op0+i)); /* ecc */
    printf ( "\n" );
  }
  printf ( "\n  number of free orbital parameters is %ld\n", nfp );

  printf ( "\n  SECTION 5: OPTIMISATION DETAILS\n\n" );
  if ( 0 == nfp ) {
    printf ( "  WARNING: *** no free orbital parameters *** \n" );
    printf ( "  optimisation of orb. parameters will not be performed\n" );
    printf ( "  all input in this section is ignored\n\n" );
  }
  GETLNG( &nruns );
  printf ( "  number of independent optimisation runs = %ld\n", nruns );
  GETLNG( &niter );
  printf ( "  number of allowed iterations per run = %ld\n", niter );
  GETDBL( &stoprat );
  printf ( "  stop when simplex shrinked by factor = %lg\n", stoprat );

  printf ( "\n  SECTION 6: OUTPUT FILES\n\n" );

  printf ( "  write model spectra to file = " );
  GETSTR ( modfn ); printf ( "%s\n", modfn );
  printf ( "  write residuals to file = " );
  GETSTR ( resfn ); printf ( "%s\n", resfn );
  printf ( "  write radial velocities to file = " );
  GETSTR ( rvsfn ); printf ( "%s\n", rvsfn );
  printf ( "  write optimisation log to file = " );
  scanf ( "%s", logfn );
  if ( strlen ( logfn ) ) {
    printf ( "%s\n", logfn );
    if ( NULL == ( logfp = fopen ( logfn, "w" ) ) )
      DIE ( "could not open log file" );
    setbuf ( logfp, NULL );
  } else {
    printf ( "[no log file will be written]\n" );
  }

  printf ( "\n  SECTION 7: OPTIMISATION\n\n" );

  if ( 0 == nfp ) {                                            /* separation */

    double chi2;

    printf ( "  WARNING: *** no free orbital parameters *** \n" );
    chi2 = meritfn ( op0 );
    printf ( "  separation at the starting point:  chi2=%lg  gof=%.2lf\n",
      chi2, gsl_sf_gamma_inc_Q ( N*(M-K)/2.0, chi2/2.0 ) );
    dft_bck ( K, N, dftmod, mod+1 ); MxWrite ( mod, K+1, N, modfn );
    dft_bck ( M, N, dftres, res+1 ); MxWrite ( res, M+1, N, resfn );
    MxWrite( rvm, K, M, rvsfn );

  } else {                                                  /* disentangling */

    const gsl_multimin_fminimizer_type *T = gsl_multimin_fminimizer_nmsimplex;
    gsl_multimin_fminimizer *s = NULL; gsl_multimin_function minex_func;
    gsl_vector *ss, *x;
    size_t iter = 0, irun = 0, i; int status;
    double size0 = 1.0, size, chi2;
    const gsl_rng_type * R; gsl_rng * r;

    gsl_rng_env_setup (); R = gsl_rng_default; r = gsl_rng_alloc (R);

    minex_func.f = &meritfngsl;
    minex_func.n = nfp;
    minex_func.params = (void *) NULL;

    s = gsl_multimin_fminimizer_alloc (T, nfp);
    x  = gsl_vector_alloc (nfp); /* starting point */
    ss = gsl_vector_alloc (nfp); /* initial vertex size vector */

    for ( i = 0 ; i < nfp ; i++ ) {
      gsl_vector_set ( x, i, op0[opsw[i]] );
      gsl_vector_set ( ss, i, dop0[opsw[i]] );
      size0 *= dop0[opsw[i]];
    }

    chi2 = meritfngsl ( x, NULL );
    printf ( "  separation at the starting point:  chi2=%lg  gof=%.2lf\n",
      chi2, gsl_sf_gamma_inc_Q ( N*(M-K)/2.0, chi2/2.0 ) );
    dft_bck ( K, N, dftmod, mod+1 ); MxWrite ( mod, K+1, N, modfn );
    dft_bck ( M, N, dftres, res+1 ); MxWrite ( res, M+1, N, resfn );
    MxWrite( rvm, K, M, rvsfn );

    printf
      ( "  converged disentangling runs reported only if chi2 decreases\n" );

    for ( irun = 0 ; irun < nruns ; irun++ )  {

      for ( i = 0 ; i < nfp ; i++ ) {
        double tmp = 2.0 * ( gsl_rng_uniform (r) - 0.5 );
        gsl_vector_set ( x,  i, op0[opsw[i]] + tmp * dop0[opsw[i]] );
        gsl_vector_set ( ss, i, dop0[opsw[i]] );
      }

      gsl_multimin_fminimizer_set ( s, &minex_func, x, ss );
      status = GSL_FAILURE; size = size0;

      for ( iter = 0 ; iter < niter ; iter++ ) {
        if ( gsl_multimin_fminimizer_iterate (s) ) break;
        size = gsl_multimin_fminimizer_size (s);
        status = gsl_multimin_test_size (size, stoprat*size0);
        if ( GSL_SUCCESS == status ) break;
      }

      if ( (GSL_SUCCESS == status) && (s->fval < chi2) ) {
        printf ( "\n  irun=%d  iter=%d  sxshrnkf=%.2lg  chi2=%lg  ",
          irun+1, iter, size/size0, chi2 = s->fval );
        printf ( "gof=%.2lf\n",
          gsl_sf_gamma_inc_Q ( (N*(M-K)-nfp)/2.0, chi2/2.0 ) );
        for ( i = 0 ; i < nfp ; i++ ) {
          double xi = gsl_vector_get (s->x, i);
          if ( 2 == opsw[i] || 8 == opsw[i] ) xi = fabs ( xi / (1 + xi) );
          printf ( "  %40s = ", triorb_pname(opsw[i]) );
          printf ( triorb_pformat(opsw[i]), xi );
          if ( strcmp ( triorb_punit(opsw[i]), "1" ) )
            printf ( " %s", triorb_punit(opsw[i]) );
          printf ( "\n" );
        }
        dft_bck ( K, N, dftmod, mod+1 ); MxWrite ( mod, K+1, N, modfn );
        dft_bck ( M, N, dftres, res+1 ); MxWrite ( res, M+1, N, resfn );
        MxWrite( rvm, K, M, rvsfn );
      }

      if ( (GSL_SUCCESS == status) && strlen ( logfn ) ) {
        fprintf ( logfp, "%d  %d  %.2lg  %lg  ",
          irun+1, iter, size/size0, s->fval );
        fprintf ( logfp, "%.2lf  ",
          gsl_sf_gamma_inc_Q ( (N*(M-K)-nfp)/2.0, (s->fval)/2.0 ) );
        for ( i = 0 ; i < nfp ; i++ ) {
          double xi = gsl_vector_get (s->x, i);
          if ( 2 == opsw[i] || 8 == opsw[i] ) xi = fabs ( xi / (1 + xi) );
          fprintf ( logfp, "  " );
          fprintf ( logfp, triorb_pformat(i), xi );
        }
        fprintf ( logfp, "\n" );
      }

    }

    gsl_multimin_fminimizer_free (s);
    gsl_vector_free(x);
    gsl_vector_free(ss);

    printf ( "\n  completed %ld optimisation runs\n\n", nruns );

  }

  printf ( "  EXITING REGULARLY\n\n" );

  return EXIT_SUCCESS;

}

/*****************************************************************************/

double meritfngsl ( const gsl_vector *v, void *params )
{

  long i, j;
  double op[TRIORB_NP];

  for ( i = j = 0 ; i < TRIORB_NP ; i++ )
    op[i] = dop0[i] ? gsl_vector_get ( v, j++ ) : op0[i];

  return meritfn ( op );
}

/*****************************************************************************/

double meritfn ( double *opin )
{

  long j, k;
  double op[TRIORB_NP], rv[3];

  op[ 0] = opin[ 0];
  op[ 1] = opin[ 1];
  op[ 2] = fabs ( opin[2] / (1 + opin[2]) );
  op[ 3] = opin[ 3] * (M_PI/180);
  op[ 4] = opin[ 4];
  op[ 5] = opin[ 5];
  op[ 6] = opin[ 6];
  op[ 7] = opin[ 7];
  op[ 8] = fabs ( opin[8] / (1 + opin[8]) );
  op[ 9] = opin[ 9] * (M_PI/180);
  op[10] = opin[10] / rvstep;
  op[11] = opin[11] / rvstep;
  op[12] = opin[12] * (M_PI/180);

  for ( j = 0 ; j < M ; j++ ) {
    triorb_rv ( op, otimes[j], rv );
    for ( k = 0 ; k < K ; k++ )
      *(*(rvm+k)+j) = rv[ksw[k]] + *(rvcorr+j) / rvstep;
  }

  return fd3sep ( K, M, N, dftobs, sig, rvm, lfm, dftmod, dftres );
}

/*****************************************************************************/


#define _USE_MATH_DEFINES
#include <cmath>
#include <Eigen/Core>
#include <gsl/gsl_sf_gamma.h>
#include <gsl/gsl_sf_zeta.h>
#include <gsl/gsl_sf_fermi_dirac.h>
#include <vector>
//#include <iostream>

using Eigen::Matrix;
using Eigen::VectorXd;

void li(size_t k, double x, Matrix<double, -1, 1> &ans);

template<typename T>
class Distribution {
private:
    T loc;
    T scale;
    Matrix<T, -1, 1> param;
    size_t k;
    //
    Matrix<T, -1, 1> C;
    Matrix<T, -1, 1> C_inf;
    Matrix<T, -1, 1> Gamma;
    Matrix<T, -1, -1> C_li;
    Matrix<T, -1, -1> Col;
    Matrix<T, -1, -1> Col_arg;
    Matrix<T, -1, -1> Alpha;
    Matrix<T, -1, 1> M;
    Matrix<T, -1, 1> C_y;
    std::vector<T> y_abs_pow;
    Matrix<T, -1, 1> X_pow;
    Matrix<T, -1, 1> Li;
    Matrix<T, -1, 1> I;
    Matrix<T, -1, 1> I_u;
    Matrix<T, -1, -1> I_t;
public:
    Distribution(T loc_, T scale_, Matrix<T, -1, 1> param_)
            : loc(loc_), scale(scale_), param(std::move(param_)), k(param.size()),
              C_y(Matrix<T, -1, 1>::Ones(2 * k - 2)), y_abs_pow(2 * k - 1), X_pow(Matrix<T, -1, 1>(2 * k - 2)),
              Li(Matrix<T, -1, 1>(2 * k - 2)), I_u(Matrix<T, -1, 1>(2 * k - 1)) {
        // Vector of i * Gamma(i) * (1 - 2^{1 - i}) * zeta(i)
        C = Matrix<T, -1, 1>(2 * k - 2);
        C[0] = M_LN2;

        {
            double curr = 0.5;
            for (size_t i = 2; i <= 2 * k - 2; i++) {
                C[i - 1] = i * gsl_sf_fact(i - 1) * (1 - curr) * gsl_sf_zeta_int(i);
                curr /= 2;
            }
        }

        C_inf = Matrix<T, -1, 1>::Ones(2 * k - 2);
        for (size_t i = 0; i < 2 * k - 2; i += 2) {
            C_inf[i] = -1;
        }

        // Vector of i * Gamma(i)
        Gamma = Matrix<T, -1, 1>(2 * k - 2);

        C_li = Matrix<T, -1, -1>::Zero(2 * k - 2, 2 * k - 2);
        for (size_t i = 0; i < 2 * k - 2; i++) {
            Gamma[i] = gsl_sf_fact(i + 1);
            if (i != 0) {
                C_li(i, 0) = 1 / Gamma[i - 1];
            } else {
                C_li(i, 0) = 1;
            }
        }

        // Matrix of collocation coefficients
        Col = Matrix<T, -1, -1>::Zero(2 * k - 1, 2 * k - 1);
        for (size_t i = 0; i < 2 * k - 1; i++) {
            for (size_t j = 0; j <= i; j++) {
                Col(i, j) = gsl_sf_choose(i, j);
            }
        }

        M = Matrix<T, -1, 1>(2 * k - 1);
        M[0] = 1;
        for (size_t i = 1; i < 2 * k - 1; i++) {
            M[i] = i % 2 == 1 ? 0 : 2 * C[i - 1];
        }

        precalc();
    }

    void precalc() {
        // Matrix of collocation coefficients
        std::vector<T> locs(2 * k - 1), scales(2 * k - 1);
        locs[0] = 1;
        scales[0] = 1;
        for (size_t i = 1; i < 2 * k - 1; i++) {
            locs[i] = locs[i - 1] * loc;
            scales[i] = scales[i - 1] * scale;
        }

        Col_arg = Col;
        for (size_t i = 0; i < 2 * k - 1; i++) {
            for (size_t j = 0; j <= i; j++) {
                Col_arg(i, j) *= locs[i - j] * scales[j];
            }
        }

        // Matrix of shifted alpha parameter vector
        Alpha = Matrix<T, -1, -1>::Zero(k, 2 * k - 1);
        for (size_t i = 0; i < k; i++) {
            for (size_t j = i; j < i + k; j++) {
                Alpha(i, j) = param[j - i];
            }
        }
        Alpha = param.transpose() * Alpha;
    }

    void reset_param(T loc_, T scale_, Eigen::Matrix<T, -1, 1> &param_) {
        loc = loc_;
        scale = scale_;
        param = param_;
        precalc();
    }

    double cdf(T x) {
//        std::cout << "x = " << x << std::endl;
        T y = (x - loc) / scale;
        T y_abs = abs(y);

        if (y < 0) {
            for (size_t i = 1; i < 2 * k - 2; i += 2) {
                C_y[i] = -1;
            }
        } else {
            for (size_t i = 1; i < 2 * k - 2; i += 2) {
                C_y[i] = 1;
            }
        }

        // Vector of pow(arg, i)
        y_abs_pow[0] = 1;
        for (size_t i = 1; i < 2 * k - 1; i++) {
            y_abs_pow[i] = y_abs_pow[i - 1] * y_abs;
        }

        for (size_t i = 0; i < 2 * k - 2; i++) {
            X_pow[i] = y_abs_pow[i + 1]; // (y_abs, i + 1);
        }
        X_pow = (0.5 * tanh(y_abs / 2) - 0.5) * X_pow;

        // Polylogarithms vector

//        for (size_t i = 0; i < 2 * k - 2; i++) {
//            if (y_abs >= 690) Li[i] = 0;
//            else Li[i] = -gsl_sf_fermi_dirac_int(i, -y_abs);
//        }
        double z = -std::exp(-y_abs);
        li(2 * k - 2, z, Li);
//        std::cout << "li = " << Li.transpose() << std::endl;

        // Polylogarithm coefficients
        Matrix<T, -1, -1> C_li_arg = C_li;
        for (size_t i = 0; i < 2 * k - 2; i++) {
            C_li_arg(i, 0) *= y_abs_pow[i];  // pow(y_abs, i);
            for (size_t j = 1; j <= i; j++) {
                C_li_arg(i, j) = C_li_arg(i - 1, j - 1);
            }
        }

        C_li_arg = C_li_arg * Li;
        C_li_arg = Gamma.array() * C_li_arg.array();

        I = (C_inf + C_y).array() * C.array() + C_y.array() * X_pow.array() + C_y.array() * C_li_arg.array();

        // from -infty to 0 + 0.5 * tanh
        T C_y0 = y >= 0 ? 1 : -1;
        I_u << 0.5 + C_y0 * 0.5 * tanh(y_abs / 2), I;

        // Vector I^t
        I_t = Col_arg * I_u;
        T fx = (Alpha * I_t)(0);
        T C_1 = 1 / (Alpha * Col_arg * M).sum();
//        std::cout << "fx = " << C_1 * fx << std::endl;
        return C_1 * fx;
    }
};


#ifndef __MATH_GSL_VECTOR__
#define __MATH_GSL_VECTOR__

#include <assert.h>

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

#include <gsl/gsl_matrix.h>
#include "math/vectorops.h"

class GslVectorItem {
 public:
  GslVectorItem(gsl_vector* ptr, size_t index) :
    ptr_(ptr),
    index_(index) { }

  operator const double() {
    return gsl_vector_get(ptr_, index_);
  }

  double operator =(const double v) {
    gsl_vector_set(ptr_, index_, v);
    return v;
  }

  double operator +=(const double v) {
    double oldv = gsl_vector_get(ptr_, index_);
    gsl_vector_set(ptr_, index_, oldv + v);
    return oldv + v;
  }

  double operator -=(const double v) {
    double oldv = gsl_vector_get(ptr_, index_);
    gsl_vector_set(ptr_, index_, oldv - v);
    return oldv - v;
  }
 private:
  gsl_vector* ptr_;
  size_t index_;
};

class GslVectorBase {
 public:
  /*
  double operator[](const size_t index) const {
    assert(ptr_ != NULL);
    return gsl_vector_get(ptr_, index);
  }
  */

  GslVectorItem operator[](const size_t index) const {
    assert(ptr_ != NULL);
    return GslVectorItem(ptr_, index);
  }

  GslVectorBase& operator+=(const gsl_vector* x) {
    assert(ptr_ != NULL);
    gsl_vector_add(ptr_, x);
    return *this;
  }

  GslVectorBase& operator-=(const gsl_vector* x) {
    assert(ptr_ != NULL);
    gsl_vector_sub(ptr_, x);
    return *this;
  }

  GslVectorBase& operator+=(const GslVectorBase& x) {
    assert(ptr_ != NULL);
    gsl_vector_add(ptr_, x.ptr());
    return *this;
  }

  GslVectorBase& operator-=(const GslVectorBase& x) {
    assert(ptr_ != NULL);
    gsl_vector_sub(ptr_, x.ptr());
    return *this;
  }

  GslVectorBase& operator*=(const gsl_vector* x) {
    assert(ptr_ != NULL);
    gsl_vector_mul(ptr_, x);
    return *this;
  }

  GslVectorBase& operator*=(const GslVectorBase& x) {
    assert(ptr_ != NULL);
    gsl_vector_mul(ptr_, x.ptr());
    return *this;
  }

  GslVectorBase& operator/=(const GslVectorBase& x) {
    assert(ptr_ != NULL);
    gsl_vector_div(ptr_, x.ptr());
    return *this;
  }

  double Sum() const {
    return gsl_blas_dsum(ptr_);
  }

  double L2Norm() const {
    return gsl_blas_dnrm2(ptr_);
  }

  void Normalize() const {
    assert(ptr_ != NULL);
    double s = Sum();
    gsl_vector_scale(ptr_, 1. / s);
  }

  size_t size() const {
    return ptr_->size;
  }

  GslVectorBase& operator*=(double v) {
    assert(ptr_ != NULL);
    gsl_vector_scale(ptr_, v);
    return *this;
  }

  GslVectorBase& operator/=(double v) {
    assert(ptr_ != NULL);
    gsl_vector_scale(ptr_, 1. / v);
    return *this;
  }

  // Note that the standalone product is a dot product!
  const double operator*(const gsl_vector* x) const {
    double result;
    assert(ptr_ != NULL);
    gsl_blas_ddot(ptr_, x, &result);
    return result;
  }

  const double operator*(const GslVectorBase& x) const {
    double result;
    assert(ptr_ != NULL);
    gsl_blas_ddot(ptr_, x.ptr(), &result);
    return result;
  }

  GslVectorBase& operator+=(const double x) {
    for (size_t ii = 0; ii < ptr_->size; ++ii) {
      gsl_vector_set(ptr_, ii, gsl_vector_get(ptr_, ii) + x);
    }
    return *this;
  }

  GslVectorBase& operator-=(const double x) {
    for (size_t ii = 0; ii < ptr_->size; ++ii) {
      gsl_vector_set(ptr_, ii, gsl_vector_get(ptr_, ii) - x);
    }
    return *this;
  }

  GslVectorBase& operator=(const gsl_vector* x) {
    assert(ptr_ != NULL);
    gsl_vector_memcpy(ptr_, x);

    return *this;
  }

  GslVectorBase& operator=(const GslVectorBase& x) {
    assert(ptr_ != NULL);
    return *this = x.ptr();
  }

  GslVectorBase& operator=(const double v) {
    if (v == 0.0) {
      SetZero();
    } else {
      SetAll(v);
    }
    return *this;
  }

  void SetZero() {
    assert(ptr_ != NULL);
    gsl_vector_set_zero(ptr_);
  }

  void SetAll(const double v) {
    assert(ptr_ != NULL);
    gsl_vector_set_all(ptr_, v);
  }

  int Fprintf(FILE* stream, const char* format) const {
    assert(ptr_ != NULL);
    return gsl_vector_fprintf(stream, ptr_, format);
  }

  int Fscanf(FILE* stream) {
    assert(ptr_ != NULL);
    return gsl_vector_fscanf(stream, ptr_);
  }

  const gsl_vector* ptr() const { return ptr_; }
  gsl_vector* mutable_ptr() { return ptr_; }

 protected:
  GslVectorBase() : ptr_(NULL) { }
  gsl_vector* ptr_;

 private:
  GslVectorBase(const GslVectorBase&) { }
};

class GslVector : public GslVectorBase {
 public:
  GslVector(const size_t size) : GslVectorBase() {
    Allocate(size);
  }

  GslVector() : GslVectorBase() {
  }

  ~GslVector() {
    if (ptr_ != NULL) {
      gsl_vector_free(ptr_);
    }
  }

  void Reset(gsl_vector* val) {
    if (ptr_ != NULL) {
      gsl_vector_free(ptr_);
    }
    ptr_ = val;
  }

  void Allocate(const size_t size) {
    assert(ptr_ == NULL);
    ptr_ = gsl_vector_alloc(size);
  }

  GslVectorBase& operator=(const gsl_vector* x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const GslVectorBase& x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const double v) {
    GslVectorBase::operator=(v);
    return *this;
  }
 private:
  GslVector(const GslVector&) { }
};

class GslMatrixRow : public GslVectorBase {
 public:
  GslMatrixRow(GslMatrix& matrix, const size_t row) :
   view_(gsl_matrix_row(matrix.mutable_ptr(), row)) {
     ptr_ = &view_.vector;
  }

  GslMatrixRow(gsl_matrix* matrix, const size_t row) :
   view_(gsl_matrix_row(matrix, row)) {
     ptr_ = &view_.vector;
  }

  GslVectorBase& operator=(const gsl_vector* x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const GslVectorBase& x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const double v) {
    GslVectorBase::operator=(v);
    return *this;
  }
 private:
  gsl_vector_view view_;
  GslMatrixRow(const GslMatrixRow&) { }
};

class GslMatrixColumn : public GslVectorBase {
 public:
  GslMatrixColumn(GslMatrix& matrix, const size_t col) :
   view_(gsl_matrix_column(matrix.mutable_ptr(), col)) {
     ptr_ = &view_.vector;
  }

  GslMatrixColumn(gsl_matrix* matrix, const size_t col) :
   view_(gsl_matrix_column(matrix, col)) {
     ptr_ = &view_.vector;
  }

  GslVectorBase& operator=(const gsl_vector* x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const GslVectorBase& x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const double v) {
    GslVectorBase::operator=(v);
    return *this;
  }
 private:
  gsl_vector_view view_;
  GslMatrixColumn(const GslMatrixColumn&) { }
};

class GslMatrixDiagonal : public GslVectorBase {
 public:
  GslMatrixDiagonal(GslMatrix& matrix) :
   view_(gsl_matrix_diagonal(matrix.mutable_ptr())) {
     ptr_ = &view_.vector;
  }

  GslMatrixDiagonal(gsl_matrix* matrix) :
   view_(gsl_matrix_diagonal(matrix)) {
     ptr_ = &view_.vector;
  }

  GslVectorBase& operator=(const gsl_vector* x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const GslVectorBase& x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const double v) {
    GslVectorBase::operator=(v);
    return *this;
  }
 private:
  gsl_vector_view view_;
  GslMatrixDiagonal(const GslMatrixDiagonal&) { }
};

class GslSubvector : public GslVectorBase {
 public:
 GslSubvector(GslVectorBase& vector, size_t i, size_t n) :
  view_(gsl_vector_subvector(vector.mutable_ptr(), i, n)) {
     ptr_ = &view_.vector;
  }

  GslVectorBase& operator=(const gsl_vector* x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const GslVectorBase& x) {
    GslVectorBase::operator=(x);
    return *this;
  }

  GslVectorBase& operator=(const double v) {
    GslVectorBase::operator=(v);
    return *this;
  }
 private:
  gsl_vector_view view_;
  GslSubvector(const GslSubvector&) { }
};

#endif  // __MATH_GSL_VECTOR__


#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "parmt_utils.h"
#ifdef PARMT_USE_INTEL
#include <mkl_cblas.h>
#else
#include <cblas.h>
#endif
#include "iscl/memory/memory.h"

/*!
 * @brief Converts the NED based fundamental faults to Green's functions which
 *        scale the NED based moment tensor terms s.t. \f$ u = G m \f$.
 *        where \f$ m = \{ m_{xx}, m_{yy}, m_{zz}, m_{xy}, m_{xz}, m_{yz} \} \f$
 *        are the NED moment tensor terms and \f$ u \f$ is the estimate for
 *        the icomp'th component with source receiver azimuth, az, and given
 *        channel orientation encapsulated by cmpaz and cmpinc.
 *
 * @param[in] npgrns     number of points in greens functions
 * @param[in] ldg        leading dimension of G (>= 6)
 * @param[in] icomp      =1 for vertical Green's functions.
 *                       =2 for north (channel 2) Green's functions.
 *                       =3 for east (channel 3) Green's functions.
 * @param[in] azin       source to receiver azimuth (degrees)
 * @param[in] bazin      receiver to source azimuth (degrees).  if you
 *                       do not know then set this to:
 *                         fmod(az + 180.0, 360.0)
 * @param[in] cmpaz      component azimuth (0 north, +90 east)
 * @param[in] cmpinc     component inclinantion (-90 up, 0 east/north, +90 down)
 * @param[in] ZDS        vertical greens fn for 90 degree dip slip [npgrns]
 * @param[in] ZSS        vertical greens fn for vertical strike slip [npgrns]
 * @param[in] ZDD        vertical greens fn for 45 degree dip slip [npgrns]
 * @param[in] ZEX        vertical greens fn for an explosion [npgrns]
 * @param[in] RDS        radial greens fn for 90 degree dip slip [npgrns]
 * @param[in] RSS        radial greens fn for vertical strike slip [npgrns]
 * @param[in] RDD        radial greens fn for 45 degree dip slip [npgrns]
 * @param[in] REX        radial greens fn for an explosion [npgrns]
 * @param[in] TDS        transverse greens fn for 90 degree dip slip [npgrns]
 * @param[in] TSS        transverse greens fn for vertical strike slip [npgrns]
 *
 * @param[out] G         row major Green's functions matrix [npgrns x ldg].
 *                       the columns are ordered:
 *                        \f$ \{ G_{xx}, G_{yy}, G_{zz},
 *                               G_{xy}, G_{xz}, G_{yz} \}  \f$
 *
 * @result 0 indicates success
 *
 * @author Ben Baker
 *
 * @copyright ISTI distributed under Apache 2
 *
 */
int parmt_utils_ff2mtGreensMatrix64f(const int npgrns, const int ldg,
                                     const int icomp,
                                     const double azin, const double bazin,
                                     const double cmpaz, const double cmpinc,
                                     const double *__restrict__ ZDS,
                                     const double *__restrict__ ZSS,
                                     const double *__restrict__ ZDD,
                                     const double *__restrict__ ZEX,
                                     const double *__restrict__ RDS,
                                     const double *__restrict__ RSS,
                                     const double *__restrict__ RDD,
                                     const double *__restrict__ REX,
                                     const double *__restrict__ TDS,
                                     const double *__restrict__ TSS,
                                     double *__restrict__ G)
{
    double *Gxx, *Gyy, *Gzz, *Gxy, *Gxz, *Gyz;
    int ierr;
    ierr = 0;
    if (npgrns < 1 || ldg < 6 || G == NULL)
    {
        if (npgrns < 1)
        {
            fprintf(stderr, "%s: No points in greens fns\n", __func__);
        }
        if (ldg < 6){fprintf(stderr, "%s: ldg must be >= 6\n", __func__);}
        if (G == NULL){fprintf(stderr, "%s: Error G is NULL\n", __func__);}
        return -1;
    }
    Gxx = memory_calloc64f(npgrns);
    Gyy = memory_calloc64f(npgrns);
    Gzz = memory_calloc64f(npgrns);
    Gxy = memory_calloc64f(npgrns);
    Gxz = memory_calloc64f(npgrns);
    Gyz = memory_calloc64f(npgrns);
    ierr = parmt_utils_ff2mtGreens64f(npgrns, icomp, azin, bazin,
                                      cmpaz, cmpinc,
                                      ZDS, ZSS, ZDD, ZEX,
                                      RDS, RSS, RDD, REX,
                                      TDS, TSS,
                                      Gxx, Gyy, Gzz,
                                      Gxy, Gxz, Gyz);
    if (ierr != 0)
    {
        fprintf(stderr, "%s: Failed to compute greens functions columns\n",
                __func__);
    }
    else
    {
        cblas_dcopy(npgrns, Gxx, 1, &G[0], ldg);
        cblas_dcopy(npgrns, Gyy, 1, &G[1], ldg);
        cblas_dcopy(npgrns, Gzz, 1, &G[2], ldg);
        cblas_dcopy(npgrns, Gxy, 1, &G[3], ldg);
        cblas_dcopy(npgrns, Gxz, 1, &G[4], ldg);
        cblas_dcopy(npgrns, Gyz, 1, &G[5], ldg);
    }
    memory_free64f(&Gxx);
    memory_free64f(&Gyy);
    memory_free64f(&Gzz);
    memory_free64f(&Gxy);
    memory_free64f(&Gxz);
    memory_free64f(&Gyz);
    return ierr;
}
//============================================================================//
/*!
 * @brief Converts the NED based fundamental faults to Green's functions which
 *        scale the NED based moment tensor terms s.t. \f$ u = G m \f$.
 *        where \f$ m = \{ m_{xx}, m_{yy}, m_{zz}, m_{xy}, m_{xz}, m_{yz} \} \f$
 *        are the NED moment tensor terms and \f$ u \f$ is the estimate for
 *        the icomp'th component with source receiver azimuth, az, and given
 *        channel orientation encapsulated by cmpaz and cmpinc.
 *
 * @param[in] npgrns     number of points in greens functions
 * @param[in] icomp      =1 for vertical Green's functions.
 *                       =2 for north (channel 2) Green's functions.
 *                       =3 for east (channel 3) Green's functions.
 * @param[in] azin       source to receiver azimuth (degrees)
 * @param[in] bazin      receiver to source azimuth (degrees).  if you
 *                       do not know then set this to:
 *                         fmod(az + 180.0, 360.0)
 * @param[in] cmpaz      component azimuth (0 north, +90 east)
 * @param[in] cmpinc     component inclinantion (-90 up, 0 east/north, +90 down)
 * @param[in] ZDS        vertical greens fn for 90 degree dip slip [npgrns]
 * @param[in] ZSS        vertical greens fn for vertical strike slip [npgrns]
 * @param[in] ZDD        vertical greens fn for 45 degree dip slip [npgrns]
 * @param[in] ZEX        vertical greens fn for an explosion [npgrns]
 * @param[in] RDS        radial greens fn for 90 degree dip slip [npgrns]
 * @param[in] RSS        radial greens fn for vertical strike slip [npgrns]
 * @param[in] RDD        radial greens fn for 45 degree dip slip [npgrns]
 * @param[in] REX        radial greens fn for an explosion [npgrns]
 * @param[in] TDS        transverse greens fn for 90 degree dip slip [npgrns]
 * @param[in] TSS        transverse greens fn for vertical strike slip [npgrns]
 *
 * @param[out] Gxx       greens function scaling mxx moment tensor term [npgrns]
 * @param[out] Gyy       greens function scaling myy moment tensor term [npgrns]
 * @param[out] Gzz       greens function scaling mzz moment tensor term [npgrns]
 * @param[out] Gxy       greens function scaling mxy moment tensor term [npgrns]
 * @param[out] Gxz       greens function scaling mxz moment tensor term [npgrns]
 * @param[out] Gyz       greens function scaling myz moment tensor term [npgrns]
 *
 * @result 0 indicates success
 *
 * @author Ben Baker
 *
 * @copyright ISTI distributed under Apache 2
 *
 */
int parmt_utils_ff2mtGreens64f(const int npgrns, const int icomp,
                               const double azin, const double bazin,
                               const double cmpaz, const double cmpinc,
                               const double *__restrict__ ZDS,
                               const double *__restrict__ ZSS,
                               const double *__restrict__ ZDD,
                               const double *__restrict__ ZEX,
                               const double *__restrict__ RDS,
                               const double *__restrict__ RSS,
                               const double *__restrict__ RDD,
                               const double *__restrict__ REX,
                               const double *__restrict__ TDS,
                               const double *__restrict__ TSS,
                               double *__restrict__ Gxx,
                               double *__restrict__ Gyy,
                               double *__restrict__ Gzz,
                               double *__restrict__ Gxy,
                               double *__restrict__ Gxz,
                               double *__restrict__ Gyz)
{
    double az, baz, c2a, ca, caz, cos_caz, cost,
           gxx_e, gxy_e, gxz_e, gyy_e, gyz_e, gzz_e,
           gxx_n, gxy_n, gxz_n, gyy_n, gyz_n, gzz_n,
           gxx_r, gxy_r, gxz_r, gyy_r, gyz_r, gzz_r,
           gxx_t, gxy_t, gxz_t, gyy_t, gyz_t, gzz_t,
           gxx_z, gxy_z, gxz_z, gyy_z, gyz_z, gzz_z,
           rex, rdd, rds, rss, s2a, sa, sin_caz, sint,
           tds, tss, xscal,
           zex, zdd, zds, zss;
    int i;
    const double pi180 = M_PI/180.0;
    const double half = 0.5;
    const double sixth = 1.0/6.0;
    const double third = 1.0/3.0;
    //------------------------------------------------------------------------//
    if (icomp < 1 || icomp > 3)
    {
        fprintf(stderr, "%s: Invalid component: %d\n", __func__, icomp);
        return -1;
    }
    if (icomp == 1)
    {
        if (ZDS == NULL || ZSS == NULL || ZDD == NULL || ZEX == NULL)
        {
            if (ZDS == NULL){fprintf(stderr, "%s: zds is NULL\n", __func__);}
            if (ZSS == NULL){fprintf(stderr, "%s: zss is NULL\n", __func__);}
            if (ZDD == NULL){fprintf(stderr, "%s: zdd is NULL\n", __func__);}
            if (ZEX == NULL){fprintf(stderr, "%s: zex is NULL\n", __func__);}
            return -1;
        }
    }
    else
    {
        if (RDS == NULL || RSS == NULL || RDD == NULL || REX == NULL ||
            TDS == NULL || TSS == NULL)
        {
            if (RDS == NULL){fprintf(stderr, "%s: rds is NULL\n", __func__);}
            if (RSS == NULL){fprintf(stderr, "%s: rss is NULL\n", __func__);}
            if (RDD == NULL){fprintf(stderr, "%s: rdd is NULL\n", __func__);}
            if (REX == NULL){fprintf(stderr, "%s: rex is NULL\n", __func__);}
            if (TDS == NULL){fprintf(stderr, "%s: tds is NULL\n", __func__);}
            if (TSS == NULL){fprintf(stderr, "%s: tss is NULL\n", __func__);}
            return -1;
        }
    }
    // Get the backazimuth
    az = azin;
    baz = bazin; //fmod(az + 180.0, 360.0);
    // Convert to radians
    az = az*pi180;
    baz = baz*pi180;
    caz = cmpaz*pi180;
    // Compute the geometric factors
    sa = sin(az);
    ca = cos(az);
    s2a = sin(2.*az);
    c2a = cos(2.*az);
    sint = sin(baz);
    cost = cos(baz);
    cos_caz = cos(caz);
    sin_caz = sin(caz);
    // the cmpaz will come from the metadata so fix that
    if (icomp == 3)
    {
       cos_caz = cos(caz - M_PI/2.0);
       sin_caz = sin(caz - M_PI/2.0);
    }
    // Set the columns (Minson et. al. 2008)
    xscal = 1.0;
    if (fabs(cmpinc - 90.0) < 1.e-4){xscal =-1.0;}
    if (icomp == 1)
    {
        #pragma omp simd
        for (i=0; i<npgrns; i++)
        {
            zss = ZSS[i];
            zdd = ZDD[i];
            zds = ZDS[i];
            zex = ZEX[i];
            // compute the greens function corresponding to the mt
            gxx_z = half*(zss*c2a) - sixth*zdd + third*zex;
            gyy_z =-half*(zss*c2a) - sixth*zdd + third*zex;
            gzz_z = third*(zdd + zex);
            gxy_z = zss*s2a;
            gxz_z = zds*ca;
            gyz_z = zds*sa;
            // copy and include polarity so it matches the observation
            Gxx[i] = xscal*gxx_z;
            Gyy[i] = xscal*gyy_z;
            Gzz[i] = xscal*gzz_z;
            Gxy[i] = xscal*gxy_z;
            Gxz[i] = xscal*gxz_z;
            Gyz[i] = xscal*gyz_z;
        } // Loop on points
    }
    else
    {
        // North or 2 component
        if (icomp == 2)
        {
            // Loop on rows (data points)
            //#pragma omp simd
            for (i=0; i<npgrns; i++)
            {
                rss = RSS[i];
                rdd = RDD[i];
                rds = RDS[i];
                rex = REX[i];
                tds = TDS[i];
                tss = TSS[i];

                gxx_r = half*(rss*c2a) - sixth*rdd + third*rex;
                gyy_r =-half*(rss*c2a) - sixth*rdd + third*rex;
                gzz_r = third*(rdd + rex);
                gxy_r = rss*s2a;
                gxz_r = rds*ca;
                gyz_r = rds*sa;

                gxx_t = half*tss*s2a;
                gyy_t =-half*tss*s2a;
                gzz_t = 0.0;
                gxy_t =-tss*c2a;
                gxz_t = tds*sa;
                gyz_t =-tds*ca;

                // Rotate from (r, t) -> (n, e)
                //   n = cos(baz-180)*r - sin(baz-180)*t
                //     =-cos(baz)*r + sin(baz)*t
                //   e = sin(baz-180)*r + cos(baz-180)*t
                //     =-sin(baz)*r - cos(baz)*t
                gxx_n =-cost*gxx_r + sint*gxx_t;
                gxx_e =-sint*gxx_r - cost*gxx_t;
                gyy_n =-cost*gyy_r + sint*gyy_t;
                gyy_e =-sint*gyy_r - cost*gyy_t;
                gzz_n =-cost*gzz_r + sint*gzz_t;
                gzz_e =-sint*gzz_r - cost*gzz_t;
                gxy_n =-cost*gxy_r + sint*gxy_t;
                gxy_e =-sint*gxy_r - cost*gxy_t;
                gxz_n =-cost*gxz_r + sint*gxz_t;
                gxz_e =-sint*gxz_r - cost*gxz_t;
                gyz_n =-cost*gyz_r + sint*gyz_t;
                gyz_e =-sint*gyz_r - cost*gyz_t;
                // Rotate from (n, e) -> (1, 2) around az
                //   1 = cos(comp_az)*n + sin(comp_az)*e
                Gxx[i] = cos_caz*gxx_n + sin_caz*gxx_e;
                Gyy[i] = cos_caz*gyy_n + sin_caz*gyy_e;
                Gzz[i] = cos_caz*gzz_n + sin_caz*gzz_e;
                Gxy[i] = cos_caz*gxy_n + sin_caz*gxy_e;
                Gxz[i] = cos_caz*gxz_n + sin_caz*gxz_e;
                Gyz[i] = cos_caz*gyz_n + sin_caz*gyz_e;

            } // Loop on points
        }
        // East or 3 component
        else
        {
            // Loop on rows (data points)
            #pragma omp simd
            for (i=0; i<npgrns; i++)
            {
                rss = RSS[i];
                rdd = RDD[i];
                rds = RDS[i];
                rex = REX[i];
                tds = TDS[i];
                tss = TSS[i];

                gxx_r = half*(rss*c2a) - sixth*rdd + third*rex;
                gyy_r =-half*(rss*c2a) - sixth*rdd + third*rex;
                gzz_r = third*(rdd + rex);
                gxy_r = rss*s2a;
                gxz_r = rds*ca;
                gyz_r = rds*sa;

                gxx_t = half*tss*s2a;
                gyy_t =-half*tss*s2a;
                gzz_t = 0.0;
                gxy_t =-tss*c2a;
                gxz_t = tds*sa;
                gyz_t =-tds*ca;

                // Rotate from (r, t) -> (n, e)
                //   n = cos(baz-180)*r - sin(baz-180)*t
                //     =-cos(baz)*r + sin(baz)*t
                //   e = sin(baz-180)*r + cos(baz-180)*t
                //     =-sin(baz)*r - cos(baz)*t
                gxx_n =-cost*gxx_r + sint*gxx_t;
                gxx_e =-sint*gxx_r - cost*gxx_t;
                gyy_n =-cost*gyy_r + sint*gyy_t;
                gyy_e =-sint*gyy_r - cost*gyy_t;
                gzz_n =-cost*gzz_r + sint*gzz_t;
                gzz_e =-sint*gzz_r - cost*gzz_t;
                gxy_n =-cost*gxy_r + sint*gxy_t;
                gxy_e =-sint*gxy_r - cost*gxy_t;
                gxz_n =-cost*gxz_r + sint*gxz_t;
                gxz_e =-sint*gxz_r - cost*gxz_t;
                gyz_n =-cost*gyz_r + sint*gyz_t;
                gyz_e =-sint*gyz_r - cost*gyz_t;
                // Rotate from (n, e) -> (1, 2) around az
                //   2 =-sin(comp_az)*n + cos(comp_az)*e
                Gxx[i] =-sin_caz*gxx_n + cos_caz*gxx_e;
                Gyy[i] =-sin_caz*gyy_n + cos_caz*gyy_e;
                Gzz[i] =-sin_caz*gzz_n + cos_caz*gzz_e;
                Gxy[i] =-sin_caz*gxy_n + cos_caz*gxy_e;
                Gxz[i] =-sin_caz*gxz_n + cos_caz*gxz_e;
                Gyz[i] =-sin_caz*gyz_n + cos_caz*gyz_e;
            } // Loop on points
        } // End check on 1 or 2 component
    } // End check on component
    return 0;
}


/* These tests are based on the NIST Statistical Reference Datasets
   See http://www.nist.gov/itl/div898/strd/index.html for more
   information. */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_test.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_multifit_nlin.h>
#include <gsl/gsl_blas.h>

#include <gsl/gsl_ieee_utils.h>

#include "test_longley.c"
#include "test_filip.c"
#include "test_pontius.c"
#include "test_brown.c"
#include "test_enso.c"
#include "test_kirby2.c"
#include "test_hahn1.c"
#include "test_nelson.c"
#include "test_fn.c"
#include "test_estimator.c"

void
test_lmder (gsl_multifit_function_fdf * f, double x0[],
            double * X, double F[], double * cov);

void
test_fdf (const char * name, gsl_multifit_function_fdf * f,
          double x0[], double x[], double sumsq,
          double sigma[]);

int
main (void)
{
  gsl_ieee_env_setup();

  test_longley();
  test_filip();
  test_pontius();
  test_estimator();

  {
    gsl_multifit_function_fdf f = make_fdf (&brown_f, &brown_df, &brown_fdf,
                                            brown_N, brown_P, 0);

    test_lmder(&f, brown_x0, &brown_X[0][0], brown_F, &brown_cov[0][0]);
  }

  {
    gsl_multifit_function_fdf f = make_fdf (&enso_f, &enso_df, &enso_fdf,
                                            enso_N, enso_P, 0);

    test_fdf("nist-ENSO", &f, enso_x0, enso_x, enso_sumsq, enso_sigma);
  }

  {
    gsl_multifit_function_fdf f = make_fdf (&kirby2_f, &kirby2_df, &kirby2_fdf,
                                            kirby2_N, kirby2_P, 0);

    test_fdf("nist-kirby2", &f, kirby2_x0, kirby2_x, kirby2_sumsq, kirby2_sigma);
  }

  {
    gsl_multifit_function_fdf f = make_fdf (&hahn1_f, &hahn1_df, &hahn1_fdf,
                                            hahn1_N, hahn1_P, 0);

    test_fdf("nist-hahn1", &f, hahn1_x0, hahn1_x, hahn1_sumsq, hahn1_sigma);
  }

#ifdef JUNK
  {
    gsl_multifit_function_fdf f = make_fdf (&nelson_f, &nelson_df, &nelson_fdf,
                                            nelson_N, nelson_P, 0);

    test_fdf("nist-nelson", &f, nelson_x0, nelson_x, nelson_sumsq, nelson_sigma);
  }
#endif

  /* now summarize the results */

  exit (gsl_test_summary ());
}

void
test_lmder (gsl_multifit_function_fdf * f, double x0[],
            double * X, double F[], double * cov)
{
  const gsl_multifit_fdfsolver_type *T;
  gsl_multifit_fdfsolver *s;

  const size_t n = f->n;
  const size_t p = f->p;

  int status;
  size_t iter = 0, i;

  gsl_vector_view x = gsl_vector_view_array (x0, p);

  T = gsl_multifit_fdfsolver_lmsder;
  s = gsl_multifit_fdfsolver_alloc (T, n, p);
  gsl_multifit_fdfsolver_set (s, f, &x.vector);

  do
    {
      status = gsl_multifit_fdfsolver_iterate (s);

      for (i = 0 ; i < p; i++)
        {
          gsl_test_rel (gsl_vector_get (s->x, i), X[p*iter+i], 1e-5,
                        "lmsder, iter=%u, x%u", iter, i);
        }

      gsl_test_rel (gsl_blas_dnrm2 (s->f), F[iter], 1e-5,
                    "lmsder, iter=%u, f", iter);

      iter++;
    }
  while (iter < 20);

  {
    size_t i, j;
    gsl_matrix * covar = gsl_matrix_alloc (4, 4);
    gsl_multifit_covar (s->J, 0.0, covar);

    for (i = 0; i < 4; i++)
      {
        for (j = 0; j < 4; j++)
          {
            gsl_test_rel (gsl_matrix_get(covar,i,j), cov[i*p + j], 1e-7,
                          "gsl_multifit_covar cov(%d,%d)", i, j) ;
          }
      }

    gsl_matrix_free (covar);
  }

  gsl_multifit_fdfsolver_free (s);

}

void
test_fdf (const char * name, gsl_multifit_function_fdf * f,
          double x0[], double x_final[],
          double f_sumsq, double sigma[])
{
  const gsl_multifit_fdfsolver_type *T;
  gsl_multifit_fdfsolver *s;

  const size_t n = f->n;
  const size_t p = f->p;

  int status;
  size_t iter = 0;

  gsl_vector_view x = gsl_vector_view_array (x0, p);

  T = gsl_multifit_fdfsolver_lmsder;
  s = gsl_multifit_fdfsolver_alloc (T, n, p);
  gsl_multifit_fdfsolver_set (s, f, &x.vector);

  do
    {
      status = gsl_multifit_fdfsolver_iterate (s);

#ifdef DEBUG
       printf("iter = %d  status = %d  |f| = %.18e x = \n",
         iter, status, gsl_blas_dnrm2 (s->f));

         gsl_vector_fprintf(stdout, s->x, "%.8e");
#endif
      status = gsl_multifit_test_delta (s->dx, s->x, 0.0, 1e-7);

      iter++;
    }
  while (status == GSL_CONTINUE && iter < 1000);

  {
    size_t i;
    gsl_matrix * covar = gsl_matrix_alloc (p, p);
    gsl_multifit_covar (s->J, 0.0, covar);

    for (i = 0 ; i < p; i++)
      {
        gsl_test_rel (gsl_vector_get (s->x, i), x_final[i], 1e-5,
                      "%s, lmsder, x%u", name, i);
      }

    {
      double s2 = pow(gsl_blas_dnrm2 (s->f), 2.0);

      gsl_test_rel (s2, f_sumsq, 1e-5, "%s, lmsder, |f|^2", name);

      for (i = 0; i < p; i++)
        {
          double ei = sqrt(s2/(n-p))*sqrt(gsl_matrix_get(covar,i,i));
          gsl_test_rel (ei, sigma[i], 1e-4,
                        "%s, sigma(%d)", name, i) ;
        }
    }

    gsl_matrix_free (covar);
  }

  gsl_multifit_fdfsolver_free (s);
}