# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 numpy as np
import scipy


class PrefillInterpolator:
    """
    Takes input from results of pre-deployment performance profiling to interpolate
    throughput/gpu and TTFT for a given ISL.
    """

    def __init__(self, profile_results_dir: str):
        prefill_npz_fn = (
            f"{profile_results_dir}/selected_prefill_interpolation/raw_data.npz"
        )

        with np.load(prefill_npz_fn) as raw_data:
            self.prefill_isl = raw_data["prefill_isl"]
            self.prefill_ttft = raw_data["prefill_ttft"] / 1000  # convert ms to s
            self.prefill_thpt_per_gpu = raw_data["prefill_thpt_per_gpu"]

        self.min_isl = min(self.prefill_isl)
        self.max_isl = max(self.prefill_isl)

        # perform 1d interpolation
        self.ttft_interpolator = scipy.interpolate.interp1d(
            self.prefill_isl, self.prefill_ttft, kind="cubic"
        )
        self.thpt_interpolator = scipy.interpolate.interp1d(
            self.prefill_isl, self.prefill_thpt_per_gpu, kind="cubic"
        )

    def interpolate_ttft(self, isl: float) -> float:
        isl = max(self.min_isl, min(isl, self.max_isl))
        return self.ttft_interpolator(isl)

    def interpolate_thpt_per_gpu(self, isl: float) -> float:
        isl = max(self.min_isl, min(isl, self.max_isl))
        return self.thpt_interpolator(isl)


