from chimerax.core.tools import ToolInstance
from chimerax.core.commands import run
from chimerax.ui import MainToolWindow
from Qt.QtWidgets import QVBoxLayout, QLabel, QPushButton, QWidget, QHBoxLayout, QFrame
from Qt.QtGui import QFont
import webbrowser

class AlphaFold2(ToolInstance):

    SESSION_ENDURING = False
    SESSION_SAVE = True
    help = "https://www.ebi.ac.uk/training/online/courses/alphafold/"

    def __init__(self, session, tool_name):
        super().__init__(session, tool_name)

        # Run the command automatically when the tool is initialized
        run(session, "ui tool show AlphaFold")

        # Set name displayed on title bar
        self.display_name = "AlphaFold Info"
        self.tool_window = MainToolWindow(self)
        self._build_ui()
        self.tool_window.manage('side')

    def _build_ui(self):
        layout = QVBoxLayout()

        # Title
        layout.addWidget(QLabel("AlphaFold Confidence Score Information"))

        # Explanation of pLDDT
        bold_font = QFont()
        bold_font.setBold(True)
        plddt_label = QLabel("pLDDT (per-residue confidence):")
        plddt_label.setFont(bold_font)
        layout.addWidget(plddt_label)
        layout.addWidget(QLabel("Confidence level of individual residues in a predicted structure."))


        # pLDDT color coding
        layout.addWidget(QLabel("Color Coding:"))
        color_bar = QVBoxLayout()
        colors = [
            ("0-50: Very Low", "darkorange"),
            ("50-70: Low", "yellow"),
            ("70-90: Confident", "deepskyblue"),
            ("90-100: Very High", "blue")
        ]
        for label, color in colors:
            color_item = QHBoxLayout()
            color_box = QFrame()
            color_box.setFixedSize(20, 20)
            color_box.setStyleSheet(f"background-color: {color}; border: 1px solid black;")
            color_label = QLabel(label)
            color_item.addWidget(color_box)
            color_item.addWidget(color_label)
            color_bar.addLayout(color_item)
        layout.addLayout(color_bar)

        # Explanation of PAE
        pae_label = QLabel("PAE (Predicted Aligned Error):")
        pae_label.setFont(bold_font)
        layout.addWidget(pae_label)
        layout.addWidget(QLabel(
            "PAE provides confidence in the relative positions of residues."
        ))

        # Link to PAE Viewer
        pae_viewer_button = QPushButton("Open PAE Viewer")
        pae_viewer_button.clicked.connect(
            lambda: webbrowser.open("https://subtiwiki.uni-goettingen.de/v4/paeViewerDemo")
        )
        layout.addWidget(pae_viewer_button)

        # Set layout to the tool window
        container = QWidget()
        container.setLayout(layout)
        self.tool_window.ui_area.setLayout(layout)

# Register the tool in ChimeraX
from chimerax.core.toolshed import BundleAPI

class _AlphaFoldInfoBundleAPI(BundleAPI):
    
    @staticmethod
    def start_tool(session, tool_name):
        return AlphaFoldInfo(session, tool_name)

bundle_api = _AlphaFoldInfoBundleAPI()

