#ifndef CHARGE_REGULATOR_H
#define CHARGE_REGULATOR_H

#include "AtomicRegulator.h"
#include "PoissonSolver.h"
#include <map>
#include <utility>
#include <string>
#include <vector>

namespace ATC {

  /**
   *  @class  ChargeRegulator
   *  @brief  Manager class for atom-continuum control of charge and potential
   */

  class ChargeRegulatorMethod;
  class ChargeRegulator : public AtomicRegulator {

  public:

    /** charge regulator types */
    enum ChargeRegulatorType {
      NONE=0,
      FEEDBACK,
      IMAGE_CHARGE,
      EFFECTIVE_CHARGE
    };
    enum ChargedSurfaceType {
      INSULATOR=0,
      DIELECTRIC,
      CONDUCTOR
    };

    /** parser data */
    struct ChargeRegulatorParameters {
      ChargeRegulatorParameters():
        method(NONE),
        value(0),
        groupBit(0),
        depth(0),
        permittivity(0),
        surfaceType(INSULATOR) { };
      ChargeRegulatorType method;
      double value;
      int groupBit;
      std::string groupTag;
      double depth;
      double permittivity;
      ChargedSurfaceType surfaceType;
      FSET faceset;
    };

    /** constructor */
    ChargeRegulator(ATC_Coupling *atc);

    /** destructor */
    ~ChargeRegulator();

    /** parser/modifier */
    virtual bool modify(int narg, char **arg);
    virtual void construct_methods();
    virtual void initialize();
    virtual void apply_pre_force(double dt);
    virtual void apply_post_force(double dt);
    virtual void output(OUTPUT_LIST & outputData) const;
    virtual double compute_vector(int /* n */) const {return 0;} // TODO

    void assign_poisson_solver(PoissonSolver * solver) { poissonSolver_ = solver;}
    PoissonSolver * poisson_solver(void) { return poissonSolver_;}

  protected:

    /** poisson solver */
    PoissonSolver * poissonSolver_;


    /** registry charge regulators */
    std::map<std::string,ChargeRegulatorMethod *> regulators_;
    std::map<std::string,ChargeRegulatorParameters> parameters_;

  private:
    ChargeRegulator(); // DO NOT define this
  };

  /**
   *  @class  ChargeRegulatorMethod
   *  @brief  Base class for implementation of ChargeRegulator algorithms
   */

  class ChargeRegulatorMethod : public RegulatorShapeFunction {

  public:
    ChargeRegulatorMethod(ChargeRegulator *chargeRegulator,
      ChargeRegulator::ChargeRegulatorParameters & parameters);
    ~ChargeRegulatorMethod(){};
    virtual void initialize(void);
    void set_greens_functions();
    virtual void apply_pre_force(double /* dt */){};
    virtual void apply_post_force(double /* dt */){};
    virtual void set_weights() {};
    const DENS_VEC & total_influence() const { return sum_;}
    virtual void output(OUTPUT_LIST & outputData);

  protected:
    int nlocal();
    /** pointer to charge regulator object for data */
    ChargeRegulator * chargeRegulator_;
    /** interscale manager */
    class InterscaleManager * interscaleManager_;
    /** direct interface to lammps */
    class LammpsInterface * lammpsInterface_;
    /** poisson solver */
    PoissonSolver * poissonSolver_;
    /** (subset) of Green's functions */
    std::vector<SparseVector<double> > greensFunctions_;
    /** cutoff radius */
    double rC_,rCsq_;
    /** conversion constants */
    double qV2e_, qqrd2e_;
    /** member data */
    XT_Function * targetValue_;
    double targetPhi_;
    // controlling
    FSET surface_;
    NSET nodes_;
    int atomGroupBit_;
    bool boundary_; // atoms on boundary
    DENS_VEC point_;
    DENS_VEC normal_;
    double depth_;
    ChargeRegulator::ChargedSurfaceType surfaceType_;
    double permittivity_;
    bool initialized_;
    DENS_VEC sum_;
    /** workspace */
    mutable DENS_MAT _atomElectricalForce_;
  private:
    ChargeRegulatorMethod(); // DO NOT define this
  };

