"""
 
* Copyright (c) 2024, The vLLM team.
*
* Licensed 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.
"""

import os
import random
from pathlib import Path

import aiter
import pandas as pd
from aiter import dtypes
import torch
import torch.nn.functional as F
from aiter.test_common import perftest
from aiter.utility.mp_tuner import mp_tuner
from functools import lru_cache
import json
from aiter.ops.awq_gemm_asm import *
from op_tests.test_awq_gemm import awq_post_dequant_torch

aiter.rocb_create_extension()
aiter.hipb_create_extension()


@lru_cache(maxsize=1)
def init_hipblas():
    aiter.hipb_create_extension()


@lru_cache(maxsize=1)
def init_rocblas():
    aiter.rocb_create_extension()


def call_hipb_mm(input, weight, solidx, bias, out_dtype, scale_a=None, scale_b=None, scale_type=None):
    init_hipblas()
    return aiter.hipb_mm(
        input,
        weight,
        solidx,
        bias=bias,
        out_dtype=out_dtype,
        scaleA=scale_a,
        scaleB=scale_b,
        scaleType=scale_type,
    )


def call_rocb_mm(inp, w, solidx):
    init_rocblas()
    return aiter.rocb_mm(inp, w, solidx)

def call_hsaco_awq_mm(out, iweights, input, zeros, scales, solidx, inputSols_file):
    #init_rocblas()
    aiter.awq_gemm_asm_tuning(out, iweights, input, zeros, scales, solidx, inputSols_file)
    return out



rtol = 1e-5
atol = 1

CACHE_INVALIDATE_BUFFERS = int(os.getenv("CACHE_INVALIDATE_BUFFERS", "37"))
ONE = torch.ones(1, dtype=dtypes.fp32, device="cuda")
HALF = torch.tensor(0.5, dtype=dtypes.fp32, device="cuda")


