OriginDex/Utilities/NewDBVisualiser.py

import sys
import os
import sqlite3
from PyQt5.QtWidgets import (QApplication, QMainWindow, QTableWidget, QTableWidgetItem, QPushButton, QVBoxLayout, 
                             QHBoxLayout, QWidget, QLineEdit, QLabel, QMessageBox, QTabWidget, QScrollArea, QFrame, QGridLayout)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap

class PokemonEvolutionWidget(QWidget):
    def __init__(self, conn, pfic):
        super().__init__()
        self.conn = conn
        self.pfic = pfic
        self.layout = QVBoxLayout()
        self.setLayout(self.layout)
        self.stage_width = 200  # Fixed width for each evolution stage
        self.stage_height = 250  # Fixed height for each evolution stage
        self.build_evolution_chain()

    def build_evolution_chain(self):
        chain = self.get_full_evolution_chain(self.pfic)
        self.display_evolution_chain(chain)

    def get_full_evolution_chain(self, pfic):
        cursor = self.conn.cursor()
        chain = []
        visited = set()

        def build_chain(current_pfic):
            if current_pfic in visited:
                return None
            visited.add(current_pfic)

            cursor.execute('SELECT name, form_name FROM pokemon_forms WHERE PFIC = ?', (current_pfic,))
            pokemon = cursor.fetchone()
            if not pokemon:
                return None

            name, form_name = pokemon
            node = {"pfic": current_pfic, "name": name, "form_name": form_name, "evolutions": []}

            cursor.execute('''
            SELECT ec.to_pfic, pf.name, pf.form_name, ec.method
            FROM evolution_chains ec
            JOIN pokemon_forms pf ON ec.to_pfic = pf.PFIC
            WHERE ec.from_pfic = ?
            ''', (current_pfic,))
            evolutions = cursor.fetchall()

            for evo_pfic, evo_name, evo_form, method in evolutions:
                evo_node = build_chain(evo_pfic)
                if evo_node:
                    node["evolutions"].append({"node": evo_node, "method": method})

            return node

        chain = build_chain(pfic)
        return chain

    def display_evolution_chain(self, chain):
        if not chain:
            return

        main_layout = QHBoxLayout()
        self.layout.addLayout(main_layout)

        stages = self.split_into_stages(chain)
        for stage_index, stage in enumerate(stages):
            stage_widget = QWidget()
            stage_layout = QGridLayout()
            stage_widget.setLayout(stage_layout)
            stage_widget.setFixedSize(self.stage_width, self.stage_height * len(stage))

            for row, pokemon in enumerate(stage):
                pokemon_widget = self.create_pokemon_widget(pokemon["pfic"], pokemon["name"], pokemon["form_name"])
                stage_layout.addWidget(pokemon_widget, row, 0, Qt.AlignCenter)

                if stage_index < len(stages) - 1 and pokemon.get("method"):
                    arrow_label = QLabel("→")
                    arrow_label.setStyleSheet("font-size: 24px;")
                    stage_layout.addWidget(arrow_label, row, 1, Qt.AlignCenter)

                    method_label = QLabel(pokemon["method"])
                    method_label.setWordWrap(True)
                    method_label.setFixedWidth(80)
                    stage_layout.addWidget(method_label, row, 2, Qt.AlignCenter)

            main_layout.addWidget(stage_widget)

    def split_into_stages(self, chain):
        stages = []
        current_stage = [{"pfic": chain["pfic"], "name": chain["name"], "form_name": chain["form_name"]}]
        stages.append(current_stage)

        while current_stage:
            next_stage = []
            for pokemon in current_stage:
                evolutions = self.get_evolutions(pokemon["pfic"])
                for evolution in evolutions:
                    next_stage.append({
                        "pfic": evolution["to_pfic"],
                        "name": evolution["name"],
                        "form_name": evolution["form_name"],
                        "method": evolution["method"]
                    })
            if next_stage:
                stages.append(next_stage)
            current_stage = next_stage

        return stages

    def get_evolutions(self, pfic):
        cursor = self.conn.cursor()
        cursor.execute('''
        SELECT ec.to_pfic, pf.name, pf.form_name, ec.method
        FROM evolution_chains ec
        JOIN pokemon_forms pf ON ec.to_pfic = pf.PFIC
        WHERE ec.from_pfic = ?
        ''', (pfic,))
        evolutions = cursor.fetchall()
        return [{"to_pfic": to_pfic, "name": name, "form_name": form_name, "method": method} for to_pfic, name, form_name, method in evolutions]

    def create_pokemon_widget(self, pfic, name, form_name):
        widget = QWidget()
        layout = QVBoxLayout()
        widget.setLayout(layout)

        image_path = f"images-new/{pfic}.png"
        if os.path.exists(image_path):
            pixmap = QPixmap(image_path)
            image_label = QLabel()
            image_label.setPixmap(pixmap.scaled(96, 96, Qt.KeepAspectRatio, Qt.SmoothTransformation))
            layout.addWidget(image_label, alignment=Qt.AlignCenter)

        name_label = QLabel(name)
        name_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(name_label)

        if form_name:
            form_label = QLabel(form_name)
            form_label.setAlignment(Qt.AlignCenter)
            layout.addWidget(form_label)

        return widget

