/*
 * Copyright (C) 2008-2010 Sébastien Villemot <sebastien.villemot@ens.fr>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/*
  Country indices j go from 0 to n-1

  Variables is a vector of size 5*n+1:
  - c^j at j
  - l^j at n+j
  - i^j at 2*n+j
  - k^j at 3*n+j
  - a^j at 4*n+j
  - lambda at 5*n

  Shocks is a vector of size n+1:
  - e^j at j
  - e at n

  Errors is a vector of size 5*n+1 (same indices than Michel's May 2007 specif, except that error of aggregate resource constraint is put at the end (in position 5*n).
*/

#ifndef _MODEL_SPEC_HH
#define _MODEL_SPEC_HH

#include <ostream>

#include <gsl/gsl_vector.h>

//! Possible specifications of residual of aggregate resource constraint
enum arc_resid_t
  {
    MichelARC,
    PaulARC,
    BenARC
  };

class ModelSpec
{
  friend class SmolSolution; // For accessing marginal_utilities()
  friend class MRGalSolution;
  friend class MMJSolution;
private:
  //! Which residual specification to be used in aggregate resource constraint?
  const arc_resid_t arc_resid;
protected:
  //! Will be filled by constructor of derived class
  double *taus;
  void write_common_params(std::ostream &output) const;
  double weighted_mass(double min, double max, int j) const;
  virtual void marginal_utilities(const double y_curr[], double uc[], double ul[]) const = 0;
  /*! We need previous and current variables because production is done using capital of previous period and labor of current period */
  virtual void production(const double y_prev[], const double y_curr[], double f[], double fk[], double fl[]) const = 0;
public:
  //! Number of countries
  const int n;
  const double alpha, beta, delta, rho, sigma, phi, A;

  ModelSpec(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, arc_resid_t arc_resid_arg);
  virtual ~ModelSpec();
  void steady_state(gsl_vector *y_ss) const;
  //! Computes all relative errors
  void errors(const gsl_vector *y_prev, const gsl_vector *y_curr, const gsl_vector *fwd, const gsl_vector *shocks, gsl_vector *errors) const;
  //! Computes the maximum error (in absolute value)
  double max_error(const gsl_vector *y_prev, const gsl_vector *y_curr, const gsl_vector *fwd, const gsl_vector *shocks) const;
  //! Computes the terms inside the expectancy operator of Euler equation (see Michel's May 2007 note)
  /*! fwd is therefore an array of size n */
  void forward_part(const gsl_vector *y_prev, const gsl_vector *y_curr, const gsl_vector *y_next, const gsl_vector *shocks, gsl_vector *fwd) const;
  virtual void write_output(std::ostream &output) const = 0;
  void write_state(std::ostream &output, const gsl_vector *y) const;
  void write_errors(std::ostream &output, const gsl_vector *errs) const;
  virtual int spec_no() const = 0;
  virtual bool labor() const { return true; };
};

//! Parts common to A1 and A5
class ModelSpecA1A5 : public ModelSpec
{
  friend class MRGalSolution;
  friend class MMJSolution;
protected:
  double *gammas;
  virtual void marginal_utilities(const double y_curr[], double uc[], double ul[]) const;
  virtual void production(const double y_prev[], const double y_curr[], double f[], double fk[], double fl[]) const;
public:
  ModelSpecA1A5(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, arc_resid_t arc_resid_arg);
  ~ModelSpecA1A5();
  virtual bool labor() const { return false; };
};

//! Symmetric Capital-labour Cobb-Douglas production function (common to A2, A3, A6, A7)
class ModelSpecSKLCB : public ModelSpec
{
protected:
  virtual void production(const double y_prev[], const double y_curr[], double f[], double fk[], double fl[]) const;
public:
  ModelSpecSKLCB(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, arc_resid_t arc_resid_arg);
};

//! Parts common to A2 and A6
class ModelSpecA2A6 : public ModelSpecSKLCB
{
  friend class MRGalSolution;
  friend class MMJSolution;
protected:
  double *b, *gammas, *etas;
  virtual void marginal_utilities(const double y_curr[], double uc[], double ul[]) const;
public:
  ModelSpecA2A6(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, arc_resid_t arc_resid_arg);
  virtual ~ModelSpecA2A6();
};

