import re
import mss
from PIL import Image, ImageEnhance
import pytesseract
import speech_recognition as sr
import pyttsx3
import keyboard
import threading
import os
import logging
import tkinter as tk
import ctypes
import sys
import asyncio
from datetime import datetime
from word2number import w2n
from langchain.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings  # Updated import per deprecation
import warnings

# -------------------- Configuration --------------------

# Suppress specific FutureWarning from transformers
warnings.filterwarnings(
    "ignore",
    category=FutureWarning,
    module='transformers.tokenization_utils_base'
)

# Define a custom logging filter to exclude specific messages
class ExcludeFilter(logging.Filter):
    def __init__(self, exclude_keywords):
        super().__init__()
        self.exclude_keywords = exclude_keywords

    def filter(self, record):
        return not any(keyword in record.getMessage() for keyword in self.exclude_keywords)

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # Set to INFO to include INFO and higher level logs

# Create a stream handler
stream_handler = logging.StreamHandler()
stream_handler.setLevel(logging.INFO)

# Define the log format
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s')
stream_handler.setFormatter(formatter)

# Define keywords to exclude from logging
excluded_keywords = [
    "Mouse moved to",
    "Width:",
    "Height:",
    "Monitor for capture:"
]

# Add the custom filter to the handler
stream_handler.addFilter(ExcludeFilter(excluded_keywords))

# Add the handler to the logger
if not logger.handlers:
    logger.addHandler(stream_handler)
else:
    # Remove existing handlers to prevent duplicate logs
    logger.handlers = []
    logger.addHandler(stream_handler)

# Initialize TTS engine
tts_engine = pyttsx3.init()
is_listening = False
ocr_text = ""

# Threading locks
lock = threading.Lock()
tts_lock = threading.Lock()

# Output directory for captured and preprocessed images
OUTPUT_DIR = r"D:\AI\EXP2\TempS"

# Initialize RAG components
embedding_model_name = "bennegeek/stella_en_1.5B_v5"  # New embedding model
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)
try:
    vector_store = FAISS.load_local(r"D:\AI\EXP2\faiss_index", embeddings, allow_dangerous_deserialization=True)
    logging.info("FAISS index loaded successfully.")
except Exception as e:
    logging.error(f"Failed to load FAISS index: {e}")
    sys.exit(1)

def set_dpi_awareness():
    try:
        if sys.platform.startswith('win'):
            ctypes.windll.shcore.SetProcessDpiAwareness(1)
            logging.info("Process DPI awareness set to SYSTEM_AWARE.")
    except Exception as e:
        logging.warning(f"Could not set DPI awareness: {e}. Continuing without setting DPI awareness.")

def get_scaling_factor():
    try:
        if sys.platform.startswith('win'):
            user32 = ctypes.windll.user32
            try:
                user32.SetProcessDPIAware()
                dpi = user32.GetDpiForWindow(user32.GetForegroundWindow())
                scaling_factor = dpi / 96
                logging.info(f"DPI Scaling Factor detected: {scaling_factor}")
                return scaling_factor
            except AttributeError:
                user32.SetProcessDPIAware()
                hdc = user32.GetDC(0)
                LOGPIXELSX = 88
                dpi = ctypes.windll.gdi32.GetDeviceCaps(hdc, LOGPIXELSX)
                scaling_factor = dpi / 96
                ctypes.windll.user32.ReleaseDC(0, hdc)
                logging.info(f"DPI Scaling Factor detected via fallback: {scaling_factor}")
                return scaling_factor
        else:
            logging.info("Non-Windows OS detected. Defaulting DPI scaling factor to 1.0.")
            return 1.0
    except Exception as e:
        logging.warning(f"Could not retrieve DPI scaling factor: {e}. Defaulting to 1.0.")
        return 1.0

def convert_number_words(text):
    """
    Converts number words in the speech text to digits.
    Example: "five hundred ninety-seven" -> "597"
    """
    try:
        # Attempt to convert the entire text
        number = w2n.word_to_num(text)
        return str(number)
    except:
        # Attempt to extract digits directly if conversion fails
        digits = re.findall(r'\d+', text)
        return ''.join(digits) if digits else text  # Return the number if found, else original text

