#!/usr/bin/env python3
"""
Script to predict PAM sequences for protein sequences and print the PAM distribution.

Edit the SEQUENCES dictionary below with your protein sequences.
This script shows the full softmaxed probability distribution for each position in the PAM.
"""

from collections import Counter
import torch

# Import PAMMatching from cas9.objectives
from cas9.objectives import PAMMatching

# ============================================================================
# EDIT THIS DICTIONARY WITH YOUR SEQUENCES
# Key = sequence name, Value = protein sequence string
# ============================================================================

SEQUENCES = {
    "sacas9": "MKRNYILGLDIGITSVGYGIIDYETRDVIDAGVRLFKEANVENNEGRRSKRGARRLKRRRRHRIQRVKKLLFDYNLLTDHSELSGINPYEARVKGLSQKLSEEEFSAALLHLAKRRGVHNVNEVEEDTGNELSTKEQISRNSKALEEKYVAELQLERLKKDGEVRGSINRFKTSDYVKEAKQLLKVQKAYHQLDQSFIDTYIDLLETRRTYYEGPGEGSPFGWKDIKEWYEMLMGHCTYFPEELRSVKYAYNADLYNALNDLNNLVITRDENEKLEYYEKFQIIENVFKQKKKPTLKQIAKEILVNEEDIKGYRVTSTGKPEFTNLKVYHDIKDITARKEIIENAELLDQIAKILTIYQSSEDIQEELTNLNSELTQEEIEQISNLKGYTGTHNLSLKAINLILDELWHTNDNQIAIFNRLKLVPKKVDLSQQKEIPTTLVDDFILSPVVKRSFIQSIKVINAIIKKYGLPNDIIIELAREKNSKDAQKMINEMQKRNRQTNERIEEIIRTTGKENAKYLIEKIKLHDMQEGKCLYSLEAIPLEDLLNNPFNYEVDHIIPRSVSFDNSFNNKVLVKQEENSKKGNRTPFQYLSSSDSKISYETFKKHILNLAKGKGRISKTKKEYLLEERDINRFSVQKDFINRNLVDTRYATRGLMNLLRSYFRVNNLDVKVKSINGGFTSFLRRKWKFKKERNKGYKHHAEDALIIANADFIFKEWKKLDKAKKVMENQMFEEKQAESMPEIETEQEYKEIFITPHQIKHIKDFKDYKYSHRVDKKPNRELINDTLYSTRKDDKGNTLIVNNLNGLYDKDNDKLKKLINKSPEKLLMYHHDPQTYQKLKLIMEQYGDEKNPLYKYYEETGNYLTKYSKKDNGPVIKKIKYYGNKLNAHLDITDDYPNSRNKVVKLSLKPYRFDVYLDNGVYKFVTVKNLDVIKKENYYEVNSKCYEEAKKLKKISNQAEFIASFYNNDLIKINGELYRVIGVNNDLLNRIEVNMIDITYREYLENMNDKRPPRIIKTIASKTQSIKKYSTDILGNLYEVKSKKHPQIIKKG",
    
    "cjcas9": "MARILAFDIGISSIGWAFSENDELKDCGVRIFTKVENPKTGESLALPRRLARSARKRLARRKARLNHLKHLIANEFKLNYEDYQSFDESLAKAYKGSLISPYELRFRALNELLSKQDFARVILHIAKRRGYDDIKNSDDKEKGAILKAIKQNEEKLANYQSVGEYLYKEYFQKFKENSKEFTNVRNKKESYERCIAQSFLKDELKLIFKKQREFGFSFSKKFEEEVLSVAFYKRALKDFSHLVGNCSFFTDEKRAPKNSPLAFMFVALTRIINLLNNLKNTEGILYTKDDLNALLNEVLKNGTLTYKQTKKLLGLSDDYEFKGEKGTYFIEFKKYKEFIKALGEHNLSQDDLNEIAKDITLIKDEIKLKKALAKYDLNQNQIDSLSKLEFKDHLNISFKALKLVTPLMLEGKKYDEACNELNLKVAINEDKKDFLPAFNETYYKDEVTNPVVLRAIKEYRKVLNALLKKYGKVHKINIELAREVGKNHSQRAKIEKEQNENYKAKKDAELECEKLGLKINSKNILKLRLFKEQKEFCAYSGEKIKISDLQDEKMLEIDHIYPYSRSFDDSYMNKVLVFTKQNQEKLNQTPFEAFGNDSAKWQKIEVLAKNLPTKKQKRILDKNYKDKEQKNFKDRNLNDTRYIARLVLNYTKDYLDFLPLSDDENTKLNDTQKGSKVHVEAKSGMLTSALRHTWGFSAKDRNNHLHHAIDAVIIAYANNSIVKAFSDFKKEQESNSAELYAKKISELDYKNKRKFFEPFSGFRQKVLDKIDEIFVSKPERKKPSGALHEETFRKEEEFYQSYGGKEGVLKALELGKIRKVNGKIVKNGDMFRVDIFKHKKTNKFYAVPIYTMDFALKVLPNKAVARSKKGEIKDWILMDENYEFCFSLYKDSLILIQTKDMQEPEFVYYNAFTSSTVSLIVSKHDNKFETLSKNQKILFKNANEKEVIAKSIGIQNLKVFEKYIVSALGEVTKAEFRQREDFKK",
    
    "spcas9": "MDKKYSIGLDIGTNSVGWAVITDEYKVPSKKFKVLGNTDRHSIKKNLIGALLFDSGETAEATRLKRTARRRYTRRKNRICYLQEIFSNEMAKVDDSFFHRLEESFLVEEDKKHERHPIFGNIVDEVAYHEKYPTIYHLRKKLVDSTDKADLRLIYLALAHMIKFRGHFLIEGDLNPDNSDVDKLFIQLVQTYNQLFEENPINASGVDAKAILSARLSKSRRLENLIAQLPGEKKNGLFGNLIALSLGLTPNFKSNFDLAEDAKLQLSKDTYDDDLDNLLAQIGDQYADLFLAAKNLSDAILLSDILRVNTEITKAPLSASMIKRYDEHHQDLTLLKALVRQQLPEKYKEIFFDQSKNGYAGYIDGGASQEEFYKFIKPILEKMDGTEELLVKLNREDLLRKQRTFDNGSIPHQIHLGELHAILRRQEDFYPFLKDNREKIEKILTFRIPYYVGPLARGNSRFAWMTRKSEETITPWNFEEVVDKGASAQSFIERMTNFDKNLPNEKVLPKHSLLYEYFTVYNELTKVKYVTEGMRKPAFLSGEQKKAIVDLLFKTNRKVTVKQLKEDYFKKIECFDSVEISGVEDRFNASLGTYHDLLKIIKDKDFLDNEENEDILEDIVLTLTLFEDREMIEERLKTYAHLFDDKVMKQLKRRRYTGWGRLSRKLINGIRDKQSGKTILDFLKSDGFANRNFMQLIHDDSLTFKEDIQKAQVSGQGDSLHEHIANLAGSPAIKKGILQTVKVVDELVKVMGRHKPENIVIEMARENQTTQKGQKNSRERMKRIEEGIKELGSQILKEHPVENTQLQNEKLYLYYLQNGRDMYVDQELDINRLSDYDVDHIVPQSFLKDDSIDNKVLTRSDKNRGKSDNVPSEEVVKKMKNYWRQLLNAKLITQRKFDNLTKAERGGLSELDKAGFIKRQLVETRQITKHVAQILDSRMNTKYDENDKLIREVKVITLKSKLVSDFRKDFQFYKVREINNYHHAHDAYLNAVVGTALIKKYPKLESEFVYGDYKVYDVRKMIAKSEQEIGKATAKYFFYSNIMNFFKTEITLANGEIRKRPLIETNGETGEIVWDKGRDFATVRKVLSMPQVNIVKKTEVQTGGFSKESILPKRNSDKLIARKKDWDPKKYGGFDSPTVAYSVLVVAKVEKGKSKKLKSVKELLGITIMERSSFEKNPIDFLEAKGYKEVKKDLIIKLPKYSLFELENGRKRMLASAGELQKGNELALPSKYVNFLYLASHYEKLKGSPEDNEQKQLFVEQHKHYLDEIIEQISEFSKRVILADANLDKVLSAYNKHRDKPIREQAENIIHLFTLTNLGAPAAFKYFDTTIDRKRYTSTKEVLDATLIHQSITGLYETRIDLSQLGGD",
    
    "nmcas9": "MAAFKPNPINYILGLDIGIASVGWAMVEIDEEENPIRLIDLGVRVFERAEVPKTGDSLAMVRRLARSVRRLTRRRAHRLLRARRLLKREGVLQAADFDENGLIKSLPNTPWQLRAAALDRKLTPLEWSAVLLHLIKHRGYLSQRKNEGETADKELGALLKGVADNAHALQTGDFRTPAELALNKFEKESGHIRNQRGDYSHTFSRKDLQAELILLFEKQKEFGNPHISGGLKEGIETLLMTQRPALSGDAVQKMLGHCTFEPAEPKAAKNTYTAERFIWLTKLNNLRILEQGSERPLTDTERATLMDEPYRKSKLTYAQARKLLGLEDTAFFKGLRYGKDNAEASTLMEMKAYHAISRALEKEGLKDKKSPLNLSPELQDEIGTAFSLFKTDEDITGRLKDRIQPEILEALLKHISFDKFVQISLKALRRIVPLMEQGKRYDEACAEIYGDHYGKKNTEEKIYLPPIPADEIRNPVVLRALSQARKVINGVVRRYGSPARIHIETAREVGKSFKDRKEIEKRQEENRKDREKAAAKFREYFPNFVGEPKSKDILKLRLYEQQHGKCLYSGKEINLGRLNEKGYVEIDHALPFSRTWDDSFNNKVLVLGSENQNKGNQTPYEYFNGKDNSREWQEFKARVETSRFPRSKKQRILLQKFDEDGFKERNLNDTRYVNRFLCQFVADRMRLTGKGKKRVFASNGQITNLLRGFWGLRKVRAENDRHHALDAVVVACSTVAMQQKITRFVRYKEMNAFDGKTIDKETGEVLHQKTHFPQPWEFFAQEVMIRVFGKPDGKPEFEEADTPEKLRTLLAEKLSSRPEAVHEYVTPLFVSRAPNRKMSGQGHMETVKSAKRLDEGVSVLRVPLTQLKLKDLEKMVNREREPKLYEALKARLEAHKDDPAKAFAEPFYKYDKAGNRTQQVKAVRVEQVQKTGVWVRNHNGIADNATMVRVDVFEKGDKYYLVPIYSWQVAKGILPDRAVVQGKDEEDWQLIDDSFNFKFSLHPNDLVEVITKKARMFGYFASCHRGTGNINIRIHDLDHKIGKNGILEGIGVKTALSFQKYQIDELGKEIRPCRLKKRPPVR",
    
    "st1cas9": "MSDLVLGLDIGIGSVGVGILNKVTGEIIHKNSRIFPAAQAENNLVRRTNRQGRRLARRKKHRRVRLNRLFEESGLITDFTKISINLNPYQLRVKGLTDELSNEELFIALKNMVKHRGISYLDDASDDGNSSVGDYAQIVKENSKQLETKTPGQIQLERYQTYGQLRGDFTVEKDGKKHRLINVFPTSAYRSEALRILQTQQEFNPQITDEFINRYLEILTGKRKYYHGPGNEKSRTDYGRYRTSGETLDNIFGILIGKCTFYPDEFRAAKASYTAQEFNLLNDLNNLTVPTETKKLSKEQKNQIINYVKNEKAMGPAKLFKYIAKLLSCDVADIKGYRIDKSGKAEIHTFEAYRKMKTLETLDIEQMDRETLDKLAYVLTLNTEREGIQEALEHEFADGSFSQKQVDELVQFRKANSSIFGKGWHNFSVKLMMELIPELYETSEEQMTILTRLGKQKTTSSSNKTKYIDEKLLTEEIYNPVVAKSVRQAIKIVNAAIKEYGDFDNIVIEMARETNEDDEKKAIQKIQKANKDEKDAAMLKAANQYNGKAELPHSVFHGHKQLATKIRLWHQQGERCLYTGKTISIHDLINNSNQFEVDHILPLSITFDDSLANKVLVYATANQEKGQRTPYQALDSMDDAWSFRELKAFVRESKTLSNKKKEYLLTEEDISKFDVRKKFIERNLVDTRYASRVVLNALQEHFRAHKIDTKVSVVRGQFTSQLRRHWGIEKTRDTYHHHAVDALIIAASSQLNLWKKQKNTLVSYSEDQLLDIETGELISDDEYKESVFKAPYQHFVDTLKSKEFEDSILFSYQVDSKFNRKISDATIYATRQAKVGKDKADETYVLGKIKDIYTQDGYDAFMKIYKKDKSKFLMYRHDPQTFEKVIEPILENYPNKQINEKGKEVPCNPFLKYKEEHGYIRKYSKKGNGPEIKSLKYYDSKLGNHIDITPKDSNNKVVLQSVSPWRADVYFNKTTGKYEILGLKYADLQFEKGTGTYKISQEKYNDIKKKEGVDSDSEFKFTLYKNDLLLVKDTETKEQQLFRFLSRTMPKQKHYVELKPYDKQKFEGGEALIKVLGNVANSGQCKKGLGKSNISIYKVRTDVLGNQHIIKNEGDKPKLDF",
    "nmecas9":
    "MAAFKPNPINYILGLDIGIASVGWAMVEIDEEENPIRLIDLGVRVFERAEVPKTGDSLAMVRRLARSVRRLTRRRAHRLLRARRLLKREGVLQAADFDENGLIKSLPNTPWQLRAAALDRKLTPLEWSAVLLHLIKHRGYLSQRKNEGETADKELGALLKGVADNAHALQTGDFRTPAELALNKFEKESGHIRNQRGDYSHTFSRKDLQAELILLFEKQKEFGNPHISGGLKEGIETLLMTQRPALSGDAVQKMLGHCTFEPAEPKAAKNTYTAERFIWLTKLNNLRILEQGSERPLTDTERATLMDEPYRKSKLTYAQARKLLGLEDTAFFKGLRYGKDNAEASTLMEMKAYHAISRALEKEGLKDKKSPLNLSPELQDEIGTAFSLFKTDEDITGRLKDRIQPEILEALLKHISFDKFVQISLKALRRIVPLMEQGKRYDEACAEIYGDHYGKKNTEEKIYLPPIPADEIRNPVVLRALSQARKVINGVVRRYGSPARIHIETAREVGKSFKDRKEIEKRQEENRKDREKAAAKFREYFPNFVGEPKSKDILKLRLYEQQHGKCLYSGKEINLGRLNEKGYVEIDHALPFSRTWDDSFNNKVLVLGSENQNKGNQTPYEYFNGKDNSREWQEFKARVETSRFPRSKKQRILLQKFDEDGFKERNLNDTRYVNRFLCQFVADRMRLTGKGKKRVFASNGQITNLLRGFWGLRKVRAENDRHHALDAVVVACSTVAMQQKITRFVRYKEMNAFDGKTIDKETGEVLHQKTHFPQPWEFFAQEVMIRVFGKPDGKPEFEEADTPEKLRTLLAEKLSSRPEAVHEYVTPLFVSRAPNRKMSGQGHMETVKSAKRLDEGVSVLRVPLTQLKLKDLEKMVNREREPKLYEALKARLEAHKDDPAKAFAEPFYKYDKAGNRTQQVKAVRVEQVQKTGVWVRNHNGIADNATMVRVDVFEKGDKYYLVPIYSWQVAKGILPDRAVVQGKDEEDWQLIDDSFNFKFSLHPNDLVEVITKKARMFGYFASCHRGTGNINIRIHDLDHKIGKNGILEGIGVKTALSFQKYQIDELGKEIRPCRLKKRPPVR"
}

