#include "AdvectionDiffusionSolverFE.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 <CompuCell3D/Field3D/Field3DIO.h>
#include <BasicUtils/BasicClassGroup.h>

#include <CompuCell3D/plugins/NeighborTracker/NeighborTrackerPlugin.h>



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




using namespace CompuCell3D;
using namespace std;



//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
AdvectionDiffusionSolverFE::AdvectionDiffusionSolverFE()
 : DiffusableGraph<float>()
{
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
AdvectionDiffusionSolverFE::~AdvectionDiffusionSolverFE()
{

}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::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();
  
  NeighborTrackerPlugin * neighborTrackerPlugin=(NeighborTrackerPlugin*)(Simulator::pluginManager.get("NeighborTracker"));
  neighborTrackerAccessorPtr= neighborTrackerPlugin->getNeighborTrackerAccessorPtr();
  
  ///setting member function pointers
  diffusePtr=&AdvectionDiffusionSolverFE::diffuse;
  secretePtr=&AdvectionDiffusionSolverFE::secrete;
  
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::extraInit(Simulator *simulator){

   
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
double AdvectionDiffusionSolverFE::computeAverageCellRadius(){
   double sumRadius=0.0;
   double radical=1.0/3.0;
   int numberOfCells=0;
   CellG *cell;
   CellInventory::cellInventoryIterator cInvItr;
   
   for(cInvItr=cellInventoryPtr->cellInventoryBegin() ; cInvItr !=cellInventoryPtr->cellInventoryEnd() ;++cInvItr ){
         
      cell=*cInvItr;
      sumRadius+=pow(cell->volume,radical);
      ++numberOfCells;
      
   }

   return 1.33*sumRadius/numberOfCells;
   
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::update(){

   for(unsigned int i = 0 ; i <= diffDataVec.size() ; ++i ){
      updateLocalCellInventory(i);
   }

}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::updateLocalCellInventory(unsigned int idx){

   CellG *cell;
   CellInventory::cellInventoryIterator cInvItr;
   
   map<CellG*,float> * concentrationField=concentrationFieldMapVector[idx];
   map<CellG*,float>::iterator concentrationItr;
   map<CellG*,float>::iterator concentrationEndItr=concentrationField->end();
   
   //first eliminate from concentration field all the cells that do not appear any longer in cell inventory
   for(concentrationItr = concentrationField->begin() ; concentrationItr != concentrationField->end() ; /*++ concentrationItr*/ ){
      cInvItr=cellInventoryPtr->find(concentrationItr->first);
      if(cInvItr==cellInventoryPtr->cellInventoryEnd()){
         concentrationField->erase(concentrationItr++);//erasing cell from concentrationField
      }else{
         concentrationItr++;
      }
   }
      
   for(cInvItr=cellInventoryPtr->cellInventoryBegin() ; cInvItr !=cellInventoryPtr->cellInventoryEnd() ;++cInvItr ){
         
      cell=*cInvItr;
      concentrationItr=concentrationField->find(cell);
         if(concentrationItr==concentrationEndItr){//no such cell addfress int he map
            concentrationField->insert(make_pair(cell,0.0));//inserting new cell with default concentration in the field 0
         }
      
      
   }


}



//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::start() {
   
   update();   
   cerr<<"GOT HERE BEFORE INITIALIZE FIELD"<<endl;
   initializeConcentration();
   
   cerr<<"GOT HERE AFTER INITIALIZE FIELD"<<endl;
}

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

void AdvectionDiffusionSolverFE::initializeConcentration(){


   for(unsigned int i = 0 ; i <diffDataVec.size() ; ++i){
      if(diffDataVec[i].concentrationFileName.empty()) continue;

      readConcentrationField(diffDataVec[i].concentrationFileName,concentrationFieldVector[i]);
      field2CellMap(concentrationFieldVector[i] , concentrationFieldMapVector[i] );
      cellMap2Field(concentrationFieldMapVector[i] , concentrationFieldVector[i]);
   }


   


}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::readConcentrationField(std::string fileName,Field3DImpl<float> *concentrationField){

   ifstream in(fileName.c_str());

   ASSERT_OR_THROW(string("Could not open chemical concentration file '") +
                      fileName + "'!", in.is_open());

   Point3D pt;
   float c;
   while(!in.eof()){
      in>>pt.x>>pt.y>>pt.z>>c;
      if(!in.fail())
         concentrationField->set(pt,c);
   }

}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::cellMap2Field(std::map<CellG *, float> *concentrationMapField , Field3DImpl<float> *concentrationField){

CellG *currentCellPtr=0;
Point3D pt;
float currentConcentration;
std::map<CellG *, float>::iterator mitr;



   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);
            
            mitr=concentrationMapField->find(currentCellPtr);
            if(mitr != concentrationMapField->end() ){
               currentConcentration = mitr->second;
//                cerr<<"curent concentration="<<currentConcentration <<endl;
            }else{
               currentConcentration =0.0;
//                cerr<<"no such cell"<<endl;
            }
            concentrationField->set(pt,currentConcentration);
      }
   
}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::field2CellMap(Field3DImpl<float> *concentrationField , std::map<CellG *, float> *concentrationMapField ){

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

   std::map<CellG *, float>::iterator mitr;

   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);
            
            mitr=concentrationMapField->find(currentCellPtr);
            currentConcentration=concentrationField->get(pt);

            
            if(mitr != concentrationMapField->end() ){
                
                mitr->second=currentConcentration;
            }else{
               //cerr<<"No such cell"<<endl;
            }
            
      }




}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::step(const unsigned int _currentStep) {
   currentStep=_currentStep;
   update(); //updates cell inventories of particular fields
   cellMap2Field(concentrationFieldMapVector[0] ,concentrationFieldVector[0] );
   (this->*secretePtr)();
   (this->*diffusePtr)();
   
   
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////


//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::secreteSingleField(unsigned int idx){

   SecretionDataFlexAD & secrData=secrDataVec[idx];
   map<CellG*,float> * concentrationField=concentrationFieldMapVector[idx];
   map<CellG*,float>::iterator concentrationItr;
   map<CellG*,float>::iterator concentrationEndItr=concentrationField->end();
   
   std::map<unsigned char,float>::iterator mitr;
   std::map<unsigned char,float>::iterator end_mitr=secrData.typeIdSecrConstMap.end();

   float currentConcentration;
   float secrConst;

   

   //the assumption is that medium has type ID 0
   mitr=secrData.typeIdSecrConstMap.find(automaton->getTypeId("Medium"));


   CellG *cell;
   CellInventory::cellInventoryIterator cInvItr;
   
   for(cInvItr=cellInventoryPtr->cellInventoryBegin() ; cInvItr !=cellInventoryPtr->cellInventoryEnd() ;++cInvItr ){
         
      cell=*cInvItr;
      
      concentrationItr = concentrationField->find(cell);
      

         
         mitr=secrData.typeIdSecrConstMap.find(cell->type);
         if(mitr!=end_mitr){
            
            secrConst=mitr->second;
            currentConcentration=concentrationItr->second;
            concentrationItr->second=currentConcentration+secrConst;
            
         }

            
   }
      

}

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