def strip_markdown(text):
    """Function to remove Markdown syntax from text."""
    # Remove asterisks, underscores, backticks
    text = re.sub(r'[*_`]', '', text)
    # Remove links [text](url)
    text = re.sub(r'\[([^\]]+)\]\([^\)]+\)', r'\1', text)
    # Remove headings ## Heading
    text = re.sub(r'#+\s?', '', text)
    # Remove Markdown inline code blocks
    text = re.sub(r'```[\s\S]*?```', '', text)
    # Remove any remaining backslashes
    text = text.replace("\\", "")
    # Add more patterns as necessary
    return text

def speak_text(text):
    global tts_engine
    logging.info(f"Speaking out the response: {text}")
    with tts_lock:
        try:
            tts_engine.say(text)

            # Add hotkeys for pausing (F4) and stopping (F5)
            keyboard.add_hotkey('f4', pause_speaking)
            keyboard.add_hotkey('f5', stop_speaking)

            tts_engine.runAndWait()
            logging.info("Text-to-Speech completed.")
        except RuntimeError as e:
            logging.error(f"TTS Error: {e}")
        finally:
            keyboard.remove_hotkey('f4')
            keyboard.remove_hotkey('f5')

def pause_speaking():
    global tts_engine
    logging.info("Paused speaking...")
    # Note: pyttsx3 does not natively support pausing. Placeholder for future implementation.

def stop_speaking():
    global tts_engine
    logging.info("Stopped speaking.")
    tts_engine.stop()

def save_captured_image(image, output_dir, prefix="captured_region"):
    try:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = os.path.join(output_dir, f"{prefix}_{timestamp}.png")
        image.save(filename)
        logging.info(f"Captured image saved as {filename}")
    except Exception as e:
        logging.error(f"Failed to save captured image: {e}")

def save_preprocessed_image(image, output_dir, prefix="preprocessed_image"):
    try:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = os.path.join(output_dir, f"{prefix}_{timestamp}.png")
        image.save(filename)
        logging.info(f"Preprocessed image saved as {filename}")
    except Exception as e:
        logging.error(f"Failed to save preprocessed image: {e}")

def perform_ocr(image):
    global ocr_text
    try:
        logging.info("Starting OCR process with image preprocessing...")

        save_captured_image(image, OUTPUT_DIR)

        gray = image.convert('L')
        # Removed DEBUG logs related to image processing

        enhancer = ImageEnhance.Contrast(gray)
        gray_enhanced = enhancer.enhance(1.0)
        # Removed DEBUG logs related to image processing

        threshold = 190
        gray_threshold = gray_enhanced.point(lambda x: 0 if x < threshold else 255, '1')
        # Removed DEBUG logs related to image processing

        final_image = gray_threshold

        save_preprocessed_image(final_image, OUTPUT_DIR)

        logging.info("Image preprocessing complete. Performing OCR...")

        custom_config = r'--oem 3 --psm 6'
        raw_text = pytesseract.image_to_string(final_image, lang='eng', config=custom_config)

        if not raw_text.strip():
            logging.warning("OCR returned empty text.")
            speak_text("I couldn't detect any text in the selected area. Please try again with a different region.")
        else:
            ocr_text = raw_text.strip()
            logging.info(f"OCR Text: {ocr_text}")  # Log the extracted OCR text
            logging.info("OCR process completed successfully.")
    except pytesseract.TesseractNotFoundError:
        logging.error("Tesseract executable not found. Please ensure Tesseract is installed and PATH is set correctly.")
        speak_text("Tesseract is not installed or not found. Please install Tesseract and try again.")
        ocr_text = ""
    except Exception as e:
        logging.error(f"OCR Error: {e}")
        speak_text("An error occurred during OCR processing. Please try again.")
        ocr_text = ""