# Model configuration
MODEL_NAME = "Profluent-Bio/protein2pam-cas9_full"  # HuggingFace model name
DEVICE = None  # None = auto-detect (cuda if available, else cpu)
MIN_CONFIDENCE = 0.7  # Minimum probability threshold for predicting a specific nucleotide (vs. 'N')
                       # If max probability < MIN_CONFIDENCE, predict 'N'


def get_pam_probability_distributions(sequences_list, device=None, model_name="Profluent-Bio/protein2pam-cas9_full", min_confidence=0.55):
    """
    Get detailed probability distributions for PAM predictions.
    
    Args:
        sequences_list: List of protein sequence strings
        device: torch device (default: auto-detect)
        model_name: HuggingFace model name for PAM prediction
        min_confidence: Minimum probability threshold for predicting a specific nucleotide
    
    Returns:
        List of dictionaries containing probability distributions and predicted PAMs
    """
    if not sequences_list:
        print("Error: No sequences provided.")
        return []
    
    # Auto-detect device if not provided
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    print(f"Using device: {device}")
    print(f"Loading PAM prediction model: {model_name}...")
    
    # Create a dummy PAMMatching object (target_pam doesn't matter for prediction)
    pam_predictor = PAMMatching(
        device=device,
        target_pam="NNNNNNNNNN",  # Dummy target, we only use get_pam_probability_distributions()
        model_name=model_name,
        use_entropy_for_n_positions=True,
        use_ce_loss=False
    )
    
    print(f"Computing PAM probability distributions for {len(sequences_list)} sequences...")
    
    # Get probability distributions
    prob_distributions = pam_predictor.get_pam_probability_distributions(sequences_list)
    
    # Recompute predicted PAMs with custom min_confidence
    nucleotides = ['A', 'C', 'G', 'T']
    for dist in prob_distributions:
        predicted_pam = []
        for pos_info in dist['per_position']:
            max_prob = pos_info['max_probability']
            if max_prob < min_confidence:
                predicted_pam.append('N')
            else:
                predicted_pam.append(pos_info['predicted'])
        dist['predicted_pam_custom'] = ''.join(predicted_pam)
    
    return prob_distributions