void AdvectionDiffusionSolverFE::secreteOnContactSingleField(unsigned int idx){

   SecretionDataFlexAD & secrData=secrDataVec[idx];
   map<CellG*,float> * concentrationField=concentrationFieldMapVector[idx];
   map<CellG*,float>::iterator concentrationItr;
   map<CellG*,float>::iterator concentrationEndItr=concentrationField->end();

   std::map<unsigned char,SecretionOnContactData>::iterator mitr;
   std::map<unsigned char,SecretionOnContactData>::iterator end_mitr=secrData.typeIdSecrOnContactDataMap.end();
   map<unsigned char, float> * contactCellMapPtr;
   std::map<unsigned char, float>::iterator mitrTypeConst;
   
   set<NeighborSurfaceData>::iterator neighborItr;
   NeighborTracker * neighborTracker;
   
   
   float currentConcentration;
   float secrConst;

   

   CellG *cell;
   CellG *nCellPtr;
   
   CellInventory::cellInventoryIterator cInvItr;
   
   for(cInvItr=cellInventoryPtr->cellInventoryBegin() ; cInvItr !=cellInventoryPtr->cellInventoryEnd() ;++cInvItr ){
         
      cell=*cInvItr;
      
      concentrationItr = concentrationField->find(cell);
      

         mitr=secrData.typeIdSecrOnContactDataMap.find(cell->type);
         
         //cerr<<"cell->type="<<(int)cell->type<<endl;
         if(mitr!=end_mitr){
            contactCellMapPtr = &(mitr->second.contactCellMap);
            neighborTracker = neighborTrackerAccessorPtr->get(cell->extraAttribPtr);
            for(
                  neighborItr = neighborTracker->cellNeighbors.begin() ;
                  neighborItr != neighborTracker->cellNeighbors.end() ;
                  ++neighborItr
            ){
               nCellPtr = neighborItr->neighborAddress;
               if(!nCellPtr) continue;
               
               mitrTypeConst=contactCellMapPtr->find(nCellPtr->type);
               if(mitrTypeConst != contactCellMapPtr->end()){//OK to secrete, contact detected
                  secrConst = mitrTypeConst->second;
                  currentConcentration=concentrationItr->second;
                  concentrationItr->second=currentConcentration+secrConst;
               }

                              
            }
            
         }

            
   }
      

}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::diffuseSingleField(unsigned int idx){

   map<CellG*,float> * concentrationField=concentrationFieldMapVector[idx];
   map<CellG*,float>::iterator concentrationItr;
   map<CellG*,float>::iterator nConcentrationItr;
   map<CellG*,float>::iterator scratchConcentrationItr;
   map<CellG*,float>::iterator concentrationEndItr=concentrationField->end();
   map<CellG*,float> * scratchField=concentrationFieldMapVector[diffDataVec.size()];

   set<unsigned char>::iterator sitr;
   set<unsigned char>::iterator end_sitr=diffDataVec[idx].avoidTypeIdSet.end();

   
   NeighborTracker * neighborTracker;
   set<NeighborSurfaceData>::iterator neighborItr;
   
   DiffusionData & diffData = diffDataVec[idx];
   CellG * currentCellPtr;
   CellG * nCellPtr;

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

   float diffConst=diffData.diffConst;
   float decayConst=diffData.decayConst;
   float deltaT=diffData.deltaT;
   float dt_dx2=deltaT/( averageRadius*averageRadius );

   float currentConcentration=0.0;
   short neighborCounter=0;

   CellInventory::cellInventoryIterator cInvItr;
   
   for(concentrationItr = concentrationField->begin() ; concentrationItr != concentrationField->end() ; ++ concentrationItr ){
   
      currentCellPtr = concentrationItr->first;
      currentConcentration = concentrationItr->second;
      
      scratchConcentrationItr = scratchField->find(currentCellPtr);
      
      if(diffData.avoidTypeIdSet.find(currentCellPtr->type) != end_sitr){
         scratchConcentrationItr->second = currentConcentration;
         continue;
      }

      
      //find concentrationSum
      concentrationSum=0.0;
      neighborCounter=0;

//       cInvItr=cellInventoryPtr->find(currentCellPtr);
//          if(cInvItr == cellInventoryPtr->cellInventoryEnd()){
//             cerr<<"ERROR - current cell cell: "<<*cInvItr<<" not in the inventory"<<endl;
//             cerr<<"probably internal module error"<<endl;
//             exit(0);
//          }
// 

            
      neighborTracker = neighborTrackerAccessorPtr->get(currentCellPtr->extraAttribPtr);
         
      

            
      for(neighborItr = neighborTracker->cellNeighbors.begin() ; neighborItr != neighborTracker->cellNeighbors.end() ; ++neighborItr){

         nCellPtr=neighborItr->neighborAddress;
         //check if it is still in neighbor boundary
/*         cInvItr=cellInventoryPtr->find(nCellPtr);
         if(cInvItr == cellInventoryPtr->cellInventoryEnd()){
            cerr<<"ERROR - cell: "<<*cInvItr<<" not in the inventory"<<endl;
            exit(0);
         }            */
         
         if(nCellPtr && diffData.avoidTypeIdSet.find(nCellPtr->type) != end_sitr){
            continue;
         }

         nConcentrationItr=concentrationField->find(nCellPtr);
         if(nConcentrationItr != concentrationEndItr){
            concentrationSum += nConcentrationItr->second;
            ++neighborCounter;
         }
      }
      updatedConcentration =  dt_dx2*diffConst*(concentrationSum - neighborCounter*currentConcentration)
                           -deltaT*(decayConst*currentConcentration)
                           +currentConcentration;

      scratchConcentrationItr->second = updatedConcentration;
      
   }
   
   scrarch2Concentration(scratchField,concentrationField);
   cellMap2Field(concentrationField ,concentrationFieldVector[idx] );
      
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::scrarch2Concentration( map<CellG *, float> * scratchField , map<CellG *, float> *concentrationField){
   
   map<CellG*,float>::iterator concentrationItr;
   map<CellG*,float>::iterator scratchConcentrationItr;

   for(concentrationItr = concentrationField->begin() ; concentrationItr != concentrationField->end() ; ++ concentrationItr ){
      scratchConcentrationItr=scratchField->find(concentrationItr->first);
      concentrationItr->second = scratchConcentrationItr->second;
      scratchConcentrationItr->second=0.0;
   }

}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::diffuse() {
   averageRadius = computeAverageCellRadius();
   for(unsigned int i = 0 ; i < diffDataVec.size() ; ++i ){
      diffuseSingleField(i);
   }

      

}


void AdvectionDiffusionSolverFE::secrete() {



   for(unsigned int i = 0 ; i < diffDataVec.size() ; ++i ){
      for(unsigned int j = 0 ; j <secrDataVec[i].secretionFcnPtrVec.size() ; ++j){
         (this->*secrDataVec[i].secretionFcnPtrVec[j])(i);
      }
   
   }
      

}

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



//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void AdvectionDiffusionSolverFE::readXML(XMLPullParser &in) {
  in.skip(TEXT);
  string concentrationFieldName;
  unsigned int concentrationFieldCounter=0;

  float diffConst;
  float decayConst;
  float secrConst;
  float deltaX;
  float deltaT;

   while(in.check(START_ELEMENT)) {
      if (in.getName() == "DiffusionField") {
         cerr<<"Got Diffusion field tag"<<endl;
         diffDataVec.push_back(DiffusionData());
         secrDataVec.push_back(SecretionDataFlexAD());
         DiffusionData & diffData=diffDataVec[diffDataVec.size()-1];
         SecretionData & secrData=secrDataVec[secrDataVec.size()-1];
         diffData.setAutomaton(potts->getAutomaton());
         secrData.setAutomaton(potts->getAutomaton());
         
         in.match(START_ELEMENT);
         in.skip(TEXT);
       
         while(in.check(START_ELEMENT)){
            cerr<<"AFTER INNER LOOP MATCH START_ELEMENT"<<endl;
            if(in.getName() == "DiffusionData"){
               cerr<<"Will read diffusionData"<<endl;
               
               in.match(START_ELEMENT);
               diffData.readXML(in);
               in.match(END_ELEMENT);
               //in.skip(TEXT);
               cerr<<"FINISHED reading diffusionData"<<endl;
            }
            else if (in.getName() == "SecretionData"){
               cerr<<"Will read secrData"<<endl;
               in.match(START_ELEMENT);
               secrData.readXML(in);
               in.match(END_ELEMENT);
               cerr<<"FINISHED read secrData"<<endl;
            }
            else{
               throw BasicException(string("Unexpected element '") + in.getName() +
                                 "'!", in.getLocation());
            }
           in.skip(TEXT);
         }
         in.skip(TEXT);
         in.match(END_ELEMENT);

         
      }
      in.skip(TEXT);
   }
   


   cerr<<"AdvectionDiffusionSolverFE: extra Init in read XML"<<endl;

   ///allocate fields including scrartch field
   allocateDiffusableFieldVector(diffDataVec.size()+1,fieldDim); //+1 is for additional scratch field
   
   ///first "translate" type names to type id's
//    Automaton *automaton=potts->getAutomaton();
//    for(unsigned int i = 0 ; i < diffDataVec.size() ; ++i){
//       for(unsigned int j = 0 ; j < diffDataVec[i].avoidTypeVec.size() ; ++j){
//          diffDataVec[i].avoidTypeIdSet.insert(automaton->getTypeId(diffDataVec[i].avoidTypeVec[j]));
//       }
//       for(unsigned int j = 0 ; j < diffDataVec[i].secreteTypeVec.size() ; ++j){
//          diffDataVec[i].secreteTypeIdSet.insert(automaton->getTypeId(diffDataVec[i].secreteTypeVec[j]));
//       }
//    }

   ///assign vector of field names
   concentrationFieldNameVector.assign(diffDataVec.size(),string(""));

//    cerr<<"diffDataVec.size()="<<diffDataVec.size()<<endl;
   
   for(unsigned int i = 0 ; i < diffDataVec.size() ; ++i){
      concentrationFieldNameVector[i] = diffDataVec[i].fieldName;
//       cerr<<" concentrationFieldNameVector[i]="<<concentrationFieldNameVector[i]<<endl;
   }
   
   //register fields once they have been allocated
   for(unsigned int i = 0 ; i < diffDataVec.size() ; ++i){
      simPtr->registerConcentrationField(concentrationFieldNameVector[i] , concentrationFieldVector[i]);
      cerr<<"registring field: "<<concentrationFieldNameVector[i]<<" field address="<<concentrationFieldVector[i]<<endl;
   }

   
   ///assigning member method ptrs to the vector
   
   for(unsigned int i = 0 ; i < secrDataVec.size() ; ++i){
       
       secrDataVec[i].secretionFcnPtrVec.assign(secrDataVec[i].secrTypesNameSet.size(),0);
      unsigned int j=0;
      for(set<string>::iterator sitr=secrDataVec[i].secrTypesNameSet.begin() ; sitr != secrDataVec[i].secrTypesNameSet.end()  ; ++sitr){
         if((*sitr)=="Secretion"){
             secrDataVec[i].secretionFcnPtrVec[j]=&AdvectionDiffusionSolverFE::secreteSingleField;
            ++j;
         }
         else if((*sitr)=="SecretionOnContact"){
             secrDataVec[i].secretionFcnPtrVec[j]=&AdvectionDiffusionSolverFE::secreteOnContactSingleField;
            ++j;
         }
      }      
   }
   

   cerr<<"ALLOCATED ALL FIELDS"<<endl;


      

   
   //exit(0);
    

}