def get_selected_region(scaling_factor):
    root = tk.Tk()
    root.title("Select Region (Press Esc to Cancel)")
    root.attributes('-alpha', 0.3)
    root.attributes("-topmost", True)
    root.overrideredirect(True)

    screen_width = root.winfo_screenwidth()
    screen_height = root.winfo_screenheight()
    root.geometry(f"{screen_width}x{screen_height}")

    canvas = tk.Canvas(root, cursor="cross", bg="grey")
    canvas.pack(fill=tk.BOTH, expand=True)

    rect = None
    start_x = start_y = end_x = end_y = 0

    def on_button_press(event):
        nonlocal start_x, start_y, rect
        start_x = int(event.x * scaling_factor)
        start_y = int(event.y * scaling_factor)
        rect = canvas.create_rectangle(event.x, event.y, event.x, event.y, outline='red', width=2)
        logging.info("Please select the region by clicking and dragging the mouse.")

    def on_move_press(event):
        nonlocal end_x, end_y
        end_x = int(event.x * scaling_factor)
        end_y = int(event.y * scaling_factor)
        canvas.coords(rect, start_x / scaling_factor, start_y / scaling_factor, event.x, event.y)

    def on_button_release(event):
        root.quit()
        root.destroy()

    def on_escape(event):
        logging.info("Capture cancelled by user.")
        root.quit()
        root.destroy()
        # Do not exit the program to allow continuous operation

    canvas.bind("<ButtonPress-1>", on_button_press)
    canvas.bind("<B1-Motion>", on_move_press)
    canvas.bind("<ButtonRelease-1>", on_button_release)
    root.bind("<Escape>", on_escape)

    root.mainloop()

    left = min(start_x, end_x)
    top = min(start_y, end_y)
    right = max(start_x, end_x)
    bottom = max(start_y, end_y)

    logging.info(f"Selected region coordinates (scaled): Left={left}, Top={top}, Right={right}, Bottom={bottom}")
    return (left, top, right, bottom)

def generate_final_response(ocr_text):
    logging.info("Generating final response...")

    # Retrieve relevant documents from the knowledge base (RAG)
    query = ocr_text
    retrieved_docs = vector_store.similarity_search(query, k=5)  # Adjust 'k' as needed
    retrieved_text = "\n".join([doc.page_content for doc in retrieved_docs])

    # Combine OCR text with retrieved context
    context = f"""Assume the role of a strategist who excels at logical problem-solving and planning based on the following information.

Retrieved Information:
{retrieved_text}

OCR Text:
{ocr_text}

Instructions:
1. Understand the OCR text, which presents a multiple-choice question with options (A, B, C, D).
2. Analyze the retrieved information to determine the correct answer.
3. Respond in plain text without any Markdown formatting.
4. Eliminate all Markdown syntax such as asterisks (*), underscores (_), backticks (`), and brackets ([]).
5. Provide the answer in the following format: "The answer is ((X)Answer) because (reason)."
6. Ensure the response is clear, concise, and free from unnecessary syntax.
7. Keep your answer strictly to 100 words.
8. If calculations are involved, follow this structure:
   - We have been given (calculation_data),
   - The formula to be used is (formula),
   - We need to find (answer_variable).
9. Perform calculations attentively.

Provide the answer below:
"""

    try:
        # Generate response from the LLM with the enriched context
        response = ollama.generate(model='mistral-nemo', prompt=context)
        raw_response = response.get('response', '').strip()

        # Strip Markdown formatting
        plain_text_response = strip_markdown(raw_response)

        # Validate and format the response
        if "because" in plain_text_response:
            formatted_answer = plain_text_response
            logging.info("Generated formatted response.")
        else:
            formatted_answer = "The response could not be structured correctly."
            logging.warning(formatted_answer)

        if formatted_answer:
            speak_text(formatted_answer)
        else:
            logging.warning("No valid response could be generated.")

        return formatted_answer
    except Exception as e:
        logging.error(f"LLM Error during final response generation: {e}")
        return ""

def reset_state():
    global ocr_text
    ocr_text = ""

# -------------------- Helper Functions --------------------


def transcribe_speech(short_timeout=5, phrase_time_limit=5):
    """
    Transcribes speech from the microphone with a short timeout.
    If no speech is detected within the timeout, proceeds without user input.

    Args:
        short_timeout (int): Maximum number of seconds to wait for a phrase to start.
        phrase_time_limit (int): Maximum number of seconds a phrase can last.
    """
    global is_listening, speech_text
    recognizer = sr.Recognizer()
    with sr.Microphone() as source:
        logging.info("Listening for voice input... Please say your query or question.")
        recognizer.adjust_for_ambient_noise(source, duration=1)  # Adjust for noise for 1 second
        try:
            # Listen for speech with shorter timeout and phrase time limit
            audio = recognizer.listen(source, timeout=short_timeout, phrase_time_limit=phrase_time_limit)
            logging.info("Processing the audio with Google Speech Recognition...")
            text = recognizer.recognize_google(audio)
            # Convert number words to digits if necessary
            speech_text = convert_number_words(text)
            logging.info(f"Speech Transcribed: {speech_text}")

            # Provide confirmation
            speak_text(f"I heard: {speech_text}.")
        except sr.WaitTimeoutError:
            logging.warning("Listening timed out, no speech detected. Proceeding without voice input.")
            speech_text = ""  # Ensure speech_text is empty
        except sr.UnknownValueError:
            logging.warning("Speech recognition could not understand audio. Proceeding without voice input.")
            speech_text = ""  # Ensure speech_text is empty
        except sr.RequestError as e:
            logging.error(f"Could not request results from speech recognition service; {e}")
            speech_text = ""  # Ensure speech_text is empty
        finally:
            with lock:
                is_listening = False  # Ensure listening flag is reset

