#pragma once

#include "logging.h"
#include "device_functions-generated.h"
#include "model_interface.h"
#include "raii_wrapper.h"
#include "model.h"
#include "macros.h"
#include "jagged.h"
#include <algorithm>
#include <deque>
#include <fstream>
#include <iostream>
#include <string>
#include <unordered_map>
#include <math.h>
#include <iomanip>


void expand_415_constant_folding (
  const int64_t* const src, // input tensor
  int64_t* const dst, // output tensor
  const int64_t head_size, // how many times to repeat the first part of the tensor.
  cudaStream_t stream);

void gemm_rcr_bias_416_constant_folding(
  void*,
  void*,
  void*,
  void*,
  uint8_t*,

    int,


  int64_t*,

  int64_t*,

  int64_t*,


  int64_t*,

  int64_t*,


  int64_t*,

  int64_t*,

  int64_t*,

  cudaStream_t
);

namespace ait {

// Model is the class that actually performs inference. It owns memory for
// intermediate tensors and dynamic dimensions. Constants are owned by
// the model's owning container object, and input/output memory is owned
// by the user.
// Once an inference run has started, it is not safe to re-use the Model
// until the run has finished!
class ConstantFolder : public ModelBase<ConstantFolder> {
  
  

  public:
    ConstantFolder(
        size_t blob_size,
        size_t workspace_size,
        size_t unique_workspace_size,
        size_t num_inputs,
        size_t num_outputs,
        size_t num_unbound_constants,
        uint8_t* constants,
        AITemplateAllocator& allocator)
        : ModelBase(
            blob_size,
            workspace_size,
            unique_workspace_size,
            num_inputs,
            num_outputs,
            num_unbound_constants,
            constants,
            allocator) {
         constant_name_to_ptr_["pool_probe"] = const_cast<const void**>(reinterpret_cast<void**>(&pool_probe));
     constant_name_to_ptr_["pool_q_weight"] = const_cast<const void**>(reinterpret_cast<void**>(&pool_q_weight));
     constant_name_to_ptr_["pool_q_bias"] = const_cast<const void**>(reinterpret_cast<void**>(&pool_q_bias));
    auto* blob_ptr = static_cast<uint8_t*>(blob_.get());
        expand_415_0 = reinterpret_cast<decltype(expand_415_0)>(blob_ptr + 0);
    
         params_[1].shape_ptrs = {ParamDim(1, 1, &pool_probe_dim_0), ParamDim(1, 1, &pool_probe_dim_1), ParamDim(1152, 1152, &pool_probe_dim_2)};
     params_[2].shape_ptrs = {ParamDim(1152, 1152, &pool_q_weight_dim_0), ParamDim(1152, 1152, &pool_q_weight_dim_1)};
     params_[3].shape_ptrs = {ParamDim(1152, 1152, &pool_q_bias_dim_0)};
     params_[0].shape_ptrs = {ParamDim(8, 8, &reshape_417_0_dim_0), ParamDim(16, 16, &reshape_417_0_dim_1), ParamDim(1, 1, &reshape_417_0_dim_2), ParamDim(72, 72, &reshape_417_0_dim_3)};

      
      
    }

    ~ConstantFolder() {
      
      
    }

    void SetUpInputsOutputs() {
        
if (pool_probe == nullptr) {
    throw std::runtime_error("Constant pool_probe was not set! Set the value with set_constant.");
}
    

if (pool_q_weight == nullptr) {
    throw std::runtime_error("Constant pool_q_weight was not set! Set the value with set_constant.");
}
    

if (pool_q_bias == nullptr) {
    throw std::runtime_error("Constant pool_q_bias was not set! Set the value with set_constant.");
}
    
     reshape_417_0 = static_cast<decltype(reshape_417_0)>(params_[0].ptr);

if (reshape_417_0 == nullptr) {
    throw std::runtime_error("Constant reshape_417_0 was not set! Set the value with set_constant.");
}
    
    }

    void ResetConstants(uint8_t* constants) {
        /*
         * This can be called if we want to use a different piece of memory
         * for the constants to be consumed.
         */
        
    }

    void DeviceToDeviceCopies(StreamType stream) {
  
    }


