OriginDex/DBEditor/plan2.py

import sqlite3
import os
from collections import defaultdict

class EfficiencyOriginDexPlanner:
    def __init__(self, db_path):
        self.conn = sqlite3.connect(db_path)
        self.conn.row_factory = sqlite3.Row

    def get_all_data(self):
        cursor = self.conn.cursor()
        cursor.execute("SELECT * FROM games ORDER BY generation, id")
        games = cursor.fetchall()
        
        cursor.execute("SELECT * FROM pokemon_forms ORDER BY generation, national_dex")
        pokemon_forms = cursor.fetchall()
        
        cursor.execute("SELECT * FROM encounters")
        encounters = cursor.fetchall()
        
        cursor.execute("SELECT * FROM evolution_chains")
        evolutions = cursor.fetchall()
        
        return games, pokemon_forms, encounters, evolutions

    def build_data_structures(self):
        games, pokemon_forms, encounters, evolutions = self.get_all_data()
        
        pokemon_by_gen = defaultdict(list)
        game_by_gen = defaultdict(list)
        encounter_data = defaultdict(list)
        evolution_map = defaultdict(list)
        
        for game in games:
            game_by_gen[game['generation']].append(game)
        
        for pokemon in pokemon_forms:
            pokemon_by_gen[pokemon['generation']].append(pokemon)
        
        for encounter in encounters:
            encounter_data[encounter['game_id']].append(encounter)
        
        for evolution in evolutions:
            evolution_map[evolution['from_pfic']].append(evolution['to_pfic'])
        
        return pokemon_by_gen, game_by_gen, encounter_data, evolution_map

    def generate_efficient_plan(self, generation_groups):
        pokemon_by_gen, game_by_gen, encounter_data, evolution_map = self.build_data_structures()
        plan = []
        caught_pokemon = set()

        for group in generation_groups:
            group_plan = self.plan_for_group(group, pokemon_by_gen, game_by_gen, encounter_data, evolution_map, caught_pokemon)
            if group_plan:
                plan.extend(group_plan)

        return plan

    def plan_for_group(self, generations, pokemon_by_gen, game_by_gen, encounter_data, evolution_map, caught_pokemon):
        group_plan = []
        needed_pokemon = set()
        games_in_group = []

        for gen in generations:
            needed_pokemon.update(pokemon['PFIC'] for pokemon in pokemon_by_gen[gen])
            games_in_group.extend(game_by_gen[gen])

        print(f"Initial needed_pokemon count: {len(needed_pokemon)}")
        print(f"Games in group: {[game['name'] for game in games_in_group]}")

        # Check for new evolutions of already caught Pokémon
        new_evolutions = set()
        for pfic in caught_pokemon:
            if pfic in evolution_map:
                for evolved_pfic in evolution_map[pfic]:
                    if self.get_pokemon_generation(evolved_pfic) in generations and evolved_pfic not in caught_pokemon:
                        new_evolutions.add(pfic)
                        needed_pokemon.add(evolved_pfic)

        print(f"Needed pokemon after adding new evolutions: {len(needed_pokemon)}")

        # Create a dictionary to store the Pokémon available in each game, including evolutions
        game_pokemon = {}
        for game in games_in_group:
            game_pokemon[game['id']] = set()
            for encounter in encounter_data[game['id']]:
                game_pokemon[game['id']].add(encounter['pfic'])
                # Add all possible evolutions
                current_pfic = encounter['pfic']
                while current_pfic in evolution_map:
                    for evolved_pfic in evolution_map[current_pfic]:
                        if self.get_pokemon_generation(evolved_pfic) in generations:
                            game_pokemon[game['id']].add(evolved_pfic)
                            current_pfic = evolved_pfic
                            break
                    else:
                        break

        selected_games = []
        remaining_pokemon = needed_pokemon.copy()

        while remaining_pokemon:
            # Add debugging information
            print("\nDebug: Game selection process")
            for game in games_in_group:
                exclusive_groups = self.get_exclusive_groups(game['id'])
                pokemon_covered = set()
                
                for pfic in remaining_pokemon & game_pokemon[game['id']]:
                    is_exclusive = False
                    for group in exclusive_groups.values():
                        if any(ge['pfic'] == pfic for ge in group):
                            is_exclusive = True
                            # If it's exclusive, only add it if no other Pokémon from its group is already covered
                            if not any(other_ge['pfic'] in pokemon_covered for other_ge in group):
                                pokemon_covered.add(pfic)
                            break
                    if not is_exclusive:
                        pokemon_covered.add(pfic)

                print(f"  {game['name']}: {len(pokemon_covered)} Pokemon covered")
                print(f"    First 15 Pokemon: {[self.get_pokemon_name(pfic) for pfic in list(pokemon_covered)[:15]]}")

            best_game = max(games_in_group, key=lambda g: len(self.get_covered_pokemon(g, remaining_pokemon, game_pokemon[g['id']])))
            pokemon_covered = self.get_covered_pokemon(best_game, remaining_pokemon, game_pokemon[best_game['id']])
            print(f"\nBest game: {best_game['name']}, Pokemon covered: {len(pokemon_covered)}")
            print(f"First 10 Pokemon covered by {best_game['name']}: {[self.get_pokemon_name(pfic) for pfic in list(pokemon_covered)[:10]]}")
            
            if not pokemon_covered:
                print("No more Pokémon can be covered. Breaking loop.")
                break
            
            selected_games.append(best_game)
            remaining_pokemon -= pokemon_covered

            # Check if we can remove any previously selected games
            for game in selected_games[:-1]:
                if all(pokemon in self.get_covered_pokemon(best_game, needed_pokemon, game_pokemon[best_game['id']]) 
                       for pokemon in self.get_covered_pokemon(game, needed_pokemon, game_pokemon[game['id']])):
                    selected_games.remove(game)
                    print(f"Removed {game['name']} as it's covered by {best_game['name']}")

        print(f"Selected games: {[game['name'] for game in selected_games]}")

        for game in selected_games:
            game_plan = self.plan_for_game(game, needed_pokemon, encounter_data, evolution_map, caught_pokemon, set(generations), new_evolutions)
            if game_plan:
                # Update caught_pokemon and game_pokemon based on the new JSON structure
                for plan_entry in game_plan:
                    for pokemon in plan_entry["pokemon"]:
                        pfic = pokemon["pfic"]
                        caught_pokemon.add(pfic)
                        
                        # Handle evolutions
                        for evolution in pokemon["evolve_to"]:
                            caught_pokemon.add(evolution["pfic"])
                        
                        # Handle breeding
                        for breeding in pokemon["breed_for"]:
                            caught_pokemon.add(breeding["pfic"])
                    
                    # Update game_pokemon to remove exclusive encounters that weren't used
                    exclusive_groups = self.get_exclusive_groups(game['id'])
                    used_exclusive_groups = set()
                    
                    # Track used exclusive groups based on the new structure
                    for pokemon in plan_entry["pokemon"]:
                        pfic = pokemon["pfic"]
                        for group_id, group_encounters in exclusive_groups.items():
                            if any(ge['pfic'] == pfic for ge in group_encounters):
                                used_exclusive_groups.add(group_id)
                
                    # Remove unused exclusive encounters
                    for group_id, group_encounters in exclusive_groups.items():
                        if group_id not in used_exclusive_groups:
                            for ge in group_encounters:
                                if ge['pfic'] in game_pokemon[game['id']]:
                                    game_pokemon[game['id']].remove(ge['pfic'])
                
                group_plan.extend(game_plan)

        if not group_plan:
            print("Warning: No games were selected or no Pokémon were planned to be caught.")

        return group_plan

    def plan_for_game(self, game, needed_pokemon, encounter_data, evolution_map, caught_pokemon, target_generations, new_evolutions):
        game_plan = {
            "game_name": game['name'],
            "pokemon": []
        }
        to_catch = defaultdict(int)
        to_evolve = defaultdict(list)
        to_breed = defaultdict(list)
        planned_evolutions = defaultdict(int)
        processed_chains = set()

        exclusive_groups = self.get_exclusive_groups(game['id'])
        used_exclusive_groups = set()

        # First, handle new evolutions for already caught Pokémon
        for pfic in new_evolutions:
            evolution_chains, visit_counts, baby_forms = self.get_evolution_chains(pfic, evolution_map, target_generations)
            
            new_evolutions_count = sum(1 for chain in evolution_chains for pokemon in chain if pokemon not in caught_pokemon)
            
            if new_evolutions_count > 0:
                if any(encounter['pfic'] == pfic for encounter in encounter_data[game['id']]):
                    to_catch[pfic] += new_evolutions_count - planned_evolutions[pfic]
                    planned_evolutions[pfic] += new_evolutions_count
                
                    for chain in evolution_chains:
                        if len(chain) > 1:
                            for i in range(len(chain) - 1):
                                from_pfic, to_pfic = chain[i], chain[i+1]
                                if to_pfic not in caught_pokemon:
                                    to_evolve[from_pfic].append(to_pfic)
                                    caught_pokemon.add(to_pfic)
                                    needed_pokemon.discard(to_pfic)

        # Then proceed with the regular planning
        for encounter in encounter_data[game['id']]:
            if encounter['pfic'] in needed_pokemon and encounter['pfic'] not in caught_pokemon:
                # Check if this encounter is part of an exclusive group
                skip_encounter = False
                for group_id, group_encounters in exclusive_groups.items():
                    if group_id in used_exclusive_groups:
                        if any(ge['encounter_id'] == encounter['id'] for ge in group_encounters):
                            skip_encounter = True
                            break
                    else:
                        for group_encounter in group_encounters:
                            if group_encounter['encounter_id'] == encounter['id']:
                                if not group_encounter['starter']:
                                    used_exclusive_groups.add(group_id)
                                break
                        else:
                            continue
                        break

                if skip_encounter:
                    continue

                evolution_chains, visit_counts, baby_forms = self.get_evolution_chains(encounter['pfic'], evolution_map, target_generations)
                base_form = evolution_chains[0][0]
                
                if base_form not in caught_pokemon:
                    chain_key = tuple(sorted(set(pokemon for chain in evolution_chains for pokemon in chain)))
                    if chain_key not in processed_chains:
                        processed_chains.add(chain_key)
                        
                        num_to_catch = max(0, visit_counts[base_form] - planned_evolutions[base_form])
                        if num_to_catch > 0:
                            to_catch[encounter['pfic']] += num_to_catch
                            planned_evolutions[base_form] += num_to_catch

                        for chain in evolution_chains:
                            if len(chain) > 1:
                                for i in range(len(chain) - 1):
                                    from_pfic, to_pfic = chain[i], chain[i+1]
                                    if from_pfic not in baby_forms:
                                        if not self.is_final_evolution(to_pfic, evolution_map) or (self.is_final_evolution(to_pfic, evolution_map) and not self.is_in_to_evolve(to_pfic, to_evolve)):
                                            to_evolve[from_pfic].append(to_pfic)
                                        else:
                                            if encounter['pfic'] in to_catch:
                                                to_catch[encounter['pfic']] += -1
                                            if base_form in planned_evolutions:
                                                planned_evolutions[base_form] += -1

                                    if from_pfic in baby_forms and from_pfic not in encounter_data[game['id']]:
                                        if to_pfic not in to_breed or from_pfic not in to_breed[to_pfic]:
                                            to_breed[to_pfic].append(from_pfic)

                                        if from_pfic in to_catch:
                                            del to_catch[from_pfic]

                        caught_pokemon.update(pfic for chain in evolution_chains for pfic in chain)
                        needed_pokemon.difference_update(pfic for chain in evolution_chains for pfic in chain)

        # Remove encounters from exclusive groups that weren't used
        for group_id, group_encounters in exclusive_groups.items():
            if group_id not in used_exclusive_groups:
                for group_encounter in group_encounters:
                    if group_encounter['pfic'] in to_catch:
                        del to_catch[group_encounter['pfic']]
                    if group_encounter['pfic'] in to_evolve:
                        del to_evolve[group_encounter['pfic']]
                    if group_encounter['pfic'] in to_breed:
                        del to_breed[group_encounter['pfic']]

        if to_catch or to_evolve or to_breed:
            for pfic, count in to_catch.items():
                pokemon_entry = {
                    "pfic": pfic,
                    "name": self.get_pokemon_name(pfic).split(" (")[0],
                    "form_name": self.get_form_name(pfic),
                    "catch_count": count,
                    "evolve_to": [],
                    "breed_for": []
                }
                
                # Add evolution information as flat structure
                if pfic in to_evolve:
                    processed_evolutions = set()
                    evolution_queue = [(pfic, to_pfic, 1) for to_pfic in to_evolve[pfic]]
                    
                    while evolution_queue:
                        from_pfic, to_pfic, count = evolution_queue.pop(0)
                        if (from_pfic, to_pfic) not in processed_evolutions:
                            processed_evolutions.add((from_pfic, to_pfic))
                            
                            evolution = {
                                "pfic": to_pfic,
                                "name": self.get_pokemon_name(to_pfic).split(" (")[0],
                                "form_name": self.get_form_name(to_pfic),
                                "count": count
                            }
                            
                            pokemon_entry["evolve_to"].append(evolution)
                            
                            # Add next evolution stage to queue if it exists
                            if to_pfic in evolution_map:
                                for next_pfic in evolution_map[to_pfic]:
                                    if next_pfic in to_evolve[to_pfic]:
                                        evolution_queue.append((to_pfic, next_pfic, count))
                
                # Add breeding information
                if pfic in to_breed:
                    for baby_pfic in to_breed[pfic]:
                        breeding = {
                            "pfic": baby_pfic,
                            "name": self.get_pokemon_name(baby_pfic).split(" (")[0],
                            "form_name": self.get_form_name(baby_pfic),
                            "count": 1
                        }
                        pokemon_entry["breed_for"].append(breeding)
                
                game_plan["pokemon"].append(pokemon_entry)

        return [game_plan] if game_plan["pokemon"] else []

    def plan_evolutions(self, evolution_chains, evolution_map, caught_pokemon):
        chains, visit_counts = self.get_evolution_chains(evolution_chains[0][0], evolution_map)
        evolution_plan = []
        needed_counts = defaultdict(int)
        
        # Count how many of each Pokémon we need based on visit counts
        for pfic, count in visit_counts.items():
            if pfic not in caught_pokemon:
                needed_counts[pfic] = count
        
        # Plan evolutions
        for chain in chains:
            for i in range(len(chain) - 1):
                from_pfic, to_pfic = chain[i], chain[i+1]
                if needed_counts[to_pfic] > 0:
                    evolution_plan.append((from_pfic, to_pfic))
                    needed_counts[to_pfic] -= 1
                    needed_counts[from_pfic] -= 1
        
        return evolution_plan, needed_counts

    def get_evolution_chains(self, pfic, evolution_map, target_generations):
        visited = defaultdict(int)
        unique_paths = set()
        baby_forms = set()
        
        def dfs(current_pfic, path):
            path_tuple = tuple(path)
            if path_tuple in unique_paths:
                return
            
            unique_paths.add(path_tuple)
            for pfic in path:
                visited[pfic] += 1
                if self.is_baby_pokemon(pfic):
                    baby_forms.add(pfic)
            
            if current_pfic in evolution_map:
                for next_pfic in evolution_map[current_pfic]:
                    if self.get_pokemon_generation(next_pfic) in target_generations:
                        dfs(next_pfic, path + [next_pfic])

        dfs(pfic, [pfic])
        chains = [list(path) for path in unique_paths]
        
        # Adjust visit counts for baby Pokémon
        for baby in baby_forms:
            visited[baby] = 1
        
        return chains, visited, baby_forms

    def get_pokemon_name(self, pfic):
        cursor = self.conn.cursor()
        cursor.execute("SELECT name, form_name FROM pokemon_forms WHERE PFIC = ?", (pfic,))
        result = cursor.fetchone()
        if result['form_name']:
            return f"{result['name']} ({result['form_name']})"
        return result['name']

    def get_form_name(self, pfic):
        cursor = self.conn.cursor()
        cursor.execute("SELECT form_name FROM pokemon_forms WHERE PFIC = ?", (pfic,))
        result = cursor.fetchone()
        return result['form_name'] if result and result['form_name'] else ""

    def display_plan(self, generation_groups, output_file='efficiency_plan.json'):
        plan = self.generate_efficient_plan(generation_groups)
        
        # Write to JSON file
        import json
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(plan, f, indent=4, ensure_ascii=False)
        
        print(f"\nPlan has been written to {os.path.abspath(output_file)}")

    def is_baby_pokemon(self, pfic):
        cursor = self.conn.cursor()
        cursor.execute("SELECT is_baby_form FROM pokemon_forms WHERE PFIC = ?", (pfic,))
        result = cursor.fetchone()
        return result['is_baby_form'] if result else False

    def get_pokemon_generation(self, pfic):
        cursor = self.conn.cursor()
        cursor.execute("SELECT generation FROM pokemon_forms WHERE PFIC = ?", (pfic,))
        result = cursor.fetchone()
        return result['generation'] if result else None

    def is_final_evolution(self, pfic, evolution_map):
        return pfic not in evolution_map

    def is_in_to_evolve(self, pfic, to_evolve):
        return any(pfic in to_pfics for to_pfics in to_evolve.values())

    def get_exclusive_groups(self, game_id):
        cursor = self.conn.cursor()
        cursor.execute('''
            SELECT eeg.group_id, e.id as encounter_id, e.pfic, e.starter
            FROM encounter_exclusive_groups eeg
            JOIN encounters e ON eeg.encounter_id = e.id
            JOIN exclusive_encounter_groups eeg2 ON eeg.group_id = eeg2.id
            WHERE eeg2.game_id = ?
        ''', (game_id,))
        groups = defaultdict(list)
        for row in cursor.fetchall():
            groups[row['group_id']].append({
                'encounter_id': row['encounter_id'],
                'pfic': row['pfic'],
                'starter': row['starter']
            })
        return groups

    def get_covered_pokemon(self, game, needed_pokemon, available_pokemon):
        exclusive_groups = self.get_exclusive_groups(game['id'])
        pokemon_covered = set()
        
        for pfic in needed_pokemon & available_pokemon:
            is_exclusive = False
            for group in exclusive_groups.values():
                if any(ge['pfic'] == pfic for ge in group):
                    is_exclusive = True
                    # If it's exclusive, only add it if no other Pokémon from its group is already covered
                    if not any(other_ge['pfic'] in pokemon_covered for other_ge in group):
                        pokemon_covered.add(pfic)
                    break
            if not is_exclusive:
                pokemon_covered.add(pfic)
        
        return pokemon_covered

planner = EfficiencyOriginDexPlanner('pokemon_forms.db')
planner.display_plan([[1, 2], [3, 4, 5], [6], [7], [8], [9]])