#include "basin.hpp"

void
Basin::compute_leak_direction(){
  switch(position){
    case Middle:
    if(abs(hleft-hright)<delta_both_leak) leak_direction=Both;
    else leak_direction=(hleft>hright)?Right:Left;
    break;
    case BorderLeft:
    leak_direction=Right;
    break;
    case BorderRight:
    leak_direction=Left;
    break;
    default:
    leak_direction=None;
    break;
  };
}

void Basin::compute_maxs(double* hsoil){
  switch(leak_direction){
    case Left:
    hmax=hleft;
    xmax_left=xleft;
    {
      int x;
      for(x=xleft+1;x<xright;++x){
        if(hsoil[x]+Physics::minimal_height_to_runoff>hmax){
          xmax_right=x;
          break;
        }
      }
      xmax_right=min(x,(int)xright);
    }
    break;
    case Right:
    hmax=hright;
    xmax_right=xright;
    {
      int x;
      for(x=xright-1;x>xleft;--x){
        if(hsoil[x]+Physics::minimal_height_to_runoff>hmax){
          xmax_left=x;
          break;
        }
      }
      xmax_left=max(x,(int)xleft);
    }
    break;
    case Both:
    hmax=min(hleft,hright);
    xmax_left=xleft;
    xmax_right=xright;
    break;
    case None:
    default:
    hmax=numeric_limits<double>::infinity();
    xmax_left=xleft;
    xmax_right=xright;
    break;
  };
}

Basin::Basin(Basin* left_,Basin* right_){
  left=left_;
  right=right_;
  xleft=left->xleft;
  xright=right->xright;
  hleft=left->hleft;
  hright=right->hright;
  hmin=left->hmax;
  xacc=left->xacc;
  if(left->position==BorderLeft){
    if(right->position==BorderRight) position=BorderBoth;
    else position=BorderLeft;
  }
  else if(right->position==BorderRight) position=BorderRight;
  else position=Middle;
  compute_leak_direction();

}

void
Basin::display(string indent){
  cout<<indent<<'['<<xleft<<','<<xright<<"] : "<<hmax<<" : ["<<xmax_left<<","<<xmax_right<<"] and "<<hmin<<endl;
  if(left==nullptr) return;
  left->display(indent+' ');
  right->display(indent+' ');
}

void
Basin::compute_vflow(double* hsoil,double* hov){
  if(is_primitive()){
    vflow=max(0.0,hov[xacc]-hsoil[xacc]-Physics::minimal_height_to_runoff);
  }
  else{
    left->compute_vflow(hsoil,hov);
    right->compute_vflow(hsoil,hov);
    vflow=left->vflow+right->vflow;
  }
}

void
Basin::repartition(double* hsoil,double* hov,Direction* rundir){

  overflow=false;
  vmin=0;
  if(not is_primitive()){
    for(size_t ix=xleft;ix<=xright;++ix){
      vmin+=max(0.0,hmin-min(hov[ix],hsoil[ix]+Physics::minimal_height_to_runoff));
    }
  }
  if(vflow!=0 and vflow>=vmin){
    vmax=0;
    for(size_t ix=xmax_left;ix<=xmax_right;++ix){
      vmax+=max(0.0,hmax-min(hov[ix],hsoil[ix]+Physics::minimal_height_to_runoff));
    }
    if(vflow-vmax>1e-10){
      overflow=true;
      vleak=vflow-vmax;
      if(not is_primitive()){
        double delta=hov[right->xacc]-hsoil[right->xacc]-Physics::minimal_height_to_runoff;
        hov[left->xacc]+=delta;
        hov[right->xacc]-=delta;
      }
      hov[xacc]-=vleak;
      switch(leak_direction){
        case Left:
        hov[xleft-1]+=vleak;
        runoff(xleft-1,Left,rundir,hsoil,hov);
        break;
        case Right:
        hov[xright+1]+=vleak;
        runoff(xright+1,Right,rundir,hsoil,hov);
        break;
        case Both:
        hov[xleft-1]+=vleak/2;
        runoff(xleft-1,Left,rundir,hsoil,hov);
        hov[xright+1]+=vleak/2;
        runoff(xright+1,Right,rundir,hsoil,hov);
        break;
        default:
        cerr<<"[Bug] Impossible case ! ... in theory :)"<<endl;
        exit(-1);
        break;
      }
    }
  }
  else{
    if(not is_primitive()){
      left->repartition(hsoil,hov,rundir);
      right->repartition(hsoil,hov,rundir);
      overflow=left->overflow|right->overflow;
    }
  }
}

void
Basin::runoff(size_t xstart,Direction dir,Direction* rundir,double* hsoil,double* hov){
  double r=max(0.0,hov[xstart]-hsoil[xstart]-Physics::minimal_height_to_runoff);
  int x=xstart;
  int xstep=(dir==Left)?-1:1;
  while(r>0 and rundir[x]!=None){
    hov[x]-=r;
    x+=xstep;
    hov[x]+=r;
    r=max(0.0,hov[x]-hsoil[x]-Physics::minimal_height_to_runoff);
  }

}

double
Basin::f(double r,double* hsoil,double* hov,double vext){
  double res=-vext;
  for(size_t ix=xleft;ix<=xright;++ix){
    res+=max(r-min(hov[ix],hsoil[ix]+Physics::minimal_height_to_runoff),0.0);
  }
  return res;
}

double
Basin::df(double r,double* hsoil,double* hov,double vext){
  double res=0;
  for(size_t ix=xleft;ix<=xright;++ix){
    if(r>min(hov[ix],hsoil[ix]+Physics::minimal_height_to_runoff)) ++res;
  }
  return res;
}

void
Basin::flatten(double* hsoil,double* hov){
  double crit=1;
  size_t count=0;
  if(vflow>1e-9){
    double r=hsoil[xacc]+1;//Physics::minimal_height_to_runoff;
    while(crit>1e-10 and count<100){
      ++count;
      double rold=r;
      r-=f(r,hsoil,hov,vflow)/df(r,hsoil,hov,vflow);
      crit=abs(r-rold)+abs(f(r,hsoil,hov,vflow));
    }
    if(crit>1e-10){
      cerr<<" Overland coverge + [TODO] on fait quoi ?"<<endl;
    }
    if(r>=hmin or is_primitive()){
      for(size_t ix=xleft;ix<=xright;++ix){
        hov[ix]=max(r,min(hov[ix],hsoil[ix]+Physics::minimal_height_to_runoff));
      }


    }
    else{
      left->flatten(hsoil,hov);
      right->flatten(hsoil,hov);
    }
  }
}