class Gemm:

    def __init__(
        self,
        m,
        n,
        k,
        g,
        bias,
        indtype,
        outdtype,
        scaleAB=False,
        rocblas_decode=False,
        mp=1,
        inputSols_file=None,
        awqgemm=False,
        fastNoCheck=1
    ):
        self.m = m
        self.k = k
        self.n = n
        self.bias = torch.randn(n, device="cuda").to(outdtype) if bias else None
        self.indtype = indtype
        self.outdtype = outdtype
        self.scaleAB = scaleAB
        self.use_rocblas = indtype == outdtype and str(indtype) != "dtypes.fp8"
        self.nb = CACHE_INVALIDATE_BUFFERS

        self.awqgemm = awqgemm  
        self.g = g
        if self.awqgemm:
            self.awq_w4a16_init()
            dequantized_weights, iweights, zeros, iscales = awq_post_dequant_torch(self.weights, self.scales, self.qzeros, self.g)
            if fastNoCheck == 1:
                self.ref  = None
            else:
                self.ref  = torch.matmul(self.inp, dequantized_weights)

        else:
            self.inp = torch.randn((self.m, self.k), device="cuda").to(self.indtype)
            self.weights = torch.randn((self.n, self.k), device="cuda").to(self.indtype)
            if fastNoCheck == 1:
                self.ref  = None
            else:
                self.ref = self.get_gemm_ref()

        self.blob = torch.ones(128 * 1024 * 1024, dtype=dtypes.fp32, device="cuda")
        self.topn = 20  # number of top solutions from each source
        self.hipb_sols = []
        self.rocb_sols = []
        self.rtol = 1e-2
        self.atol = 1e-2
        self.check_err_ratio = 0.01
        self.start = torch.cuda.Event(enable_timing=True)
        self.end = torch.cuda.Event(enable_timing=True)
        # prefer hipblaslt unless rocblas time is less than this
        # ratio of hipblaslt time
        self.hipb_prefer_ratio = 0.995
        self.rocblas_decode = rocblas_decode
        self.mp = mp
        self.inputSols_file = inputSols_file

    def awq_w4a16_init(self):

        dtype=dtypes.fp16
        G = self.g if self.g != -1 else self.k
        device = "cuda"
        input_rows = self.m
        input_cols = self.k
        input_dtype = dtype
        qweight_rows = input_cols
        qweight_cols = self.n // 2
        scales_rows = qweight_rows // G
        scales_cols = self.n
        scales_dtype = dtype
        qzeros_rows = scales_rows
        qzeros_cols = qweight_cols

        self.inp = torch.rand((input_rows, input_cols),
                dtype=input_dtype,
                device=device)
        
        self.weights = torch.randint(0,
                            torch.iinfo(torch.int8).max,
                            (qweight_rows, qweight_cols),
                            dtype=torch.int8,
                            device=device)
        self.qzeros = torch.randint(0,
                            torch.iinfo(torch.int8).max,
                            (qzeros_rows, qzeros_cols),
                            dtype=torch.int8,
                            device=device)
        self.scales = torch.rand((scales_rows, scales_cols),
                        dtype=scales_dtype,
                        device=device)

        self.out_asm = torch.rand((self.m, self.n),
                        dtype=input_dtype,
                        device=device)

    def find_hipblas_sols(self):
        sols = aiter.hipb_findallsols(
            self.inp,
            self.weights.t(),
            bias=self.bias,
            out_dtype=self.outdtype,
            scaleA=HALF if self.scaleAB else None,
            scaleB=HALF if self.scaleAB else None,
        )
        print(
            "M N K bias dtype outdtype",
            self.m,
            self.n,
            self.k,
            self.bias is not None,
            self.indtype,
            self.outdtype,
            self.scaleAB,
            ">>> Total hipb solutions",
            len(sols),
            flush=True,
        )
        # print(sols)
        self.hipb_sols = sols

    def get_gemm_ref(self):
        scaleA = HALF if self.scaleAB else ONE
        scaleB = HALF if self.scaleAB else ONE
        if self.indtype == dtypes.fp8:
            try:
                ref = torch._scaled_mm(
                    self.inp,
                    self.weights.t(),
                    bias=self.bias,
                    scale_a=scaleA,
                    scale_b=scaleB,
                    out_dtype=self.outdtype,
                )
            except RuntimeError:
                ref = (
                    F.linear(self.inp.to(dtypes.fp32), self.weights.to(dtypes.fp32))
                    * scaleA
                    * scaleB
                )
                ref = (
                    (ref.to(self.outdtype) + self.bias)
                    if self.bias is not None
                    else ref.to(self.outdtype)
                )
            if type(ref) is tuple and len(ref) == 2:
                ref = ref[0]
        else:
            ref = F.linear(self.inp, self.weights, self.bias).to(self.outdtype)
        return ref

    def hipb_time_all_sols(self, fast_mode=0, top_sols=0):
        coldi = 20
        warmi = 20
        if fast_mode:
            coldi = 2
            warmi = 5
        solutions = self.hipb_sols
        if top_sols:
            solutions = self.hipb_top_sols
        task = []
        scaleA = HALF if self.scaleAB else None
        scaleB = HALF if self.scaleAB else None

        gtimes = {}
        for solidx in solutions:
            task.append(
                (
                    solidx,
                    call_hipb_mm,
                    (
                        self.inp,
                        self.weights.t(),
                        solidx,
                        self.bias if self.bias is not None else None,
                        self.outdtype,
                        scaleA,
                        scaleB,
                    ),
                    {
                        "num_warmup": warmi,
                        "num_iters": coldi,
                    },
                    self.ref if fast_mode == 0 else None,
                    self.rtol,
                    self.atol,
                )
            )
        ret = mp_tuner(task, self.mp)
        for solidx, us, err_ratio in ret:
            if fast_mode == 0:
                if err_ratio > self.check_err_ratio:
                    continue
            gtimes[solidx] = us / 1000.0
        self.hipb_gtimedf = pd.DataFrame.from_dict(
            gtimes, orient="index", columns=["gtimems"]
        ).sort_values(by="gtimems")
        self.hipb_gtimedf.to_csv("/tmp/hipb_gtimedf.csv")
        print(">>> HipBlasLt top solutions, Fast Mode", fast_mode)
        print(self.hipb_gtimedf.head(self.topn))

    def find_rocblas_sols(self):
        if self.scaleAB or self.bias is not None:
            sols = []
        else:
            sols = aiter.rocb_findallsols(self.inp, self.weights.t())
        print(
            "M N K dtype",
            self.m,
            self.n,
            self.k,
            self.indtype,
            self.outdtype,
            ">>> Total rocb solutions",
            len(sols),
            flush=True,
        )
        # print(sols)
        self.rocb_sols = sols

    def rocb_time_all_sols(self, fast_mode=0, top_sols=0):
        coldi = 20
        warmi = 20
        if fast_mode:
            coldi = 2
            warmi = 5
        solutions = self.rocb_sols
        if top_sols:
            solutions = self.rocb_top_sols
        task = []
        gtimes = {}
        for solidx in solutions:
            task.append(
                (
                    solidx,
                    call_rocb_mm,
                    (
                        self.inp,
                        self.weights.t(),
                        solidx,
                    ),
                    {
                        "num_warmup": warmi,
                        "num_iters": coldi,
                    },
                )
            )
        ret = mp_tuner(task, self.mp)
        for solidx, us, err_ratio in ret:
            if fast_mode == 0:
                if err_ratio > self.check_err_ratio:
                    continue
            gtimes[solidx] = us / 1000.0
        self.rocb_gtimedf = pd.DataFrame.from_dict(
            gtimes, orient="index", columns=["gtimems"]
        ).sort_values(by="gtimems")
        self.rocb_gtimedf.to_csv("/tmp/rocb_gtimedf.csv")
        print(">>> Rocblas top solutions, Fast Mode", fast_mode, flush=True)
        print(self.rocb_gtimedf.head(self.topn), flush=True)

    def find_hsaco_sols(self):

        sols = []
        sols_names = []
        with open(self.inputSols_file, 'r', encoding='utf-8') as file:
            data = json.load(file)

        for solution in data["kernels"]:
            sols.append(solution["solutionId"])
            sols_names.append(solution["kernel_name"])
            #print(">>> hsaco solutions ",solution, flush=True)

        print(
            "M N K dtype",
            self.m,
            self.n,
            self.k,
            self.indtype,
            self.outdtype,
            ">>> Total hsaco solutions",
            len(sols),
            flush=True,
        )
        # print(sols)
        self.hsaco_sols = sols
        self.hsaco_sols_name = sols_names


    def hsaco_time_all_sols(self, fast_mode=0, top_sols=0):

        warmi = self.warmupIters
        coldi = self.runIters

        solutions = self.hsaco_sols
        if top_sols:
            solutions = self.hsaco_top_sols
        task = []
        gtimes = {}
        for solidx in solutions:
            task.append(
                (
                    solidx,
                    call_hsaco_awq_mm,
                    (
                        self.out_asm,
                        self.weights,
                        self.inp,
                        self.qzeros,
                        self.scales,
                        solidx,
                        self.inputSols_file,
                    ),
                    {
                        "num_warmup": warmi,
                        "num_iters": coldi,
                    },
                    self.ref if fast_mode == 0 else None,
                )
            )
        ret = mp_tuner(task, self.mp)
        for solidx, us, err_ratio in ret:
            check = "NO_CHECK"
            if fast_mode == 0:
                if err_ratio > self.check_err_ratio:
                    check = "FAIL"
                    continue
                else:
                    check = "PASS"
            gtimes[solidx] = [us / 1000.0, check]
        self.hsaco_gtimedf = pd.DataFrame.from_dict(
            gtimes, orient="index", columns=["gtimems", "status"]
        ).sort_values(by="gtimems")
        self.hsaco_gtimedf.to_csv("/tmp/hsaco_gtimedf.csv")
        self.kernel_name = self.hsaco_sols_name[solutions[self.hsaco_gtimedf.index[0]]]
        print(">>> hsaco top solutions, Fast Mode", fast_mode, flush=True)
        print(self.hsaco_gtimedf.head(self.topn), flush=True)
        
    def warmup(self, warmi=500):
        for i in range(warmi):
            self.blob = self.blob + 0.00001

    def functional_get_topn_fastest(self):
        rocb_topn = []
        for solidx in self.rocb_gtimedf.index[: self.topn]:
            rocb_topn.append(solidx)
        self.rocb_top_sols = rocb_topn
        hipb_topn = []
        for solidx in self.hipb_gtimedf.index[: self.topn]:
            hipb_topn.append(solidx)
        self.hipb_top_sols = hipb_topn

    def find_fastest_solution(self):
        if self.hsacoOnly == 1:
            self.find_hsaco_sols()
            self.warmup()
            if self.fastNoCheck == 0:
                self.hsaco_time_all_sols(fast_mode=0)
            else:
                self.hsaco_time_all_sols(fast_mode=1)

            if self.fastNoCheck == 2:  #run one more time to check pass or fail
                hsaco_topn = []
                for solidx in self.hsaco_gtimedf.index[: 1]:
                    hsaco_topn.append(solidx)
                self.hsaco_top_sols = hsaco_topn
                self.hsaco_time_all_sols(fast_mode=0, top_sols=1)

            if len(self.hsaco_gtimedf) > 0:
                print(">>> Hsaco solutions found!", flush=True)
                self.best_libtype = "hsaco"
                self.best_solidx = self.hsaco_gtimedf.index[0]
                self.best_soltime = self.hsaco_gtimedf.gtimems.iloc[0]
            print(
                ">>> Fastest Solution is",
                self.best_libtype,
                self.best_solidx,
                self.best_soltime,
                flush=True,
            )
            return

        if self.use_rocblas:
            self.find_rocblas_sols()
        if not (self.rocblas_decode and self.m == 1):
            self.find_hipblas_sols()
        self.warmup()
        self.rocb_time_all_sols(fast_mode=1)
        self.warmup()
        self.hipb_time_all_sols(fast_mode=1)
        self.functional_get_topn_fastest()
        self.warmup()
        self.rocb_time_all_sols(fast_mode=0, top_sols=1)
        self.warmup()
        self.hipb_time_all_sols(fast_mode=0, top_sols=1)
        if len(self.rocb_gtimedf) > 0 and len(self.hipb_gtimedf) > 0:
            best_rocb_time = self.rocb_gtimedf.gtimems.iloc[0]
            best_hipb_time = self.hipb_gtimedf.gtimems.iloc[0]
            if best_rocb_time < best_hipb_time * self.hipb_prefer_ratio:
                self.best_libtype = "rocblas"
                self.best_solidx = self.rocb_gtimedf.index[0]
                self.best_soltime = best_rocb_time
            else:
                self.best_libtype = "hipblaslt"
                self.best_solidx = self.hipb_gtimedf.index[0]
                self.best_soltime = best_hipb_time
            # self.check_gemm_ref(self.best_libtype,self.best_solidx)
        elif len(self.hipb_gtimedf) > 0:
            print(">>> Only hipblas solutions found!", flush=True)
            best_hipb_time = self.hipb_gtimedf.gtimems.iloc[0]
            self.best_libtype = "hipblaslt"
            self.best_solidx = self.hipb_gtimedf.index[0]
            self.best_soltime = best_hipb_time
        elif len(self.rocb_gtimedf) > 0:
            print(">>> Only rocblas solutions found!", flush=True)
            best_rocb_time = self.rocb_gtimedf.gtimems.iloc[0]
            self.best_libtype = "rocblas"
            self.best_solidx = self.rocb_gtimedf.index[0]
            self.best_soltime = best_rocb_time
        else:
            print(">>> No rocblas or hipblas solutions found!", flush=True)
            self.best_libtype = "rocblas"
            self.best_solidx = 0
            self.best_soltime = 0
        print(
            ">>> Fastest Solution is",
            self.best_libtype,
            self.best_solidx,
            self.best_soltime,
            flush=True,
        )


