app.py

import re
import streamlit as st
import pandas as pd
from functools import lru_cache
from ortools.sat.python import cp_model

# Precompile common regex patterns
NOISE_PATTERNS = [
    re.compile(r'\bthe person\b', re.IGNORECASE),
    re.compile(r'\bperson\b', re.IGNORECASE),
    re.compile(r'\bsmoker\b', re.IGNORECASE),
    re.compile(r'\buses?\b', re.IGNORECASE),
    re.compile(r'\bloves?\b', re.IGNORECASE),
    re.compile(r'\bpartial to\b', re.IGNORECASE),
    re.compile(r'\bbouquet of\b', re.IGNORECASE),
    re.compile(r'\bboquet of\b', re.IGNORECASE),
    re.compile(r'\bvase of\b', re.IGNORECASE),
    re.compile(r'\bmany unique\b', re.IGNORECASE),
    re.compile(r'\bbouquet\b', re.IGNORECASE),
    re.compile(r'\bvase\b', re.IGNORECASE),
    re.compile(r'\barrangement\b', re.IGNORECASE)
]
ARTICLE_PATTERN = re.compile(r'\b(a|an|the)\b', re.IGNORECASE)
EXTRA_WORDS_PATTERN = re.compile(r'\b(owner|lover|enthusiast)\b', re.IGNORECASE)
NON_ALNUM_PATTERN = re.compile(r'[^a-z0-9 ]')
MULTISPACE_PATTERN = re.compile(r'\s+')

# Mapping for ordinal words.
ordinal_map = {
    "first": 1,
    "second": 2,
    "third": 3,
    "fourth": 4,
    "fifth": 5,
    "sixth": 6
}

# Mapping for number words in between–clues.
word_to_num = {
    "one": 1,
    "two": 2,
    "three": 3,
    "four": 4,
    "five": 5,
    "six": 6
}

def sanitize_token(text):
    text = text.lower()
    for pattern in NOISE_PATTERNS:
        text = pattern.sub('', text)
    text = ARTICLE_PATTERN.sub('', text)
    text = EXTRA_WORDS_PATTERN.sub('', text)
    text = NON_ALNUM_PATTERN.sub(' ', text)
    text = MULTISPACE_PATTERN.sub(' ', text)
    return text.strip()

def normalize_token(token, candidate_key=None):
    token_norm = token.lower()
    if candidate_key == "month":
        month_map = {
            "january": "jan",
            "february": "feb",
            "march": "mar",
            "april": "april",
            "may": "may",
            "june": "jun",
            "july": "jul",
            "august": "aug",
            "september": "sept",
            "october": "oct",
            "november": "nov",
            "december": "dec",
        }
        for full, abbr in month_map.items():
            token_norm = token_norm.replace(full, abbr)
    elif candidate_key == "nationalities":
        nat_map = {
            "swedish": "swede",
            "british": "brit",
            "danish": "dane"
        }
        for full, abbr in nat_map.items():
            token_norm = token_norm.replace(full, abbr)
    return token_norm

@lru_cache(maxsize=1024)
def lemmatize_text_cached(text):
    if nlp is not None:
        doc = nlp(text)
        return " ".join(token.lemma_ for token in doc)
    return text

def lemmatize_text(text):
    return lemmatize_text_cached(text)

def get_category_key(category):
    cat_lower = category.lower()
    if "favorite" in cat_lower and "color" in cat_lower:
        return "favorite_color"
    if "hair" in cat_lower:
        return "hair_color"
    if "name" in cat_lower:
        return "name"
    if "vacation" in cat_lower:
        return "vacation"
    if "occupation" in cat_lower:
        return "occupation"
    if "flower" in cat_lower:
        return "flower"
    if "lunch" in cat_lower:
        return "lunch"
    if "smoothie" in cat_lower:
        return "smoothie"
    if "hobby" in cat_lower:
        return "hobby"
    if "pet" in cat_lower or "animal" in cat_lower:
        return "animals"
    if "birthday" in cat_lower or "month" in cat_lower:
        return "month"
    if "nationalities" in cat_lower:
        return "nationalities"
    tokens = cat_lower.split()
    return tokens[-1] if tokens else cat_lower

