#include "pooling.hpp"

namespace Layer{

  Vector
  Pooling::feed_forward(Vector x){
    x_in=x;
    for(size_t f=0;f<nf;++f){
      for(size_t k=0;k<mi;++k){
        for(size_t l=0;l<mi;++l){
          size_t i=p*k;
          size_t j=q*l;
          Real temp=x[indice3(f,i,j,ni,nj)];
          for(;i<min(p*k+p,ni);++i){
            for(;j<min(q*l+q,nj);++j){
              temp=max(temp,x[indice3(f,i,j,ni,nj)]);
            }
          }
          x_out[indice3(f,k,l,mi,mj)]=temp;
        }
      }
    }
    return x_out;
  }

  Vector
  Pooling::back_propagation(Vector d){
    for(size_t f=0;f<nf;++f){
      for(size_t i=0;i<ni;++i){
        size_t k=i/p;
        for(size_t j=0;j<nj;++j){
          size_t l=j/q;
          bool is_max=true;
          Real val=x_in[indice3(f,i,j,ni,nj)];
          for(size_t i2=k*p;i2<min(k*p+p,ni) and is_max;++i2){
            for(size_t j2=l*q;j2<min(l*q+q,nj) and is_max;++j2){
              if(x_in[indice3(f,i2,j2,ni,nj)]>val) is_max=false;
            }
          }
          delta[indice3(f,i,j,ni,nj)]=(is_max)?d[indice3(f,k,l,mi,mj)]:0;
        }
      }
    }
    return delta;
  }
}