class GemmTuner:

    def __init__(self, indtype, outdtype, tuned_file=None, rocblas_decode=False, mp=1, inputSols_file=None):
        self.gemm_problems = pd.DataFrame(columns=["M", "N", "K", "G", "bias"])
        self.indtype = indtype
        self.outdtype = outdtype
        self.rocblas_decode = rocblas_decode
        self.tuned_file = tuned_file
        self.mp = mp
        if Path(tuned_file).is_file():
            self.tuned_shapes = pd.read_csv(tuned_file)
        else:
            self.tuned_shapes = None
        self.inputSols_file = inputSols_file

    def add_gemm(self, m, n, k, g, indtype, bias=False, outdtype=None, scaleAB=False, awq=False):
        assert indtype is not None
        outdtype = outdtype if outdtype is not None else indtype
        assert outdtype is not None
        if self.tuned_shapes is None or (
            self.tuned_shapes[
                (self.tuned_shapes["M"] == m)
                & (self.tuned_shapes["N"] == n)
                & (self.tuned_shapes["K"] == k)
                & (self.tuned_shapes["bias"] == bias)
                & (self.tuned_shapes["dtype"] == str(indtype))
                & (self.tuned_shapes["outdtype"] == str(outdtype))
            ].empty
        ):
            entry = {
                "M": [m],
                "N": [n],
                "K": [k],
                "G": [g],
                "bias": [bias],
                "dtype": [indtype],
                "outdtype": [outdtype],
                "scaleAB": [scaleAB],
                "awqgemm": [awq],
            }
            df = pd.DataFrame(entry)
            self.gemm_problems = pd.concat([self.gemm_problems, df], ignore_index=True)
        else:
            print(
                f">>>Info: Found Duplicate shape(M:{m},"
                f" N:{n}, K:{k} bias:{bias}), skipping"
            )

    def find_best_sols(self):
        df = self.gemm_problems
        soldf = pd.DataFrame(columns=["libtype", "solidx", "soltimes", "kernelName"])
        for i in range(len(df)):
            ds = df.loc[i, :]
            indtype = ds["dtype"]
            outdtype = ds["outdtype"]
            gemmobj = Gemm(
                ds["M"],
                ds["N"],
                ds["K"],
                ds["G"],
                ds["bias"],
                indtype=indtype,
                outdtype=outdtype,
                scaleAB=ds["scaleAB"],
                rocblas_decode=self.rocblas_decode,
                mp=self.mp,
                inputSols_file=self.inputSols_file,
                awqgemm=ds["awqgemm"],
                fastNoCheck = self.fastNoCheck
            )
            gemmobj.warmupIters = self.warmupIters
            gemmobj.runIters = self.runIters
            gemmobj.fastNoCheck = self.fastNoCheck
            gemmobj.hsacoOnly = self.hsacoOnly

            gemmobj.find_fastest_solution()
            soldf.loc[i, "libtype"] = gemmobj.best_libtype
            soldf.loc[i, "solidx"] = gemmobj.best_solidx
            soldf.loc[i, "soltimes"] = round(gemmobj.best_soltime * 1000, 2)
            soldf.loc[i, "kernelName"] = (
                aiter.getHipblasltKernelName(int(gemmobj.best_solidx))
                if gemmobj.best_libtype == "hipblaslt"
                else ""
            )
            if gemmobj.best_libtype == "hsaco":
                soldf.loc[i, "kernelName"] = gemmobj.kernel_name

            del gemmobj
            torch.cuda.empty_cache()

        finaldf = pd.concat([self.gemm_problems, soldf], axis=1)
        if self.tuned_shapes is not None:
            finaldf = pd.concat([finaldf, self.tuned_shapes])
        finaldf["solidx"] = finaldf["solidx"].convert_dtypes("int64")
        finaldf.to_csv(self.tuned_file, index=False)
        pd.set_option('display.width', None)
        pd.set_option('display.max_colwidth', None)
        print(finaldf)