def shorten_category(category):
    key = get_category_key(category)
    return key.replace('_', ' ')

# Try loading spaCy with the transformer-based model.
try:
    import spacy
    nlp = spacy.load("en_core_web_trf")
except Exception as e:
    st.warning("spaCy model could not be loaded; proceeding without it: " + str(e))
    nlp = None

class PuzzleSolver:
    def __init__(self, puzzle_text, debug=False):
        self.puzzle_text = puzzle_text
        self.num_houses = None
        self.categories = {}
        self.category_keys = {}
        self.clues = []
        self.var = {}
        self.model = cp_model.CpModel()
        self.debug = debug
        self.category_keywords = {
            "nationalities": ["swede", "norwegian", "german", "chinese", "dane", "brit", "danish", "swedish", "british"],
            "name": ["name"],
            "vacation": ["vacation", "trip", "break"],
            "occupation": ["occupation", "job"],
            "lunch": ["lunch", "soup", "stew", "grilled", "cheese", "spaghetti", "pizza", "stir"],
            "smoothie": ["smoothie", "cherry", "dragonfruit", "watermelon", "lime", "blueberry", "desert"],
            "models": ["phone", "model", "iphone", "pixel", "oneplus", "samsung", "xiaomi", "huawei"],
            "hair_color": ["hair"],
            "month": ["month", "birthday", "birth"],
            "hobby": ["photography", "cooking", "knitting", "woodworking", "paints", "painting", "gardening"],
            "pet": ["rabbit", "hamster", "fish", "cat", "bird", "dog"],
            "animals": ["rabbit", "dog", "horse", "fish", "bird", "cat"]
        }

    def parse_puzzle(self):
        m = re.search(r"There are (\d+) houses", self.puzzle_text, re.IGNORECASE)
        self.num_houses = int(m.group(1)) if m else 6
        cat_pattern = re.compile(r"^[-*]\s*(.*?):\s*(.+)$")
        for line in self.puzzle_text.splitlines():
            line = line.strip()
            m = cat_pattern.match(line)
            if m:
                cat_label = m.group(1).strip()
                attr_line = m.group(2).strip()
                attrs = [x.strip() for x in attr_line.split(",") if x.strip()]
                self.categories[cat_label] = attrs
                self.category_keys[cat_label] = get_category_key(cat_label)
                if self.debug:
                    st.write(f"Parsed category: '{cat_label}' with attributes {attrs}")
                    st.write(f"Assigned key for category: {self.category_keys[cat_label]}")
        clues_section = False
        for line in self.puzzle_text.splitlines():
            if "### Clues:" in line:
                clues_section = True
                continue
            if clues_section:
                clean = line.strip()
                if clean:
                    self.clues.append(clean)
                    if self.debug:
                        st.write(f"Parsed clue: {clean}")

    def build_variables(self):
        for cat, attrs in self.categories.items():
            self.var[cat] = {}
            for attr in attrs:
                self.var[cat][attr] = self.model.NewIntVar(1, self.num_houses, f"{cat}_{attr}")
            self.model.AddAllDifferent(list(self.var[cat].values()))
            if self.debug:
                st.write(f"Added all-different constraint for category '{cat}'.")

    def find_attribute(self, token):
        token_san = sanitize_token(token)
        candidate_key = None
        for key, kws in self.category_keywords.items():
            if any(kw in token_san for kw in kws):
                candidate_key = key
                if self.debug:
                    st.write(f"Debug: Token '{token}' suggests category key '{candidate_key}' based on keywords {kws}.")
                break
        if candidate_key == "pet":
            candidate_key = "animals"
        token_lemmatized = lemmatize_text(token_san)
        if self.debug:
            st.write(f"Debug: Lemmatized token for '{token}': '{token_lemmatized}'")
        if candidate_key == "hobby" and "paint" in token_lemmatized:
            token_lemmatized = token_lemmatized.replace("paint", "painting")
            if self.debug:
                st.write(f"Debug: Adjusted hobby token to '{token_lemmatized}' for proper matching.")
        if candidate_key in ["month", "nationalities"]:
            token_san = normalize_token(token_san, candidate_key)
            if self.debug:
                st.write(f"Debug: Normalized token for {candidate_key}: '{token_san}'")
        if candidate_key:
            categories_to_search = [(cat, attrs) for cat, attrs in self.categories.items() if self.category_keys.get(cat) == candidate_key]
            if self.debug:
                st.write(f"Debug: Restricted search to categories: {[cat for cat, _ in categories_to_search]}")
        else:
            categories_to_search = self.categories.items()
        best = None
        best_len = 0
        for cat, attrs in categories_to_search:
            for attr in attrs:
                attr_san = sanitize_token(attr)
                if candidate_key in ["month", "nationalities"]:
                    attr_san = normalize_token(attr_san, candidate_key)
                pattern = rf'\b{re.escape(attr_san)}\b'
                if re.search(pattern, token_san) or re.search(pattern, token_lemmatized):
                    if len(attr_san) > best_len:
                        best = (cat, attr)
                        best_len = len(attr_san)
                else:
                    alt = attr_san[:-1] if attr_san.endswith('s') else attr_san + 's'
                    if re.search(rf'\b{re.escape(alt)}\b', token_san) or re.search(rf'\b{re.escape(alt)}\b', token_lemmatized):
                        if len(attr_san) > best_len:
                            best = (cat, attr)
                            best_len = len(attr_san)
        if best is None and candidate_key in ["month", "nationalities"]:
            if self.debug:
                st.write(f"Debug: Fallback for {candidate_key}: no match found in token '{token_san}'. Trying explicit substrings.")
            mapping = {}
            if candidate_key == "month":
                mapping = {"jan": "jan", "feb": "feb", "mar": "mar",
                           "april": "april", "may": "may", "jun": "jun",
                           "jul": "jul", "aug": "aug", "sept": "sept", "oct": "oct", "nov": "nov", "dec": "dec"}
            elif candidate_key == "nationalities":
                mapping = {"swede": "swede", "norwegian": "norwegian", "german": "german",
                           "chinese": "chinese", "dane": "dane", "brit": "brit"}
            for key_abbr in mapping.values():
                if re.search(rf'\b{re.escape(key_abbr)}\b', token_san):
                    for cat, attrs in categories_to_search:
                        for attr in attrs:
                            attr_san = normalize_token(sanitize_token(attr), candidate_key)
                            if attr_san == key_abbr:
                                best = (cat, attr)
                                best_len = len(attr_san)
                                if self.debug:
                                    st.write(f"Debug: Found fallback match: '{attr_san}' in token '{token_san}'.")
                                break
                        if best is not None:
                            break
                if best is not None:
                    break
        if best is None and self.debug:
            st.write(f"DEBUG: No attribute found for token '{token}' (sanitized: '{token_san}', lemmatized: '{token_lemmatized}').")
        return best

    def find_all_attributes_in_text(self, text):
        found = []
        text_san = sanitize_token(text)
        for cat, attrs in self.categories.items():
            for attr in attrs:
                attr_san = sanitize_token(attr)
                if re.search(rf'\b{re.escape(attr_san)}\b', text_san):
                    found.append((cat, attr))
        unique = []
        seen = set()
        for pair in found:
            if pair not in seen:
                unique.append(pair)
                seen.add(pair)
        return unique

    def spacy_equality_extraction(self, text):
        if nlp is None:
            return None, None
        doc = nlp(text)
        for token in doc:
            if token.lemma_ == "be" and token.dep_ == "ROOT":
                subj = None
                attr = None
                for child in token.children:
                    if child.dep_ in ["nsubj", "nsubjpass"]:
                        subj = child
                    if child.dep_ in ["attr", "acomp"]:
                        attr = child
                if subj and attr:
                    subject_span = doc[subj.left_edge.i : subj.right_edge.i+1].text
                    attr_span = doc[attr.left_edge.i : attr.right_edge.i+1].text
                    return subject_span, attr_span
        ents = list(doc.ents)
        if len(ents) >= 2:
            return ents[0].text, ents[1].text
        return None, None

    def apply_constraint_equality(self, token1, token2):
        a1 = self.find_attribute(token1)
        a2 = self.find_attribute(token2)
        if a1 and a2:
            cat1, attr1 = a1
            cat2, attr2 = a2
            self.model.Add(self.var[cat1][attr1] == self.var[cat2][attr2])
            if self.debug:
                st.write(f"Added constraint: [{cat1}][{attr1}] == [{cat2}][{attr2}]")
        else:
            if self.debug:
                st.write(f"Warning: could not apply equality between '{token1}' and '{token2}'")

    def apply_constraint_inequality(self, token, house_number):
        a1 = self.find_attribute(token)
        if a1:
            cat, attr = a1
            self.model.Add(self.var[cat][attr] != house_number)
            if self.debug:
                st.write(f"Added constraint: [{cat}][{attr}] != {house_number}")
        else:
            if self.debug:
                st.write(f"Warning: could not apply inequality for '{token}' at house {house_number}")

    def apply_constraint_position(self, token1, op, token2):
        a1 = self.find_attribute(token1)
        a2 = self.find_attribute(token2)
        if a1 and a2:
            cat1, attr1 = a1
            cat2, attr2 = a2
            if op == "==":
                self.model.Add(self.var[cat1][attr1] == self.var[cat2][attr2])
                if self.debug:
                    st.write(f"Added constraint: [{cat1}][{attr1}] == [{cat2}][{attr2}]")
            elif op == "<":
                self.model.Add(self.var[cat1][attr1] < self.var[cat2][attr2])
                if self.debug:
                    st.write(f"Added constraint: [{cat1}][{attr1}] < [{cat2}][{attr2}]")
            elif op == ">":
                self.model.Add(self.var[cat1][attr1] > self.var[cat2][attr2])
                if self.debug:
                    st.write(f"Added constraint: [{cat1}][{attr1}] > [{cat2}][{attr2}]")
            elif op == "+1":
                self.model.Add(self.var[cat1][attr1] + 1 == self.var[cat2][attr2])
                if self.debug:
                    st.write(f"Added constraint: [{cat1}][{attr1}] + 1 == [{cat2}][{attr2}]")
            elif op == "-1":
                self.model.Add(self.var[cat1][attr1] - 1 == self.var[cat2][attr2])
                if self.debug:
                    st.write(f"Added constraint: [{cat1}][{attr1}] - 1 == [{cat2}][{attr2}]")
        else:
            if self.debug:
                st.write(f"Warning: could not apply position constraint between '{token1}' and '{token2}' with op '{op}'")

    def apply_constraint_next_to(self, token1, token2):
        a1 = self.find_attribute(token1)
        a2 = self.find_attribute(token2)
        if a1 and a2:
            cat1, attr1 = a1
            cat2, attr2 = a2
            diff = self.model.NewIntVar(0, self.num_houses, f"diff_{attr1}_{attr2}")
            self.model.AddAbsEquality(diff, self.var[cat1][attr1] - self.var[cat2][attr2])
            self.model.Add(diff == 1)
            if self.debug:
                st.write(f"Added next-to constraint: |[{cat1}][{attr1}] - [{cat2}][{attr2}]| == 1")
        else:
            if self.debug:
                st.write(f"Warning: could not apply next-to constraint between '{token1}' and '{token2}'")

    def apply_constraint_between(self, token1, token2, houses_between):
        a1 = self.find_attribute(token1)
        a2 = self.find_attribute(token2)
        if a1 and a2:
            cat1, attr1 = a1
            cat2, attr2 = a2
            diff = self.model.NewIntVar(0, self.num_houses, f"between_{attr1}_{attr2}")
            self.model.AddAbsEquality(diff, self.var[cat1][attr1] - self.var[cat2][attr2])
            self.model.Add(diff == houses_between + 1)
            if self.debug:
                st.write(f"Added between constraint: |[{cat1}][{attr1}] - [{cat2}][{attr2}]| == {houses_between + 1}")
        else:
            if self.debug:
                st.write(f"Warning: could not apply between constraint for '{token1}' and '{token2}' with {houses_between} houses in between")

    def apply_constraint_fixed(self, token, house_number):
        a1 = self.find_attribute(token)
        if a1:
            cat, attr = a1
            self.model.Add(self.var[cat][attr] == house_number)
            if self.debug:
                st.write(f"Added fixed constraint: [{cat}][{attr}] == {house_number}")
        else:
            if self.debug:
                st.write(f"Warning: could not apply fixed constraint for '{token}' at house {house_number}")

    def process_clue(self, clue):
        text = re.sub(r'^\d+\.\s*', '', clue).strip()
        if self.debug:
            st.write(f"Processing clue: {text}")
        ordinal_numbers = r"(?:\d+|first|second|third|fourth|fifth|sixth)"
        m_fixed = re.search(rf"(.+?) is in the ({ordinal_numbers}) house", text, re.IGNORECASE)
        if m_fixed:
            token = m_fixed.group(1).strip()
            num_str = m_fixed.group(2).strip().lower()
            house_num = int(num_str) if num_str.isdigit() else ordinal_map.get(num_str)
            if house_num is not None:
                self.apply_constraint_fixed(token, house_num)
                return
        m_not = re.search(rf"(.+?) is not in the ({ordinal_numbers}) house", text, re.IGNORECASE)
        if m_not:
            token = m_not.group(1).strip()
            num_str = m_not.group(2).strip().lower()
            house_num = int(num_str) if num_str.isdigit() else ordinal_map.get(num_str)
            if house_num is not None:
                self.apply_constraint_inequality(token, house_num)
                return
        m_left = re.search(r"(.+?) is directly left of (.+)", text, re.IGNORECASE)
        if m_left:
            token1 = m_left.group(1).strip()
            token2 = m_left.group(2).strip()
            self.apply_constraint_position(token1, "+1", token2)
            return
        m_right = re.search(r"(.+?) is directly right of (.+)", text, re.IGNORECASE)
        if m_right:
            token1 = m_right.group(1).strip()
            token2 = m_right.group(2).strip()
            self.apply_constraint_position(token1, "-1", token2)
            return
        m_sl = re.search(r"(.+?) is somewhere to the left of (.+)", text, re.IGNORECASE)
        if m_sl:
            token1 = m_sl.group(1).strip()
            token2 = m_sl.group(2).strip()
            self.apply_constraint_position(token1, "<", token2)
            return
        m_sr = re.search(r"(.+?) is somewhere to the right of (.+)", text, re.IGNORECASE)
        if m_sr:
            token1 = m_sr.group(1).strip()
            token2 = m_sr.group(2).strip()
            self.apply_constraint_position(token1, ">", token2)
            return
        m_next = re.search(r"(.+?) and (.+?) are next to each other", text, re.IGNORECASE)
        if m_next:
            token1 = m_next.group(1).strip()
            token2 = m_next.group(2).strip()
            self.apply_constraint_next_to(token1, token2)
            return
        m_between = re.search(rf"There (?:are|is) (\d+|one|two|three|four|five|six) house(?:s)? between (.+?) and (.+)", text, re.IGNORECASE)
        if m_between:
            num_str = m_between.group(1).strip().lower()
            houses_between = int(num_str) if num_str.isdigit() else word_to_num.get(num_str)
            token1 = m_between.group(2).strip()
            token2 = m_between.group(3).strip()
            self.apply_constraint_between(token1, token2, houses_between)
            return
        m_eq = re.search(r"(.+)\sis(?: the)?\s(.+)", text, re.IGNORECASE)
        if m_eq:
            token1 = m_eq.group(1).strip()
            token2 = m_eq.group(2).strip()
            token1 = re.sub(r"^(the person who\s+|who\s+)", "", token1, flags=re.IGNORECASE).strip()
            token2 = re.sub(r"^(a\s+|an\s+|the\s+)", "", token2, flags=re.IGNORECASE).strip()
            a1 = self.find_attribute(token1)
            a2 = self.find_attribute(token2)
            if a1 and a2:
                self.apply_constraint_equality(token1, token2)
                return
            else:
                if self.debug:
                    st.write("Equality regex failed to extract valid attributes using token cleaning.")
        if nlp is not None:
            left, right = self.spacy_equality_extraction(text)
            if left and right:
                if self.debug:
                    st.write(f"spaCy extracted equality: '{left}' == '{right}'")
                self.apply_constraint_equality(left, right)
                return
        if self.debug:
            st.write(f"Unprocessed clue: {text}")

    def process_all_clues(self):
        for clue in self.clues:
            self.process_clue(clue)

    def solve(self):
        solver = cp_model.CpSolver()
        # Use all available cores (0 means all available, 1 means single core for deployment to streamlit community cloud)
        solver.parameters.num_search_workers = 1
        status = solver.Solve(self.model)
        if status in (cp_model.OPTIMAL, cp_model.FEASIBLE):
            solution = {}
            for house in range(1, self.num_houses + 1):
                solution[house] = {}
                for cat, attr_dict in self.var.items():
                    for attr, var in attr_dict.items():
                        if solver.Value(var) == house:
                            solution[house][cat] = attr
            return solution
        else:
            if self.debug:
                st.write("No solution found. The clues may be contradictory or incomplete.")
            return None

    def print_solution(self, solution):
        if solution:
            headers = ["House"] + [shorten_category(cat) for cat in self.categories.keys()]
            table = []
            for house in sorted(solution.keys()):
                row = [str(house)]
                for cat in self.categories.keys():
                    row.append(solution[house].get(cat, ""))
                table.append(row)
            df = pd.DataFrame(table, columns=headers)
            return df
        else:
            return None

