#include "FlexibleDiffusionSolverFE.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 <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>


//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// std::ostream & operator<<(std::ostream & out,CompuCell3D::DiffusionData & diffData){
//    
//    
// }


using namespace CompuCell3D;
using namespace std;



//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
FlexibleDiffusionSolverFE::FlexibleDiffusionSolverFE()
 : DiffusableVector<float>(),deltaX(1.0),deltaT(1.0)
{
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
FlexibleDiffusionSolverFE::~FlexibleDiffusionSolverFE()
{

}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::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=&FlexibleDiffusionSolverFE::diffuse;
  secretePtr=&FlexibleDiffusionSolverFE::secrete;
  
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::extraInit(Simulator *simulator){

   
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::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);
// 
//    }

   
//cerr<<"GOT HERE BEFORE INITIALIZE FIELD"<<endl;
   dt_dx2=deltaT/(deltaX*deltaX);
   initializeConcentration();
   
//   cerr<<"GOT HERE AFTER INITIALIZE FIELD"<<endl;
}

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

void FlexibleDiffusionSolverFE::initializeConcentration(){

//    for(unsigned int i = 0 ; i <diffDataVec.size() ; ++i){
//       if(diffDataVec[i].concentrationFileName.empty()) continue;
//          
//       ifstream cFile(diffDataVec[i].concentrationFileName.c_str());
//       ASSERT_OR_THROW(string("Could not open chemical concentration file '") +
//                       diffDataVec[i].concentrationFileName + "'!", cFile.is_open());
//       cFile>>*concentrationFieldVector[i];      
//    }


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

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


   

//    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++){
// 
// 
//             c=p.x*p.y;
// 
// 
//             for(unsigned int i = 0 ; i < diffDataVec.size() ; ++i){
//                concentrationFieldVector[i]->set(p,c);
//             }
//       }
// 
//    return;



   

//    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++){
//          
//             
//             c=exp((-(p.x-fieldDim.x/2)*(p.x-fieldDim.x/2)-(p.y-fieldDim.y/2)*(p.y-fieldDim.y/2)) / 1000.0)*p.z*0.04;
//          
//            
//             for(unsigned int i = 0 ; i < diffDataVec.size() ; ++i){
//                concentrationFieldVector[i]->set(p,c);
//             }
//       }
// 
//    return;
//    
// 

}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::step(const unsigned int _currentStep) {
   //cout << "IN" << endl;
    currentStep=_currentStep;
   //cout << "SECRETE" << endl;  
   (this->*secretePtr)();
   //cout << "DIFFUSE" << endl;
   (this->*diffusePtr)();
   
   //cout << "OUT" << endl;
}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::secreteOnContactSingleField(unsigned int idx){
//cerr<<"\n\n\tsecrete on contact "<<idx<<"\n\n"<<endl;

SecretionData & secrData=secrDataVec[idx];

// if(!secrData.active)
//    return;

std::map<unsigned char,SecretionOnContactData>::iterator mitr;
std::map<unsigned char,SecretionOnContactData>::iterator end_mitr=secrData.typeIdSecrOnContactDataMap.end();

CellG *currentCellPtr;
Field3DImpl<float> * concentrationField=concentrationFieldVector[idx];
float currentConcentration;
float secrConst;
float secrConstMedium=0.0;
std::map<unsigned char, float> * contactCellMapMediumPtr;
std::map<unsigned char, float> * contactCellMapPtr;
std::map<unsigned char, float>::iterator mitrTypeConst;

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

if( mitr != end_mitr ){
   secreteInMedium=true;
   contactCellMapMediumPtr = &(mitr->second.contactCellMap);
}



Point3D pt;

unsigned int token = 0;
double distance = 0;
Point3D n;
CellG *nCell=0;
Field3D<CellG *> *fieldG = potts->getCellFieldG();
unsigned char type;



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

                        
            if(secreteInMedium && ! currentCellPtr){
                  token = 0;
                  distance = 0;
                  //set changePt so that it is accesible from other methods
                  while (true) {
                     n = fieldG->getNeighbor(pt, token, distance, false);
                     if (distance > 1) break;//only nearest neighbors
                     nCell = fieldG->get(n);
                     
                     if(nCell)
                        type=nCell->type;
                     else
                        type=0;
                        

                     mitrTypeConst=contactCellMapMediumPtr->find(type);
                     
                     if(mitrTypeConst != contactCellMapMediumPtr->end()){//OK to secrete, contact detected
                        secrConstMedium = mitrTypeConst->second;
                        concentrationField->set(pt,currentConcentration+secrConstMedium);
                     }
                  }
                  continue;
            }

             if(currentCellPtr){
                mitr=secrData.typeIdSecrOnContactDataMap.find(currentCellPtr->type);
                if(mitr!=end_mitr){
                  
                  contactCellMapPtr = &(mitr->second.contactCellMap);
                  
                  token = 0;
                  distance = 0;
                  //set changePt so that it is accesible from other methods
                  while (true) {
                     n = fieldG->getNeighbor(pt, token, distance, false);
                     if (distance > 1) break;//only nearest neighbors
                     nCell = fieldG->get(n);
                     if(nCell)
                        type=nCell->type;
                     else
                        type=0;
                        
                     mitrTypeConst=contactCellMapPtr->find(type);
                     if(mitrTypeConst != contactCellMapPtr->end()){//OK to secrete, contact detected
                        secrConst=mitrTypeConst->second;
                        concentrationField->set(pt,currentConcentration+secrConst);
                     }
                     
                  }
                  
                }
             }
         }

}

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

