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, QComboBox, QWidget
import os
import requests
import pandas as pd
import zipfile

class ChopChopMissense(ToolInstance):
    """
    A ChimeraX tool for fetching AlphaMissense data, aligning sequences, and coloring structures.
    """
    SESSION_ENDURING = False
    SESSION_SAVE = True

    def __init__(self, session, tool_name):
        super().__init__(session, tool_name)
        self.display_name = "Missense"
        self.tool_window = MainToolWindow(self)
        self._build_ui()

    def _build_ui(self):
        all_models, all_chains, *_ = self.select_models_and_chains(self.session)

        layout = QVBoxLayout()

        self.combobox1 = QComboBox()
        for chain in all_chains:
            self.combobox1.addItem(chain)
        layout.addWidget(QLabel("Choose model ID and chain ID for alignment"))
        layout.addWidget(self.combobox1)

        self.text1 = QLineEdit()
        layout.addWidget(QLabel("Enter Human Protein UniProt ID to fetch AlphaMissense data"))
        layout.addWidget(self.text1)

        button1 = QPushButton("ChopChop Missense")
        layout.addWidget(button1)
        button1.clicked.connect(self.fetch_align_and_color)

        container = QWidget()
        container.setLayout(layout)
        self.tool_window.ui_area.setLayout(layout)
        self.tool_window.manage('side')

    def select_models_and_chains(self, session):
        self.session.logger.info("Selecting models and chains.")
        all_chains = []
        all_models = []

        for model in session.models:
            all_models.append(f'{model.id_string}')
            for chain in model.chains:
                all_chains.append(f'{model.id_string}:{chain.chain_id}')

        self.session.logger.info(f"Found models: {all_models}, chains: {all_chains}")
        return all_models, all_chains

    def fetch_align_and_color(self):
        self.session.logger.info("Starting fetch, align, and color process.")
        selected_chain = self.combobox1.currentText()
        uniprot_id = self.text1.text().strip()

        if not selected_chain or not uniprot_id:
            self.session.logger.warning("Please provide both a model/chain and a UniProt ID.")
            return

        try:
            model_id, chain_id = selected_chain.split(":")
        except ValueError:
            self.session.logger.warning("Invalid chain selection format. Ensure correct selection.")
            return

        self.session.logger.info(f"Selected model: {model_id}, chain: {chain_id}")

        sequence1 = self.get_chain_sequence_as_string(model_id, chain_id)

        if not sequence1:
            self.session.logger.warning("Failed to get chain sequence.")
            return

        zip_file_path, extracted_folder = self.download_hotspot_file(uniprot_id, Path.home() / "Downloads")
        if not extracted_folder:
            self.session.logger.warning("Failed to download and extract AlphaMissense data.")
            return

        csv_file = extracted_folder / f"AlphaMissense-Hotspot-{uniprot_id}.tsv"
        sequence2 = self.extract_tsv_sequence_as_string(csv_file)

        if not sequence2:
            self.session.logger.warning("Failed to extract sequence from AlphaMissense TSV.")
            return

        alignment_file = self.perform_alignment(sequence1, sequence2, extracted_folder)
        if not alignment_file:
            self.session.logger.warning("Failed to perform sequence alignment.")
            return

        aligned_sequence, aligned_tsv, residue_numbers = self.parse_alignment_with_residues(alignment_file)

        if aligned_sequence is None or aligned_tsv is None:
            self.session.logger.warning("Failed to parse alignment file.")
            return

        self.color_model(aligned_sequence, aligned_tsv, residue_numbers, csv_file, model_id, chain_id, extracted_folder)

    def get_chain_sequence_as_string(self, model_id, chain_id):
        self.session.logger.info(f"Fetching sequence for model {model_id}, chain {chain_id}.")
        output_dir = Path("/tmp")
        output_dir.mkdir(parents=True, exist_ok=True)
        sequence1_path = output_dir / "chopchop_sequence1.fasta"

        try:
            run(self.session, f"sequence chain #{model_id}/{chain_id}")
            self.session.logger.info("Sequence alignment performed.")
            run(self.session, f"save {sequence1_path} format fasta")
            self.session.logger.info(f"Sequence saved to {sequence1_path}.")

            with open(sequence1_path, "r", encoding="utf-8-sig") as file:
                lines = file.readlines()

            header = True
            sequence1 = ""
            for line in lines:
                if header:
                    header = False
                else:
                    sequence1 += line

            sequence1 = "".join(sequence1.split())
            return sequence1

        except Exception as e:
            self.session.logger.warning(f"Error fetching chain sequence: {e}")
            return None

    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:
            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}")

            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 extract_tsv_sequence_as_string(self, csv_file):
        self.session.logger.info(f"Extracting sequence from AlphaMissense CSV file: {csv_file}.")
        try:
            df = pd.read_csv(csv_file, sep="\t", engine="python")
            sequence_column = "a.a."

            if sequence_column not in df.columns:
                raise ValueError(f"Column '{sequence_column}' not found in TSV file.")

            sequence2 = "".join(df[sequence_column].dropna().astype(str))
            sequence2 = "".join(sequence2.split())
            return sequence2

        except Exception as e:
            self.session.logger.warning(f"Error extracting sequence from CSV: {e}")
            return None

    def perform_alignment(self, sequence1, sequence2, extracted_folder):
        self.session.logger.info("Performing alignment for sequences.")

        try:
            alignment_file_path = extracted_folder / "chopchop_alignment.fasta"

            protein_name1 = "proteinA"
            protein_name2 = "proteinB"

            run(self.session, f"sequence align {sequence1},{sequence2} program muscle")

            alignment_Id = 1
            while True:
                try:
                    run(self.session, f"save {alignment_file_path} format fasta alignment {alignment_Id}")
                    alignment_Id += 1
                except:
                    break

            self.session.logger.info(f"Alignment completed and saved to {alignment_file_path}.")
            return alignment_file_path

        except Exception as e:
            self.session.logger.warning(f"Error during alignment: {e}")
            return None

    def parse_alignment_with_residues(self, alignment_file):
        self.session.logger.info(f"Parsing alignment file: {alignment_file}.")
        try:
            with open(alignment_file, "r") as f:
                lines = f.readlines()

            aligned_sequence = ""
            aligned_tsv = ""
            residue_numbers = []

            in_sequence = False
            current_residue = 1

            for line in lines:
                if line.startswith(">"):
                    if in_sequence:
                        break
                    in_sequence = True
                    continue

                if in_sequence:
                    aligned_sequence += line.strip()

                    for char in line.strip():
                        if char != "-":
                            residue_numbers.append(current_residue)
                            current_residue += 1
                        else:
                            residue_numbers.append(None)

            self.session.logger.info("Alignment parsed successfully with residue mapping.")
            return aligned_sequence, aligned_tsv, residue_numbers
        except Exception as e:
            self.session.logger.warning(f"Error parsing alignment file: {e}")
            return None, None, None

    def color_model(self, aligned_sequence, aligned_tsv, residue_numbers, csv_file, model_id, chain_id, extracted_folder):
        self.session.logger.info("Coloring model based on alignment and AlphaMissense data.")
        attribute_lines = [
            "attribute: MissenseScores",
            "match mode: 1-to-1",
            "recipient: residues"
        ]

        try:
            df = pd.read_csv(csv_file, sep="\t", engine="python")
            score_col = "mean AM score from all"

            if score_col not in df.columns:
                raise ValueError(f"Column '{score_col}' not found in TSV file.")

            scores = dict(zip(df.index + 1, df[score_col]))

            for res, res_num in zip(aligned_sequence, residue_numbers):
                if res == "-" or res_num is None:
                    attribute_lines.append(f"\t:{res_num}\t5")  # No score for gaps
                    continue

                score = scores.get(res_num, None)
                if score is None:
                    attribute_lines.append(f"\t:{res_num}\t5")  # No score available
                elif score > 0.78:
                    attribute_lines.append(f"\t:{res_num}\t4")  # Pathogenic high
                elif score > 0.564:
                    attribute_lines.append(f"\t:{res_num}\t3")  # Pathogenic low
                elif score > 0.34:
                    attribute_lines.append(f"\t:{res_num}\t2")  # Ambiguous
                else:
                    attribute_lines.append(f"\t:{res_num}\t1")  # Benign

            attribute_file = extracted_folder / "MissenseScores.defattr"
            with open(attribute_file, "w") as f:
                f.write("\n".join(attribute_lines))
                f.write("\n")  # Ensure the last line is empty

            self.session.logger.info(f"Attributes saved to {attribute_file}.")

            run(self.session, f"open {attribute_file}")
            run(self.session, "color byattribute MissenseScores palette 1,navy:2,cornflowerblue:3,firebrick:4,darkred:5,white")
            self.session.logger.info("Coloring completed successfully.")
        except Exception as e:
            self.session.logger.warning(f"Error during coloring: {e}")