# Streamlit UI
st.title("Zebra Logic Puzzle Solver")
st.subheader("🦓 ZebraLogic: Benchmarking the Logical Reasoning Ability of Language Models")
st.markdown("""
Copy the Zebra Logic Puzzles description [from the huggingface site](https://huggingface.co/spaces/allenai/ZebraLogic), and paste it below.
""")

puzzle_text = st.text_area("Puzzle Input", height=300)
show_debug = st.checkbox("Show Debug Output", value=False)

# Use session_state to ensure the solution is computed only once per click.
if "puzzle_solved" not in st.session_state:
    st.session_state["puzzle_solved"] = False

if st.button("Solve Puzzle") or st.session_state["puzzle_solved"]:
    # Indicate that we've clicked the button
    st.session_state["puzzle_solved"] = True

    solver_instance = PuzzleSolver(puzzle_text, debug=show_debug)
    solver_instance.parse_puzzle()
    solver_instance.build_variables()
    solver_instance.process_all_clues()
    
    # st.subheader("Parsed Attributes (Categories & Their Attributes)")
    # for cat, attrs in solver_instance.categories.items():
    #     st.markdown(f"**{cat}**: {', '.join(attrs)}")

    # st.subheader("Parsed Clues")
    # for i, clue in enumerate(solver_instance.clues, start=1):
    #     st.markdown(f"{i}. {clue}")
    
    solution = solver_instance.solve()
    st.subheader("Solution Table")
    df_solution = solver_instance.print_solution(solution)
    if df_solution is not None:
        st.table(df_solution)
    else:
        st.error("No solution found. The clues may be contradictory or incomplete.")