import json
import os
import glob
import sys
import time
from pathlib import Path
from typing import Tuple

from huggingface_hub import hf_hub_download
from PIL import Image
import gradio as gr
import torch
from fairscale.nn.model_parallel.initialize import initialize_model_parallel

from llama import LLaMA, ModelArgs, Tokenizer, Transformer, VisionModel

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

PROMPT_DICT = {
    "prompt_input": (
        "Below is an instruction that describes a task, paired with an input that provides further context. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:"
    ),
    "prompt_no_input": (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        "### Instruction:\n{instruction}\n\n### Response:"
    ),
}


def setup_model_parallel() -> Tuple[int, int]:
    os.environ['RANK'] = '0'
    os.environ['WORLD_SIZE'] = '1'
    os.environ['MP'] = '1'
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '2223'
    local_rank = int(os.environ.get("LOCAL_RANK", -1))
    world_size = int(os.environ.get("WORLD_SIZE", -1))

    torch.distributed.init_process_group("nccl")
    initialize_model_parallel(world_size)
    torch.cuda.set_device(local_rank)

    # seed must be the same in all processes
    torch.manual_seed(1)
    return local_rank, world_size


def load(
    ckpt0_path: str,
    ckpt1_path: str,
    param_path: str,
    tokenizer_path: str,
    instruct_adapter_path: str,
    caption_adapter_path: str,
    local_rank: int,
    world_size: int,
    max_seq_len: int,
    max_batch_size: int,
) -> LLaMA:
    start_time = time.time()
    print("Loading")
    instruct_adapter_checkpoint = torch.load(
        instruct_adapter_path, map_location="cpu")
    caption_adapter_checkpoint = torch.load(
        caption_adapter_path, map_location="cpu")
    with open(param_path, "r") as f:
        params = json.loads(f.read())

    model_args: ModelArgs = ModelArgs(
        max_seq_len=max_seq_len, max_batch_size=max_batch_size, **params
    )
    model_args.adapter_layer = int(
        instruct_adapter_checkpoint['adapter_query.weight'].shape[0] / model_args.adapter_len)
    model_args.cap_adapter_layer = int(
        caption_adapter_checkpoint['cap_adapter_query.weight'].shape[0] / model_args.cap_adapter_len)

    tokenizer = Tokenizer(model_path=tokenizer_path)
    model_args.vocab_size = tokenizer.n_words
    torch.set_default_tensor_type(torch.cuda.HalfTensor)
    model = Transformer(model_args)

    # To reduce memory usuage
    ckpt0 = torch.load(ckpt0_path, map_location='cuda')
    model.load_state_dict(ckpt0, strict=False)
    del ckpt0
    torch.cuda.empty_cache()

    ckpt1 = torch.load(ckpt1_path, map_location='cuda')
    model.load_state_dict(ckpt1, strict=False)
    del ckpt1
    torch.cuda.empty_cache()

    vision_model = VisionModel(model_args)

    torch.set_default_tensor_type(torch.FloatTensor)
    model.load_state_dict(instruct_adapter_checkpoint, strict=False)
    model.load_state_dict(caption_adapter_checkpoint, strict=False)
    vision_model.load_state_dict(caption_adapter_checkpoint, strict=False)

    generator = LLaMA(model, tokenizer, vision_model)
    print(f"Loaded in {time.time() - start_time:.2f} seconds")
    return generator


def instruct_generate(
    instruct: str,
    input: str = 'none',
    max_gen_len=512,
    temperature: float = 0.1,
    top_p: float = 0.75,
):
    if input == 'none':
        prompt = PROMPT_DICT['prompt_no_input'].format_map(
            {'instruction': instruct, 'input': ''})
    else:
        prompt = PROMPT_DICT['prompt_input'].format_map(
            {'instruction': instruct, 'input': input})

    results = generator.generate(
        [prompt], max_gen_len=max_gen_len, temperature=temperature, top_p=top_p
    )
    result = results[0].strip()
    print(result)
    return result


def download_llama_adapter(instruct_adapter_path, caption_adapter_path):
    if not os.path.exists(instruct_adapter_path):
        os.system(
            f"wget -q -O {instruct_adapter_path} https://github.com/ZrrSkywalker/LLaMA-Adapter/releases/download/v.1.0.0/llama_adapter_len10_layer30_release.pth")

    if not os.path.exists(caption_adapter_path):
        os.system(
            f"wget -q -O {caption_adapter_path} https://github.com/ZrrSkywalker/LLaMA-Adapter/releases/download/v.1.0.0/llama_adapter_len10_layer30_caption_vit_l.pth")


# ckpt_path = "/data1/llma/7B/consolidated.00.pth"
# param_path = "/data1/llma/7B/params.json"
# tokenizer_path = "/data1/llma/tokenizer.model"
ckpt0_path = hf_hub_download(
    repo_id="csuhan/llama_storage", filename="consolidated.00_part0.pth")
ckpt1_path = hf_hub_download(
    repo_id="csuhan/llama_storage", filename="consolidated.00_part1.pth")
param_path = hf_hub_download(
    repo_id="nyanko7/LLaMA-7B", filename="params.json")
tokenizer_path = hf_hub_download(
    repo_id="nyanko7/LLaMA-7B", filename="tokenizer.model")
instruct_adapter_path = "llama_adapter_len10_layer30_release.pth"
caption_adapter_path = "llama_adapter_len10_layer30_caption_vit_l.pth"
max_seq_len = 512
max_batch_size = 1

# download models
# download_llama_adapter(instruct_adapter_path, caption_adapter_path)

local_rank, world_size = setup_model_parallel()
if local_rank > 0:
    sys.stdout = open(os.devnull, "w")

generator = load(
    ckpt0_path, ckpt1_path, param_path, tokenizer_path, instruct_adapter_path, caption_adapter_path, local_rank, world_size, max_seq_len, max_batch_size
)


def create_instruct_demo():
    with gr.Blocks() as instruct_demo:
        with gr.Row():
            with gr.Column():
                instruction = gr.Textbox(lines=2, label="Instruction")
                input = gr.Textbox(
                    lines=2, label="Context input", placeholder='none')
                max_len = gr.Slider(minimum=1, maximum=512,
                                    value=128, label="Max length")
                with gr.Accordion(label='Advanced options', open=False):
                    temp = gr.Slider(minimum=0, maximum=1,
                                     value=0.1, label="Temperature")
                    top_p = gr.Slider(minimum=0, maximum=1,
                                      value=0.75, label="Top p")

                run_botton = gr.Button("Run")

            with gr.Column():
                outputs = gr.Textbox(lines=10, label="Output")

        inputs = [instruction, input, max_len, temp, top_p]

        examples = [
            "Tell me about alpacas.",
            "Write a Python program that prints the first 10 Fibonacci numbers.",
            "Write a conversation between the sun and pluto.",
            "Write a theory to explain why cat never existed",
        ]
        examples = [
            [x, "none", 128, 0.1, 0.75]
            for x in examples]

        gr.Examples(
            examples=examples,
            inputs=inputs,
            outputs=outputs,
            fn=instruct_generate,
            cache_examples=os.getenv('SYSTEM') == 'spaces'
        )
        run_botton.click(fn=instruct_generate, inputs=inputs, outputs=outputs)
    return instruct_demo


description = """
# TAPA: xxx
"""

with gr.Blocks(css='style.css') as demo:
    gr.Markdown(description)
    with gr.TabItem("Instruction-Following"):
        create_instruct_demo()

demo.queue(api_open=True, concurrency_count=1).launch()