class DecodeInterpolator:
    """
    Takes input from results of pre-deployment performance profiling to interpolate
    throughput/gpu and ITL for a given decode context length.
    """

    def __init__(self, profile_results_dir: str, resolution: int = 100):
        decode_npz_fn = (
            f"{profile_results_dir}/selected_decode_interpolation/raw_data.npz"
        )

        with np.load(decode_npz_fn) as raw_data:
            self.x_kv_usage = raw_data["x_kv_usage"]
            self.y_context_length = raw_data["y_context_length"]
            self.z_itl = raw_data["z_itl"]
            self.z_thpt_per_gpu = raw_data["z_thpt_per_gpu"]
            self.max_kv_tokens = raw_data["max_kv_tokens"][0]

        # pre-compute the interpolation grid for fast lookup
        self.resolution = resolution
        self.xi = np.linspace(0, 1, resolution)
        self.yi = np.linspace(0, max(self.y_context_length), resolution)
        self.X, self.Y = np.meshgrid(self.xi, self.yi)

        # perform 2d interpolation with fallback for NaN values
        self.itl_interpolator = scipy.interpolate.griddata(
            (self.x_kv_usage, self.y_context_length),
            self.z_itl,
            (self.X, self.Y),
            method="cubic",
        )
        # Fill NaN values using nearest neighbor interpolation
        nan_mask = np.isnan(self.itl_interpolator)
        if np.any(nan_mask):
            itl_nearest = scipy.interpolate.griddata(
                (self.x_kv_usage, self.y_context_length),
                self.z_itl,
                (self.X, self.Y),
                method="nearest",
            )
            self.itl_interpolator[nan_mask] = itl_nearest[nan_mask]
        self.itl_interpolator /= 1000  # convert ms to s

        self.thpt_interpolator = scipy.interpolate.griddata(
            (self.x_kv_usage, self.y_context_length),
            self.z_thpt_per_gpu,
            (self.X, self.Y),
            method="cubic",
        )
        # Fill NaN values using nearest neighbor interpolation
        nan_mask = np.isnan(self.thpt_interpolator)
        if np.any(nan_mask):
            thpt_nearest = scipy.interpolate.griddata(
                (self.x_kv_usage, self.y_context_length),
                self.z_thpt_per_gpu,
                (self.X, self.Y),
                method="nearest",
            )
            self.thpt_interpolator[nan_mask] = thpt_nearest[nan_mask]

    def compute_idx(self, concurrency: float, context_length: float) -> tuple[int, int]:
        kv_usage = concurrency * context_length / self.max_kv_tokens
        # Calculate x index (kv_usage)
        ix = int(
            np.clip(
                np.round((kv_usage - self.xi[0]) / (self.xi[1] - self.xi[0])),
                0,
                self.resolution - 1,
            )
        )
        # Calculate y index (context_length)
        iy = int(
            np.clip(
                np.round((context_length - self.yi[0]) / (self.yi[1] - self.yi[0])),
                0,
                self.resolution - 1,
            )
        )
        return ix, iy

    def interpolate_itl(self, concurrency: float, context_length: float) -> float:
        ix, iy = self.compute_idx(concurrency, context_length)
        return self.itl_interpolator[iy, ix]

    def interpolate_thpt_per_gpu(
        self, concurrency: float, context_length: float
    ) -> float:
        ix, iy = self.compute_idx(concurrency, context_length)
        return self.thpt_interpolator[iy, ix]

    def find_best_throughput_per_gpu(
        self, itl: float, context_length: float
    ) -> tuple[float, float, float]:
        # find the max kv_load that has itl <= target itl
        # here we cannot use binary search as interpolated itl might not be monotonic
        iy = int(
            np.clip(
                np.round((context_length - self.yi[0]) / (self.yi[1] - self.yi[0])),
                0,
                self.resolution - 1,
            )
        )
        iy = max(0, min(iy, self.resolution - 1))

        for ix in range(self.resolution - 1, -1, -1):
            if self.itl_interpolator[iy, ix] <= itl:
                return (
                    self.thpt_interpolator[iy, ix],
                    self.itl_interpolator[iy, ix],
                    self.xi[ix],
                )
        return self.thpt_interpolator[iy, 0], self.itl_interpolator[iy, 0], self.xi[0]


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--profile_results_dir", type=str, required=True)
    parser.add_argument("--isl", type=int, default=3000)
    parser.add_argument("--osl", type=int, default=150)
    parser.add_argument("--ttft", type=float, default=0.1, help="in s")
    parser.add_argument("--itl", type=float, default=0.01, help="in s")
    args = parser.parse_args()

    print(f"ISL={args.isl}, OSL={args.osl}")
    print(f"TTFT={args.ttft}s, ITL={args.itl}s")
    print(f"Using profile results from {args.profile_results_dir}")
    print("")

    # first interpolate prefill
    print("Interpolating prefill performance ...")
    prefill_interpolator = PrefillInterpolator(args.profile_results_dir)

    est_ttft = prefill_interpolator.interpolate_ttft(args.isl)
    est_thpt_per_gpu = prefill_interpolator.interpolate_thpt_per_gpu(args.isl)

    if est_ttft <= args.ttft:
        print(
            f"\tEstimated TTFT={est_ttft:.3f}s <= target TTFT={args.ttft:.3f}s. Requests can queue {args.ttft - est_ttft:.3f}s maximally while meeting TTFT SLA."
        )
    else:
        print(
            f"\tEstimated TTFT={est_ttft:.3f}s > target TTFT={args.ttft:.3f}s. Cannot meet TTFT SLA."
        )

    print(
        f"\tEstimated throughput: {est_thpt_per_gpu:.2f} tokens/s/gpu. Request rate at {est_thpt_per_gpu / args.isl:.2f} requests/s will saturate one GPU."
    )

    print("")

    # then interpolate decode
    decode_interpolator = DecodeInterpolator(args.profile_results_dir)

    print("Interpolating decode performance ...")
    context_length = args.isl + args.osl // 2
    print(f"\tAverage context length: isl + osl/2 = {context_length}.")
    (
        est_thpt_per_gpu,
        est_itl,
        est_kv_usage,
    ) = decode_interpolator.find_best_throughput_per_gpu(args.itl, context_length)
    if est_itl <= args.itl:
        print(
            f"\tEstimated ITL={est_itl:.4f}s <= target ITL={args.itl:.4f}s at {est_kv_usage*100:.2f}% active kv usage."
        )
        print(
            f"\tEstimated throughput: {est_thpt_per_gpu:.2f} token/s/gpu. Request rate at {est_thpt_per_gpu / args.osl:.2f} requests/s will saturate one GPU."
        )
    else:
        print(
            f"\tEstimated ITL={est_itl:.4f}s > target ITL={args.itl:.4f}s. Cannot meet ITL SLA."
        )
