#include "network.hpp"

Network::Network(){
  C=Quadratic;
}

void
Network::push_layer(Layer::Layer& l){
  if(layers.empty()){
    n_in=l.get_input_size();
    layers.push_back(&l);
    cout<<"In size = "<<n_in<<endl;
  }
  else{
    assert(l.get_input_size()==layers.back()->get_output_size());
    layers.push_back(&l);
  }
  n_out=l.get_output_size();
}

void
Network::is_done(){
  last_delta=init_vector(n_out);
}

Vector
Network::feed_forward(Vector x_in){
  Vector x=x_in;
  for(auto it=layers.begin();it!=layers.end();++it){
    //cout<<" - Try feed_forward on layer "<<(*it)->name<<endl;
    x=(*it)->feed_forward(x);
  }
  a=x;
  return a;
}

Real
Network::eval(Dataset* dataset){
  size_t n=dataset->get_test_size();
  size_t nb=0;
  for(size_t i=0;i<n;++i){
    pair<Vector,Vector> t=dataset->get_test(i);
    Vector a=feed_forward(t.first);
    if(argmax(a,n_out)==argmax(t.second,n_out)) ++nb;
  }
  Real res=Real(nb)/Real(n)*100;
  cout<<"> Res = "<<res<<"%"<<endl;
  return res;
}

void
Network::shuffle(size_t* tab,size_t size){
  default_random_engine generator;
  uniform_int_distribution<int> distribution(0,size-1);
  for(size_t k=0;k<size;++k){
    size_t i=distribution(generator);
    size_t j=distribution(generator);
    swap(tab[i],tab[j]);
  }
}

void
Network::train(Dataset* dataset,size_t nb_epochs,size_t batch_size,Real eta){
  size_t train_size=dataset->get_train_size();
  size_t nb_batchs=(train_size-1)/batch_size+1;
  size_t* indices=new size_t[train_size];
  for(size_t i=0;i<train_size;++i){
    indices[i]=i;
  }
  for(size_t epoch=0;epoch<nb_epochs;++epoch){
    cout<<"Epoch "<<epoch<<endl;
    shuffle(indices,train_size);
    for(size_t batch=0;batch<nb_batchs;++batch){
      size_t begin=batch*batch_size;
      size_t end=min(train_size,begin+batch_size);
      update_batch(dataset,indices,begin,end,eta);
    }
    eval(dataset);
  }
  delete[] indices;
}

void
Network::update_batch(Dataset* dataset,size_t* indices,size_t begin,size_t end,Real eta){
  Real batch_size=end-begin;
  for(auto it=layers.begin();it!=layers.end();++it){
    (*it)->init_nabla();
  }
  for(size_t i=begin;i<end;++i){
    pair<Vector,Vector> data=dataset->get_train(indices[i]);
    //cout<<"Call back_propagation on batch data "<<i-begin<<"/"<<batch_size<<endl;
    back_propagation(data.first,data.second,eta);
  }
  Real eta_batch=eta/batch_size;
  for(auto it=layers.begin();it!=layers.end();++it){
    (*it)->update(eta_batch);
  }
}

void
Network::back_propagation(Vector x,Vector y,Real eta){
  Vector z=feed_forward(x);
  //cout<<" - Feed forward done"<<endl;
  compute_last_delta(y);
  Vector delta=last_delta;
  //cout<<" - Last_delta computed"<<endl;
  for(auto it=layers.rbegin();it!=layers.rend();++it){
    //cout<<" - Try back_propagation on layer "<<(*it)->name<<endl;
    delta=(*it)->back_propagation(delta);
    //cout<<"    - Done"<<endl;
  }
}

void
Network::compute_last_delta(Vector y){
  switch(C){
    case Quadratic:
    case CrossEntropy:
    for(size_t i=0;i<n_out;++i){
      last_delta[i]=a[i]-y[i];
    }
    break;
    default:
    assert(false);
    break;
  }
}