import subprocess 

subprocess.run(
    'pip install flash-attn --no-build-isolation',
    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
    shell=True
)
import os
import re
import time
import torch
import spaces
import gradio as gr
from threading import Thread
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    BitsAndBytesConfig, 
    TextIteratorStreamer
)

# Configuration Constants
MODEL_ID = "Daemontatox/PathFinderAi3.0"


# Understand]: Analyze the question to identify key details and clarify the goal.
# [Plan]: Outline a logical, step-by-step approach to address the question or problem.
# [Reason]: Execute the plan, applying logical reasoning, calculations, or analysis to reach a conclusion. Document each step clearly.
# [Reflect]: Review the reasoning and the final answer to ensure it is accurate, complete, and adheres to the principle of openness.
# [Respond]: Present a well-structured and transparent answer, enriched with supporting details as needed.
# Use these tags as headers in your response to make your thought process easy to follow and aligned with the principle of openness.

DEFAULT_SYSTEM_PROMPT = """You are a highly skilled and meticulous reasoning engine, capable of breaking down complex problems into a series of logical, step-by-step inferences. Your primary goal is to arrive at accurate and well-justified conclusions by explicitly showing your thought process.

**Reasoning Principles:**

* **Engage in Chain of Thought Reasoning:** For every query, explicitly outline your reasoning process in a step-by-step manner. Think aloud, explaining each inference you make.
* **Focus on Clarity and Detail:** Each step in your thought process should be clear, concise, and logically connected to the previous and subsequent steps. Don't skip any crucial reasoning.
* **Verify Each Step:** Before moving to the next step, briefly verify the validity of the current step based on the information provided or your existing knowledge.
* **Consider Multiple Perspectives (if applicable):** If the query allows for different interpretations or approaches, briefly acknowledge and evaluate alternative possibilities before settling on the most likely path.
* **Explicitly State Assumptions:**  If you need to make any assumptions to proceed, clearly state them before using them in your reasoning.
* **Identify Information Gaps:** If you lack the necessary information to definitively answer a query, acknowledge the gap and explain why it prevents you from reaching a conclusion. Suggest potential avenues for obtaining the missing information.
* **Prioritize Accuracy over Speed:** Take your time and focus on producing accurate and well-reasoned responses, even if it takes more steps.
* **Self-Correction:** If you realize you've made an error in your reasoning, explicitly acknowledge the mistake, explain why it was incorrect, and correct your path.
* **Conclude with a Clear Answer:** After outlining your thought process, provide a clear and concise answer to the original query, based on the reasoning you presented.

**Output Format:**

Present your reasoning in a structured format, using clear headings or numbered steps for each stage of your thought process. For example:

**Thought Process:**

1. **Step 1: Analyze the Input:** [Explain your understanding of the input query.]
2. **Step 2: Identify Key Information:** [Highlight the crucial pieces of information relevant to the problem.]
3. **Step 3: Apply Relevant Knowledge/Rules:** [Explain the relevant knowledge or rules you are applying.]
4. **Step 4: Make an Inference:** [Clearly state the inference you are making based on the previous steps.]
5. **Step 5: Verify the Inference:** [Briefly explain why this inference is valid.]
6. **Step 6: Repeat Steps 3-5 as needed...**
7. **Step N: Synthesize and Conclude:** [Combine the results of your reasoning to reach a conclusion.]

**Final Answer:** [Your concise answer to the original query.]

**Important Considerations:**

* **Context is Key:** Pay close attention to the specific details and nuances of the user's prompt.
* **Follow Instructions Carefully:** Adhere to any specific instructions or constraints provided in the prompt.
* **Be Explicit, Not Implicit:** Don't assume the user understands your reasoning unless you explicitly state it.

**Example:**

**User Prompt:**  What is the capital of France and when was the Eiffel Tower completed?

**Your Response:**

**Thought Process:**

1. **Step 1: Analyze the Input:** The user is asking for two pieces of information: the capital of France and the completion date of the Eiffel Tower.
2. **Step 2: Identify Key Information needed:** I need to access my knowledge base about geography and historical landmarks.
3. **Step 3: Recall Information - Capital of France:**  My internal knowledge base indicates that Paris is the capital of France.
4. **Step 4: Verify the Inference:**  Paris is widely recognized and documented as the capital of France.
5. **Step 5: Recall Information - Eiffel Tower Completion:** My internal knowledge base indicates that the Eiffel Tower was completed in 1889.
6. **Step 6: Verify the Inference:** Historical records and reliable sources confirm that the Eiffel Tower was completed in 1889.
7. **Step 7: Synthesize and Conclude:** I have retrieved both pieces of information requested by the user.

**Final Answer:** The capital of France is Paris, and the Eiffel Tower was completed in 1889.

**Adaptability:**

You can adapt this system prompt by:

* **Specifying the domain of expertise:** "You are a highly skilled **mathematical** reasoning engine..." or "You are a meticulous **historical** reasoning engine..."
* **Adjusting the complexity of reasoning:** For simpler tasks, you might remove the "Consider Multiple Perspectives" section.
* **Adding specific constraints:**  "You must only use information provided in the prompt."
* **Requesting specific reasoning styles:** "Focus on deductive reasoning" or "Emphasize causal relationships."

By using this detailed system prompt as a foundation and tailoring it to your specific needs, you can significantly improve the LLM's ability to engage in effective Chain of Thought reasoning and provide more accurate and transparent answers. Remember to experiment and refine the prompt based on the performance you observe."""
# UI Configuration
TITLE = "<h1><center>AI Reasoning Assistant</center></h1>"
PLACEHOLDER = "Ask me anything! I'll think through it step by step."