class DatabaseVisualizer(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Pokémon Database Visualizer")
        self.setGeometry(100, 100, 1200, 800)

        self.conn = sqlite3.connect('pokemon_forms.db')
        self.cursor = self.conn.cursor()

        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)

        self.layout = QVBoxLayout()
        self.central_widget.setLayout(self.layout)

        self.tab_widget = QTabWidget()
        self.layout.addWidget(self.tab_widget)

        self.forms_tab = QWidget()
        self.evolutions_tab = QWidget()
        self.tab_widget.addTab(self.forms_tab, "Pokémon Forms")
        self.tab_widget.addTab(self.evolutions_tab, "Evolution Chains")

        self.setup_forms_tab()
        self.setup_evolutions_tab()

    def setup_forms_tab(self):
        layout = QVBoxLayout()
        self.forms_tab.setLayout(layout)

        self.search_layout = QHBoxLayout()
        self.search_label = QLabel("Search:")
        self.search_input = QLineEdit()
        self.search_button = QPushButton("Search")
        self.search_button.clicked.connect(self.search_pokemon)
        self.search_layout.addWidget(self.search_label)
        self.search_layout.addWidget(self.search_input)
        self.search_layout.addWidget(self.search_button)
        layout.addLayout(self.search_layout)

        self.table = QTableWidget()
        layout.addWidget(self.table)

        self.save_button = QPushButton("Save Changes")
        self.save_button.clicked.connect(self.save_changes)
        layout.addWidget(self.save_button)

        self.load_forms_data()

    def setup_evolutions_tab(self):
        layout = QVBoxLayout()
        self.evolutions_tab.setLayout(layout)

        self.evolution_search_layout = QHBoxLayout()
        self.evolution_search_label = QLabel("Search Pokémon:")
        self.evolution_search_input = QLineEdit()
        self.evolution_search_button = QPushButton("Search")
        self.evolution_search_button.clicked.connect(self.search_evolution)
        self.evolution_search_layout.addWidget(self.evolution_search_label)
        self.evolution_search_layout.addWidget(self.evolution_search_input)
        self.evolution_search_layout.addWidget(self.evolution_search_button)
        layout.addLayout(self.evolution_search_layout)

        self.evolution_scroll_area = QScrollArea()
        self.evolution_scroll_area.setWidgetResizable(True)
        layout.addWidget(self.evolution_scroll_area)

    def load_forms_data(self):
        self.cursor.execute('''
        SELECT pf.PFIC, pf.name, pf.form_name, pf.national_dex, pf.generation, ps.storable_in_home
        FROM pokemon_forms pf
        LEFT JOIN pokemon_storage ps ON pf.PFIC = ps.PFIC
        ''')
        data = self.cursor.fetchall()

        self.table.setColumnCount(6)
        self.table.setHorizontalHeaderLabels(["PFIC", "Name", "Form Name", "National Dex", "Generation", "Storable in Home"])
        self.table.setRowCount(len(data))

        for row, record in enumerate(data):
            for col, value in enumerate(record):
                item = QTableWidgetItem(str(value))
                if col == 0:  # PFIC column
                    item.setFlags(item.flags() & ~Qt.ItemIsEditable)  # Make PFIC non-editable
                self.table.setItem(row, col, item)

        self.table.resizeColumnsToContents()

    def search_pokemon(self):
        search_term = self.search_input.text().lower()
        for row in range(self.table.rowCount()):
            match = False
            for col in range(self.table.columnCount()):
                item = self.table.item(row, col)
                if item and search_term in item.text().lower():
                    match = True
                    break
            self.table.setRowHidden(row, not match)

    def save_changes(self):
        try:
            for row in range(self.table.rowCount()):
                pfic = self.table.item(row, 0).text()
                name = self.table.item(row, 1).text()
                form_name = self.table.item(row, 2).text() or None
                national_dex = int(self.table.item(row, 3).text())
                generation = int(self.table.item(row, 4).text())
                storable_in_home = self.table.item(row, 5).text().lower() == 'true'

                self.cursor.execute('''
                UPDATE pokemon_forms
                SET name = ?, form_name = ?, national_dex = ?, generation = ?
                WHERE PFIC = ?
                ''', (name, form_name, national_dex, generation, pfic))

                self.cursor.execute('''
                INSERT OR REPLACE INTO pokemon_storage (PFIC, storable_in_home)
                VALUES (?, ?)
                ''', (pfic, storable_in_home))

            self.conn.commit()
            QMessageBox.information(self, "Success", "Changes saved successfully!")
        except Exception as e:
            QMessageBox.critical(self, "Error", f"An error occurred while saving changes: {str(e)}")

    def search_evolution(self):
        search_term = self.evolution_search_input.text().lower()

        self.cursor.execute('''
        SELECT DISTINCT name, PFIC
        FROM pokemon_forms
        WHERE LOWER(name) LIKE ?
        ''', (f'%{search_term}%',))

        matching_pokemon = self.cursor.fetchall()

        if matching_pokemon:
            pokemon_name, pfic = matching_pokemon[0]
            evolution_widget = PokemonEvolutionWidget(self.conn, pfic)
            self.evolution_scroll_area.setWidget(evolution_widget)
        else:
            QMessageBox.information(self, "No Results", "No Pokémon found matching the search term.")

    def closeEvent(self, event):
        self.conn.close()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = DatabaseVisualizer()
    window.show()
    sys.exit(app.exec_())