#include "convolution.hpp"

namespace Layer{

  void
  Convolution::init(Real m,Real d){
    default_random_engine generator;
    normal_distribution<Real> distribution(m,d);
    for(size_t i=0;i<mf*nf*i*q;++i){
      K[i]=distribution(generator);
    }
    for(size_t i=0;i<mf*nf;++i){
      b[i]=distribution(generator);
    }
  }

  Vector
  Convolution::feed_forward(Vector x){
    x_in=x;
    for(size_t g=0;g<mf;++g){
      for(size_t k=0;k<mi;++k){
        for(size_t l=0;l<mj;++l){
          Real temp=0;
          for(size_t f=0;f<nf;++f){
            for(size_t r=0;r<p;++r){
              for(size_t s=0;s<q;++s){
                temp+=x[indice3(f,k+r,l+s,ni,nj)]*K[indice4(g,f,r,s,nf,p,q)];
              }
            }
            temp+=b[indice2(g,f,nf)];
          }
          x_out[indice3(g,k,l,mi,mj)]=temp;
        }
      }
    }
    return x_out;
  }

  void
  Convolution::init_nabla(){
    for(size_t i=0;i<mf*nf;++i){
      nabla_b[i]=0;
    }
    for(size_t i=0;i<mf*nf*p*q;++i){
      nabla_K[i]=0;
    }
  }

  Vector
  Convolution::back_propagation(Vector d){
    for(size_t f=0;f<nf;++f){
      for(size_t i=0;i<ni;++i){
        for(size_t j=0;j<nj;++j){
          Real temp=0;
          for(size_t g=0;g<mf;++g){
            size_t r=(i>=mi-1)?i-mi+1:0;
            for(;r<min(i,p);++r){
              size_t s=(j>=mj-1)?j-mj+1:0;
              for(;s<min(j,q);++s){
                temp+=K[indice4(g,f,r,s,nf,p,q)]*d[indice3(g,i-r,j-s,mi,mj)];
              }//s
            }//r
          }//g
          delta[indice3(f,i,j,ni,nj)]=temp;
        }//j
      }//i
    }//f
    //display(delta,nf*ni*nj);
    //char a;cin>>a;
    //cout<<"  - Update nabla_b"<<endl;
    //Update nabla_b<<
    for(size_t g=0;g<mf;++g){
      for(size_t f=0;f<nf;++f){
        Real temp=0;
        for(size_t k=0;k<mi;++k){
          for(size_t l=0;l<mj;++l){
            temp+=d[indice3(g,k,l,mi,mj)];
          }//l
        }//k
        nabla_b[indice2(g,f,nf)]+=temp;
      }
    }
    //Update nabla_w
    for(size_t g=0;g<mf;++g){
      for(size_t f=0;f<nf;++f){
        for(size_t r=0;r<p;++r){
          for(size_t s=0;s<q;++s){
            Real temp=0;
            for(size_t k=0;k<mi;++k){
              for(size_t l=0;l<mj;++l){
                temp+=d[indice3(g,k,l,mi,mj)]*x_in[indice3(f,k+r,l+s,ni,nj)];
              }//l
            }//k
            nabla_K[indice4(g,f,r,s,nf,p,q)]+=temp;
          }
        }
      }
    }
    return delta;
  }

  void
  Convolution::update(Real eta){
    //Update b
    for(size_t i=0;i<mf*nf;++i){
      b[i]-=eta*nabla_b[i];
    }
    //Update K
    for(size_t i=0;i<mf*nf*p*q;++i){
      K[i]-=eta*nabla_K[i];
    }
  }
}