/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#include <catch2/catch.hpp>
#include <vector>
#include <cstring>
#include <sstream>
#include <fstream>

#include "count_min.hpp"
#include "common_defs.hpp"
#include "test_allocator.hpp"

namespace datasketches {

using count_min_sketch_test_alloc = count_min_sketch<uint64_t, test_allocator<uint64_t>>;
using alloc = test_allocator<uint64_t>;

TEST_CASE("CountMin sketch test allocator: serialize-deserialize empty", "[cm_sketch_alloc]"){
  test_allocator_total_bytes = 0;
  test_allocator_net_allocations = 0;
  {
    uint8_t n_hashes = 1 ;
    uint32_t n_buckets = 5 ;
    std::stringstream s(std::ios::in | std::ios::out | std::ios::binary);
    count_min_sketch_test_alloc c(n_hashes, n_buckets, DEFAULT_SEED, alloc(0)) ;
    c.serialize(s);
    count_min_sketch_test_alloc d = count_min_sketch_test_alloc::deserialize(s, DEFAULT_SEED, alloc(0)) ;
    REQUIRE(c.get_num_hashes() == d.get_num_hashes()) ;
    REQUIRE(c.get_num_buckets() == d.get_num_buckets()) ;
    REQUIRE(c.get_seed() == d.get_seed()) ;
    uint64_t zero = 0;
    REQUIRE(c.get_estimate(zero) == d.get_estimate(zero)) ;
    REQUIRE(c.get_total_weight() == d.get_total_weight()) ;

    // Check that all entries are equal and 0
    for(auto di: d){
      REQUIRE(di == 0) ;
    }
  }
  REQUIRE(test_allocator_total_bytes == 0);
  REQUIRE(test_allocator_net_allocations == 0);
}

TEST_CASE("CountMin sketch test allocator: serialize-deserialize non-empty", "[cm_sketch_alloc]"){
  test_allocator_total_bytes = 0;
  test_allocator_net_allocations = 0;
  {
    uint8_t n_hashes = 3 ;
    uint32_t n_buckets = 1024 ;
    std::stringstream s(std::ios::in | std::ios::out | std::ios::binary);
    count_min_sketch_test_alloc c(n_hashes, n_buckets, DEFAULT_SEED, alloc(0)) ;
    for(uint64_t i=0 ; i < 10; ++i) c.update(i,10*i*i) ;
    c.serialize(s);
    count_min_sketch_test_alloc d = count_min_sketch_test_alloc::deserialize(s, DEFAULT_SEED, alloc(0)) ;
    REQUIRE(c.get_num_hashes() == d.get_num_hashes()) ;
    REQUIRE(c.get_num_buckets() == d.get_num_buckets()) ;
    REQUIRE(c.get_seed() == d.get_seed()) ;
    REQUIRE(c.get_total_weight() == d.get_total_weight()) ;
    for(uint64_t i=0 ; i < 10; ++i){
      REQUIRE(c.get_estimate(i) == d.get_estimate(i)) ;
    }

    auto c_it = c.begin() ;
    auto d_it = d.begin() ;
    while(c_it != c.end()){
      REQUIRE(*c_it == *d_it) ;
      ++c_it ;
      ++d_it ;
    }
  }
  REQUIRE(test_allocator_total_bytes == 0);
  REQUIRE(test_allocator_net_allocations == 0);
}

TEST_CASE("CountMin sketch test allocator: bytes serialize-deserialize empty", "[cm_sketch_alloc]"){
  test_allocator_total_bytes = 0;
  test_allocator_net_allocations = 0;
  {
    uint8_t n_hashes = 3 ;
    uint32_t n_buckets = 32 ;
    count_min_sketch_test_alloc c(n_hashes, n_buckets, DEFAULT_SEED, alloc(0)) ;
    auto bytes = c.serialize() ;

    REQUIRE_THROWS_AS(count_min_sketch_test_alloc::deserialize(bytes.data(), bytes.size(), DEFAULT_SEED-1, alloc(0)), std::invalid_argument);
    auto d = count_min_sketch_test_alloc::deserialize(bytes.data(), bytes.size(), DEFAULT_SEED, alloc(0)) ;
    REQUIRE(c.get_num_hashes() == d.get_num_hashes()) ;
    REQUIRE(c.get_num_buckets() == d.get_num_buckets()) ;
    REQUIRE(c.get_seed() == d.get_seed()) ;
    uint64_t zero = 0;
    REQUIRE(c.get_estimate(zero) == d.get_estimate(zero)) ;
    REQUIRE(c.get_total_weight() == d.get_total_weight()) ;

    // Check that all entries are equal and 0
    for(auto di: d){
      REQUIRE(di == 0) ;
    }
  }
  REQUIRE(test_allocator_total_bytes == 0);
  REQUIRE(test_allocator_net_allocations == 0);
}

TEST_CASE("CountMin sketch test allocator: bytes serialize-deserialize non-empty", "[cm_sketch_alloc]"){
  test_allocator_total_bytes = 0;
  test_allocator_net_allocations = 0;
  {
    uint8_t n_hashes = 5 ;
    uint32_t n_buckets = 64 ;
    count_min_sketch_test_alloc c(n_hashes, n_buckets, DEFAULT_SEED, alloc(0)) ;
    for(uint64_t i=0 ; i < 10; ++i) c.update(i,10*i*i) ;

    auto bytes = c.serialize() ;
    REQUIRE_THROWS_AS(count_min_sketch_test_alloc::deserialize(bytes.data(), bytes.size(), DEFAULT_SEED-1, alloc(0)), std::invalid_argument);
    auto d = count_min_sketch_test_alloc::deserialize(bytes.data(), bytes.size(), DEFAULT_SEED, alloc(0)) ;

    REQUIRE(c.get_num_hashes() == d.get_num_hashes()) ;
    REQUIRE(c.get_num_buckets() == d.get_num_buckets()) ;
    REQUIRE(c.get_seed() == d.get_seed()) ;
    REQUIRE(c.get_total_weight() == d.get_total_weight()) ;

    // Check that all entries are equal
    auto c_it = c.begin() ;
    auto d_it = d.begin() ;
    while(c_it != c.end()){
      REQUIRE(*c_it == *d_it) ;
      ++c_it ;
      ++d_it ;
    }

    // Check that the estimates agree
    for(uint64_t i=0 ; i < 10; ++i){
      REQUIRE(c.get_estimate(i) == d.get_estimate(i)) ;
    }
  }
  REQUIRE(test_allocator_total_bytes == 0);
  REQUIRE(test_allocator_net_allocations == 0);
}

} // namespace datasketches