#include "kernel.hpp"

Kernel::Kernel(string filename){
  fstream file;
  file.open(filename.c_str(),fstream::in|fstream::binary);
  InputData::load(file,true);

  initSolutions();

  Solution& S0=*solution[0];

  //Init pressure
  for(size_t x=0;x<geometry.nX;++x){
    for(size_t z=0;z<geometry.nZ[x];++z){
      S0.P[x][z]=initial_state->Pinit[x][z];
    }
  }
  //Init hydraulic head
  for(size_t x=0;x<geometry.nX;++x){
    S0.hydr[x]=initial_state->hsat[x]+Psat/(Physics::rho*Physics::g);
  }

  //Init overland level
  for(size_t x=0;x<geometry.nX;++x){
    S0.hov[x]=Physics::invPsol(geometry.hsoil[x],S0.P[x][geometry.nZ[x]-1]); //TODO[Chrisophe] à initiliser
  }

  //Init l
  for(size_t x=0;x<geometry.nX;++x){
    double t=initial_state->hsat[x];
    S0.l[x]=t;
    S0.hsat[x]=t;
  }

  //Init Pl
  for(size_t x=0;x<geometry.nX;++x){
    S0.Pl[x]=Psat;
  }


  indice_x_Richards=new bool[geometry.nX];

  div_w=new double*[geometry.nX];
  P_temp=new double*[geometry.nX];
  for(size_t i=0;i<geometry.nX;++i){
    div_w[i]=new double[geometry.nZ[i]];
    P_temp[i]=new double[geometry.nZ[i]];
  }
  t_sub_A=new double*[geometry.nX];
  t_sub_B=new double*[geometry.nX];
  t_sub_C=new double*[geometry.nX];
  t_diag_A=new double*[geometry.nX];
  t_diag_B=new double*[geometry.nX];
  t_diag_C=new double*[geometry.nX];
  t_sup_A=new double*[geometry.nX];
  t_sup_B=new double*[geometry.nX];
  t_sup_C=new double*[geometry.nX];
  t_F=new double*[geometry.nX];
  t_G=new double*[geometry.nX];
  t_S=new double*[geometry.nX];
  t_H=new double*[geometry.nX];
  t_R=new double*[geometry.nX];
  t_I=new double*[geometry.nX];
  t_BB=new double*[geometry.nX];

  for(size_t i=0;i<geometry.nX;++i){
    size_t nZ=geometry.nZ[i];
    t_sub_A[i]=new double[nZ-2];
    t_sub_B[i]=new double[nZ-2];
    t_sub_C[i]=new double[nZ-2];
    t_diag_A[i]=new double[nZ-1];
    t_diag_B[i]=new double[nZ-1];
    t_diag_C[i]=new double[nZ-1];
    t_sup_A[i]=new double[nZ-2];
    t_sup_B[i]=new double[nZ-2];
    t_sup_C[i]=new double[nZ-2];
    t_F[i]=new double[nZ-1];
    t_G[i]=new double[nZ-1];
    t_S[i]=new double[nZ-1];
    t_H[i]=new double[nZ-1];
    t_R[i]=new double[nZ-1];
    t_I[i]=new double[nZ-1];
    t_BB[i]=new double[nZ-1];
  }

  delete initial_state;
  initial_state=nullptr;

  file.close();
  step=0;
  m=1;
  scheme_error=0;
}

void
Kernel::initSolutions(){
  Solution* S0=new Solution();
  S0->init(geometry);
  solution[0]=S0;
  for(size_t i=1;i<=Time::nT;++i) solution[i]=nullptr;
  for(size_t i=0;i<Time::nT+2;++i){
    Solution* S=new Solution();
    S->init(geometry);
    pool_solutions.push(S);
  }
}

bool
Kernel::next(){
  //cout<<"[Kernel::next]"<<endl;
  if(m>1) m=m/2;
  spaceSolution();
  if(scheme_error>0) return false;
  ++step;
  //cout<<"[next] done"<<endl;
  return true;
}