CSS = """
.duplicate-button {
    margin: auto !important;
    color: white !important;
    background: black !important;
    border-radius: 100vh !important;
}
h3 {
    text-align: center;
}
.message-wrap {
    overflow-x: auto;
}
.message-wrap p {
    margin-bottom: 1em;
}
.message-wrap pre {
    background-color: #f6f8fa;
    border-radius: 3px;
    padding: 16px;
    overflow-x: auto;
}
.message-wrap code {
    background-color: rgba(175,184,193,0.2);
    border-radius: 3px;
    padding: 0.2em 0.4em;
    font-family: monospace;
}
.custom-tag {
    color: #0066cc;
    font-weight: bold;
}
.chat-area {
    height: 500px !important;
    overflow-y: auto !important;
}
"""

def initialize_model():
    """Initialize the model with appropriate configurations"""
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True
    )

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        device_map="cuda",
        attn_implementation="flash_attention_2",
        quantization_config=quantization_config

    )

    return model, tokenizer

def format_text(text):
    """Format text with proper spacing and tag highlighting (but keep tags visible)"""
    tag_patterns = [
        (r'<Thinking>', '\n<Thinking>\n'),
        (r'</Thinking>', '\n</Thinking>\n'),
        (r'<Critique>', '\n<Critique>\n'),
        (r'</Critique>', '\n</Critique>\n'),
        (r'<Revising>', '\n<Revising>\n'),
        (r'</Revising>', '\n</Revising>\n'),
        (r'<Final>', '\n<Final>\n'),
        (r'</Final>', '\n</Final>\n')
    ]
    
    formatted = text
    for pattern, replacement in tag_patterns:
        formatted = re.sub(pattern, replacement, formatted)
    
    formatted = '\n'.join(line for line in formatted.split('\n') if line.strip())
    
    return formatted

def format_chat_history(history):
    """Format chat history for display, keeping tags visible"""
    formatted = []
    for user_msg, assistant_msg in history:
        formatted.append(f"User: {user_msg}")
        if assistant_msg:
            formatted.append(f"Assistant: {assistant_msg}")
    return "\n\n".join(formatted)
    
def create_examples():
    """Create example queries for the UI"""
    return [
        "Explain the concept of artificial intelligence.",
        "How does photosynthesis work?",
        "What are the main causes of climate change?",
        "Describe the process of protein synthesis.",
        "What are the key features of a democratic government?",
        "Explain the theory of relativity.",
        "How do vaccines work to prevent diseases?",
        "What are the major events of World War II?",
        "Describe the structure of a human cell.",
        "What is the role of DNA in genetics?"
    ]

