# vim: set expandtab shiftwidth=4 softtabstop=4:

# === UCSF ChimeraX Copyright ===
# Copyright 2016 Regents of the University of California.
# All rights reserved.  This software provided pursuant to a
# license agreement containing restrictions on its disclosure,
# duplication and use.  For details see:
# http://www.rbvi.ucsf.edu/chimerax/docs/licensing.html
# This notice must be embedded in or attached to all copies,
# including partial copies, of the software or any revisions
# or derivations thereof.
# === UCSF ChimeraX Copyright ===

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 QLineEdit, QComboBox, QMainWindow, QApplication, QWidget, QVBoxLayout, QLabel, QPushButton, QHBoxLayout, QFileDialog
from Qt.QtGui import QColor
from Qt.QtWidgets import QFrame
import csv
import os
import requests
import distutils.dir_util
import shutil  
from Qt.QtGui import QAction

class ChopChopMFalignment(ToolInstance):

    # Inheriting from ToolInstance makes us known to the ChimeraX tool mangager,
    # so we can be notified and take appropriate action when sessions are closed,
    # saved, or restored, and we will be listed among running tools and so on.
    #
    # If cleaning up is needed on finish, override the 'delete' method
    # but be sure to call 'delete' from the superclass at the end.

    SESSION_ENDURING = False    # Does this instance persist when session closes
    SESSION_SAVE = True         # We do save/restore in sessions
    help = "help:user/tools/tutorial.html"
                                # Let ChimeraX know about our help page

    def __init__(self, session, tool_name):
        # 'session'   - chimerax.core.session.Session instance
        # 'tool_name' - string

        # Initialize base class.
        super().__init__(session, tool_name)

        # Set name displayed on title bar (defaults to tool_name)
        # Must be after the superclass init, which would override it.
        self.display_name = "Sequence"

        # Create the main window for our tool.  The window object will have
        # a 'ui_area' where we place the widgets composing our interface.
        # The window isn't shown until we call its 'manage' method.
        #
        # Note that by default, tool windows are only hidden rather than
        # destroyed when the user clicks the window's close button.  To change
        # this behavior, specify 'close_destroys=True' in the MainToolWindow
        # constructor.
        #from chimerax.ui import MainToolWindow
        self.tool_window = MainToolWindow(self)

        # We will be adding an item to the tool's context menu, so override
        # the default MainToolWindow fill_context_menu method
        self.tool_window.fill_context_menu = self.fill_context_menu

        # Our user interface is simple enough that we could probably inline
        # the code right here, but for any kind of even moderately complex
        # interface, it is probably better to put the code in a method so
        # that this __init__ method remains readable.
        self._build_ui()
    
    
    def _build_ui(self):
        """
        Build the user interface for the tool.
        Includes model/chain selection, alignment options, scoring scheme legend,
        and the new color customization section.
        """
        # Extract available model IDs and chain IDs
        all_models, all_chains, all_modl_index, all_chains_with_index = self.select_models_and_chains(self.session)

        # Define the main layout
        layout = QVBoxLayout()

        # Model and Chain Selection
        self.combobox1 = QComboBox()
        for x in all_chains:
            x.split(":")
            self.combobox1.addItem("model {}, chain {}".format(x[0], x[2]))
        layout.addWidget(QLabel("Choose model ID and chain ID for alignment"))
        layout.addWidget(self.combobox1)

        # Input field for protein sequence or UniProt ID
        self.text1 = QLineEdit()
        layout.addWidget(QLabel("Enter an amino acid sequence or UniProt ID to be aligned to your model"))
        layout.addWidget(self.text1)

        # Alignment Button
        button1 = QPushButton("ChopChop SequenceAlignment")
        layout.addWidget(button1)

        # Download folder path
        self.download_path_input = QLineEdit()
        self.download_path_input.setText(str(Path.home() / "Downloads"))
        layout.addWidget(QLabel("Enter download folder path:"))
        layout.addWidget(self.download_path_input)

        # Browse button for folder selection
        browse_button = QPushButton("Browse")
        browse_button.clicked.connect(self.browse_folder)
        layout.addWidget(browse_button)

        # Download Button
        button2 = QPushButton("Download")
        layout.addWidget(button2)

        # Connect button actions
        button1.clicked.connect(self.align)
        button2.clicked.connect(self.save)

        # Scoring Scheme Legend
        layout.addWidget(QLabel("Scoring Scheme:"))
        color_bar = QVBoxLayout()
        colors = [
            ("conserved", "teal"),
            ("semi-conserved", "palegoldenrod"),
            ("not conserved", "lightsalmon"),
            ("gap", "dimgrey"),
        ]
        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 after running ChopChp SequenceAlignment if needed:"))
        self.color_inputs = {}
        color_entries = [
            ("1 (Conserved):", "teal"),
            ("2 (Semi-Conserved):", "palegoldenrod"),
            ("3 (Not Conserved):", "lightsalmon"),
            ("4 (Gap):", "dimgrey"),
        ]
        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)

        # Show the tool window
        self.tool_window.manage('side')
    
    def apply_custom_colors(self):
        """
        Apply the new color scheme to the alignment visualization using ChimeraX.
        Collects the user-defined colors from the input fields and constructs
        the ChimeraX command to apply the color scheme.
        """
        # 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 MUSCLEalignment palette " + ":".join(palette)

        # Execute the command in ChimeraX
        try:
            run(self.session, color_command)
            print("Custom color scheme applied:", color_command)
        except Exception as e:
            print("Error applying custom color scheme:", str(e))
    

    def browse_folder(self):
        # Open a folder selection dialog
        folder_path = QFileDialog.getExistingDirectory(self.tool_window.ui_area, "Select Download Folder")
        if folder_path:
            self.download_path_input.setText(folder_path)

    def browse_file(self):
        # Placeholder for browse file logic
        pass

    # Removed the assignment to self.tool_window.fill_context_menu

    def fill_context_menu(self, menu, x, y):
        # This method should populate the context menu with any additional options you want
        # Add placeholder functionality
        action = QAction("Sample Action", menu)
        action.triggered.connect(lambda: print("Sample Action triggered"))
        menu.addAction(action)
    
    
  

    def is_semi_conserved(self, residue1, residue2):
        # Define the criteria for semi-conserved residues based on similarity of amino acid properties
        property_similarities = {
            'G': ['G'],
            'P': ['P'],
            'A': ['A'],
            'S': ['S'],
            'T': ['T'],
            'V': ['V', 'I'],
            'L': ['L', 'I', 'V'],
            'I': ['I', 'L', 'V'],
            'M': ['M'],
            'F': ['F', 'Y', 'W'],
            'Y': ['Y', 'F', 'W'],
            'W': ['W', 'F', 'Y'],
            'H': ['H', 'N', 'Q'],
            'N': ['N', 'H', 'Q'],
            'Q': ['Q', 'H', 'N'],
            'R': ['R', 'K'],
            'K': ['K', 'R'],
            'D': ['D', 'E'],
            'E': ['E', 'D'],
            'C': ['C'],
        }
        return residue1.upper() in property_similarities.get(residue2.upper(), [])
        
    def add_chain_id(self, file_location, file_name, chain_id, model_id):
        try:
            # Construct the file path
            file_path = file_location + file_name
            # Load the file contents
            with open(file_path) as file:
                lines = file.readlines()
            # Add "/" and chain ID before ":" in the first column starting from line 10
            modified_lines = []
            for i, line in enumerate(lines):
                if i >= 9 and line.startswith("\t:"):
                    modified_lines.append(f"\t#{model_id}/{chain_id}{line[1:]}")
                else:
                    modified_lines.append(line)
            # Construct the new file path
            new_file_path = file_location + "/Chain{}_{}".format(chain_id, file_name[1:])
            # Save the modified contents to the new file
            with open(new_file_path, "w") as file:
                file.writelines(modified_lines)
            print(f"Chain ID added successfully! Modified file saved as: {new_file_path}")
        except FileNotFoundError:
            print("File not found. Please enter a valid file location and name.")
        except Exception as e:
            print("An error occurred:", str(e))
    
    def align(self):
        #create a directory for output files
        mypath = Path().home()
        output_dir = "{}/ChopChop_output".format(mypath)
        if not os.path.exists(output_dir):
            os.mkdir(output_dir)
        #get input chain ID and model ID from combobox 
        model_chain = str(self.combobox1.currentText()) 
        if len(model_chain) == 16:
            model = model_chain[6]
            chain = model_chain[15]
        if len(model_chain) == 17:
            model = model_chain[6] + model_chain[7]
            chain = model_chain[16]
        if len(model_chain) == 18:
            model = model_chain[6] + model_chain[7] + model_chain[8]
            chain = model_chain[17]
        #get fasta sequence of chosen chain
        run(self.session, "sequence chain #{}/{}".format(model, chain))
        run(self.session, "save {}/chopchop_sequence1.fasta format fasta alignment {}/{}".format(output_dir, model, chain))
        file = open("{}/chopchop_sequence1.fasta".format(output_dir), mode = "r", encoding = "utf-8-sig")
        lines = file.readlines()
        file.close()
        header = True
        sequence1 = ""
        for line in lines:
            if header == False:
                sequence1 = sequence1 + line
            if header == True:
                header = False 
        sequence1 = " ".join(sequence1.split())
        sequence1 = sequence1.replace(" ", "")
        sequence1 = sequence1.replace("\n", "")       
        #format fasta sequence from input text
        sequence2 = str(self.text1.text()) 
        sequence2 = " ".join(sequence2.split())
        sequence2 = sequence2.replace(" ", "")
        sequence2 = sequence2.replace("\n", "")
        #get fasta sequence from uniprot ID if given  
        if any(char.isdigit() for char in sequence2):
            #Downloads the FASTA file for a given UniProt ID and saves it in the output directory.
            url = f"https://www.uniprot.org/uniprot/{sequence2}.fasta"
            response = requests.get(url)
            if response.status_code == 200:
                fasta_file_path = os.path.join(output_dir, f"{sequence2}.fasta")
                with open(fasta_file_path, "w") as fasta_file:
                    fasta_file.write(response.text)
                    print(f"FASTA file for {sequence2} downloaded successfully.")
            else:
                raise ValueError(f"Failed to download FASTA file for UniProt ID: {sequence2}")
            file = open(fasta_file_path, mode = "r", encoding = "utf-8-sig")
            lines = file.readlines()
            file.close()
            header = True
            sequence2 = ""
            for line in lines:
                if header == False:
                    sequence2 = sequence2 + line
                if header == True:
                    header = False 
            sequence2 = " ".join(sequence2.split())
            sequence2 = sequence2.replace(" ", "")
            sequence2 = sequence2.replace("\n", "")
        #name proteins for internal use
        protein_name1 = "proteinA"
        protein_name2 = "proteinB"
        # Perform sequence alignment
        run(self.session, "sequence align {},{} program muscle".format(sequence1, sequence2))
        #save last alignment in fasta format
        alignment_Id = 1
        while True:
            try: 
                run(self.session, "save {}/chopchop_alignment.fasta format fasta alignment {}".format(output_dir, str(alignment_Id)))
                alignment_Id += 1
            except:
                break
        # Parse the alignment results
        alignment = []
        with open("{}/chopchop_alignment.fasta".format(output_dir)) as alignment_file:
            lines = alignment_file.readlines()
            #get first alignment sequence
            start_index = 1
            end_index = lines.index("\n")
            sequence1_alignment = "".join(lines[start_index:end_index]).strip()
            sequence1_alignment = sequence1_alignment.replace("\n", "")
            #get second alignment sequence
            start_index = lines.index("\n") + 2
            sequence2_alignment = "".join(lines[start_index:]).strip()
            sequence2_alignment = sequence2_alignment.replace("\n", "")
            print("Sequence 1 Alignment:", sequence1_alignment)
            print("Sequence 2 Alignment:", sequence2_alignment)
            for residue1, residue2 in zip(sequence1_alignment, sequence2_alignment):
                row = []
                row.append(residue1)  # Residue from sequence 1
                row.append(residue2)  # Residue from sequence 2
                alignment.append(row)
        # Define scoring criteria for different residue types
        conserved_score = 1
        semi_conserved_score = 2
        non_conserved_score = 3
        gap_score = 4
        # Calculate alignment scores
        scores = []
        for residue1, residue2 in alignment:
            if residue1 == residue2:
                score = conserved_score
            elif residue1 != '-' and residue2 != '-':
                if self.is_semi_conserved(residue1, residue2):
                    score = semi_conserved_score
                else:
                    score = non_conserved_score
            elif residue1 == '-' or residue2 == '-':
                score = gap_score
            else:
                score = non_conserved_score
            scores.append(score)
        # Write alignment results to TSV file
        output_file = output_dir + f"/{protein_name1}_{protein_name2}_alignment.csv"
        with open(output_file, mode='w', newline='') as file:
            writer = csv.writer(file, delimiter='\t')
            writer.writerow([f"{protein_name1}", f"{protein_name2}", "Alignment Score"])
            writer.writerow(["", "", "1: Conserved, 2: Semi-Conserved, 3: Not Conserved, 4: Gap"])
            for (residue1, residue2), score in zip(alignment, scores):
                writer.writerow([residue1, residue2, score])
        print("Alignment results saved to:", output_file)  
        # Print alignment information for debugging
        print("Alignment scores:", scores)
        print("Alignment dictionary:")
        print(alignment)
        #save as csv file
        file_path = output_dir
        file_name = f"/{protein_name1}_{protein_name2}_alignment.csv"
        file = file_path + file_name
        # Read the CSV file
        with open(file, 'r') as inp:
            lines = inp.readlines()
        #remove lines with double quotes
        file_name = f"/{protein_name1}_{protein_name2}_alignment_final.csv"
        output_file = file_path + file_name
        with open(output_file, 'w') as out:
            for line in lines:
                if not '"' in line:
                    out.write(line)
        print("Rows with double quotes removed. Saved to", output_file)        
        #parse amino acid sequence of alignment and alignment scores for model in chimerax
        # Read the input .tsv file and skip the second row
        input_file_path = output_file
        with open(input_file_path, 'r') as tsv_file:
            # Create a CSV reader object
            tsv_reader = csv.reader(tsv_file, delimiter='\t')
            # Skip the header row
            next(tsv_reader)
            # Skip the second row
            next(tsv_reader)
            # Create separate lists for protein names and alignment scores
            protein_names_column1 = []
            protein_names_column2 = []
            alignment_scores_column3 = []
            # Iterate over the rows in the .tsv file
            for row in tsv_reader:
                protein_name_column1 = row[0]
                protein_name_column2 = row[1]
                alignment_score = row[2]
                # Append protein names and alignment scores to the respective lists
                protein_names_column1.append(protein_name_column1)
                protein_names_column2.append(protein_name_column2)
                alignment_scores_column3.append(alignment_score)
        # Create a new CSV file for column 1 protein names and alignment scores
        output_file_path_column1 = output_dir + "/column1_proteins.csv"
        with open(output_file_path_column1, 'w', newline='') as output_file_column1:
            csv_writer_column1 = csv.writer(output_file_column1, delimiter='\t')
            csv_writer_column1.writerow(['Residue', 'Alignment Score'])
            for i in range(len(protein_names_column1)):
                csv_writer_column1.writerow([protein_names_column1[i], alignment_scores_column3[i]])
            output_file_column1.write("# Scoring: 1 = Conserved, 2 = Semi-Conserved, 3 = Not Conserved, 4 = Gap\n")
        print("Extraction complete!")
        # Load the CSV file and remove rows starting with "-" in the first column
        directory = output_dir
        filename = "/column1_proteins.csv"
        filepath = output_file_path_column1
        data = []
        with open(filepath, 'r') as file:
            reader = csv.reader(file, delimiter='\t')
            for row in reader:
                if not row[0].startswith('-'):
                    data.append(row)
        # Number the residues and remove capital letters before ":"
        header = '#  \n'
        header += '#  Use this file to assign the attribute in Chimera with the Define Attribute tool.\n'
        header += '#  open .defattr then use the command: color byattribute MUSCLEalignment palette 1,teal:2,palegoldenrod:3,lightsalmon:4,dimgrey\n'
        header += '#  \n'
        header += '#  Scoring: 1 = Conserved, 2 = Semi-Conserved, 3 = Not Conserved, 4 = Gap\n'
        header += '#  \n'
        header += 'attribute: MUSCLEalignment\n'
        header += 'match mode: 1-to-1\n'
        header += 'recipient: residues\n'
        for i in range(1, len(data) - 1):
            residue = data[i][0]
            residue_parts = residue.split(':')
            residue_parts[0] = residue_parts[0].lstrip('ARNDCEQGHILKMFPSTWYV')
            data[i][0] = '\t' + residue_parts[0] + ':' + str(i).strip('"')
        # Save the updated data with the modified file name
        output_filename = f"{os.path.splitext(filename)[0]}_MUSCLEalignment4ChimeraX.defattr"
        output_filepath = directory + output_filename
        with open(output_filepath, 'w', newline='') as file:
            file.write(header)
            for row in data[1:-1]:
                row_str = '\t'.join(row).strip('"') + '\n'
                file.write(row_str)
        file.close()
        print("The updated data has been saved as", output_filename)
        #color the model in chimerax according to the defattr file
        self.add_chain_id(directory, output_filename, chain, model)
        run(self.session, "open {}/Chain{}_{}".format(directory, chain, output_filename[1:]))
        run(self.session, "color byattribute MUSCLEalignment palette 1,teal:2,palegoldenrod:3,lightsalmon:4,dimgrey")
        
        # Execute lighting commands
        run(self.session, "lighting flat")
        run(self.session, "lighting soft")
        

        
    def save(self):
        #copy output folder to download directory
        download_dir = Path(self.download_path_input.text().strip())
        mypath = Path().home()
        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
        output_dir = "{}/ChopChop_output".format(mypath)
        distutils.dir_util.copy_tree(output_dir, str(download_dir))
        #delete original output folder
        shutil.rmtree(output_dir)
        print("Alignment downloaded in {}".format(download_dir))
        

    def select_models_and_chains(self, session):
        all_chains = []
        all_chains_with_index = {}
        all_models = []
        all_modl_index = {}

        for index, model in enumerate(session.models):
            if '.' not in model.id_string:
                all_models.append(f'{model.id_string}')
                all_modl_index[f'{model.id_string}'] = index

                for chain in model.chains:
                    info = f'{model.id_string}:{chain.chain_id}'
                    all_chains.append(info)
                    all_chains_with_index[info] = index

        return all_models, all_chains, all_modl_index, all_chains_with_index
      
    def fill_context_menu(self, menu, x, y):
        # Add any tool-specific items to the given context menu (a QMenu instance).
        # The menu will then be automatically fijlled out with generic tool-related actions
        # (e.g. Hide Tool, Help, Dockable Tool, etc.) 
        #
        # The x,y args are the x() and y() values of QContextMenuEvent, in the rare case
        # where the items put in the menu depends on where in the tool interface the menu
        # was raised.
        #from Qt.QtGui import QAction
        clear_action = QAction("Clear", menu)
        clear_action.triggered.connect(lambda *args: self.line_edit.clear())
        menu.addAction(clear_action)
