import pandas as pd
import os
import subprocess
import sys
from Bio import SeqIO
import shutil
from fuson_plm.utils.logging import open_logfile, log_update
    
def ensure_mmseqs_in_path(mmseqs_dir):
    """
    Checks if MMseqs2 is in the PATH. If it's not, add it. MMseqs2 will not run if this is not done correctly.
    
    Args:
        mmseqs_dir (str): Directory containing MMseqs2 binaries
    """
    mmseqs_bin = os.path.join(mmseqs_dir, 'mmseqs')
    
    # Check if mmseqs is already in PATH
    if shutil.which('mmseqs') is None:
        # Export the MMseqs2 directory to PATH
        os.environ['PATH'] = f"{mmseqs_dir}:{os.environ['PATH']}"
        log_update(f"\tAdded {mmseqs_dir} to PATH")

def process_fasta(fasta_path):
    fasta_sequences = SeqIO.parse(open(fasta_path),'fasta')
    d = {}
    for fasta in fasta_sequences:
        id, sequence = fasta.id, str(fasta.seq)

        d[id] = sequence
    
    return d

def analyze_clustering_result(input_fasta: str, tsv_path: str):
    """
    Args:
        input_fasta (str): path to input fasta file
    """
    
    # Process input fasta
    input_d = process_fasta(input_fasta)
    
    # Process clusters.tsv
    clusters = pd.read_csv(f'{tsv_path}',sep='\t',header=None)
    clusters = clusters.rename(columns={
        0: 'representative seq_id',
        1: 'member seq_id'
    })

    clusters['representative seq'] = clusters['representative seq_id'].apply(lambda seq_id: input_d[seq_id])
    clusters['member seq'] = clusters['member seq_id'].apply(lambda seq_id: input_d[seq_id])
    
    # Sort them so that splitting results are reproducible
    clusters = clusters.sort_values(by=['representative seq_id','member seq_id'],ascending=True).reset_index(drop=True)
    
    return clusters

def make_fasta(sequences: dict, fasta_path: str):
    """
    Makes a fasta file from sequences, where the key is the header and the value is the sequence.

    Args:
        sequences (dict): A dictionary where the key is the header and the value is the sequence.

    Returns:
        str: The path to the fasta file.
    """
    with open(fasta_path, 'w') as f:
      for header, sequence in sequences.items():
        f.write(f'>{header}\n{sequence}\n')
    
    return fasta_path

def run_mmseqs_clustering(input_fasta, output_dir, min_seq_id=0.3, c=0.8, cov_mode=0, cluster_mode=0, path_to_mmseqs='fuson_plm/mmseqs'):
    """
    Runs MMSeqs2 clustering using easycluster module
    
    Args:
        input_fasta (str): path to input fasta file, formatted >header\nsequence\n>header\nsequence.... 
        output_dir (str): path to output dir for clustering results
        min_seq_id (float): number [0,1] representing --min-seq-id in cluster command
        c (float): nunber [0,1] representing -c in cluster command
        cov_mode (int): number 0, 1, 2, or 3 representing --cov-mode in cluster command 
        cluster_mode (int): number 0, 1, or 2 representing --cluster-mode in cluster command
        
    """
    # Get mmseqs dir
    log_update("\nRunning MMSeqs clustering...")
    mmseqs_dir = os.path.join(path_to_mmseqs[0:path_to_mmseqs.index('/mmseqs')], 'mmseqs/bin')

    # Ensure MMseqs2 is in the PATH
    ensure_mmseqs_in_path(mmseqs_dir)
    
    # Define paths for MMseqs2
    mmseqs_bin = "mmseqs"  # Ensure this is in your PATH or provide the full path to mmseqs binary

    # Create the output directory
    os.makedirs(output_dir, exist_ok=True)

    # Run MMseqs2 easy-cluster
    cmd_easy_cluster = [
        mmseqs_bin, "easy-cluster", input_fasta, os.path.join(output_dir, "mmseqs"), output_dir,
        "--min-seq-id", str(min_seq_id),
        "-c", str(c),
        "--cov-mode", str(cov_mode),
        "--cluster-mode", str(cluster_mode),
        "--dbtype", "1"
    ]
    
    # Write the command to a log file
    log_update("\n\tCommand entered to MMSeqs2:")
    log_update("\t" + " ".join(cmd_easy_cluster) + "\n")
        
    subprocess.run(cmd_easy_cluster, check=True)

    log_update(f"Clustering completed. Results are in {output_dir}")
    
def cluster_summary(clusters: pd.DataFrame):
    """
    Summarizes how many clusters were formed, how big they are, etc ...
    """
    grouped_clusters = clusters.groupby('representative seq_id')['member seq_id'].count().reset_index().rename(columns={'member seq_id':'member count'})
    assert len(grouped_clusters) == len(clusters['representative seq_id'].unique()) # make sure number of cluster reps = # grouped clusters
    
    total_seqs = sum(grouped_clusters['member count'])
    log_update(f"Created {len(grouped_clusters)} clusters of {total_seqs} sequences")
    log_update(f"\t{len(grouped_clusters.loc[grouped_clusters['member count']==1])} clusters of size 1")
    csize1_seqs = sum(grouped_clusters[grouped_clusters['member count']==1]['member count'])
    log_update(f"\t\tsequences: {csize1_seqs} ({round(100*csize1_seqs/total_seqs, 2)}%)")
    
    log_update(f"\t{len(grouped_clusters.loc[grouped_clusters['member count']>1])} clusters of size > 1")
    csizeg1_seqs = sum(grouped_clusters[grouped_clusters['member count']>1]['member count'])
    log_update(f"\t\tsequences: {csizeg1_seqs} ({round(100*csizeg1_seqs/total_seqs, 2)}%)")
    log_update(f"\tlargest cluster: {max(grouped_clusters['member count'])}")

    log_update("\nCluster size breakdown below...")
    
    value_counts = grouped_clusters['member count'].value_counts().reset_index().rename(columns={'index':'cluster size (n_members)','member count': 'n_clusters'})
    log_update(value_counts.sort_values(by='cluster size (n_members)',ascending=True).to_string(index=False))