OriginDex/Utilities/Update_evolution_information.py

import sqlite3
from typing import List, Optional
from dataclasses import dataclass
from fuzzywuzzy import fuzz
import re
from cache_manager import CacheManager
from DetermineOriginGame import get_evolution_data_from_bulbapedia

@dataclass
class EvolutionInfo:
    from_pfic: str
    to_pfic: str
    method: str

def create_evolution_table():
    conn = sqlite3.connect('pokemon_forms.db')
    cursor = conn.cursor()
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS evolution_chains (
        from_pfic TEXT,
        to_pfic TEXT,
        method TEXT,
        PRIMARY KEY (from_pfic, to_pfic),
        FOREIGN KEY (from_pfic) REFERENCES pokemon_forms (PFIC),
        FOREIGN KEY (to_pfic) REFERENCES pokemon_forms (PFIC)
    )
    ''')
    conn.commit()
    return conn

def insert_evolution_info(conn, evolution_info: EvolutionInfo):
    cursor = conn.cursor()
    cursor.execute('''
    INSERT OR REPLACE INTO evolution_chains 
    (from_pfic, to_pfic, method)
    VALUES (?, ?, ?)
    ''', (evolution_info.from_pfic, evolution_info.to_pfic, evolution_info.method))
    conn.commit()

def strip_pokemon_name(pokemon_name: str, form_name: str) -> str:
    """Remove the Pokémon's name from the form name if present."""
    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(form1: str, form2: str, threshold: int = 80) -> bool:
    """Perform fuzzy matching between two form names."""
    if form1 is None or form2 is None:
        return form1 == form2
    return fuzz.ratio(form1.lower(), form2.lower()) >= threshold

def get_pokemon_form_by_name(conn, name: str, form: Optional[str] = None, threshold: int = 80, gender: Optional[str] = None) -> Optional[str]:
    cursor = conn.cursor()
    cursor.execute('SELECT PFIC, name, form_name FROM pokemon_forms WHERE name = ?', (name,))
    results = cursor.fetchall()

    if not results:
        return None
    
    if form is None and gender is None:
        if len(results) > 1:
            if results[0][2] == None:
                return results[0][0]
            else:
                return get_pokemon_form_by_name(conn, name, "Male", threshold=100, gender=gender)
        else:
            return results[0][0]  # Return the PFIC of the first result if no form is specified
    
    if gender:
        gendered_form = get_pokemon_form_by_name(conn, name, gender, threshold=100)
        if gendered_form:
            return gendered_form

    stripped_form = strip_pokemon_name(name, form)
    
    for pfic, pokemon_name, db_form in results:
        stripped_db_form = strip_pokemon_name(pokemon_name, db_form)
        if fuzzy_match_form(stripped_form, stripped_db_form, threshold):
            return 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][0]

    return None

def process_evolution_chain(conn, evolution_chain, cache, gender: Optional[str] = None):
    for stage in evolution_chain:
        from_pfic = get_pokemon_form_by_name(conn, stage.pokemon, stage.form, gender=gender)
        if not from_pfic:
            print(f"Warning: Could not find PFIC for {stage.pokemon} {stage.form}")
            continue

        if stage.next_stage:
            to_pfic = get_pokemon_form_by_name(conn, stage.next_stage.pokemon, stage.next_stage.form, gender=gender)
            if to_pfic:
                evolution_info = EvolutionInfo(from_pfic, to_pfic, stage.next_stage.method)
                insert_evolution_info(conn, evolution_info)

        for branch in stage.branches:
            to_pfic = get_pokemon_form_by_name(conn, branch.pokemon, branch.form, gender=gender)
            if to_pfic:
                evolution_info = EvolutionInfo(from_pfic, to_pfic, branch.method)
                insert_evolution_info(conn, evolution_info)

def update_evolution_chains():
    cache = CacheManager()
    conn = create_evolution_table()

    cursor = conn.cursor()
    cursor.execute('SELECT DISTINCT name, form_name FROM pokemon_forms')
    pokemon_forms = cursor.fetchall()

    for name, form in pokemon_forms:
        print(f"Processing {name} {form if form else ''}")

        if form and name in form:
            form = form.replace(name, "").strip()

        gender = None
        if form and "male" in form.lower():
            gender = form
            form = None
        
        evolution_chain = get_evolution_data_from_bulbapedia(name, form, cache, gender)
        if evolution_chain:
            if name == "Tauros": # Bulbapedia has a weird formatting for Tauros.
                for stage in evolution_chain:
                    if stage.form:
                        stage.form = stage.form.replace("Paldean Form(", "").replace(")", "").strip()
            process_evolution_chain(conn, evolution_chain, cache, gender)

    conn.close()

if __name__ == "__main__":
    update_evolution_chains()