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, QFileDialog, QComboBox, QFrame, QHBoxLayout, QLineEdit
from Qt.QtGui import QFont
import xml.etree.ElementTree as ET
import webbrowser
import os

class PDBePISA(ToolInstance):

    SESSION_ENDURING = False
    SESSION_SAVE = True
    help = "https://www.ebi.ac.uk/pdbe/pisa/pistart.html"

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

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

    def _build_ui(self):
        layout = QVBoxLayout()

        # Hyperlink Button at the top
        pisa_button = QPushButton("Open PDBePISA Website")
        pisa_button.clicked.connect(lambda: webbrowser.open(self.help))
        layout.addWidget(pisa_button)

        # Title
        title = QLabel("PDBePISA: Macromolecular Interface Analysis")
        title_font = QFont()
        title_font.setBold(True)
        title_font.setPointSize(14)
        title.setFont(title_font)
        layout.addWidget(title)

        # Description
        description = QLabel(
            "How to get the XML file\n"
            "Open PDBePISA and enter your PDB ID or upload your .pdb file, analyse the structure for interfaces."
            "At the PISA Interface List, press the 'Details' button under Interfaces,"
            "scroll down to Interfacing residues at the PISA Interface page and press the 'XML' button with Display level Residues."
            "Save the .xml file and open it in ChopChopMF."
        )
        description.setWordWrap(True)
        layout.addWidget(description)

        # File Selection Button
        file_button = QPushButton("Select PDBePISA XML File")
        file_button.clicked.connect(self.select_file)
        layout.addWidget(file_button)

        # Add header above ChopChop Interfaces button
        layout.addWidget(QLabel("Select defattr file for coloring interfaces:"))

        # Add ChopChop Interfaces Button
        chop_button = QPushButton("ChopChop Pisa Interfaces")
        chop_button.clicked.connect(self.chopchop_interfaces)
        layout.addWidget(chop_button)

        # Scoring Scheme Legend
        layout.addWidget(QLabel("Scoring Scheme:"))
        color_bar = QVBoxLayout()
        colors = [
            ("Buried", "darkorange"),
            ("Hydrogen Bond", "cornflowerblue"),
            ("Salt Bridge", "purple"),
        ]
        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)

        # Color Customization Section
        layout.addWidget(QLabel("Customize Colors for Scoring Scheme:"))
        self.color_inputs = {}
        color_entries = [
            ("1 (Buried):", "darkorange"),
            ("2 (Hydrogen Bond):", "cornflowerblue"),
            ("3 (Salt Bridge):", "purple"),
        ]
        for label, default_color in color_entries:
            entry_layout = QHBoxLayout()
            color_label = QLabel(label)
            color_input = QLineEdit(default_color)  # Pre-fill with default color
            entry_layout.addWidget(color_label)
            entry_layout.addWidget(color_input)
            layout.addLayout(entry_layout)
            self.color_inputs[label.split(" ")[0]] = color_input

        # Button to Apply New Color Scheme
        apply_color_button = QPushButton("Apply New Color Scheme")
        apply_color_button.clicked.connect(self.apply_custom_colors)
        layout.addWidget(apply_color_button)

        # Set the layout in the main container
        container = QWidget()
        container.setLayout(layout)
        self.tool_window.ui_area.setLayout(layout)

    def select_file(self):
        # Open file dialog to select XML file
        xml_file, _ = QFileDialog.getOpenFileName(
            caption="Select PDBePISA XML File",
            filter="XML Files (*.xml)"
        )
        if xml_file:
            residues_by_chain, defattr_lines = self.parse_pdbepisa_xml(xml_file)
            commands = self.generate_commands(residues_by_chain)
            self.run_chimerax_commands(commands)
            self.write_defattr_file(defattr_lines, xml_file)

    def chopchop_interfaces(self):
        # Open file dialog to select the generated defattr file
        defattr_file, _ = QFileDialog.getOpenFileName(
            caption="Select Generated Defattr File",
            filter="Defattr Files (*.defattr)"
        )
        if defattr_file:
            # Run the color command using the defattr file
            run(self.session, f"open {defattr_file}")
            run(self.session, "color byattribute residue_score palette 1,darkorange:2,cornflowerblue:3,purple")
            self.session.logger.info("ChopChop Interfaces coloring applied.")

    def apply_custom_colors(self):
        # Collect the updated colors from the input fields
        palette = []
        for key, input_field in self.color_inputs.items():
            color = input_field.text().strip()  # Get the user-provided color
            palette.append(f"{key},{color}")   # Format as "1,color", "2,color", etc.

        # Construct the ChimeraX color command
        color_command = "color byattribute residue_score palette " + ":".join(palette)

        # Execute the command in ChimeraX
        try:
            run(self.session, color_command)
            self.session.logger.info("Custom color scheme applied.")
        except Exception as e:
            self.session.logger.error(f"Error applying custom color scheme: {e}")

    # Function to parse the PDBE PISA XML and return interacting residues
    def parse_pdbepisa_xml(self, xml_file):
        try:
            tree = ET.parse(xml_file)
            root = tree.getroot()
        except ET.ParseError as e:
            self.session.logger.error(f"XML parsing error: {e}")
            return {}, []

        # Dictionary to hold residues for each chain
        residues_by_chain = {}
        defattr_lines = [
            "#\n",
            "# Use this file to assign the attribute in ChimeraX with the Define Attribute tool.\n",
            "# Use the command: color byattribute residue_score palette 1,darkorange:2,cornflowerblue:3,purple\n",
            "#\n",
            "# Scoring: 1 = Buried, 2 = Buried with H-bond, 3 = Buried with Salt Bridge\n",
            "#\n",
            "attribute: residue_score\n",
            "match mode: 1-to-1\n",
            "recipient: residues\n"
        ]

        # Iterate over each 'RESIDUE' element in the XML
        for residue in root.findall('.//RESIDUE'):
            structure_element = residue.find('STRUCTURE')
            bsa_element = residue.find('BURIEDSURFACEAREA')
            hsdc_element = residue.find('HSDC')

            if structure_element is not None and bsa_element is not None and bsa_element.text:
                structure_text = structure_element.text.strip()
                bsa_value = float(bsa_element.text.strip())
                hsdc_value = hsdc_element.text.strip() if hsdc_element is not None else ""

                # Parse the structure text to extract chain ID and residue number
                chain_id, _, residue_info = structure_text.partition(':')
                chain_id = chain_id.strip()
                res_num = ''.join(filter(str.isdigit, residue_info))

                # If the buried surface area is greater than zero, add to the list
                if bsa_value > 0:
                    if chain_id not in residues_by_chain:
                        residues_by_chain[chain_id] = []
                    residues_by_chain[chain_id].append(res_num)

                    # Assign score based on HSDC value
                    score = 1  # Default score for buried residues
                    if hsdc_value == "H":
                        score = 2
                    elif hsdc_value == "HS":
                        score = 3

                    # Add to defattr lines
                    defattr_lines.append(f"/{chain_id}:{res_num}\t{score}\n")

        return residues_by_chain, defattr_lines

    # Function to generate Chimera and ChimeraX commands
    def generate_commands(self, residues_by_chain):
        if not residues_by_chain:
            self.session.logger.info("No interacting residues found or XML file is empty/incorrect.")
            return None

        chimerax_command = []

        # Iterate over chains and their residues to create the selection commands
        for chain, residues in residues_by_chain.items():
            residues_str = ','.join(residues)
            chimerax_command.append(f"#1/{chain}:{residues_str}")

        return 'select ' + ' '.join(chimerax_command)

    def run_chimerax_commands(self, commands):
        if commands:
            run(self.session, commands)
            self.session.logger.info(f"Executed ChimeraX command: {commands}")
            # Color the selected residues in cornflowerblue
            run(self.session, "color sel darkorange")
            self.session.logger.info("Colored selected residues in darkorange.")

    def write_defattr_file(self, defattr_lines, xml_file):
        output_file = os.path.splitext(xml_file)[0] + "_output.defattr"
        try:
            with open(output_file, 'w') as file:
                for line in defattr_lines:
                    if line.startswith("/"):
                        file.write(f"\t{line}")  # Ensure initial tab before residue entries
                    else:
                        file.write(line)
            self.session.logger.info(f"Defattr file written to {output_file}")
        except Exception as e:
            self.session.logger.error(f"Error writing defattr file: {e}")

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

class _PDBePISABundleAPI(BundleAPI):
    
    @staticmethod
    def start_tool(session, tool_name):
        return PDBePISA(session, tool_name)

bundle_api = _PDBePISABundleAPI()


