import json
from tqdm import tqdm
import time

import re
import torch
import os
from transformers import AutoTokenizer, AutoModel
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
os.environ['HF_HOME'] = "/home/admin/HFHOME"

#model_id = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
#model_id = "meta-llama/Llama-3.1-70B-Instruct"

# use quantization to lower GPU usage
# 4 bit:
# bnb_config = BitsAndBytesConfig(
#     load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
# )
# 8 bit:
# bnb_config = BitsAndBytesConfig(
#     load_in_8bit=True, bnb_8bit_use_double_quant=True, bnb_8bit_quant_type="nf8", bnb_8bit_compute_dtype=torch.bfloat16
# )

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    #quantization_config=bnb_config
)
terminators = [
    tokenizer.eos_token_id,
    tokenizer.convert_tokens_to_ids("<|eot_id|>")
]
model.generation_config.pad_token_id = tokenizer.eos_token_id #tokenizer.pad_token_id


def remove_think_tags(strings):
    if isinstance(strings, str):
      return re.sub(r'<think>[\s\S]*?</think>', '', strings)
    else:
      return [re.sub(r'<think>[\s\S]*?</think>', '', s) for s in strings]


SYS_PROMPT_simpler = """
Analyze the input sentence and divide it into several simpler and clearer sentences. The new sentences should retain the main meaning of the original sentence but be easier to understand. Use short and direct sentences and avoid complex terms or confusing grammatical structures.
The output should be Persian sentences only.
Each sentence should either be factual and express a minimal fact, or be conditional and express a minimal rule. A minimal fact or rule is a fact or rule that cannot be expressed in two separate and independent sentences.
start the output with this word: "HAMTAA"
"""

SYS_PROMPT_modified ="""
Edit the text below carefully for punctuation (including punctuation, spaces and half-spaces), and spelling errors, and perform a complete editorial review considering all editing points. The output should only be the corrected text without any additional explanation. 
The text should be preserved completely and only edited. Do not alter the sentence structure, add or remove words, or change the original meaning in any way. Please pay attention to correcting homophone words and fix these words based on the meaning of the sentence.
start the output with this word: "HAMTAA"
"""

SYS_PROMPT_translated ="""

translate the Persian content below into fluent and accurate English, ensuring the following criteria are met:

Meaning and Context: Preserve the original meaning and context of the text throughout the translation. These texts are all parts of legal laws in different domains.
Linguistic Accuracy: Adhere to proper English grammar, vocabulary, and syntax, demonstrating a native-level proficiency in the language.
Proper noun: do not translate those words that are proper nouns.
Cultural Sensitivity: Translate any technical terms, idiomatic expressions, and culturally specific references appropriately to ensure clarity and relevance in Persian.
The final English translation should be professional, natural, and suitable for its intended audience, which is law's audiance.

Note that the output must be nothing but the final traslation. start the output with this word: "HAMTAA"
"""

def generate(system_prompt, formatted_prompt):
  formatted_prompt = formatted_prompt[:50000] # to avoid GPU OOM
  messages = [{"role":"system","content":system_prompt},{"role":"user","content":formatted_prompt}]
  # tell the model to generate
  input_ids = tokenizer.apply_chat_template(
      messages,
      add_generation_prompt=True,
      return_tensors="pt"
  ).to(model.device)
  outputs = model.generate(
      input_ids,
      max_new_tokens=2048,
      eos_token_id=terminators,
      do_sample=True,
      temperature=0.6,
      top_p=0.9,
  )
  response = outputs[0][input_ids.shape[-1]:]
  return tokenizer.decode(response, skip_special_tokens=True)


def simpler(text, context):
    user_prompt_with_context = f"""input:
    {text}
    context:
    "{context}"
    """
    return generate(SYS_PROMPT_simpler, user_prompt_with_context)

def modified(text, context):
    user_prompt_with_context = f"""input:
    {text}
    context:
    "{context}"
    """
    return generate(SYS_PROMPT_modified, user_prompt_with_context)

def translated(text):
    user_prompt_with_context = f"""input:
    {text}
    """
    return generate(SYS_PROMPT_translated, user_prompt_with_context)

def think_remover(text):
    if 'HAMTAA' in text:
        part_after_HAMTAA = text.split('HAMTAA')[1]
        return part_after_HAMTAA
    if '</think>' in text:
        part_after_tag = text.split('</think>')[1]
        return part_after_tag
    else:
        return text

if __name__ == "__main__":
    print('start')
    start_time = time.time()
    inputfile = open('./merged_output.json', "r", encoding='utf-8')
    data = json.load(inputfile)
    inputfile.close()
    for index, item in enumerate(tqdm(data)):
        content = item.get("content","")
        context = item.get("context_sentences","")
        contextdict = item.get("context","")
        context_sentences = contextdict.get("context_sentences","")
        context = ""
        for a in context_sentences:
            if context == "":
                context = a
            else:
                context = context + " " + a
        context1_simpler= simpler(content, context)
        context1_simpler = think_remover(context1_simpler)
        item['simplified_content'] = context1_simpler
        context1_modified= modified(content, context)
        context1_modified = think_remover(context1_modified)
        item['modified_content'] = context1_modified
        context1_translated= translated(content)
        context1_translated = think_remover(context1_translated)
        item['translated_content'] = context1_translated


    outputfile = open('./merged_output_01.json', "w", encoding='utf-8')
    outputfile.write(json.dumps(data, ensure_ascii=False, indent = 4))
    outputfile.close()
    end_time = time.time()
    print(f"elapsed time:   {end_time-start_time}")
    print("end")