diff --git a/database/db_controller.py b/database/db_controller.py index 10b2fff..0eca676 100644 --- a/database/db_controller.py +++ b/database/db_controller.py @@ -1,6 +1,8 @@ import sqlite3 import threading import json +import os +import networkx as nx class DBController: def __init__(self, db_path=':memory:', max_connections=10): @@ -9,6 +11,7 @@ class DBController: self.conn = sqlite3.connect(db_path, check_same_thread=False) self.conn.row_factory = sqlite3.Row self.cursor = self.conn.cursor() + self.graph = nx.DiGraph() self.init_database() def init_database(self): @@ -27,6 +30,11 @@ class DBController: # Close the file-based database connection disk_conn.close() + if os.path.exists("pokemon_evolution_graph.json"): + with open("pokemon_evolution_graph.json", "r") as f: + data = json.load(f) + self.graph = nx.node_link_graph(data) + def save_changes(self): with self.lock: # Count the number of records before backup for verification @@ -40,6 +48,10 @@ class DBController: self.conn.backup(disk_conn) disk_conn.close() + data = nx.node_link_data(self.graph) + with open("pokemon_evolution_graph.json", "w") as f: + json.dump(data, f) + def close(self): self.save_changes() self.conn.close() @@ -95,12 +107,21 @@ class DBController: "generation", "is_baby_form", "storable_in_home", + "gender_relevant" ] query = self.craft_pokemon_json_query(fields, pfic) self.cursor.execute(query) results = self.cursor.fetchone() return dict(results) + def get_pokemon_details_by_name(self, name, fields): + query = self.craft_pokemon_json_query(fields) + name = name.replace("'", "''") + query += f" WHERE JSON_EXTRACT(data, '$.name') = '{name}'" + self.cursor.execute(query) + results = self.cursor.fetchall() + return [dict(row) for row in results] + def get_list_of_pokemon_forms(self): fields = [ "pfic", @@ -141,3 +162,55 @@ class DBController: WHERE PFIC = ? ''', (updated_data_str, pfic)) self.conn.commit() + + def update_evolution_graph(self, evolutions): + for evolution in evolutions: + from_pfic = evolution["from_pfic"] + to_pfic = evolution["to_pfic"] + method = evolution["method"] + + # Add nodes if they do not already exist + if not self.graph.has_node(from_pfic): + self.graph.add_node(from_pfic) + + if not self.graph.has_node(to_pfic): + self.graph.add_node(to_pfic) + + # Add the edge representing the evolution, with the method as an attribute + self.graph.add_edge(from_pfic, to_pfic, method=method) + + def get_evolution_graph(self, pfic): + return list(self.graph.successors(pfic)) + + def get_evolution_paths(self, start_node): + paths = [] + + # Define a recursive function to traverse the graph + def traverse(current_node, current_path): + # Add the current node to the path + current_path.append(current_node) + + # Get successors of the current node + successors = list(self.graph.successors(current_node)) + + if not successors: + # If there are no successors, add the current path to paths list + paths.append(current_path.copy()) + else: + # Traverse each successor and add edge metadata + for successor in successors: + method = self.graph[current_node][successor]["method"] + # Add the edge metadata as a tuple (to_node, method) + current_path.append((successor, method)) + + # Recur for the successor + traverse(successor, current_path) + + # Backtrack (remove the last node and edge metadata) + current_path.pop() + current_path.pop() + + # Start traversal from the start_node + traverse(start_node, []) + + return paths diff --git a/ui/main_window_controller.py b/ui/main_window_controller.py index e20ce15..8723d36 100644 --- a/ui/main_window_controller.py +++ b/ui/main_window_controller.py @@ -125,6 +125,7 @@ class MainWindowController: def on_evolutions_gathered(self, data): print("Works Done!") + db.update_evolution_graph(data) def reinitialize_database(self): pass @@ -168,8 +169,10 @@ class MainWindowController: else: self.view.image_label.setText("Image not found") - #self.load_evolution_chain(pfic) + self.load_evolution_chain(pfic) #self.load_encounter_locations(pfic) self.current_pfic = pfic - + def load_evolution_chain(self, pfic): + chain = db.get_evolution_paths(pfic) + self.view.update_evolution_tree(chain, pfic) diff --git a/ui/main_window_view.py b/ui/main_window_view.py index 12601f4..fafdc6c 100644 --- a/ui/main_window_view.py +++ b/ui/main_window_view.py @@ -221,4 +221,43 @@ class PokemonUI(QWidget): display_name = get_display_name(pokemon, not pokemon["gender_relevant"]) item = QListWidgetItem(display_name) item.setData(Qt.ItemDataRole.UserRole, pokemon["pfic"]) - self.pokemon_list.addItem(item) \ No newline at end of file + self.pokemon_list.addItem(item) + + def update_evolution_tree(self, evolution_chain, selected_pfic): + tree_items = {} + #for item in evolution_chain: + # print(item) + for pfic in evolution_chain: + pokemon_details = db.get_pokemon_details(pfic) + display_name = get_display_name(pokemon_details, not pokemon_details["gender_relevant"]) + item = QTreeWidgetItem([display_name, method if method else ""]) + item.setData(0, Qt.ItemDataRole.UserRole, current_pfic) + tree_items[current_pfic] = item + + if current_pfic == selected_pfic: + item.setBackground(0, QColor(255, 255, 0, 100)) # Highlight selected Pokémon + + # Second pass: build the tree structure + root = None + for current_pfic, name, form_name, method in evolution_chain: + item = tree_items[current_pfic] + + # Find the parent of this item + #parent_pfic = event_system.call_sync('get_evolution_parent', data=current_pfic) + parent_pfic = None + + if parent_pfic: + parent_item = tree_items.get(parent_pfic[0]) + if parent_item: + parent_item.addChild(item) + elif not root: + root = item + self.evolution_tree.addTopLevelItem(root) + + # Expand the entire tree + self.evolution_tree.expandAll() + + # Scroll to and select the current Pokémon + current_item = tree_items[selected_pfic] + self.evolution_tree.scrollToItem(current_item) + self.evolution_tree.setCurrentItem(current_item) \ No newline at end of file diff --git a/ui/workers/gather_evolutions_worker.py b/ui/workers/gather_evolutions_worker.py index 3bc5844..fad179e 100644 --- a/ui/workers/gather_evolutions_worker.py +++ b/ui/workers/gather_evolutions_worker.py @@ -1,10 +1,13 @@ from typing import Optional from PyQt6.QtCore import QObject, pyqtSignal, QRunnable from bs4 import BeautifulSoup, Tag +from fuzzywuzzy import fuzz from cache import cache from db import db -from utility.functions import get_form_name, get_display_name +import re + +from utility.functions import get_form_name, get_display_name, parse_pfic class GatherEvolutionsWorkerSignals(QObject): finished = pyqtSignal(list) @@ -22,23 +25,50 @@ class GatherEvolutions(QRunnable): except Exception as e: print(f"Error gathering Pokémon home storage status: {e}") - def gather_evolution_data(self): + def gather_evolution_data(self, force_refresh = False): all_pokemon_forms = db.get_list_of_pokemon_forms() evolutions = [] for pokemon_form in all_pokemon_forms: print(f"Processing {get_display_name(pokemon_form)}'s evolutions") - url = f"https://bulbapedia.bulbagarden.net/wiki/{pokemon_form["name"]}_(Pokémon)" + pokemon_name = pokemon_form["name"] + form = get_form_name(pokemon_form) + + cache_record_name = f"chain_{pokemon_name}_{form}" + if force_refresh: + cache.purge(cache_record_name) + + cached_entry = cache.get(cache_record_name) + if cached_entry != None: + evolutions.extend(cached_entry) + continue + + #form = get_form_name(pokemon_form, not pokemon_form["gender_relevant"]) + search_form = form + if search_form and pokemon_name in search_form: + search_form = search_form.replace(pokemon_name, "").strip() + + gender = None + if search_form and "male" in search_form.lower(): + gender = search_form + search_form = None + + if pokemon_name == "Flabébé": + # Bulbapedia doesn't detail out Flabébé's evolution chain fully. as its exactly the same for each form, but the coloured form remains constant + # through the evolution line, Red->Red->Red, Yellow->Yellow->Yellow etc. + search_form = None + + url = f"https://bulbapedia.bulbagarden.net/wiki/{pokemon_name}_(Pokémon)" page_data = cache.fetch_url(url) if not page_data: continue + soup = BeautifulSoup(page_data, 'html.parser') evolution_section = soup.find('span', id='Evolution_data') if not evolution_section: continue + evolution_table = None - form = get_form_name(pokemon_form, not pokemon_form["gender_relevant"]) - pokemon_name = pokemon_form["name"] evolution_table = evolution_section.parent.find_next('table') if form: form_without_form = form.replace('Form', '').replace('form', '').strip() @@ -51,12 +81,54 @@ class GatherEvolutions(QRunnable): if not evolution_table: continue + evolution_chain = [] if pokemon_name == "Eevee": evolution_chain = self.parse_eevee_evolution_chain(evolution_table, pokemon_form) - evolutions.append(evolution_chain) + #evolutions.append(evolution_chain) else: evolution_chain = self.parse_evolution_chain(evolution_table, pokemon_form) - evolutions.append(evolution_chain) + #evolutions.append(evolution_chain) + + chain = [] + for pokemon in evolution_chain: + from_pfic = self.get_pokemon_form_by_name(pokemon["pokemon"], pokemon["form"], gender=gender) + if not from_pfic: + #logger.warning(f"Could not find PFIC for {stage.pokemon} {stage.form}") + continue + + stage = pokemon["next_stage"] + if stage: + to_pfic = self.get_pokemon_form_by_name(stage["pokemon"], stage["form"], gender=gender) + if to_pfic: + evolution_info = { + "from_pfic": from_pfic, + "to_pfic": to_pfic, + "method": stage["method"] + } + evolutions.append(evolution_info) + chain.append(evolution_info) + + #insert_evolution_info(evolution_info) + + #if "breed" in stage["next_stage"]["method"].lower(): + # update_pokemon_baby_status(from_pfic, True) + + for branch in pokemon["branches"]: + to_pfic = self.get_pokemon_form_by_name(branch["pokemon"], branch["form"], gender=gender) + if to_pfic: + evolution_info = { + "from_pfic": from_pfic, + "to_pfic": to_pfic, + "method": branch["method"] + } + evolutions.append(evolution_info) + chain.append(evolution_info) + #EvolutionInfo(from_pfic, to_pfic, branch.method) + #insert_evolution_info(evolution_info) + + #if "breed" in branch.method.lower(): + # update_pokemon_baby_status(from_pfic, True) + cache.set(cache_record_name, chain) return evolutions @@ -233,4 +305,58 @@ class GatherEvolutions(QRunnable): pokemon_name = self.extract_pokemon_name(td) stage = self.extract_stage_form(td) return pokemon_name, stage - return None, None \ No newline at end of file + return None, None + + def get_pokemon_form_by_name(self, name: str, form: Optional[str] = None, threshold: int = 80, gender: Optional[str] = None): + fields = [ + "pfic", + "name", + "form_name" + ] + results = db.get_pokemon_details_by_name(name, fields) + #results = db_controller.execute_query('SELECT PFIC, name, form_name FROM pokemon_forms WHERE name = ?', (name,)) + + if not results: + return None + + results.sort(key=lambda x: parse_pfic(x["pfic"])) + + if form is None and gender is None: + if len(results) > 1: + if results[0]["form_name"] == None: + return results[0]["pfic"] + else: + return self.get_pokemon_form_by_name(name, "Male", threshold=100, gender=gender) + else: + return results[0]["pfic"] # Return the PFIC of the first result if no form is specified + + if gender: + gendered_form = self.get_pokemon_form_by_name(name, gender, threshold=100) + if gendered_form: + return gendered_form + + stripped_form = self.strip_pokemon_name(name, form) + + for entry in results: + stripped_db_form = self.strip_pokemon_name(entry["name"], entry["form_name"]) + if self.fuzzy_match_form(stripped_form, stripped_db_form, threshold): + return entry["pfic"] + + # Some times we get a form for a pokemon that doesn't really have one. + if len(results) > 1 and form != None: + return results[0]["pfic"] + + return None + + def strip_pokemon_name(self, pokemon_name: str, form_name: str) -> str: + if form_name: + form_name = form_name.replace("Form", "").strip() + form_name = re.sub(f'{re.escape(pokemon_name)}\\s*', '', form_name, flags=re.IGNORECASE).strip() + form_name = form_name.replace(" ", " ") + return form_name + return form_name + + def fuzzy_match_form(self, form1: str, form2: str, threshold: int = 80) -> bool: + if form1 is None or form2 is None: + return form1 == form2 + return fuzz.ratio(form1.lower(), form2.lower()) >= threshold \ No newline at end of file diff --git a/ui/workers/gather_pokemon_forms_worker.py b/ui/workers/gather_pokemon_forms_worker.py index 0a602d0..e41fec6 100644 --- a/ui/workers/gather_pokemon_forms_worker.py +++ b/ui/workers/gather_pokemon_forms_worker.py @@ -59,7 +59,7 @@ class GatherPokemonFormsWorker(QRunnable): return form_name return "None" - def process_pokemon_entry(self, national_dex_number, pokemon_soup, force_refresh = False): + def process_pokemon_entry(self, national_dex_number, pokemon_soup, force_refresh = True): found_forms = [] generation = get_generation_from_national_dex(national_dex_number) pokemon_name = pokemon_soup.get_text(strip=True)