#include "richards.hpp"

namespace Kernel{

  Richards::Richards(){
    sup_A=nullptr;
    diag_A=nullptr;
    sub_A=nullptr;
    sup_B=nullptr;
    diag_B=nullptr;
    sub_B=nullptr;
    sup_C=nullptr;
    diag_C=nullptr;
    sub_C=nullptr;
    F=nullptr;
    G=nullptr;
    S=nullptr;
    H=nullptr;
    R=nullptr;
    I=nullptr;
    BB=nullptr;
  }

  void
  Richards::init(size_t ix_,size_t nZ_,double dZ_,double* pumps_){
    pumps=pumps_;
    ix=ix_;
    nZ=nZ_;
    dZ=dZ_;
    temp_P[0]=new double[nZ];
    temp_P[1]=new double[nZ];

    sub_A=new double[nZ-2];
    sub_B=new double[nZ-2];
    sub_C=new double[nZ-2];
    diag_A=new double[nZ-1];
    diag_B=new double[nZ-1];
    diag_C=new double[nZ-1];
    sup_A=new double[nZ-2];
    sup_B=new double[nZ-2];
    sup_C=new double[nZ-2];
    F=new double[nZ-1];
    G=new double[nZ-1];
    S=new double[nZ-1];
    H=new double[nZ-1];
    R=new double[nZ-1];
    I=new double[nZ-1];
    BB=new double[nZ-1];

  }

  void
  Richards::run(){
    #ifdef MYDEBUG
    if(Debug::ix==ix) cout<<"     [Richards::run] start"<<endl;
    #endif
    norm_previous_P=norm2(previous_P,nZ);
    has_converged=weighted_run(1);
    if(!has_converged) has_converged=weighted_run(0.5);

    #ifdef MYDEBUG
    if(Debug::ix==ix) cout<<"     [Richards::run] stop"<<endl;
    #endif

  }

  bool
  Richards::weighted_run(double w){
    Debug::debug_Thomas=(Debug::ix==ix);
    #ifdef MYDEBUG
    if(Debug::ix==ix){
      cout<<"      [Richards::weighted_run] start"<<endl;
      cout<<"      [Richards::weighted_run] w = "<<w<<endl;
    }
    #endif
    double error=numeric_limits<double>::infinity();
    size_t count=0;
    in_P=near_P;
    out_P=temp_P[1];
    while(error>=tolerence_Richards and count<max_iterations_Richards){
      ++count;

      #ifdef MYDEBUG
      if(Debug::level>1 and Debug::ix==ix) cout<<"      [Richards::weighted_run] count = "<<count<<endl;
      #endif

      solve_system();

      if(w<1){
        for(size_t i=0;i<nZ;++i){
          out_P[i]=w*out_P[i]+(1-w)*in_P[i];
        }
      }
      //Computed P is in out_P=temp_P[count%2]
      /*if(Debug::ix==ix) Debug::display("in_P",in_P,0,nZ);
      if(Debug::ix==ix) Debug::display("out_P",out_P,0,nZ);*/
      error=error2(in_P,out_P,nZ)/norm_previous_P;
      if(count==1){
        in_P=temp_P[1];
        out_P=temp_P[0];
      }
      else{
        swap(in_P,out_P);
      }
      #ifdef MYDEBUG
      if(Debug::ix==ix){
        cout<<"      [Richards::weighted_run] error = "<<error<<endl;;
      }
      #endif
    }
   if(error<tolerence_Richards){
     //Last computed P is in temp_P[count%2] which is now also in_P
     assert(in_P==temp_P[count%2]);
     swap(P,temp_P[count%2]);
     #ifdef MYDEBUG
     if(Debug::ix==ix){
       if(Debug::level>1) cout<<"      [Richards::weighted_run] converge"<<endl;
       cout<<"      [Richards::weighted_run] stop"<<endl;
     }
     #endif
     return true;
   }
   #ifdef MYDEBUG
   if(Debug::ix==ix){
     cout<<"      [Richards::weighted_run] not converge"<<endl;
     cout<<"      [Richards::weighted_run] stop"<<endl;
   }
   #endif
   return false;
 }