void
Kernel::spaceSolution(){
  //cout<<"[Kernel::spaceSolution]"<<endl;
  size_t n=step+1; //We want solution at time n
  size_t k=0; //we work on time interval [n-1+k/m,n-1+(k+1)/m]
  Solution *S=solution[step];
  while(m<=max_time_subdivisions){
    //cout<<" k,m = "<<k<<','<<m<<endl;
    S=Piccard(S);
    if(S==nullptr){
      m*=2;
      k=0;
      S=solution[step];
    }
    else{
      ++k;
      if(k==m) break;
    }
  }
  if(m>max_time_subdivisions){
    cerr<<"Max time subdivision reached"<<endl;
    scheme_error=1;
  }
  //cout<<"[spaceSolution] done"<<endl;
}

Solution*
Kernel::Piccard(Solution* Sin){
  //Return a valid solution or nullptr if not converge
  //cout<<"[Kernel::Piccard]"<<endl;
  memcpy(Sin.l,Sin.hsat,sizeof(double)*geometry.nX);
  //Compute Pl with l=hsat of input solution and P=pressure of input_solution

  Solution* Spl=pool_solutions.pop();
  compute_Pl(*Spl,Sin.hsat,Sin.P);
  double dt=Time::dT/m; // util ????

  size_t count=0;

  //Specicy to compute Richards on all columns
  for(size_t i=0;i<geometry.nX;++i) indice_x_Richards[i]=true;

  double error=0;
  Solution* Snew=allVerticalRichards(Sin,Sin,Spl);
  if(Snew==nullptr) return false; //allVerti... must compute sol_out.P and sol_out.hsat
  //cout<<"[Piccard] done"<<endl;
  //Quitte de facon anticipe
  return true;
  compute_l(Snew,error);
  compute_Pl(Snew,Snew.l,Snew.P);
  if(!horizontalProblem(Sin,Snew,error)) return false; //Set hydr of Snew voir mettre hydr_in
  if(!overlandProblem(Sin,Sin,Snew,error)) return false; //Set hov of Snew

  //Comput error
  double error_prev=error;

  while(error>=tolerence_Piccard and count<max_iterations_Piccard and (abs(error_prev-error)>tolerence_Piccard/100 or error<oscilation_Piccard)){
    Solution* Sold=Snew;
    error_prev=error;
    ++count;
    error=0;
    Snew=allVerticalRichards(Sin,Sold,Sold);
    if(Snew==nullptr) return false;
    compute_l(Snew,error);
    compute_Pl(Snew,Snew.l,Snew.P);
    if(!horizontalProblem(Sold,Snew,error)) return false;
    if(!overlandProblem(Sin,Sold,Snew,error)) return false;
  }

  compute_mass(Snew);
  Snew.nstep_Piccard=count;
  //cout<<"[Piccard] done"<<endl;
  if(error>=tolerence_Piccard){
    return nullptr;
    //Faire attention à la pile -> rempiler des solutions
  }
  return Snew;
}


void
Kernel::compute_l(Solution& S,double& error) {
  //Update S.l using S.hsat
  bool e=0;
  for(size_t ix=0;ix<geometry.nX;++ix){
    double a=S.l[ix];
    double b=max(S.hsat[ix],a);
    S.l[ix]=b;
    double t=b-a;
    e+=t*t;
  }
  error+=sqrt(e);
}

void
Kernel::compute_Pl(Solution& S,const double* h,double** P){
  //cout<<"[Kernel::compute_Pl]"<<endl;
  for(size_t ix=0;ix<geometry.nX;++ix){
    double* Px=P[ix];
    if(h[ix]==geometry.hsoil[ix]){
      S.Pl[ix]=Px[geometry.nZ[ix]-1];
      //if(S.Pl[ix]!=-2000) cout<<"error "<<ix<<endl;
    }
    else{
      size_t a=(h[ix]-geometry.hbot[ix])/geometry.dZ[ix];
      double p1=Px[a];
      double p2=Px[a+1];
      S.Pl[ix]=p1+(p2-p1)/geometry.dZ[ix]*(h[ix]-geometry.Z[ix][a]);
      //cout<<ix<<" : "<<S.Pl[ix]<<endl;

    }
  }
}