//! Parts common to A3 and A7
class ModelSpecA3A7 : public ModelSpecSKLCB
{
  friend class MRGalSolution;
  friend class MMJSolution;
protected:
  double *gammas;
  const double time_endowment, psi;
  virtual void marginal_utilities(const double y_curr[], double uc[], double ul[]) const;
public:
  ModelSpecA3A7(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double time_endowment_arg, arc_resid_t arc_resid_arg);
  virtual ~ModelSpecA3A7();
};

//! Parts common to A4 and A8
class ModelSpecA4A8 : public ModelSpec
{
  friend class MRGalSolution;
  friend class MMJSolution;
protected:
  double *gammas, *mus, *xis, *b;
  virtual void marginal_utilities(const double y_curr[], double uc[], double ul[]) const;
  virtual void production(const double y_prev[], const double y_curr[], double f[], double fk[], double fl[]) const;
public:
  const double time_endowment;
  ModelSpecA4A8(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double time_endowment_arg, arc_resid_t arc_resid_arg);
  virtual ~ModelSpecA4A8();
};

class ModelSpecA1 : public ModelSpecA1A5
{
public:
  const double gamma;
  ModelSpecA1(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double gamma_arg, arc_resid_t arc_resid_arg);
  virtual void write_output(std::ostream &output) const;
  virtual int spec_no() const;
};

class ModelSpecA2 : public ModelSpecA2A6
{
public:
  const double gamma, eta;
  ModelSpecA2(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double gamma_arg, double eta_arg, arc_resid_t arc_resid_arg);
  virtual void write_output(std::ostream &output) const;
  virtual int spec_no() const;
};

class ModelSpecA3 : public ModelSpecA3A7
{
public:
  const double gamma;
  ModelSpecA3(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double time_endowment_arg, double gamma_arg, arc_resid_t arc_resid_arg);
  virtual void write_output(std::ostream &output) const;
  virtual int spec_no() const;
};

class ModelSpecA4 : public ModelSpecA4A8
{
public:
  const double gamma, mu, xi;
  ModelSpecA4(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double time_endowment_arg, double gamma_arg, double mu_arg, double xi_arg, arc_resid_t arc_resid_arg);
  virtual void write_output(std::ostream &output) const;
  virtual int spec_no() const;
};

class ModelSpecA5 : public ModelSpecA1A5
{
public:
  const double gamma_min, gamma_max;
  ModelSpecA5(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double gamma_min_arg, double gamma_max_arg, arc_resid_t arc_resid_arg);
  virtual void write_output(std::ostream &output) const;
  virtual int spec_no() const;
};

class ModelSpecA6 : public ModelSpecA2A6
{
public:
  const double gamma_min, gamma_max, eta_min, eta_max;
  ModelSpecA6(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double gamma_min_arg, double gamma_max_arg, double eta_min_arg, double eta_max_arg, arc_resid_t arc_resid_arg);
  virtual void write_output(std::ostream &output) const;
  virtual int spec_no() const;
};

class ModelSpecA7 : public ModelSpecA3A7
{
public:
  const double gamma_min, gamma_max;
  ModelSpecA7(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double time_endowment_arg, double gamma_min_arg, double gamma_max_arg, arc_resid_t arc_resid_arg);
  virtual void write_output(std::ostream &output) const;
  virtual int spec_no() const;
};

class ModelSpecA8 : public ModelSpecA4A8
{
public:
  const double gamma_min, gamma_max, mu_min, mu_max, xi_min, xi_max;
  ModelSpecA8(int n_arg, double alpha_arg, double beta_arg, double delta_arg, double rho_arg, double sigma_arg, double phi_arg, double time_endowment_arg, double gamma_min_arg, double gamma_max_arg, double mu_min_arg, double mu_max_arg, double xi_min_arg, double xi_max_arg, arc_resid_t arc_resid_arg);
  virtual void write_output(std::ostream &output) const;
  virtual int spec_no() const;
};

#endif // !_MODEL_SPEC_HH