 void
 Richards::solve_system(){
   #ifdef MYDEBUG
   if(Debug::ix==ix) cout<<"       [Richards::solve_system] start"<<endl;
   assert(nZ>=3);
   #endif
   //Compute A
   diag_A[0]=(Physics::phi*dZ*Physics::ds(in_P[0]))/(2*dt);
   for(size_t i=1;i<nZ-1;++i){
     diag_A[i]=(Physics::phi*dZ*Physics::ds(in_P[i]))/dt;
   }


   //Compute B

   //TODO : A optimiser boucle par rapport aux appels de Kr
   double alpha=Physics::k0/(2*Physics::rho*Physics::g*dZ);
   diag_B[0]=alpha*(Physics::kr(in_P[0])+Physics::kr(in_P[1]));
   sup_B[0]=-diag_B[0];
   diag_B[nZ-2]=alpha*(Physics::kr(in_P[nZ-3])+2*Physics::kr(in_P[nZ-2])+Physics::kr(in_P[nZ-1]));
   sub_B[nZ-3]=-alpha*(Physics::kr(in_P[nZ-3])+Physics::kr(in_P[nZ-2]));
   for(size_t i=1;i<nZ-2;++i){
     diag_B[i]=alpha*(Physics::kr(in_P[i-1])+2*Physics::kr(in_P[i])+Physics::kr(in_P[i+1]));
     sub_B[i-1]=-alpha*(Physics::kr(in_P[i-1])+Physics::kr(in_P[i]));
     sup_B[i]=-alpha*(Physics::kr(in_P[i])+Physics::kr(in_P[i+1]));
   }


   //Compute C
   double hk=Physics::k0/2;
   double temp=1/(dZ*Physics::rho*Physics::g);
   diag_C[0]=-hk*Physics::dkr(in_P[0])*(temp*(in_P[1]-in_P[0])+1);
   sup_C[0]=-hk*Physics::dkr(in_P[1])*(temp*(in_P[1]-in_P[0])+1);
   diag_C[nZ-2]=hk*Physics::dkr(in_P[nZ-3])*temp*(-in_P[nZ-3]+2*in_P[nZ-2]-in_P[nZ-1]);
   sub_C[nZ-3]=hk*Physics::dkr(in_P[nZ-3])*(temp*(in_P[nZ-2]-in_P[nZ-3])+1);
   for(size_t i=1;i<nZ-2;++i){
     diag_C[i]=hk*Physics::dkr(in_P[i])*temp*(-in_P[i-1]+2*in_P[i]-in_P[i+1]);
     sub_C[i-1]=hk*Physics::dkr(in_P[i-1])*(temp*(in_P[i]-in_P[i-1])+1);
     sup_C[i]=-hk*Physics::dkr(in_P[i+1])*(temp*(in_P[i+1]-in_P[i])+1);
   }

   //Compute G
   clear(G,nZ-1);
   G[nZ-2]=-alpha*(Physics::kr(in_P[nZ-2])+Physics::kr(in_P[nZ-1]))*in_P[nZ-1];

   //Compute S
   S[0]=-hk*(Physics::kr(in_P[0])+Physics::kr(in_P[1]));
   for(size_t i=1;i<nZ-1;++i){
     S[i]=hk*(Physics::kr(in_P[i-1])-Physics::kr(in_P[i+1]));
   }

   //Compute H
   clear(H,nZ-1);
   H[nZ-2]=-hk*Physics::dkr(in_P[nZ-1])*(temp*(in_P[nZ-1]-in_P[nZ-2])+1);

   //Compute R
   clear(R,nZ-1);
   R[0]=-flux_bot;

   //Compute I
   I[0]=(Physics::phi*dZ*(Physics::s(in_P[0])-Physics::s(previous_P[0])))/(2*dt);
   for(size_t i=1;i<nZ-1;++i){
     I[i]=(Physics::phi*dZ*(Physics::s(in_P[i])-Physics::s(previous_P[i])))/dt;
   }

   //Compute BB
   for(size_t i=1;i<nZ-2;++i){
     BB[i]=dZ*(pumps[i-1]/6+(2*pumps[i])/3+pumps[i+1]/6);
   }
   BB[0]=dZ*(pumps[0]/3+pumps[1]/6);
   BB[nZ-2]=dZ*(pumps[nZ-3]/6+(2*pumps[nZ-2])/3);

   //Compute F
   F[0]=R[0]-G[0]-I[0]-S[0]-H[0]+BB[0]+(diag_A[0]+diag_C[0])*in_P[0]+sup_C[0]*in_P[1];
   for(size_t i=1;i<nZ-2;++i){
     F[i]=div_w[i]*dZ+R[i]-G[i]-I[i]-S[i]-H[i]+BB[i]+(diag_A[i]+diag_C[i])*in_P[i]+sub_C[i-1]*in_P[i-1]+sup_C[i]*in_P[i+1];
   }
   F[nZ-2]=div_w[nZ-2]*dZ+R[nZ-2]-G[nZ-2]-I[nZ-2]-S[nZ-2]-H[nZ-2]+BB[nZ-2]+(diag_A[nZ-2]+diag_C[nZ-2])*in_P[nZ-2]+sub_C[nZ-3]*in_P[nZ-3];

 // /  if(ix==39)  display("F",F,nZ-1);
   //TODO Add contribution of BB in F

   for(size_t i=0;i<nZ-2;++i){
     sub_A[i]=(sub_B[i]+sub_C[i]);
     diag_A[i]+=(diag_B[i]+diag_C[i]);
     sup_A[i]=(sup_B[i]+sup_C[i]);
   }
   diag_A[nZ-2]+=(diag_B[nZ-2]+diag_C[nZ-2]);
   //Debug::display("sup_A",sup_A,0,nZ-2);
  Thomas(nZ-1,sub_A,diag_A,sup_A,F,out_P);
  out_P[nZ-1]=P[nZ-1];
  #ifdef MYDEBUG
  if(Debug::ix==ix){
    cout<<"       [Richards::solve_system] out = "<<out_P<<endl;
    cout<<"       [Richards::solve_system] stop"<<endl;
  }
  #endif

 }
}