#include "ReactionDiffusionSolverFE_SavHog.h"

#include <CompuCell3D/Simulator.h>
#include <CompuCell3D/Automaton/Automaton.h>
#include <CompuCell3D/Potts3D/Potts3D.h>
#include <CompuCell3D/Potts3D/CellInventory.h>
#include <CompuCell3D/Field3D/WatchableField3D.h>
#include <CompuCell3D/Field3D/Field3DImpl.h>
#include <CompuCell3D/Field3D/Field3D.h>
#include <BasicUtils/BasicClassGroup.h>



#include <XMLCereal/XMLPullParser.h>
#include <XMLCereal/XMLSerializer.h>
#include <BasicUtils/BasicString.h>
#include <BasicUtils/BasicException.h>
#include <BasicUtils/BasicRandomNumberGenerator.h>
#include <string>
#include <cmath>
#include <iostream>

using namespace CompuCell3D;
using namespace std;
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
ReactionDiffusionSolverFE_SavHog::ReactionDiffusionSolverFE_SavHog()
 : DiffusableVector<float>(),deltaX(1.0),deltaT(1.0)
{
   imposeDiffusionBox=false;
//    C1=20;
//    C2=3;
//    C3=15;
//    c1=0.0065;
//    c2=0.841;
//    a=0.15;
//    k=3.5;
//    b=0.35;
//    eps1=0.5;
//    eps2=0.0589;
//    eps3=0.5;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
ReactionDiffusionSolverFE_SavHog::~ReactionDiffusionSolverFE_SavHog()
{

}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void ReactionDiffusionSolverFE_SavHog::init(Simulator *simulator) {
  simPtr=simulator;
  potts = simulator->getPotts();
  automaton=potts->getAutomaton();
   
  ///getting cell inventory
  cellInventoryPtr=& potts->getCellInventory(); 
  
  ///getting field ptr from Potts3D
  cellFieldG=potts->getCellFieldG();
  fieldDim=cellFieldG->getDim();
  
  

  
  ///setting member function pointers
  diffusePtr=&ReactionDiffusionSolverFE_SavHog::diffuse;



  
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void ReactionDiffusionSolverFE_SavHog::start() {
//     if(diffConst> (1.0/6.0-0.05) ){ //hard coded condtion for stability of the solutions - assume dt=1 dx=dy=dz=1
// 
//       cerr<<"CANNOT SOLVE DIFFUSION EQUATION: STABILITY PROBLEM - DIFFUSION CONSTANT TOO LARGE. EXITING..."<<endl;
//       exit(0);
// 
//    }


   dt_dx2=deltaT/(deltaX*deltaX);
   initializeConcentration();
   
   
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

void ReactionDiffusionSolverFE_SavHog::initializeConcentration(){

/*   Point3D p;
   float c;
   for (p.z = 0; p.z < fieldDim.z; p.z++)
      for (p.y = 0; p.y < fieldDim.y; p.y++)
         for (p.x = 0; p.x < fieldDim.x; p.x++){
         if(p.z<10 || sqrt(1.0*(p.x-40)*(p.x-40)+(p.y-40)*(p.y-40))<10.0){
            
            c=exp((-(p.x-40)*(p.x-40)-(p.y-40)*(p.y-40)) / 1000.0)*p.z*0.04;
         }else{
            c=0.0;
         }
         concentrationFieldVector[0]->set(p,c);
         
      }

   return;*/
   
   BasicRandomNumberGenerator *rand = BasicRandomNumberGenerator::getInstance();
   

   CellInventory::cellInventoryIterator cInvItr;

   CellG * cell;
   map<CellG *,float> concentrationMap;
   map<CellG *,float> refractorinessMap;
/*   float concentrationMin=4.5;//0.5;
   float concentrationMax=5.0;//1.0;

   float refractorinessMin=-2.0;//1.0;
   float refractorinessMax=1.0;//2.0;*/
   float concentrationMin=0.5;
   float concentrationMax=1.0;
   
   float refractorinessMin=1.0;
   float refractorinessMax=2.0;
   
   float randomConcentration;
   float randomRefractoriness;
   
   
   float x,y,z;
   
   for(cInvItr=cellInventoryPtr->cellInventoryBegin() ; cInvItr !=cellInventoryPtr->cellInventoryEnd() ;++cInvItr ){

      cell=*cInvItr;
      if(cell->type!=automaton->getTypeId("Ground") && cell->type!=automaton->getTypeId("Wall")){
      // do not put anything in the groundCell or wall
         randomConcentration=fabs(concentrationMax-concentrationMin)*rand->getRatio()+concentrationMin;
         randomRefractoriness=fabs(refractorinessMax-refractorinessMin)*rand->getRatio()+refractorinessMin;
         
         concentrationMap.insert(make_pair(cell,randomConcentration));
         refractorinessMap.insert(make_pair(cell,randomRefractoriness));
      }
      
   }
   Point3D pt;
   CellG * currentCellPtr;
   map<CellG *,float>::iterator concentrationMapItr;
   map<CellG *,float>::iterator refractorinessMapItr;
   
   //once we have decided about the initial values we need to initialize lattice
   for (pt.z = 0; pt.z < fieldDim.z; pt.z++)
      for (pt.y = 0; pt.y < fieldDim.y; pt.y++)
         for (pt.x = 0; pt.x < fieldDim.x; pt.x++){
            
            currentCellPtr=cellFieldG->get(pt);
            if(!currentCellPtr) continue;///medium does not participate in FN equations
            
            //only autocycling cells will have concentration and refractoriness initialized
            if(currentCellPtr->type!=automaton->getTypeId("Autocycling"))continue;
            

            concentrationMapItr=concentrationMap.find(currentCellPtr);
            if(concentrationMapItr!=concentrationMap.end()){///setting all pixels belonging to given cell to same concentr.
               concentrationFieldVector[0]->set(pt,concentrationMapItr->second);
            }

            refractorinessMapItr=refractorinessMap.find(currentCellPtr);
            if(refractorinessMapItr!=refractorinessMap.end()){///setting all pixels belonging to given cell to same refract.
               concentrationFieldVector[1]->set(pt,refractorinessMapItr->second);
            }

                        
   }



}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void ReactionDiffusionSolverFE_SavHog::step(const unsigned int _currentStep) {
   currentStep=_currentStep;
   
   (this->*diffusePtr)();


//    Point3D p;
//    float c;
//    for (p.z = 0; p.z < fieldDim.z; p.z++)
//       for (p.y = 0; p.y < fieldDim.y; p.y++)
//          for (p.x = 0; p.x < fieldDim.x; p.x++){
//          if(p.z<10 || sqrt(1.0*(p.x-40)*(p.x-40)+(p.y-40)*(p.y-40))<10.0){
//             
//             c=exp((-(p.x-40)*(p.x-40)-(p.y-40)*(p.y-40)) / 1000.0)*p.z*0.04;
//          }else{
//             c=0.0;
//          }
//          concentrationFieldVector[0]->set(p,c);
//          
//       }
// 

   
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void ReactionDiffusionSolverFE_SavHog::diffuse() {
/// 'n' denotes neighbor

///this is the diffusion equation 
///C_{xx}+C_{yy}+C_{zz}=(1/a)*C_{t}
///a - diffusivity - diffConst

///Finite difference method:
///T_{0,\delta \tau}=F*\sum_{i=1}^N T_{i}+(1-N*F)*T_{0}
///N - number of neighbors
///will have to double check this formula


Point3D pt, n;
unsigned int token = 0;
double distance;
CellG * currentCellPtr=0,*nCell=0;

short currentCellType=0;
float concentrationSum=0.0;
float updatedConcentration=0.0;

float currentRefract=0.0;//refractoriness
float updatedRefract=0.0;//refractoriness

float currentConcentration=0.0;
short neighborCounter=0;


float F;/// F=a*\delta \tau / (\delta x)^2
        /// in the simplest case (\delta*\tau=1 \delta x=1) F=a;

        
   F=diffConst;
                    
   for (pt.z = 0; pt.z < fieldDim.z; pt.z++)
    for (pt.y = 0; pt.y < fieldDim.y; pt.y++)
      for (pt.x = 0; pt.x < fieldDim.x; pt.x++){

         currentCellPtr=cellFieldG->get(pt);
         /// go to the next iteration in case you get Medium cell - medium does not participate in the diffusion
         ///(model dependent may be changed if necessary
         if(!currentCellPtr){
            //continue;
            
//          if(pt.z >maxDiffusionZ) continue;
            
            if(imposeDiffusionBox)
               if(!insideDiffusionBox(pt)) continue;
               
          }     
         updatedConcentration=0.0;
         updatedRefract=0.0;
         concentrationSum=0.0;
         neighborCounter=0;
         token=0; ///have to zero the token every time we access neighbors - otherwise the neighbor distance will keep growing up
         ///loop over nearest neighbors
         while (true) {
            n = cellFieldG->getNeighbor(pt, token, distance, false);
            ///cerr<<"pt="<<pt<<" n="<<n<<" distance="<<distance<<" token="<<token<<endl;
             
            
            if (distance > 1) break;

                           
               nCell = cellFieldG->get(n); //pulls out cell from a field

               /// this has to be AFTER if (distance > 1) break; otherwise there might be performance problems
               /// because  if (distance > 1) break; could be skipped even when distance is >1
               
//             if(n.z>maxDiffusionZ && !nCell) continue; /// will not diffuse to medium only above maxDiffusionZ
               if(imposeDiffusionBox)
                  if(!insideDiffusionBox(n) && !nCell) continue;


                                             

                              
               //if(!nCell) continue; /// assume that medium does not participate in the diffusion
                                    ///(model dependent may be changed if necessary)

               
               ///cerr<<"neighborPoint="<<neighborPoint<<" nCell="<<nCell<<endl;
               concentrationSum += concentrationFieldVector[0]->get(n);
               ++neighborCounter;
               //cerr<<"neighborCounter="<<neighborCounter<<endl;      
         }
         
    
         currentConcentration = concentrationFieldVector[0]->get(pt);
         currentRefract = concentrationFieldVector[1]->get(pt);

         if(
            currentCellPtr &&
            (
               currentCellPtr->type==automaton->getTypeId("Autocycling")||
               currentCellPtr->type==automaton->getTypeId("Prespore")||
               currentCellPtr->type==automaton->getTypeId("Prestalk")
            )
            
         ){
         
            updatedConcentration =  dt_dx2*F*(concentrationSum - neighborCounter*currentConcentration)
                                    -deltaT*(f(currentConcentration)+currentRefract)
                                    +currentConcentration;
         }
         else {
            if(currentCellPtr && currentCellPtr->type==automaton->getTypeId("Wall")){
               updatedConcentration =currentConcentration;
            }else {

               updatedConcentration =  dt_dx2*F*(concentrationSum - neighborCounter*currentConcentration)
                                    +currentConcentration;
            }
         
         }
         
/*         if(updatedConcentration!=0.0){
            cerr<<"updatedConcentration="<<updatedConcentration<<endl;
         }*/
         //cerr<<" OUT neighborCounter="<<neighborCounter<<" F*cSum="<<F*concentrationSum<<" updC="<<updatedConcentration<<endl;
/*         if(updatedConcentration>2){
            cerr<<"updatedConcentration="<<updatedConcentration<<endl;
            exit(0);
         }*/
         concentrationFieldVector[2]->set(pt,updatedConcentration);///updating scratch
         //cerr<<"F="<<F<<" concentrationSum="<<concentrationSum<<" updated Concentr="<<updatedConcentration<<endl;

         ///Refractoriness
         
         if(currentCellPtr && currentCellPtr->type==automaton->getTypeId("Autocycling")){
            updatedRefract = deltaT*eps(currentConcentration)*(k*currentConcentration-b-currentRefract)+currentRefract;
            concentrationFieldVector[1]->set(pt,updatedRefract);///updating Refractoriness
         }
         else if(currentCellPtr &&
                  (currentCellPtr->type==automaton->getTypeId("Prestalk") || currentCellPtr->type==automaton->getTypeId("Prespore") ))
         {
            updatedRefract = deltaT*eps(currentConcentration)*(k*currentConcentration-currentRefract)+currentRefract;///remove -b
            concentrationFieldVector[1]->set(pt,updatedRefract);///updating Refractoriness
         }
         else{
            
         }

      }
      ///copy updated values from scratch to concentration field   
      scrarch2Concentration(concentrationFieldVector[2],concentrationFieldVector[0]);
      

}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void ReactionDiffusionSolverFE_SavHog::scrarch2Concentration( Field3DImpl<float> *scratchField, Field3DImpl<float> *concentrationField){

Point3D pt;
CellG * currentCellPtr=0;
float tempValue;

   for (pt.z = 0; pt.z < fieldDim.z; pt.z++)
    for (pt.y = 0; pt.y < fieldDim.y; pt.y++)
      for (pt.x = 0; pt.x < fieldDim.x; pt.x++){
         
         tempValue=scratchField->get(pt);
         concentrationField->set(pt,tempValue);
         scratchField->set(pt,0.0);
         currentCellPtr = cellFieldG->get(pt);
         if(!currentCellPtr && pt.z>maxDiffusionZ){
            concentrationFieldVector[1]->set(pt,0.0);
         }
      }
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void ReactionDiffusionSolverFE_SavHog::outputField( std::ostream & _out, Field3DImpl<float> *_concentrationField){
Point3D pt;
float tempValue;

   for (pt.z = 0; pt.z < fieldDim.z; pt.z++)
    for (pt.y = 0; pt.y < fieldDim.y; pt.y++)
      for (pt.x = 0; pt.x < fieldDim.x; pt.x++){
         tempValue=_concentrationField->get(pt);
         _out<<pt<<" "<<tempValue<<endl;
      }
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void ReactionDiffusionSolverFE_SavHog::readXML(XMLPullParser &in) {
  in.skip(TEXT);
  string concentrationFieldName;
  unsigned int concentrationFieldCounter=0;

    
  while (in.check(START_ELEMENT)) {
    if (in.getName() == "DiffusionConstant") {
      diffConst = BasicString::parseDouble(in.matchSimple());

    }
    else if (in.getName() == "DecayConstant") {
      decayConst = BasicString::parseDouble(in.matchSimple());

    }
    else if (in.getName() == "DeltaX") {
      deltaX = BasicString::parseDouble(in.matchSimple());
    }
    else if (in.getName() == "DeltaT") {
      deltaT = BasicString::parseDouble(in.matchSimple());
    }
    else if (in.getName() == "MaxDiffusionZ") {
       maxDiffusionZ= BasicString::parseUInteger(in.matchSimple());
    }

    else if (in.getName() == "fFunctionParameters") {
       C1 = BasicString::parseDouble(in.getAttribute("C1").value);
       C2 = BasicString::parseDouble(in.getAttribute("C2").value);
       C3 = BasicString::parseDouble(in.getAttribute("C3").value);
       a = BasicString::parseDouble(in.getAttribute("a").value);
       in.matchSimple();
    }

    else if (in.getName() == "epsFunctionParameters") {
       eps1 = BasicString::parseDouble(in.getAttribute("eps1").value);
       eps2 = BasicString::parseDouble(in.getAttribute("eps2").value);
       eps3 = BasicString::parseDouble(in.getAttribute("eps3").value);
       in.matchSimple();
    }
    else if (in.getName() == "IntervalParameters") {
       c1 = BasicString::parseDouble(in.getAttribute("c1").value);
       c2 = BasicString::parseDouble(in.getAttribute("c2").value);
       in.matchSimple();
    }
    else if (in.getName() == "RefractorinessParameters") {
       k = BasicString::parseDouble(in.getAttribute("k").value);
       b = BasicString::parseDouble(in.getAttribute("b").value);
       in.matchSimple();
    }

    else if (in.getName() == "MinDiffusionBoxCorner") {
      minDiffusionBoxCorner.readXML(in);
      in.matchSimple();
    }
    
    else if (in.getName() == "MaxDiffusionBoxCorner") {
      maxDiffusionBoxCorner.readXML(in);
      in.matchSimple();
    }

                
    else if (in.getName() == "NumberOfFields") {
       numberOfFields = BasicString::parseUInteger(in.matchSimple());
      ///constructing concentration and scratch fields
      allocateDiffusableFieldVector(3, fieldDim);/// c,r, scratch fields
    }
    else if (in.getName() == "FieldName") {
      concentrationFieldName = in.matchSimple();
      concentrationFieldNameVector[concentrationFieldCounter] = concentrationFieldName;
      ///register concentration field to make it accesible from external viewers etc.

      ASSERT_OR_THROW("You are trying to define more field that you declared!",concentrationFieldCounter<=numberOfFields);
      
      simPtr->registerConcentrationField(concentrationFieldName , concentrationFieldVector[concentrationFieldCounter]);
      ++concentrationFieldCounter;
    }
    
     else {
      throw BasicException(string("Unexpected element '") + in.getName() + 
            "'!", in.getLocation());
    }

    in.skip(TEXT);
  }

  if(minDiffusionBoxCorner==maxDiffusionBoxCorner)
      imposeDiffusionBox=false;
  else
      imposeDiffusionBox=true;

/*  cerr<<" c1="<<c1<<" c2="<<c2<<" C1="<<C1<<" C2="<<C2<<" C3="<<C3<<endl;
  cerr<<" a="<<a<<" b="<<b<<endl;
  cerr<<" eps1="<<eps1<<" eps2="<<eps2<<" eps3="<<eps3<<" k="<<k<<endl;
  exit(0);*/
  

}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