def print_sequence_pam_distribution(seq_name, prob_dist, min_confidence):
    """Print the PAM probability distribution for a single sequence."""
    print("\n" + "="*80)
    print(f"Sequence: {seq_name}")
    print("="*80)
    
    print(f"\nPredicted PAM (min_confidence={min_confidence:.2f}): {prob_dist['predicted_pam_custom']}")
    print(f"Predicted PAM (default min_confidence=0.55): {prob_dist['predicted_pam']}")
    
    print("\nPosition-wise Probability Distribution:")
    print("-" * 80)
    print(f"{'Pos':<5} {'A':<10} {'C':<10} {'G':<10} {'T':<10} {'Pred':<6} {'Max Prob':<10} {'Entropy':<10}")
    print("-" * 80)
    
    nucleotides = ['A', 'C', 'G', 'T']
    for pos_info in prob_dist['per_position']:
        pos = pos_info['position']
        probs = pos_info['probabilities']
        pred = pos_info['predicted']
        max_prob = pos_info['max_probability']
        entropy = pos_info['entropy']
        
        # Determine if this position would be 'N' with custom threshold
        final_pred = 'N' if max_prob < min_confidence else pred
        
        print(f"{pos+1:<5} "
              f"{probs['A']:<10.4f} "
              f"{probs['C']:<10.4f} "
              f"{probs['G']:<10.4f} "
              f"{probs['T']:<10.4f} "
              f"{final_pred:<6} "
              f"{max_prob:<10.4f} "
              f"{entropy:<10.4f}")
    
    print("-" * 80)


