import streamlit as st
import re

class TreeofThoughts:
    def __init__(self, model, search_method="BFS"):
        self.model = model
        self.search_method = search_method

    def solve(self, problem, k=3, T=3, b=3, vth=0.5, context="I dont know this answer"):
        """Solve a problem using Tree of Thoughts approach.

        Args:
            problem: The problem statement
            k: Number of thoughts to generate per step
            T: Maximum depth/steps to explore
            b: Beam width (number of states to maintain)
            vth: Value threshold for accepting a solution

        Returns:
            Best solution found
        """
        try:
            with st.spinner("Thinking deeply about your complex question..."):
                # Initialize with the problem
                states = [problem]
                print(context)
                # Explore for T steps
                for t in range(T):
                    # Generate k thoughts for each state
                    all_thoughts = []
                    for state in states[:b]:  # Only consider top b states
                        thoughts_prompt = f"""
                        I'm solving this problem step-by-step:

                        Problem and context: {state}

                        Based on the context you know, generate {k} different approaches or thoughts to make progress on this problem.
                        For each approach, provide clear reasoning and next steps. Also if not present in context, reply with 1. Sorry, I don't have any information for every thought.

                        Format rules:
                        -Respond ONLY with a numbered list (e.g., "1. First approach", "2. Second approach").
                        -DO NOT include any introductory text or explanations.
                        -Start directly with "1." followed by the first thought.
                        """

                        response = self.model.generate(thoughts_prompt)
                        # print("fjfjfjfj")
                        # print(response)
                        thoughts_text = response.get("answer", "No answer found.")
                        thoughts = self._parse_thoughts(thoughts_text, k)
                        all_thoughts.extend([(state, thought) for thought in thoughts])

                    # Evaluate all thoughts
                    eval_results = []
                    for parent_state, thought in all_thoughts:
                        value_prompt = f"""
                        Evaluate how promising this approach is for solving the original problem.

                        Original problem: {problem}
                        Current approach: {thought}
                        Previous steps: {parent_state}

                        Rate this approach on a scale from 0.0 to 1.0, where:
                        - 0.0 means completely unpromising, leading to a dead end
                        - 1.0 means will definitely lead to a correct solution

                        Provide your numerical rating and brief justification.
                        """

                        response = self.model.generate(value_prompt)
                        eval_text = response.get("answer", "No answer found.")
                        value = self._extract_value(eval_text)

                        st.caption(f"I am thinking about: {thought} I would rate it's relevance as {value}!!")
                        
                        # Create new state by combining parent state and thought
                        new_state = f"{parent_state}\n\nStep {t+1}: {thought}"
                        eval_results.append((new_state, value))

                    # Sort by value and keep top b states
                    eval_results.sort(key=lambda x: x[1], reverse=True)
                    states = [state for state, value in eval_results[:b]]

                    # Check if we have a satisfactory solution
                    if eval_results and eval_results[0][1] >= vth:
                        break

                # Generate final answer based on best state
                if states:
                    best_state = states[0]
                    final_prompt = f"""
                    Based on this reasoning process, and from the context:{context},
                    Provide a comprehensive and detailed, elaborate and structured answer to the original problem.

                    Original problem: {problem}

                    Reasoning process:
                    {best_state}
    
                    Final answer:
                    """

                    response = self.model.generate(final_prompt)
                    final_answer = response.get("answer", "No answer found.")
                    return final_answer
                else:
                    return "I couldn't find a satisfactory solution to this problem."

        except Exception as e:
            st.error(f"Error in Tree of Thoughts: {str(e)}")
            return "An error occurred while processing your complex request."

    def _parse_thoughts(self, thoughts_text, k):
        """Parse the generated thoughts from text."""
        thoughts = []
        lines = thoughts_text.strip().split("\n")

        #  Find the index of the first occurrence of "1."
        start_index = next((i for i, line in enumerate(lines) if re.match(r"^1\.", line.strip())), None)

        # If "1." is found, remove everything before it
        if start_index is not None:
            lines = lines[start_index:]

        current_thought = ""

        for line in lines:
            if any(line.startswith(f"{i}.") for i in range(1, k+1)):
                if current_thought:
                    thoughts.append(current_thought.strip())
                current_thought = line.split(".", 1)[1].strip()
            else:
                current_thought += " " + line.strip()

        if current_thought and current_thought not in thoughts:
            thoughts.append(current_thought.strip())

        # Ensure we have k thoughts (or at least 1)
        while len(thoughts) < k:
            if thoughts:
                thoughts.append(f"Alternative approach building on {thoughts[0]}")
            else:
                thoughts.append("Start by breaking down the problem into manageable parts")

        return thoughts[:k]

    def _extract_value(self, eval_text):
        """Extract numerical value from evaluation text."""
        try:
            # Look for numbers between 0.0 and 1.0
            for word in eval_text.split():
                if word.replace('.', '', 1).isdigit():
                    value = float(word)
                    if 0.0 <= value <= 1.0:
                        return value

            # If no value was found, use default based on presence of positive words
            positive_words = ["promising", "good", "effective", "useful", "correct"]
            negative_words = ["sorry","unpromising", "poor", "ineffective", "useless", "incorrect"]
            hallucinate= ["don't have any information"]

            positive_count = sum(eval_text.lower().count(word) for word in positive_words)
            negative_count = sum(eval_text.lower().count(word) for word in negative_words)
            hallucinate_count = sum(eval_text.lower().count(phrase) for phrase in hallucinate)

            if hallucinate_count > 0:
                return 0.9

            if positive_count > negative_count:
                return 0.7
            elif negative_count > positive_count:
                return 0.3
            else:
                return 0.5
        except:
            return 0.5  # Default fallback