#include 
#include 
#include 

#include "gtest/gtest.h"

#include "caffe/blob.hpp"
#include "caffe/common.hpp"
#include "caffe/filler.hpp"
#include "caffe/layers/lrn_layer.hpp"

#ifdef USE_CUDNN
#include "caffe/layers/cudnn_lcn_layer.hpp"
#include "caffe/layers/cudnn_lrn_layer.hpp"
#endif

#include "caffe/test/test_caffe_main.hpp"
#include "caffe/test/test_gradient_check_util.hpp"

using std::min;
using std::max;

int test_h=5;
int test_w=5;

namespace caffe {
 

template 
static void dump_blob(const Blob * blob, const char * outfile)
{
   std::ofstream os;
   os.open(outfile);

   for(int i=0;ishape(0);i++)
     for(int j=0;jshape(1);j++)
        for(int k=0;kshape(2);k++)
            for(int l=0;lshape(3);l++)
       {
          Dtype data=blob->data_at(i,j,k,l);

          osheight()*bottom->width()+
                   l*bottom->width()+k;

           ptr=bottom->mutable_cpu_data();

           if(fixed)
              ptr[offset]=val;
           else
              ptr[offset]=offset;

        }


}


template 
class LRNLayerTest : public MultiDeviceTest {
  typedef typename TypeParam::Dtype Dtype;

 protected:
  LRNLayerTest()
      : epsilon_(Dtype(1e-5)),
        blob_bottom_(new Blob()),
        blob_top_(new Blob()) {}
  virtual void SetUp() {
    Caffe::set_random_seed(1701);
    blob_bottom_->Reshape(2, 7, test_h,test_w);
    // fill the values
    FillerParameter filler_param;
    GaussianFiller filler(filler_param);
    filler.Fill(this->blob_bottom_);
    blob_bottom_vec_.push_back(blob_bottom_);
    blob_top_vec_.push_back(blob_top_);
  }
  virtual ~LRNLayerTest() { delete blob_bottom_; delete blob_top_; }
  void ReferenceLRNForward(const Blob& blob_bottom,
      const LayerParameter& layer_param, Blob* blob_top);

  Dtype epsilon_;
  Blob* const blob_bottom_;
  Blob* const blob_top_;
  vector blob_bottom_vec_;
  vector blob_top_vec_;
};

template 
void LRNLayerTest::ReferenceLRNForward(
    const Blob& blob_bottom, const LayerParameter& layer_param,
    Blob* blob_top) {
  typedef typename TypeParam::Dtype Dtype;
  blob_top->Reshape(blob_bottom.num(), blob_bottom.channels(),
      blob_bottom.height(), blob_bottom.width());
  Dtype* top_data = blob_top->mutable_cpu_data();
  LRNParameter lrn_param = layer_param.lrn_param();
  Dtype alpha = lrn_param.alpha();
  Dtype beta = lrn_param.beta();
  int size = lrn_param.local_size();
  switch (lrn_param.norm_region()) {
  case LRNParameter_NormRegion_ACROSS_CHANNELS:
    for (int n = 0; n < blob_bottom.num(); ++n) {
      for (int c = 0; c < blob_bottom.channels(); ++c) {
        for (int h = 0; h < blob_bottom.height(); ++h) {
          for (int w = 0; w < blob_bottom.width(); ++w) {
            int c_start = c - (size - 1) / 2;
            int c_end = min(c_start + size, blob_bottom.channels());
            c_start = max(c_start, 0);
            Dtype scale = 1.;
            for (int i = c_start; i < c_end; ++i) {
              Dtype value = blob_bottom.data_at(n, i, h, w);
              scale += value * value * alpha / size;
            }
            *(top_data + blob_top->offset(n, c, h, w)) =
              blob_bottom.data_at(n, c, h, w) / pow(scale, beta);
          }
        }
      }
    }
    break;
  case LRNParameter_NormRegion_WITHIN_CHANNEL:
    for (int n = 0; n < blob_bottom.num(); ++n) {
      for (int c = 0; c < blob_bottom.channels(); ++c) {
        for (int h = 0; h < blob_bottom.height(); ++h) {
          int h_start = h - (size - 1) / 2;
          int h_end = min(h_start + size, blob_bottom.height());
          h_start = max(h_start, 0);
          for (int w = 0; w < blob_bottom.width(); ++w) {
            Dtype scale = 1.;
            int w_start = w - (size - 1) / 2;
            int w_end = min(w_start + size, blob_bottom.width());
            w_start = max(w_start, 0);

//            std::coutForward(this->blob_bottom_vec_, this->blob_top_vec_);

  Blob top_reference;
  this->ReferenceLRNForward(*(this->blob_bottom_), layer_param,
      &top_reference);
  for (int i = 0; i < this->blob_bottom_->count(); ++i) {
    EXPECT_NEAR(this->blob_top_->cpu_data()[i], top_reference.cpu_data()[i],
                this->epsilon_);
  }
}


TYPED_TEST(LRNLayerTest, TestSetupWithinChannel) {
  typedef typename TypeParam::Dtype Dtype;
  LayerParameter layer_param;
  layer_param.mutable_lrn_param()->set_norm_region(
      LRNParameter_NormRegion_WITHIN_CHANNEL);
  layer_param.mutable_lrn_param()->set_local_size(3);

  
  layer_param.set_type("LRN");

  shared_ptr new_layer=
    LayerRegistry::CreateLayer(layer_param);

  shared_ptr layer=
   boost::static_pointer_cast (new_layer);


  layer->SetUp(this->blob_bottom_vec_, this->blob_top_vec_);

  EXPECT_EQ(this->blob_top_->num(), 2);
  EXPECT_EQ(this->blob_top_->channels(), 7);
  EXPECT_EQ(this->blob_top_->height(), test_h);
  EXPECT_EQ(this->blob_top_->width(), test_w);
}
#endif

#if 1

TYPED_TEST(LRNLayerTest, TestForwardWithinChannel) {
  typedef typename TypeParam::Dtype Dtype;
  LayerParameter layer_param;
  layer_param.mutable_lrn_param()->set_norm_region(
      LRNParameter_NormRegion_WITHIN_CHANNEL);
  layer_param.mutable_lrn_param()->set_local_size(3);
//  layer_param.mutable_lrn_param()->set_beta(1);

  
  layer_param.set_type("LRN");

  shared_ptr new_layer=
    LayerRegistry::CreateLayer(layer_param);

  shared_ptr layer=
   boost::static_pointer_cast (new_layer);

/* presetting bottom_vec and data */

  vector bottom_shape;
  bottom_shape.push_back(1);
  bottom_shape.push_back(1);
  bottom_shape.push_back(5);
  bottom_shape.push_back(5);


  this->blob_bottom_vec_[0]->Reshape(bottom_shape);

  fill_blob_data(this->blob_bottom_,1,1);


  layer->SetUp(this->blob_bottom_vec_, this->blob_top_vec_);
  layer->Forward(this->blob_bottom_vec_, this->blob_top_vec_);


  Blob top_reference;
  this->ReferenceLRNForward(*(this->blob_bottom_), layer_param,
      &top_reference);
//  for (int i = 0; i < this->blob_bottom_->count(); ++i) {
//    EXPECT_NEAR(this->blob_top_->cpu_data()[i], top_reference.cpu_data()[i],
//                this->epsilon_);
//  }

  dump_blob(this->blob_bottom_,"lrn.bottom.data");
  dump_blob(this->blob_top_,"lrn.top.data");
  dump_blob(&top_reference,"lrn.reftop.data");
  
}

#endif


}  // namespace caffe