//cerr<<"\n\n\t secrete usual on contact "<<idx<<"\n\n"<<endl;
SecretionData & secrData=secrDataVec[idx];

// if(!secrData.active)
//    return;

std::map<unsigned char,float>::iterator mitr;
std::map<unsigned char,float>::iterator end_mitr=secrData.typeIdSecrConstMap.end();

CellG *currentCellPtr;
Field3DImpl<float> * concentrationField=concentrationFieldVector[idx];
float currentConcentration;
float secrConst;
float secrConstMedium=0.0;

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

if( mitr != end_mitr){
   secreteInMedium=true;
   secrConstMedium=mitr->second;
   
}

Point3D pt;
   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);
            currentConcentration = concentrationField->get(pt);
            
            if(secreteInMedium && ! currentCellPtr){
               concentrationField->set(pt,currentConcentration+secrConstMedium);
            }

            if(currentCellPtr){
               mitr=secrData.typeIdSecrConstMap.find(currentCellPtr->type);
               if(mitr!=end_mitr){
                  secrConst=mitr->second;
                  concentrationField->set(pt,currentConcentration+secrConst);
               }
            }
         }


}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::secrete() {

   for(unsigned int i = 0 ; i < secrDataVec.size() ; ++i ){
/*      if(secretionFcnPtrVec[i]){//only active secreters are called
         (this->*secretionFcnPtrVec[i])(i);
      }*/

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

      

}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::diffuseSingleField(unsigned int idx){
//cout << "DIFFUSE SINLE FIELD" << endl;
/// '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 currentConcentration=0.0;
short neighborCounter=0;

DiffusionData & diffData = diffDataVec[idx];
float diffConst=diffData.diffConst;
float decayConst=diffData.decayConst;
float deltaT=diffData.deltaT;
float deltaX=diffData.deltaX;
float dt_dx2=deltaT/(deltaX*deltaX);



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

Automaton *automaton=potts->getAutomaton();

Field3DImpl<float> * concentrationField=concentrationFieldVector[idx];
Field3DImpl<float> * scratchField=concentrationFieldVector[diffDataVec.size()];


bool avoidMedium=false;
//the assumption is that medium has type ID 0
if(diffDataVec[idx].avoidTypeIdSet.find(automaton->getTypeId("Medium")) != end_sitr){
   avoidMedium=true;
}


   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);
         currentConcentration = concentrationField->get(pt);
         
         if(avoidMedium && !currentCellPtr){//if medium is to be avoided
            scratchField->set(pt,currentConcentration);
            continue;
         }

         if(currentCellPtr && diffData.avoidTypeIdSet.find(currentCellPtr->type)!=end_sitr){
         
           scratchField->set(pt,currentConcentration);
           continue; // avoid user defined types
         }
         
         updatedConcentration=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(avoidMedium && !nCell) continue; // avoid medium if specified by the user
               if(nCell && diffData.avoidTypeIdSet.find(nCell->type)!=end_sitr) continue; //avoid user specified types
                              
               //cerr<<"neighborPOint="<<n<<" nCell="<<nCell<<endl;
               concentrationSum += concentrationField->get(n);
               ++neighborCounter;
         }
         
    
         

         updatedConcentration =  dt_dx2*diffConst*(concentrationSum - neighborCounter*currentConcentration)
                                    -deltaT*(decayConst*currentConcentration)
                                    +currentConcentration;
         

         scratchField->set(pt,updatedConcentration);//updating scratch

      }
      //copy updated values from scratch to concentration field
      scrarch2Concentration(scratchField,concentrationField);
   
//cout << "DIFFUSE SINLE FIELD LEAVING" << endl;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::diffuse() {

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

      

}
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::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);
      }
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void FlexibleDiffusionSolverFE::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 FlexibleDiffusionSolverFE::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 FlexibleDiffusionSolverFE::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(SecretionDataFlex());
         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);
         cerr<<"AFTER INNER LOOP "<<endl;
         
/*         in.match(START_ELEMENT);
         diffData.readXML(in);
         in.match(END_ELEMENT);*/
      }
      in.skip(TEXT);
      cerr<<"AFTER OUTER LOOP "<<endl;
      //exit(0);
   }
   


   cerr<<"FlexibleDiffusionSolverFE: 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]=&FlexibleDiffusionSolverFE::secreteSingleField;
            ++j;
         }
         else if((*sitr)=="SecretionOnContact"){
            secrDataVec[i].secretionFcnPtrVec[j]=&FlexibleDiffusionSolverFE::secreteOnContactSingleField;
            ++j;
         }
      }      
   }
   

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


      
//   cerr<<"GOT HERE"<<endl;
   
   //exit(0);
    

}