@spaces.GPU(duration=120)
def chat_response(
    message: str,
    history: list,
    chat_display: str,
    system_prompt: str,
    temperature: float = 0.2,
    max_new_tokens: int = 8192,
    top_p: float = 0.8,
    top_k: int = 40,
    penalty: float = 1.2,
):
    """Generate chat responses, keeping tags visible in the output"""
    conversation = [
        {"role": "system", "content": system_prompt}
    ]
    
    for prompt, answer in history:
        conversation.extend([
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": answer}
        ])
    
    conversation.append({"role": "user", "content": message})
    
    input_ids = tokenizer.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)
    
    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=60.0,
        skip_prompt=True,
        skip_special_tokens=True
    )
    
    generate_kwargs = dict(
        input_ids=input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=False if temperature == 0 else True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        repetition_penalty=penalty,
        streamer=streamer,
    )
    
    buffer = ""
    
    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=generate_kwargs)
        thread.start()
        
        history = history + [[message, ""]]
        
        for new_text in streamer:
            buffer += new_text
            formatted_buffer = format_text(buffer)
            history[-1][1] = formatted_buffer
            chat_display = format_chat_history(history)
            
            yield history, chat_display

def process_example(example: str) -> tuple:
    """Process example query and return empty history and updated display"""
    return [], f"User: {example}\n\n"

def main():
    """Main function to set up and launch the Gradio interface"""
    global model, tokenizer
    model, tokenizer = initialize_model()
    
    with gr.Blocks(css=CSS, theme="soft") as demo:
        gr.HTML(TITLE)
        gr.DuplicateButton(
            value="Duplicate Space for private use",
            elem_classes="duplicate-button"
        )
        
        with gr.Row():
            with gr.Column():
                chat_history = gr.State([])
                chat_display = gr.TextArea(
                    value="",
                    label="Chat History",
                    interactive=False,
                    elem_classes=["chat-area"],
                )
                
                message = gr.TextArea(
                    placeholder=PLACEHOLDER,
                    label="Your message",
                    lines=3
                )
                
                with gr.Row():
                    submit = gr.Button("Send")
                    clear = gr.Button("Clear")
                
                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    system_prompt = gr.TextArea(
                        value=DEFAULT_SYSTEM_PROMPT,
                        label="System Prompt",
                        lines=5,
                    )
                    temperature = gr.Slider(
                        minimum=0,
                        maximum=1,
                        step=0.1,
                        value=0.2,
                        label="Temperature",
                    )
                    max_tokens = gr.Slider(
                        minimum=128,
                        maximum=32000,
                        step=128,
                        value=8192,
                        label="Max Tokens",
                    )
                    top_p = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        step=0.1,
                        value=0.8,
                        label="Top-p",
                    )
                    top_k = gr.Slider(
                        minimum=1,
                        maximum=100,
                        step=1,
                        value=40,
                        label="Top-k",
                    )
                    penalty = gr.Slider(
                        minimum=1.0,
                        maximum=2.0,
                        step=0.1,
                        value=1.2,
                        label="Repetition Penalty",
                    )
                
                examples = gr.Examples(
                    examples=create_examples(),
                    inputs=[message],
                    outputs=[chat_history, chat_display],
                    fn=process_example,
                    cache_examples=False,
                )
        
        # Set up event handlers
        submit_click = submit.click(
            chat_response,
            inputs=[
                message,
                chat_history,
                chat_display,
                system_prompt,
                temperature,
                max_tokens,
                top_p,
                top_k,
                penalty,
            ],
            outputs=[chat_history, chat_display],
            show_progress=True,
        )
        
        message.submit(
            chat_response,
            inputs=[
                message,
                chat_history,
                chat_display,
                system_prompt,
                temperature,
                max_tokens,
                top_p,
                top_k,
                penalty,
            ],
            outputs=[chat_history, chat_display],
            show_progress=True,
        )
        
        clear.click(
            lambda: ([], ""),
            outputs=[chat_history, chat_display],
            show_progress=True,
        )
        
        submit_click.then(lambda: "", outputs=message)
        message.submit(lambda: "", outputs=message)
    
    return demo

if __name__ == "__main__":
    demo = main()
    demo.launch()