
import math
from typing import Optional
import warnings
import psutil

import torch
from langport.model.compression import load_compress_model
from langport.model.executor.base import BaseModelExecutor
from langport.model.model_adapter import get_model_adapter, raise_warning_for_incompatible_cpu_offloading_configuration
from langport.model.monkey_patch_non_inplace import replace_llama_attn_with_non_inplace_operations
from langport.utils import get_gpu_memory


def load_model(
    model_path: str,
    device: str,
    num_gpus: int,
    max_gpu_memory: Optional[str] = None,
    load_8bit: bool = False,
    cpu_offloading: bool = False,
    deepspeed: bool = False,
    debug: bool = False,
):
    """Load a model from Hugging Face."""
    adapter = get_model_adapter(model_path)

    # Handle device mapping
    cpu_offloading = raise_warning_for_incompatible_cpu_offloading_configuration(
        device, load_8bit, cpu_offloading
    )
    if device == "cpu":
        kwargs = {"torch_dtype": torch.float32}
    elif device == "cuda":
        kwargs = {"torch_dtype": torch.float16}
        if num_gpus != 1:
            kwargs["device_map"] = "auto"
            if max_gpu_memory is None:
                kwargs[
                    "device_map"
                ] = "sequential"  # This is important for not the same VRAM sizes
                available_gpu_memory = get_gpu_memory(num_gpus)
                kwargs["max_memory"] = {
                    i: str(int(available_gpu_memory[i] * 0.85)) + "GiB"
                    for i in range(num_gpus)
                }
            else:
                kwargs["max_memory"] = {i: max_gpu_memory for i in range(num_gpus)}
    elif device == "mps":
        kwargs = {"torch_dtype": torch.float16}
        # Avoid bugs in mps backend by not using in-place operations.
        replace_llama_attn_with_non_inplace_operations()
    else:
        raise ValueError(f"Invalid device: {device}")

    if cpu_offloading:
        # raises an error on incompatible platforms
        from transformers import BitsAndBytesConfig

        if "max_memory" in kwargs:
            kwargs["max_memory"]["cpu"] = (
                str(math.floor(psutil.virtual_memory().available / 2**20)) + "Mib"
            )
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_8bit_fp32_cpu_offload=cpu_offloading
        )
        kwargs["load_in_8bit"] = load_8bit
        # Load model
        model, tokenizer = adapter.load_model(model_path, kwargs)
    elif load_8bit:
        if num_gpus != 1:
            warnings.warn(
                "8-bit quantization is not supported for multi-gpu inference."
            )
        else:
            model, tokenizer = load_compress_model(
                model_path=model_path, device=device, torch_dtype=kwargs["torch_dtype"]
            )
            # return adapter, model, tokenizer
    else:
        # Load model
        model, tokenizer = adapter.load_model(model_path, kwargs)

    if deepspeed:
        from transformers.deepspeed import HfDeepSpeedConfig
        import deepspeed

        dtype = torch.float16
        if load_8bit:
            dtype = torch.int8
        config = {
            # "mp_size": 1,        # Number of GPU
            "dtype": dtype, # dtype of the weights (fp16)
            "replace_method": "auto", # Lets DS autmatically identify the layer to replace
            "replace_with_kernel_inject": True, # replace the model with the kernel injector
            "enable_cuda_graph": True,
            "tensor_parallel": {
                "enabled": True,
                "tp_size": 1,
            },
        }
        ds_engine = deepspeed.init_inference(model=model, config=config)
        model = ds_engine.module
    else:
        if (device == "cuda" and num_gpus == 1 and not cpu_offloading) or device == "mps":
            model.to(device)

    if debug:
        print(model)

    return adapter, model, tokenizer