def print_summary_distribution(pam_distribution_dict):
    """Print summary distribution across all sequences."""
    print("\n" + "="*80)
    print("SUMMARY: PAM DISTRIBUTION ACROSS ALL SEQUENCES")
    print("="*80)
    
    total = len(pam_distribution_dict)
    pam_counter = Counter(pam_distribution_dict.values())
    
    print(f"\nTotal sequences: {total}")
    print(f"Unique PAMs: {len(pam_counter)}")
    
    # Sort by frequency (descending)
    sorted_pams = sorted(pam_counter.items(), key=lambda x: x[1], reverse=True)
    
    print("\nPAM Distribution (sorted by frequency):")
    print("-" * 80)
    print(f"{'PAM Sequence':<15} {'Count':<10} {'Percentage':<10}")
    print("-" * 80)
    
    for pam, count in sorted_pams:
        percentage = (count / total) * 100
        print(f"{pam:<15} {count:<10} {percentage:>6.2f}%")
    
    print("-" * 80)
    print(f"{'Total':<15} {total:<10} {'100.00%':<10}")
    print("="*80)


def main():
    # Get sequences from the dictionary at the top
    sequences_dict = SEQUENCES
    
    if not sequences_dict:
        print("Error: SEQUENCES dictionary is empty. Please add sequences to the dictionary at the top of the script.")
        return
    
    print(f"Processing {len(sequences_dict)} sequences...")
    print(f"Using min_confidence threshold: {MIN_CONFIDENCE}")
    
    # Extract sequence names and values
    sequence_names = list(sequences_dict.keys())
    sequences_list = list(sequences_dict.values())
    
    # Set device
    device = DEVICE
    if device is not None:
        device = torch.device(device)
    
    # Get probability distributions
    prob_distributions = get_pam_probability_distributions(
        sequences_list,
        device=device,
        model_name=MODEL_NAME,
        min_confidence=MIN_CONFIDENCE
    )
    
    # Create dictionary mapping sequence names to predicted PAMs
    predicted_pams_dict = {
        name: dist['predicted_pam_custom'] 
        for name, dist in zip(sequence_names, prob_distributions)
    }
    
    # Print detailed distributions for each sequence
    for seq_name, prob_dist in zip(sequence_names, prob_distributions):
        print_sequence_pam_distribution(seq_name, prob_dist, MIN_CONFIDENCE)
    
    # Print summary distribution
    print_summary_distribution(predicted_pams_dict)
    
    # Return dictionaries for potential programmatic use
    return predicted_pams_dict, prob_distributions


if __name__ == "__main__":
    main()