from chimerax.core.tools import ToolInstance
from chimerax.core.commands import run
from pathlib import Path
from chimerax.ui import MainToolWindow
from Qt.QtWidgets import QVBoxLayout, QLabel, QLineEdit, QPushButton, QWidget, QHBoxLayout, QRadioButton, QFileDialog
from Qt.QtGui import QColor
from Qt.QtWidgets import QFrame
import webbrowser
import os
import requests
import zipfile
import pandas as pd

class ChopChopGetMissense(ToolInstance):

    SESSION_ENDURING = False    # Does this instance persist when session closes
    SESSION_SAVE = True         # We do save/restore in sessions
    help = "help:user/tools/chopchop_get_missense.html"  # Help page URL

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

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

    def _build_ui(self):
        layout = QVBoxLayout()

        # Hyperlink and information
        hegelab_button = QPushButton("Hegelab AlphaMissense")
        hegelab_button.clicked.connect(lambda: webbrowser.open("https://alphamissense.hegelab.org/hotspot"))
        layout.addWidget(QLabel("AlphaMissense PDB and Scores are fetched from"))
        layout.addWidget(hegelab_button)

        # Input for UniProt ID
        self.uniprot_input = QLineEdit()
        layout.addWidget(QLabel("Enter UniProt ID:"))
        layout.addWidget(self.uniprot_input)

        # Input for download folder path
        self.download_path_input = QLineEdit()
        self.download_path_input.setText(str(Path.home() / "Downloads"))
        layout.addWidget(QLabel("Enter download folder path for fetched Files and Attribute file for Scoring Scheme:"))
        layout.addWidget(self.download_path_input)

        # Option to upload custom TSV
        self.use_uploaded_tsv_checkbox = QRadioButton("Use uploaded AlphaMissense TSV file")
        self.use_uploaded_tsv_checkbox.setChecked(False)
        layout.addWidget(self.use_uploaded_tsv_checkbox)

        self.file_path_input = QLineEdit()
        self.file_path_input.setPlaceholderText("Choose a TSV file")
        self.file_path_input.setDisabled(True)
        layout.addWidget(QLabel("Upload AlphaMissense TSV file:"))
        layout.addWidget(self.file_path_input)

        browse_button = QPushButton("Browse")
        browse_button.clicked.connect(self.browse_file)
        browse_button.setDisabled(True)
        layout.addWidget(browse_button)

        def toggle_inputs():
            self.file_path_input.setDisabled(not self.use_uploaded_tsv_checkbox.isChecked())
            browse_button.setDisabled(not self.use_uploaded_tsv_checkbox.isChecked())

        self.use_uploaded_tsv_checkbox.toggled.connect(toggle_inputs)

        # Execute button for downloading and processing
        execute_button = QPushButton("ChopChop Missense PDB")
        execute_button.clicked.connect(self.run_script)
        layout.addWidget(execute_button)

        # Add color bar legend
        layout.addWidget(QLabel("Scoring Scheme:"))
        color_bar = QVBoxLayout()

        colors = [
            ("1: Benign                  AlphaMissense Score: >0.34", "navy"),
            ("2: Ambiguous           AlphaMissense Score: 0.34-0.564", "cornflowerblue"),
            ("3: Pathogenic low     AlphaMissense Score: 0.564-0.78", "red"),
            ("4: Pathogenic high    AlphaMissense Score: 0.78-1.0", "darkred"),
            ("5: No score", "white")
        ]

        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)

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

    def browse_file(self):
        file_path, _ = QFileDialog.getOpenFileName(None, "Select AlphaMissense TSV File", str(Path.home()), "TSV Files (*.tsv)")
        if file_path:
            self.file_path_input.setText(file_path)
    

    def run_script(self):
        download_dir = Path(self.download_path_input.text().strip())

        if not download_dir.exists():
            try:
                download_dir.mkdir(parents=True, exist_ok=True)
            except Exception as e:
                self.session.logger.warning(f"Failed to create directory {download_dir}: {e}")
                return

        uniprot_id = self.uniprot_input.text().strip()
        if not uniprot_id:
            self.session.logger.warning("Please provide a UniProt ID.")
            return

        self.session.logger.info(f"Starting download for UniProt ID: {uniprot_id}")
        zip_file_path, extracted_folder = self.download_hotspot_file(uniprot_id, download_dir)

        if zip_file_path and extracted_folder:
            tsv_file = next(extracted_folder.glob("*.tsv"), None)

            # Process TSV or custom TSV
            if self.use_uploaded_tsv_checkbox.isChecked():
                tsv_file = Path(self.file_path_input.text().strip())
                if not tsv_file.exists():
                    self.session.logger.warning("Provided TSV file does not exist. Falling back to downloaded TSV file.")
                else:
                    self.session.logger.info(f"Using uploaded TSV file: {tsv_file}")
                    self.process_uploaded_tsv(tsv_file, extracted_folder / "AM_score.defattr")
                    return

            if tsv_file:
                self.generate_defattr_file(tsv_file, extracted_folder / "AM_score.defattr")

            # Step 4: Load PDB files into ChimeraX
            pdb_files = list(extracted_folder.glob("*.pdb"))
            if not pdb_files:
                self.session.logger.warning("No PDB files found in the extracted folder.")
                return

            for pdb_file in pdb_files:
                run(self.session, f"open {pdb_file}")

            # Step 5: Apply coloring using defattr file
            run(self.session, f"open {extracted_folder / 'AM_score.defattr'}")
            self.session.logger.info(f"Applied coloring using AM_score.defattr")

            # Step 6: Apply coloring with palette
            run(self.session, "color byattribute AM_score palette 1,navy:2,cornflowerblue:3,red:4,darkred:5,white")
            self.session.logger.info("Applied color byattribute with custom palette")

            self.session.logger.info(f"Process completed successfully for UniProt ID {uniprot_id}")

    def download_hotspot_file(self, uniprot_id, download_dir):
        base_url = "https://alphamissense.hegelab.org/download_hot"
        download_url = f"{base_url}?uniprot_id={uniprot_id}"
        zip_file_path = download_dir / f"{uniprot_id}_hotspots.zip"
        extracted_folder = download_dir / f"{uniprot_id}_hotspots"

        try:
            # Step 1: Download the zip file
            self.session.logger.info(f"Downloading hotspot dataset from {download_url}...")
            response = requests.get(download_url)
            response.raise_for_status()

            with open(zip_file_path, "wb") as f:
                f.write(response.content)

            self.session.logger.info(f"Downloaded to {zip_file_path}")

            # Step 2: Extract the zip file
            with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
                zip_ref.extractall(extracted_folder)

            self.session.logger.info(f"Extracted to {extracted_folder}")
            return zip_file_path, extracted_folder

        except Exception as e:
            self.session.logger.warning(f"Failed during processing: {e}")
            return None, None

    def process_uploaded_tsv(self, tsv_file, defattr_file):
        """Process the uploaded TSV file to generate a defattr file."""
        try:
            self.session.logger.info(f"Processing uploaded TSV file: {tsv_file}")

            # Load the TSV file to inspect columns
            df = pd.read_csv(tsv_file, sep="\t", engine='python')

            # Find the relevant columns by their closest match
            position_col = next((col for col in df.columns if "position" in col.lower()), None)
            score_col = next((col for col in df.columns if "mean" in col.lower() and "am score" in col.lower()), None)

            if not position_col or not score_col:
                raise ValueError("Required columns not found in the TSV file.")

            # Keep only the relevant columns
            df = df[[position_col, score_col]]
            df.columns = ["position", "mean AM score from all"]

            self.session.logger.info(f"TSV file loaded with {len(df)} rows and {len(df.columns)} columns.")

            # Generate defattr file
            self.generate_defattr_file_from_dataframe(df, defattr_file)

        except Exception as e:
            self.session.logger.warning(f"Failed to process uploaded TSV file: {e}")
    

    def generate_defattr_file(self, tsv_file, defattr_file):
        """Generate a defattr file for coloring based on AM_score from a TSV file."""
        try:
            df = pd.read_csv(tsv_file, sep="\t", usecols=["position", "mean AM score from all"], engine='python')
            self.generate_defattr_file_from_dataframe(df, defattr_file)
        except Exception as e:
            self.session.logger.warning(f"Error generating defattr file from TSV: {e}")

    def generate_defattr_file_from_dataframe(self, df, defattr_file):
        """Generate a defattr file from a DataFrame."""
        try:
            # Define ranges and associated scores
            ranges = {
                "benign": (0, 0.34, 1),
                "ambiguous": (0.34, 0.564, 2),
                "patho0": (0.564, 0.78, 3),
                "patho1": (0.78, 1.0, 4)
            }

            # Create defattr content with a header comment
            header = [
                "# Coloring scheme:",
                "# 1: Benign (navy)",
                "# 2: Ambiguous (cornflowerblue)",
                "# 3: Pathogenic low (red)",
                "# 4: Pathogenic high (darkred)",
                "# 5: No score (white)"
            ]
            lines = header + ["attribute: AM_score", "match mode: 1-to-1", "recipient: residues"]

            # Add the first residue with a default score of 5
            residue_scores = {1: 5}

            # Validate and clean data
            for index, row in df.iterrows():
                try:
                    residue_id = int(row["position"])
                    score = float(row["mean AM score from all"])  # Conversion that might fail
                    for category, (low, high, assigned_score) in ranges.items():
                        if low <= score < high:
                            residue_scores[residue_id] = assigned_score
                            break
                except ValueError as e:
                    self.session.logger.warning(
                        f"Skipping row {index} due to invalid data: position={row.get('position')}, "
                        f"score={row.get('mean AM score from all')}. Error: {e}"
                    )

            # Add all residues with assigned values or default score 5
            max_residue = max(residue_scores.keys())
            for residue_id in range(1, max_residue + 1):
                assigned_score = residue_scores.get(residue_id, 5)
                lines.append(f"\t:{residue_id}\t{assigned_score}")

            # Add an empty line at the end
            lines.append("")

            # Write to defattr file
            with open(defattr_file, "w") as f:
                f.write("\n".join(lines))

            self.session.logger.info(f"Generated defattr file: {defattr_file}")

        except Exception as e:
            self.session.logger.warning(f"Error generating defattr file: {e}")