  /**
   *  @class  ChargeRegulatorMethodFeedback
   *  @brief  ChargeRegulator with feedback and ad hoc positions
   *          can be used for conductor (1 charged sheet) or insulator (2 sheets)
   */

  class ChargeRegulatorMethodFeedback : public ChargeRegulatorMethod {

  public:

    /** constructor */
    ChargeRegulatorMethodFeedback(ChargeRegulator *chargeRegulator,
      ChargeRegulator::ChargeRegulatorParameters & parameters);

    /** destructor */
    ~ChargeRegulatorMethodFeedback(){};

    /** instantiate all needed data */
    virtual void construct_transfers();

    /** initialize */
    virtual void initialize(void);

    /** set influence nodes and atoms */
    void set_influence();

    void set_influence_matrix(void);

    /** post first verlet step */
    virtual void apply_pre_force(double dt);

    void apply_feedback_charges();

  protected:

    int nControlNodes_;
    NSET & controlNodes_;
    // measurement/controlled
    int influenceGroupBit_;
    int nInfluenceAtoms_;
    NSET influenceAtoms_, influenceAtomsIds_;
    int nInfluenceNodes_;
    NSET influenceNodes_;


    DENS_MAT invG_;
    DENS_MAT invNNT_;
    DENS_MAT NT_;
    DENS_MAT NTinvNNTinvG_;

  private:
    ChargeRegulatorMethodFeedback(); // DO NOT define this
  };

  /**
   *  @class  ChargeRegulatorMethodImageCharge
   *  @brief  ChargeRegulator with image charges
   */

  class ChargeRegulatorMethodImageCharge : public ChargeRegulatorMethod {

  public:

    /** constructor */
    ChargeRegulatorMethodImageCharge(ChargeRegulator *chargeRegulator,
      ChargeRegulator::ChargeRegulatorParameters & parameters);

    /** destructor */
    ~ChargeRegulatorMethodImageCharge(){};

    /** initialize */
    virtual void initialize(void);

    /** post first verlet step */
    virtual void apply_post_force(double dt);

  protected:
    double reflect(DENS_VEC & x) const {
      double dn = (x-point_).dot(normal_);
      x -= 2*dn*normal_;
      return dn;
    }
    // internal functions
    void apply_local_forces();
    void correct_charge_densities();
    void correct_forces();

    double layerDepth_;
    double permittivityRatio_;
    NSET & imageNodes_;
    DENS_MAT imageTransferOp_;

  private:
    ChargeRegulatorMethodImageCharge(); // DO NOT define this
  };

  /**
   *  @class  ChargeRegulatorMethodEffectiveCharge
   *  @brief  ChargeRegulator with effective charges at FE quadrature pts
   */

  class ChargeRegulatorMethodEffectiveCharge : public ChargeRegulatorMethod {

  typedef std::map<int,std::pair<DENS_VEC,double> > NODE_TO_XF_MAP;

  public:

    /** constructor */
    ChargeRegulatorMethodEffectiveCharge(ChargeRegulator *chargeRegulator,
      ChargeRegulator::ChargeRegulatorParameters & parameters);

    /** destructor */
    ~ChargeRegulatorMethodEffectiveCharge(){};

    /** initialize */
    virtual void initialize(void);

    /** post first verlet step */
    virtual void apply_post_force(double dt);

  protected:
    // internal functions
    void apply_local_forces();

    // member data
    double chargeDensity_;
    std::map<int,double> nodalChargePotential_;

    NODE_TO_XF_MAP nodeXFMap_;

    bool useSlab_;


  private:
    ChargeRegulatorMethodEffectiveCharge(); // DO NOT define this
  };

};

#endif