def configure_tts_voice():
    """
    Configures the TTS engine's voice, rate, and volume.
    """
    global tts_engine
    voices = tts_engine.getProperty('voices')

    # List available voices (optional)
    # Commented out DEBUG logs as they are now suppressed
    # for index, voice in enumerate(voices):
    #     logging.debug(f"Voice {index}: {voice.name}, {voice.gender}, {voice.languages}")

    # Change voice (select male/female, adjust the index as per your preference)
    if len(voices) > 1:
        tts_engine.setProperty('voice', voices[1].id)  # Change index to select a different voice
        logging.info(f"TTS voice set to: {voices[1].name}")
    else:
        logging.info("Only one voice available. Using default voice.")

    # Adjust speaking rate (lower is slower and more human-like)
    tts_engine.setProperty('rate', 150)  # Default is 200, lowering makes it more human-like

    # Adjust volume (0.0 to 1.0)
    tts_engine.setProperty('volume', 1.0)

def capture_and_save_image():
    global ocr_text, is_listening  # Added is_listening to globals

    region = get_selected_region(scaling_factor)
    left, top, right, bottom = region
    width = right - left
    height = bottom - top

    # Removed detailed dimension logging
    logging.info("Starting OCR process with image preprocessing...")

    if width <= 0 or height <= 0:
        logging.error("Invalid region selected. Please try again.")
        speak_text("Invalid region selected. Please try again by pressing F2.")
        reset_state()
        with lock:
            is_listening = False  # Reset the listening flag
        return

    with mss.mss() as sct:
        monitor = {
            "top": int(top / scaling_factor),
            "left": int(left / scaling_factor),
            "width": int(width / scaling_factor),
            "height": int(height / scaling_factor)
        }

        try:
            screenshot = sct.grab(monitor)
            img = Image.frombytes("RGB", screenshot.size, screenshot.bgra, "raw", "BGRX")
        except Exception as e:
            logging.error(f"Screen capture failed: {e}")
            speak_text("Screen capture failed. Please ensure the selected area is correct and try again.")
            reset_state()
            with lock:
                is_listening = False  # Reset the listening flag
            return

    perform_ocr(img)
    if not ocr_text.strip():
        logging.error("OCR returned empty text. Please ensure the selected region contains readable text.")
        speak_text("OCR did not detect any text. Please ensure the selected area contains readable text and try again.")
        reset_state()
        with lock:
            is_listening = False  # Reset the listening flag
        return

    generate_final_response(ocr_text)
    reset_state()
    with lock:
        is_listening = False  # Reset the listening flag after processing

def start_listening(scaling):
    global is_listening
    with lock:
        if not is_listening:
            is_listening = True
            logging.info("Started listening...")
            threading.Thread(target=capture_and_save_image, args=(), daemon=True).start()

def stop_listening():
    global is_listening
    with lock:
        if is_listening:
            is_listening = False
            logging.info("Stopped listening.")
            speak_text("Stopped listening.")

async def main():
    global scaling_factor

    configure_tts_voice()

    set_dpi_awareness()
    scaling_factor = get_scaling_factor()

    # Bind hotkeys
    keyboard.add_hotkey('f2', start_listening, args=(scaling_factor,))
    keyboard.add_hotkey('f3', stop_listening)

    logging.info("QA Assistant is running.")
    logging.info("Press 'F2' to select a region and answer a question, 'F3' to stop listening, 'Esc' to exit.")

    while True:
        await asyncio.sleep(0.1)
        if keyboard.is_pressed('esc'):
            logging.info("Exiting...")
            speak_text("Exiting the QA Assistant. Goodbye!")
            break

    keyboard.unhook_all()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logging.info("Program terminated by user.")
    except Exception as e:
        logging.error(f"An unexpected error occurred: {e}")
    finally:
        keyboard.unhook_all() in this i want to implement this : Usage
