#include <memory.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <vector>
using namespace std;

typedef int int32;
typedef float BaseFloat;
typedef int MatrixIndexT;

#ifdef HAS_MKL
#include "mkl_cblas.h"
#else
#include "cblas.h"
#endif

#include "FileBuf.h"
#include "PerfChecker.h"

void ref_sgemm(
    int M, int N, int K, float alpha, const float *A, int lda, const float *B, int ldb, float beta, float *C, int ldc)
{
    for (int row = 0, a_start = 0, c_start = 0; row < M; row++, a_start += lda, c_start += ldc)
    {
        for (int col = 0, c_index = c_start, b_start = 0; col < N; col++, c_index++, b_start += ldb)
        {
            float t = 0.0f;
            for (int k = 0, a_index = a_start, b_index = b_start; k < K; k++, a_index++, b_index++)
            {
                // t += A[row * lda + k] * B[col * ldb + k];
                t += A[a_index] * B[b_index];
            }

            C[c_index] = t;
        }
    }
}

void TestCorrect()
{
    const int M = 2;
    const int K = 3;
    const int N = 5;
    float A[M * K] = {1, 2, 3, 4, 5, 6};
    float B[K * N] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
    float C[M * N] = {0};

    float alpha = 1.0f;
    float beta = 0.0f;
    int lda = K;
    int ldb = N;
    int ldc = N;

    {
        PERF_CHECKER(checker);

        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc);

        checker.CheckPoint("cblas");

        for (int i = 0; i < M * N; i++) printf("%f ", C[i]);
        printf("\n");
    }

    float B2[K * N] = {1, 6, 11, 2, 7, 12, 3, 8, 13, 4, 9, 14, 5, 10, 15};
    memset(C, 0, M * N * sizeof(float));
    ldb = K;

    {
        PERF_CHECKER(checker);

        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, M, N, K, alpha, A, lda, B2, ldb, beta, C, ldc);

        checker.CheckPoint("cblas trans");

        for (int i = 0; i < M * N; i++) printf("%f ", C[i]);
        printf("\n");
    }

    memset(C, 0, M * N * sizeof(float));
    ldb = K;
    {
        PERF_CHECKER(checker);

        ref_sgemm(M, N, K, alpha, A, lda, B2, ldb, beta, C, ldc);

        checker.CheckPoint("ref");

        for (int i = 0; i < M * N; i++) printf("%f ", C[i]);
        printf("\n");
    }
}

void SaveResult(const char *filename, const float *data, int count)
{
    FILE *fpC = fopen(filename, "wb");
    fwrite(data, sizeof(float), count, fpC);
    fclose(fpC);
}

int TestData(int index)
{
    FileBuf buf_Data, buf_Result;

    char filename[256];
    sprintf(filename, "Data_%d", index);

    if (buf_Data.Load(filename) == false)
    {
        printf("open Data failed!");
        return -1;
    }

    sprintf(filename, "Result_%d", index);
    if (buf_Result.Load(filename) == false)
    {
        printf("open Result failed!");
        return -1;
    }

    int M, N, K;
    const char *data = (char *)buf_Data.buf();
    M = *((int *)data);
    data += sizeof(int);
    N = *((int *)data);
    data += sizeof(int);
    K = *((int *)data);
    data += sizeof(int);

    if (buf_Data.len() != (M * K + K * N + M * N) * sizeof(float) + 3 * sizeof(int))
    {
        printf("size failed A!");
        return -1;
    }


	printf("M: %d, N: %d, K: %d\n", M, N, K);
    float *A = (float *)data;
    data += M * K * sizeof(float);
    float *B = (float *)data;
    data += K * N * sizeof(float);
    float *C = (float *)data;

    float *save_C = (float *)malloc(M * N * sizeof(float));
    memcpy(save_C, C, M * N * sizeof(float));

    float alpha = 1.0f;
    float beta = 0.0f;
    int lda = K;
    int ldb = N;
    int ldc = N;

    {
        PERF_CHECKER(checker);
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc);

        checker.CheckPoint("normal");
        SaveResult("Result_cblas", C, M * N);
    }

    // make A trans
    float *transA = (float *)malloc(K * M * sizeof(float));

    for (int i = 0; i < K; i++)
    {
        for (int j = 0; j < M; j++)
        {
            transA[i * M + j] = A[j * K + i];
        }
    }

    memcpy(C, save_C, M * N * sizeof(float));
    lda = M;
	ldb = N;
    {
        PERF_CHECKER(checker);
        cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, M, N, K, alpha, transA, lda, B, ldb, beta, C, ldc);
        checker.CheckPoint("A trans");

        SaveResult("Result_cblas_trans_A", C, M * N);
    }


    // make B trans
    float *save_B = (float *)malloc(K * N * sizeof(float));
    memcpy(save_B, B, K * N * sizeof(float));

    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < K; j++)
        {
            B[i * K + j] = save_B[j * N + i];
        }
    }

    memcpy(C, save_C, M * N * sizeof(float));
	lda = K;
    ldb = K;
    {
        PERF_CHECKER(checker);
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc);
        checker.CheckPoint("B trans");

        SaveResult("Result_cblas_trans", C, M * N);
    }

    // 还是按照 trans的方式
    memcpy(C, save_C, M * N * sizeof(float));
    {
        PERF_CHECKER(checker);
        ref_sgemm(M, N, K, alpha, A, lda, B, ldb, beta, C, ldc);
        checker.CheckPoint("ref");

        SaveResult("Result_ref_sgemm", C, M * N);
    }

    return 0;
}

int main(int argc, char *argv[])
{
    if (argc < 2)
    {
        TestCorrect();
    }
    else
    {
        int index = atoi(argv[1]);
        TestData(index);
    }
    return 0;
}