bool
Kernel::allVerticalRichards(Solution& S_0,Solution& S_s,Solution& S_sp){
  //cout<<"[Kernel::allVerticalRichards]"<<endl;
  for(size_t ix=0;ix<geometry.nX;++ix){
    if(indice_x_Richards[ix]){
      if(!Richards1dEvolutiveTime(ix,S_0,S_s,S_sp)) return false;
    }
  }
  //cout<<"[Kernel::allVerticalRichards] done"<<endl;
  return true;
}

bool
Kernel::Richards1dEvolutiveTime(size_t ix,const Solution& S_0,const Solution& S_s,Solution& S_sp){
  //cout<<"[Kernel::Richards1dEvolutiveTime]"<<endl;
  //cout<<" ix = "<<ix<<endl;
  //Warning flux_bot and div_w are always zero
  double flux_bot=bottomFluxRichards(ix,S_s);
  compute_div_w(ix,S_s);
  size_t q=1;
  size_t l=0;
  bool conv;
  //cout<<"Time::dT "<<Time::dT<<endl;
  double dt=Time::dT/m*3600; //in second
  while(l<q and q<=max_Richards_time_subdivisions){
    //cout<<" l,q = "<<l<<','<<q<<endl;
    if(l=0){
      conv=Richards1d(ix,S_0,S_s,S_sp,dt,flux_bot);
    }
    else{
      conv=Richards1d(ix,S_0,S_sp,S_sp,dt,flux_bot);
    }
    if(!conv){
      q*=2;
      dt/=2;
      l=0;
    }
    else{
      ++l;
    }
  }
  //{cout<<"Continue ?"<<endl;char a;cin>>a;}

  if(q>max_Richards_time_subdivisions) return false;
  return true;
}

void
Kernel::compute_mass(Solution& S){

}

bool
Kernel::horizontalProblem(Solution& S,double& error){
  //Compute hydr from P, l, Pl, sources, time
  //return error otherwise
  return true;
}

bool
Kernel::overlandProblem(const Solution& S_in,Solution& S_out,double& error){
  //Compute hov
  //return error otherwise
  return true;
}

void
Kernel::saturationInterface(const Solution& S_in,Solution& S_out){
/*  Solution& S_in=*solution[sol_in];
  Solution& S_out=*solution[sol_out];
  for(size_t ix=0;ix<geometry.nX;++ix){
    double* Px_in=S_in.P[ix];
    size_t iz=0;
    for(;iz<geometry.nZ[ix] and Px_in[iz]>Psat;++iz);
    if(iz>0 and Px_in[iz]<=Psat){
      S_out.hsat[ix]=(Psat-Px_in[iz-1])*geometry.dZ[ix]/(Px_in[iz]-Px_in[iz-1])+geometry.Z[ix][iz-1];
    }
    else{
      S_out.hsat[ix]=(iz==0)?geometry.hbot[ix]:geometry.hsoil[ix];
    }
  }*/
}

double
Kernel::bottomFluxRichards(size_t ix,const Solution& S){
  return 0;
}

void
Kernel::compute_div_w(size_t ix,const Solution& S){
  double* div_w_ix=div_w[ix];
  for(size_t iz=0;iz<geometry.nZ[ix];++iz){
    div_w_ix[iz]=0;
  }
}

bool
Kernel::Richards1d(size_t ix,const Solution& S_0,const Solution& S_s,Solution& S_sp,double dt,double flux_bot){
  //cout<<"[Richards1d]"<<endl;
  size_t nZ=geometry.nZ[ix];
  double* P0=S_0.P[ix];
  double* Ps=S_s.P[ix];
  double* Psp=S_sp.P[ix];
  double n_P0=norm2(P0,nZ);
  if(weightedRichards1d(ix,P0,Ps,Psp,dt,flux_bot,1,n_P0)) return true;
  return weightedRichards1d(ix,P0,Ps,Psp,dt,flux_bot,0.5,n_P0);
}

