from transformers import LlamaForCausalLM, LlamaTokenizer, BitsAndBytesConfig, GenerationConfig
from utils import setup_device
import torch
import tqdm
import os

model_name = os.environ.get('LLM_MODEL')

model_path = "models/CRYSTAL-instruct" if model_name == None else model_name

device = setup_device()

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float32 if device == "cpu" else torch.bfloat16
)

model = LlamaForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32 if device == "cpu" else torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
    offload_folder="offloads",
    quantization_config=bnb_config if str(device) != "cpu" else None,
)

tokenizer = LlamaTokenizer.from_pretrained(
    model_name,
    trust_remote_code=True,
    use_fast=True,
)

PROMPT = '''### Instruction:
{}
### Input:
{}
### Response:'''

if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

tokenizer.padding_side = "left"
tokenizer = tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.eval()


def evaluate(
    prompt='',
    temperature=0.4,
    top_p=0.65,
    top_k=35,
    repetition_penalty=1.1,
    max_new_tokens=512,
    stream_output=False,
    **kwargs,
):
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to(device)
    generation_config = GenerationConfig(
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        repetition_penalty=repetition_penalty,
        **kwargs,
    )

    with torch.no_grad():
        generation_output = model.generate(
            input_ids=input_ids,
            generation_config=generation_config,
            return_dict_in_generate=True,
            output_scores=True,
            max_new_tokens=max_new_tokens,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )
    s = generation_output.sequences[0]
    output = tokenizer.decode(s, skip_special_tokens=True)
    yield output.split("### Response:")[-1].strip()


def run_instruction(
    instructions,
    inputs,
    temperature=0.4,
    top_p=0.65,
    top_k=35,
    repetition_penalty=1.1,
    max_new_tokens=512,
    stream_output=False,
):
    now_prompt = PROMPT.format(instructions+'\n', inputs)

    response = evaluate(
        now_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens, stream_output, do_sample=True
    )

    if stream_output:
        response = tqdm.tqdm(response, unit='token')

    for i in response:
        print(i)
        response = i

    return response


def search_keyword(prompt):
    instructions = """Prompt:Time: Fri, 23 August 2023 2:30PM\nWeather: 73F\nHow many friends have I told you about?
Search Keyword:Friends
Prompt:Time: Thu, 27 September 2023 3:41PM\nWeather: 62F\nWhat was our very first conversation
Chat Index:0
Prompt:Time: Tue, 21 September 2023 2:30PM\nWeather: 67F\nWhat was the last thing I said to you
Chat Index:-1
Prompt:Time: Sun, 31 October 2023 7:33AM\nWeather: 59F\nWhat was the last thing I said to you before that
Chat Index:-2
Prompt:Time: Sat, 30 October 2023 8:21PM\nWeather: 65F\nDid I ever tell you about my math class?
Search Keyword:math
Prompt:Time: Mon, 13 November 2023 4:52PM\nWeather: 55F\nWhat was my 7th grade English teacher's name?
Search Keyword:English
Prompt:Time: Wed, 15 May 2023 6:19PM\nWeather: 80F\nWhere did I say my wallet was?
Search Keyword:Wallet
Prompt:Time: Fri, 24 June 2023 1:52PM\nWeather: 92F\nWhat did Alex tell you?
Search Keyword:Alex
Prompt:Time: Sat, 19 July 2023 2:44PM\nWeather: 91F\nWhat was my first conversation today
Search Keyword:24 June"""
    answer = ''.join(run_instruction(
        instructions,
        "Prompt:"+prompt+"\n",
        temperature=0.5,
        top_p=0.5,
        top_k=200,
        repetition_penalty=1.1,
        max_new_tokens=256,
        stream_output=False,
    ))
    return answer


def identify_objects_from_text(prompt):
    instructions = """Input:The object that flies in the air from this picture is a toy helicopter
Output:Toy helicopter
Input:For the robot to be able to achieve the task, the robot needs to look for a white shirt
Output:White shirt
Input:To complete the task, the robot needs to remove the fruits from the wooden basket.
Output:fruits, wooden basket
Input:To clean up your desk, you need to gather and organize the various items scattered around it. This includes the laptop, cell phone, scissors, pens, and other objects. By putting these items back in their designated spaces or containers, you can create a more organized and clutter-free workspace.
Output:Laptop, cell phone, scissors, pens, containers
Input:The tree with a colorful sky background is the one to be looking for.
Output:Tree"""
    answer = ''.join(run_instruction(
        instructions,
        prompt,
        temperature=0.5,
        top_p=0.5,
        top_k=200,
        repetition_penalty=1.1,
        max_new_tokens=256,
        stream_output=False,
    ))
    return answer


def robotix(prompt):
    instructions = """#Get me some water
objects = [['water: 57%', (781, 592)]]
robot.target((781, 592))
object_distance = distance()
if object_distance > 10:
    robot.go("forward", object_distance, track="water")
robot.grab()
if object_distance > 10:
    robot.go("back", object_distance)
robot.release("here")
### Input:
#Stand by the table
objects = [['table: 81%', (1489, 1173)], ['table: 75%', (1971, 1293)]]
### Response:
robot.target((1489, 1173))
if distance() > 10:
    robot.go(forward, distance())
### Input:
#Put the apples in the basket
objects = [['basket: 77%', (89, 112)], ['apples: 72%', (222, 182)]]
### Response:
robot.target((281, 189))
if distance() > 10:
    robot.go("forward", distance(), track="apples")
robot.grab()
robot.target(robot.find("basket"))
robot.release(distance())
### Input:
#Go to the sofa
objects=[['sofa: 81%', (1060, 931)]]
### Response:
robot.target((1060, 931))
if distance() > 10:
    robot.go("forward", distance())
### Input:
#Go to that person over there and then come back
objects=[['person: 85%', (331, 354)]]
### Response:
robot.target((331, 354))
object_distance = distance()
if object_distance > 10:
    robot.go("forward", object_distance)
    robot.go("backward", object_distance)"""

    answer = ''.join(run_instruction(
        instructions,
        prompt,
        temperature=0.2,
        top_p=0.5,
        top_k=300,
        repetition_penalty=1.1,
        max_new_tokens=256,
        stream_output=False,
    ))
    return answer