"""
THIS IS ONLY TO SHOWCASE THE ALLAM LLM SEGMENT  
"""
import re
#PRIVATE LIBRARY
#PRIVATE LIBRARY
import numpy as np
#private libraries
#PRIVATE LIBRARY
from transformers import AutoTokenizer, AutoModelForCausalLM


#PRIVATE Clients

from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "humain-ai/ALLaM-7B-Instruct-preview"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto"
)
model.eval()


def check_if_alz_or_general(query: str) -> str:

    prompt = f"""
    You are a classifier. Classify the query into ONE of these labels:
    - Alzheimer
    - General


    #private 

    Examples:
        1. Query: "مين أنا؟" Expected answer: Alzheimer
        2. Query: "أين تقع المملكة؟" Expected answer: General

       private further queries

    Now classify:
    Query: {query}
    
    """  

    #PRIVATE HELPER FUNCTIONS
    return response['content'].strip().lower() #response is from a private varibale 


## PRIVATE Similarity checking algorithm
    
## PRIVATE arabic normalization (depends on your own way of implementation)


## Intents (depends on your own way of implementation)

##Inent detection


def is_time_query(query: str) -> bool:
    time_patterns = [
        r"كم.*الساعة",
        r"الساعة كم",
        r"ايش الوقت",
        r"وش الوقت",
        r"الوقت الحين",
        r"الوقت الآن",
        r"الحين كم الساعة",
        r"كم الوقت",
    ]
    
    #private code 

    return False


#PRIVATE
#PRIVATE

ARABIC_DAYS = {
    "Monday": "الإثنين",
    "Tuesday": "الثلاثاء",
    "Wednesday": "الأربعاء",
    "Thursday": "الخميس",
    "Friday": "الجمعة",
    "Saturday": "السبت",
    "Sunday": "الأحد"
}


#PRIVATE
#PRIVATE
#PRIVATE



"""
PRIVATE RETRIEVAL FUNCTION 


"""

"""
PRIVATE RERANKER FUNCTION 

"""


def generate_with_allam(question, context):

    messages = [
        {
            "role": "system",
            "content": (
                "أنت مساعد ذكي لمرضى الزهايمر. "
                "أجب باستخدام المعلومات الموجودة في السياق فقط."
            ),
        },
        {
            "role": "user",
                "content": f"""
    السياق:
    {context}

    السؤال:
    {question}
    """
            }
        ]

    
    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt"
    ).to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=150,
        do_sample=False,
        #temperature=0.2   ,used if do_sample is true
    )

    return tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True
    )


"""
PRIVATE ANSWER QUERY FUNCTION

it returns generate_with_allam function result
"""

if __name__ == "__main__":
    queries = [
        "مين انا؟",
        "كيف وضعي انا الصحي؟",
        "هل انا بخير؟"
    ]
    
    for query in queries:
        answer = #private function
        print(f"Query: {query}\nAnswer : {answer}\n")