You can use SentenceTransformers or transformers library to encode text.

Sentence Transformers
from sentence_transformers import SentenceTransformer

# This model supports two prompts: "s2p_query" and "s2s_query" for sentence-to-passage and sentence-to-sentence tasks, respectively.
# They are defined in `config_sentence_transformers.json`
query_prompt_name = "s2p_query"
queries = [
    "What are some ways to reduce stress?",
    "What are the benefits of drinking green tea?",
]
# docs do not need any prompts
docs = [
    "There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent stress from building up.",
    "Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.",
]

# ！The default dimension is 1024, if you need other dimensions, please clone the model and modify `modules.json` to replace `2_Dense_1024` with another dimension, e.g. `2_Dense_256` or `2_Dense_8192` !
model = SentenceTransformer("dunzhang/stella_en_1.5B_v5", trust_remote_code=True).cuda()
query_embeddings = model.encode(queries, prompt_name=query_prompt_name)
doc_embeddings = model.encode(docs)
print(query_embeddings.shape, doc_embeddings.shape)
# (2, 1024) (2, 1024)

similarities = model.similarity(query_embeddings, doc_embeddings)
print(similarities)
# tensor([[0.8179, 0.2958],
#         [0.3194, 0.7854]])

Transformers
import os
import torch
from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize

query_prompt = "Instruct: Given a web search query, retrieve relevant passages that answer the query.\nQuery: "
queries = [
    "What are some ways to reduce stress?",
    "What are the benefits of drinking green tea?",
]
queries = [query_prompt + query for query in queries]
# docs do not need any prompts
docs = [
    "There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent stress from building up.",
    "Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.",
]

# The path of your model after cloning it
model_dir = "{Your MODEL_PATH}"

vector_dim = 1024
vector_linear_directory = f"2_Dense_{vector_dim}"
model = AutoModel.from_pretrained(model_dir, trust_remote_code=True).cuda().eval()
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
vector_linear = torch.nn.Linear(in_features=model.config.hidden_size, out_features=vector_dim)
vector_linear_dict = {
    k.replace("linear.", ""): v for k, v in
    torch.load(os.path.join(model_dir, f"{vector_linear_directory}/pytorch_model.bin")).items()
}
vector_linear.load_state_dict(vector_linear_dict)
vector_linear.cuda()

# Embed the queries
with torch.no_grad():
    input_data = tokenizer(queries, padding="longest", truncation=True, max_length=512, return_tensors="pt")
    input_data = {k: v.cuda() for k, v in input_data.items()}
    attention_mask = input_data["attention_mask"]
    last_hidden_state = model(**input_data)[0]
    last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
    query_vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
    query_vectors = normalize(vector_linear(query_vectors).cpu().numpy())

# Embed the documents
with torch.no_grad():
    input_data = tokenizer(docs, padding="longest", truncation=True, max_length=512, return_tensors="pt")
    input_data = {k: v.cuda() for k, v in input_data.items()}
    attention_mask = input_data["attention_mask"]
    last_hidden_state = model(**input_data)[0]
    last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
    docs_vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
    docs_vectors = normalize(vector_linear(docs_vectors).cpu().numpy())

print(query_vectors.shape, docs_vectors.shape)
# (2, 1024) (2, 1024)

similarities = query_vectors @ docs_vectors.T
print(similarities)
# [[0.8178789  0.2958377 ]
#  [0.31938642 0.7853526 ]]

FAQ
Q: The details of training?

A: The training method and datasets will be released in the future. (specific time unknown, may be provided in a paper)

Q: How to choose a suitable prompt for my own task?

A: In most cases, please use the s2p and s2s prompts. These two prompts account for the vast majority of the training data.

Q: How to reproduce MTEB results?

A: Please use evaluation scripts in Alibaba-NLP/gte-Qwen2-1.5B-instruct or intfloat/e5-mistral-7b-instruct

Q: Why each dimension has a linear weight?

A: MRL has multiple training methods, we choose this method which has the best performance.

Q: What is the sequence length of models?

A: 512 is recommended, in our experiments, almost all models perform poorly on specialized long text retrieval datasets. Besides, the model is trained on datasets of 512 length. This may be an optimization term.

If you have any questions, please start a discussion on community.