#include <iostream>
#include <vector>
#include <cuda_runtime_api.h>
#include "thrust/host_vector.h"
#include "thrust/device_vector.h"
#include <cublasdx.hpp>
#include "common.hpp"
#include "reference.hpp"
#include "arch_runner.hpp"
#include <thrust/complex.h>

typedef commondx::detail::complex<float> complex;
const unsigned int size=80;
const unsigned int M_TILE = 32; const unsigned int N_TILE =32;const unsigned int K_TILE = 32;  //由于gemm_kernel_shared是小核

void print_device_vector( thrust::device_vector<complex>& vec)
 {
     std::cout << "num"<<vec.size()<< " "<<std::endl;
    for (size_t i = 0; i < vec.size(); ++i) {
            complex  dd=vec[i];
            std::cout << dd.real()<<"i"<<dd.imag() << " "; }
    std::cout << std::endl;
}


template<class GEMM>     //这是小核
__global__ void gemm_kernel_shared(const typename GEMM::c_value_type  alpha,const typename GEMM::a_value_type* a, const typename GEMM::b_value_type* b,const typename GEMM::c_value_type  beta, typename GEMM::c_value_type* c,
                                     int lda, int ldb, int ldc  )
 {
     extern __shared__ __align__(16) char smem[];

      // 确定本 block 负责处理矩阵哪一块 tile
    int block_row = blockIdx.y;
    int block_col = blockIdx.x;

    int row_start = block_row *80;
    int col_start = block_col *80;

    // 边界判断：如果 tile 超过矩阵范围，直接返回（防止越界）
    if (row_start >= 80 || col_start >= 80)
        return;

    // 取 A, B, C 对应 tile 起点
    auto a_tile = a + row_start;
    auto b_tile = b + col_start * ldb;
    auto c_tile = c + row_start + col_start * ldc;

    // Make global memory tensor
    auto a_global_tensor = cublasdx::make_tensor(a_tile, GEMM::get_layout_gmem_a(lda));
    auto b_global_tensor = cublasdx::make_tensor(b_tile, GEMM::get_layout_gmem_b(ldb));
    auto c_global_tensor = cublasdx::make_tensor(c_tile , GEMM::get_layout_gmem_c(ldc));

    // Make shared memory tensor
    auto [smem_a, smem_b, smem_c] = cublasdx::slice_shared_memory<GEMM>(smem);
    auto a_shared_tensor = cublasdx::make_tensor(smem_a, GEMM::get_layout_smem_a());
    auto b_shared_tensor = cublasdx::make_tensor(smem_b, GEMM::get_layout_smem_b());
    auto c_shared_tensor = cublasdx::make_tensor(smem_c, GEMM::get_layout_smem_c());

    // Load data from global memory tensor to shared memory tensor
    using alignment = cublasdx::alignment_of<GEMM>;
    cublasdx::copy<GEMM, alignment::a>(a_global_tensor, a_shared_tensor);
    cublasdx::copy<GEMM, alignment::b>(b_global_tensor, b_shared_tensor);
    cublasdx::copy<GEMM, alignment::c>(c_global_tensor, c_shared_tensor);
    cublasdx::copy_wait();

    // Execute GEMM
    GEMM().execute(alpha, a_shared_tensor, b_shared_tensor, beta, c_shared_tensor);
    __syncthreads();

    // Store data from shared memory tensor to global memory tensor
    cublasdx::copy<GEMM, alignment::c>(c_shared_tensor, c_global_tensor);
}

template<unsigned int Arch>
void tiled_gemm_example(complex alpha,const complex* A,const complex*  B, complex   beta, complex*  C,  cudaStream_t   stream) {
    using GEMM = decltype(cublasdx::Size<M_TILE, N_TILE, K_TILE>()
                  + cublasdx::Precision<float>()
                  + cublasdx::Type<cublasdx::type::complex>()
                  + cublasdx::Arrangement<cublasdx::col_major, cublasdx::col_major>()
                  + cublasdx::Function<cublasdx::function::MM>()
                  + cublasdx::SM<860>()
                  + cublasdx::Block()
                  + cublasdx::BlockDim<512>());

    dim3 grid_dim((80 + N_TILE - 1) / N_TILE, (80 + M_TILE - 1) / M_TILE);
    dim3 block_dim(GEMM::block_dim);
    gemm_kernel_shared<GEMM><<<grid_dim , block_dim, cublasdx::get_shared_storage_size<GEMM>(),stream>>>(alpha, A, B, beta,  C,80,80,80);  //
    cudaPeekAtLastError();
   cudaStreamSynchronize(stream); 
}


void  dss(int i,  cudaStream_t   stream ,complex alpha,complex* A,complex* B,complex beta,complex* C)
{
       // tiled_gemm_example<86><<<1, 1, 0,  reinterpret_cast<CUstream_st*>(stream)>>>(   alpha,A  ,B,beta, C,  stream   );
        tiled_gemm_example<86>(   alpha,A  ,B,beta, C,  stream   );
        cudaError_t err = cudaGetLastError();
        if (err != cudaSuccess) {  printf("CUDA Kernel launch error: %s\n", cudaGetErrorString(err)); }
}


int main(int, char**) {
     complex a=complex {1.0, 1.0};
     complex b={2.0, 2.0};
     complex c={0.0, 0.0};
     int N_Heri=1;
     thrust::device_vector<complex > A(size*size*N_Heri,a); 
     thrust::device_vector<complex > B(size*size*N_Heri,b);
     thrust::device_vector<complex > C(size*size*N_Heri,c);
     complex alpha{1.0, 1.0};  complex beta{1.0, 1.0};


     cudaStream_t streams[N_Heri];
     for (int i=0;i<N_Heri;i++){  cudaStreamCreate(&streams[i]); }
     for (int i=0;i<N_Heri;i++){
               dss(  i,  streams[i],      alpha,thrust::raw_pointer_cast(A.data())+i*size*size,thrust::raw_pointer_cast(B.data())+i*size*size ,beta,thrust::raw_pointer_cast(C.data())+i*size*size );

          }


     for (int i = 0; i <N_Heri; ++i) {cudaStreamDestroy(streams[i]);}

      cudaDeviceSynchronize();

     
    //print_device_vector(A);
    //print_device_vector(B);
    print_device_vector(C);
}