1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4Logo Design Core - BM25 search engine for logo design guidelines
5"""
6 
7import csv
8import re
9from pathlib import Path
10from math import log
11from collections import defaultdict
12 
13# ============ CONFIGURATION ============
14DATA_DIR = Path(__file__).parent.parent.parent / "data" / "logo"
15MAX_RESULTS = 3
16 
17CSV_CONFIG = {
18    "style": {
19        "file": "styles.csv",
20        "search_cols": ["Style Name", "Category", "Keywords", "Best For"],
21        "output_cols": ["Style Name", "Category", "Keywords", "Primary Colors", "Secondary Colors", "Typography", "Effects", "Best For", "Avoid For", "Complexity", "Era"]
22    },
23    "color": {
24        "file": "colors.csv",
25        "search_cols": ["Palette Name", "Category", "Keywords", "Psychology", "Best For"],
26        "output_cols": ["Palette Name", "Category", "Keywords", "Primary Hex", "Secondary Hex", "Accent Hex", "Background Hex", "Text Hex", "Psychology", "Best For", "Avoid For"]
27    },
28    "industry": {
29        "file": "industries.csv",
30        "search_cols": ["Industry", "Keywords", "Recommended Styles", "Mood"],
31        "output_cols": ["Industry", "Keywords", "Recommended Styles", "Primary Colors", "Typography", "Common Symbols", "Mood", "Best Practices", "Avoid"]
32    }
33}
34 
35 
36# ============ BM25 IMPLEMENTATION ============
37class BM25:
38    """BM25 ranking algorithm for text search"""
39 
40    def __init__(self, k1=1.5, b=0.75):
41        self.k1 = k1
42        self.b = b
43        self.corpus = []
44        self.doc_lengths = []
45        self.avgdl = 0
46        self.idf = {}
47        self.doc_freqs = defaultdict(int)
48        self.N = 0
49 
50    def tokenize(self, text):
51        """Lowercase, split, remove punctuation, filter short words"""
52        text = re.sub(r'[^\w\s]', ' ', str(text).lower())
53        return [w for w in text.split() if len(w) > 2]
54 
55    def fit(self, documents):
56        """Build BM25 index from documents"""
57        self.corpus = [self.tokenize(doc) for doc in documents]
58        self.N = len(self.corpus)
59        if self.N == 0:
60            return
61        self.doc_lengths = [len(doc) for doc in self.corpus]
62        self.avgdl = sum(self.doc_lengths) / self.N
63 
64        for doc in self.corpus:
65            seen = set()
66            for word in doc:
67                if word not in seen:
68                    self.doc_freqs[word] += 1
69                    seen.add(word)
70 
71        for word, freq in self.doc_freqs.items():
72            self.idf[word] = log((self.N - freq + 0.5) / (freq + 0.5) + 1)
73 
74    def score(self, query):
75        """Score all documents against query"""
76        query_tokens = self.tokenize(query)
77        scores = []
78 
79        for idx, doc in enumerate(self.corpus):
80            score = 0
81            doc_len = self.doc_lengths[idx]
82            term_freqs = defaultdict(int)
83            for word in doc:
84                term_freqs[word] += 1
85 
86            for token in query_tokens:
87                if token in self.idf:
88                    tf = term_freqs[token]
89                    idf = self.idf[token]
90                    numerator = tf * (self.k1 + 1)
91                    denominator = tf + self.k1 * (1 - self.b + self.b * doc_len / self.avgdl)
92                    score += idf * numerator / denominator
93 
94            scores.append((idx, score))
95 
96        return sorted(scores, key=lambda x: x[1], reverse=True)
97 
98 
99# ============ SEARCH FUNCTIONS ============
100def _load_csv(filepath):
101    """Load CSV and return list of dicts"""
102    with open(filepath, 'r', encoding='utf-8') as f:
103        return list(csv.DictReader(f))
104 
105 
106def _search_csv(filepath, search_cols, output_cols, query, max_results):
107    """Core search function using BM25"""
108    if not filepath.exists():
109        return []
110 
111    data = _load_csv(filepath)
112 
113    # Build documents from search columns
114    documents = [" ".join(str(row.get(col, "")) for col in search_cols) for row in data]
115 
116    # BM25 search
117    bm25 = BM25()
118    bm25.fit(documents)
119    ranked = bm25.score(query)
120 
121    # Get top results with score > 0
122    results = []
123    for idx, score in ranked[:max_results]:
124        if score > 0:
125            row = data[idx]
126            results.append({col: row.get(col, "") for col in output_cols if col in row})
127 
128    return results
129 
130 
131def detect_domain(query):
132    """Auto-detect the most relevant domain from query"""
133    query_lower = query.lower()
134 
135    domain_keywords = {
136        "style": ["style", "minimalist", "vintage", "modern", "retro", "geometric", "abstract", "emblem", "badge", "wordmark", "mascot", "luxury", "playful", "corporate"],
137        "color": ["color", "palette", "hex", "#", "rgb", "blue", "red", "green", "gold", "warm", "cool", "vibrant", "pastel"],
138        "industry": ["tech", "healthcare", "finance", "legal", "restaurant", "food", "fashion", "beauty", "education", "sports", "fitness", "real estate", "crypto", "gaming"]
139    }
140 
141    scores = {domain: sum(1 for kw in keywords if kw in query_lower) for domain, keywords in domain_keywords.items()}
142    best = max(scores, key=scores.get)
143    return best if scores[best] > 0 else "style"
144 
145 
146def search(query, domain=None, max_results=MAX_RESULTS):
147    """Main search function with auto-domain detection"""
148    if domain is None:
149        domain = detect_domain(query)
150 
151    config = CSV_CONFIG.get(domain, CSV_CONFIG["style"])
152    filepath = DATA_DIR / config["file"]
153 
154    if not filepath.exists():
155        return {"error": f"File not found: {filepath}", "domain": domain}
156 
157    results = _search_csv(filepath, config["search_cols"], config["output_cols"], query, max_results)
158 
159    return {
160        "domain": domain,
161        "query": query,
162        "file": config["file"],
163        "count": len(results),
164        "results": results
165    }
166 
167 
168def search_all(query, max_results=2):
169    """Search across all domains and combine results"""
170    all_results = {}
171    for domain in CSV_CONFIG.keys():
172        result = search(query, domain, max_results)
173        if result.get("results"):
174            all_results[domain] = result["results"]
175    return all_results
176