bool
Kernel::weightedRichards1d(size_t ix,double* P0,double* Ps,double* Psp,double dt,double flux_bot,double r,double n_P0){
//  cout<<"[weightedRichards1d]"<<endl;
  //cout<<"r = "<<r<<endl;
  //cout<<"ix = "<<ix<<endl;
  size_t nZ=geometry.nZ[ix];
  double n0=norm2(P0,nZ);
  double* Pi=Ps;
  double* Po=P_temp[ix];
  size_t count=0;
  double e=+inf;
  while(e>=tolerence_Richards and count<max_iterations_Richards){
    ++count;
    solveSystemRichards(ix,P0,Pi,Po,flux_bot,dt,count);
    e=error2(Pi,Po,nZ)/n0;
    //cout<<count<<" : "<<e<<" vs "<<tolerence_Richards<<endl;
    if(count==1){
      Pi=Po;
      Po=Psp;
    }
    else{
      swap(Pi,Po);
    }
  }
  if(e<tolerence_Richards){
    //cout<<"[converge]"<<endl;
    memcpy(Psp,Pi,nZ*sizeof(double));
    return true;
  }
  //cout<<"[don't converge]"<<endl;
  return false;
}

void
Kernel::solveSystemRichards(size_t ix,double* P0,double* P,double* Pout,double flux_bot,double dt,size_t count){
  //TODO treat fpump
  /*cout<<"[solveSystemRichards]"<<endl;
  cout<<"P0 = "<<P0<<endl;
  cout<<"Pin = "<<P<<endl;
  cout<<"Pout = "<<Pout<<endl;*/

  size_t nZ=geometry.nZ[ix];
  assert(nZ>=3);
  double dZ=geometry.dZ[ix];



  double* sup_A=t_sup_A[ix];
  double* diag_A=t_diag_A[ix];
  double* sub_A=t_sub_A[ix];
  double* sup_B=t_sup_B[ix];
  double* diag_B=t_diag_B[ix];
  double* sub_B=t_sub_B[ix];
  double* sup_C=t_sup_C[ix];
  double* diag_C=t_diag_C[ix];
  double* sub_C=t_sub_C[ix];
  double* F=t_F[ix];
  double* G=t_G[ix];
  double* S=t_S[ix];
  double* H=t_H[ix];
  double* R=t_R[ix];
  double* I=t_I[ix];
  double* BB=t_BB[ix];

  //Compute A
  diag_A[0]=(Physics::phi*dZ*Physics::ds(P[0]))/(2*dt);
  for(size_t i=1;i<nZ-1;++i){
    diag_A[i]=(Physics::phi*dZ*Physics::ds(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(P[0])+Physics::kr(P[1]));
  sup_B[0]=-diag_B[0];
  diag_B[nZ-2]=alpha*(Physics::kr(P[nZ-3])+2*Physics::kr(P[nZ-2])+Physics::kr(P[nZ-1]));
  sub_B[nZ-3]=-alpha*(Physics::kr(P[nZ-3])+Physics::kr(P[nZ-2]));
  for(size_t i=1;i<nZ-2;++i){
    diag_B[i]=alpha*(Physics::kr(P[i-1])+2*Physics::kr(P[i])+Physics::kr(P[i+1]));
    sub_B[i-1]=-alpha*(Physics::kr(P[i-1])+Physics::kr(P[i]));
    sup_B[i]=-alpha*(Physics::kr(P[i])+Physics::kr(P[i+1]));
  }


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

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

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

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

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

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

  //Compute BB
  clear(BB,nZ-1);
  //TODO : Add BB computation from fpump

  //Compute F
  F[0]=div_w[ix][0]*dZ+R[0]-G[0]-I[0]-S[0]-H[0]+(diag_A[0]+diag_C[0])*P[0]+sup_C[0]*P[1];
  for(size_t i=1;i<nZ-2;++i){
    F[i]=div_w[ix][i]*dZ+R[i]-G[i]-I[i]-S[i]-H[i]+(diag_A[i]+diag_C[i])*P[i]+sub_C[i-1]*P[i-1]+sup_C[i]*P[i+1];
  }
  F[nZ-2]=div_w[ix][nZ-2]*dZ+R[nZ-2]-G[nZ-2]-I[nZ-2]-S[nZ-2]-H[nZ-2]+(diag_A[nZ-2]+diag_C[nZ-2])*P[nZ-2]+sub_C[nZ-3]*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]);



 Thomas(nZ-1,sub_A,diag_A,sup_A,F,Pout);
 Pout[nZ-1]=P[nZ-1];
}