    ///////////////////////////////////////////////////////////////////////////
    // default RunImpl implemenation
    void RunImpl(StreamType stream) {
        
  
  
    {
        expand_415_constant_folding(
            static_cast<int64_t*>(pool_probe),
            static_cast<int64_t*>(expand_415_0),
            8l,
            stream);
    }
    
      DeviceCheckLastError(__FILE__, __LINE__);
  
  
    {
    

    gemm_rcr_bias_416_constant_folding(

        expand_415_0,
        pool_q_weight,

        pool_q_bias,

        reshape_417_0,
        global_workspace_,
        1,

        &expand_415_0_dim_0,

        &pool_probe_dim_1,

        &pool_probe_dim_2,


        &pool_q_weight_dim_0,

        &pool_q_weight_dim_1,


        &expand_415_0_dim_0,

        &pool_probe_dim_1,

        &pool_q_weight_dim_0,

        stream
    );
    }
      DeviceCheckLastError(__FILE__, __LINE__);
  
    }




    void ProfileImpl(StreamType stream, size_t iters, const std::string& filename) {
#ifdef OPTIMIZE_FOR_COMPILATION_TIME
      throw std::runtime_error("Profile is disabled, please recompile without OPTIMIZE_FOR_COMPILE_TIME flag");
#else
      std::ofstream ss(filename);
      if (!ss) {
        throw std::runtime_error(std::string("Could not open file ") + filename);
      }

      int deviceId;
      char* L2CacheSlab = nullptr;
      DevicePropertyType deviceProperties;
      GetDevice(&deviceId);
      GetDeviceProperties(&deviceProperties, deviceId);
      const size_t L2SizeInBytes = deviceProperties.l2CacheSize;
      DeviceMalloc((void**) &L2CacheSlab, L2SizeInBytes);

      ss << "{\n";
      
      {
        std::cout << "Profiling: " << "expand_415" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
    {
        expand_415_constant_folding(
            static_cast<int64_t*>(pool_probe),
            static_cast<int64_t*>(expand_415_0),
            8l,
            stream);
    }
    
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "expand_415" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"1\", \"1\", \"1152\"]]"
           << ", \"output_sizes\": " << "[[\"8\", \"1\", \"1152\"]]"
        
           << " } ";
        
          ss << ",\n";
        
      }
      
      {
        std::cout << "Profiling: " << "gemm_rcr_bias_416" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
    {
    

    gemm_rcr_bias_416_constant_folding(

        expand_415_0,
        pool_q_weight,

        pool_q_bias,

        reshape_417_0,
        global_workspace_,
        1,

        &expand_415_0_dim_0,

        &pool_probe_dim_1,

        &pool_probe_dim_2,


        &pool_q_weight_dim_0,

        &pool_q_weight_dim_1,


        &expand_415_0_dim_0,

        &pool_probe_dim_1,

        &pool_q_weight_dim_0,

        stream
    );
    }
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "gemm_rcr_bias_416" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"8\", \"1\", \"1152\"], [\"1152\", \"1152\"], [\"1152\"]]"
           << ", \"output_sizes\": " << "[[\"8\", \"1\", \"1152\"]]"
        
           << " } ";
        
          ss << "\n";
        
      }
      
      ss << "}\n";

      DeviceToDeviceCopies(stream);
      std::cout << "AIT per op profiling finished." << std::endl;
      FreeDeviceMemory(L2CacheSlab);
#endif
    }

    static std::unique_ptr<ConstantFolder> Create(
      AITemplateAllocator& allocator,
      uint8_t* constants
    ) {
      return std::make_unique<ConstantFolder>(
          36864,
          0 * (1 + 0),
          0 * (1 + 0),
          0,
          1,
          3,
          constants,
          allocator
      );
    }

  private:
   void* pool_probe {nullptr};
   void* pool_q_weight {nullptr};
   void* pool_q_bias {nullptr};
   void* expand_415_0 {nullptr};
   void* reshape_417_0 {nullptr};
   int64_t pool_probe_dim_0 { 1 };
   int64_t pool_probe_dim_1 { 1 };
   int64_t pool_probe_dim_2 { 1152 };
   int64_t pool_q_weight_dim_0 { 1152 };
   int64_t pool_q_weight_dim_1 { 1152 };
   int64_t pool_q_bias_dim_0 { 1152 };
   int64_t expand_415_0_dim_0 { 8 };
   int64_t reshape_417_0_dim_0 { 8 };
   int64_t reshape_417_0_dim_1 { 16 };
   int64_t reshape_417_0_dim_2 { 1 };
   int64_t reshape_417_0_dim_3 { 72 };


};
} // namespace ait