#!/usr/bin/env python3
"""
Chapter 2: Python for Next Gen Biologist
Script: ch02_data_structures.py
Description: Genomic data structures in Python: Lists, Tuples, Dictionaries, and Sets.
             Includes universal genetic codon table, restriction enzyme motifs, and k-mer sets.
Usage:
    python ch02_data_structures.py --seq ATGCGATCGATCGATCGATAGCTAGCTA
"""

import argparse

CODON_TABLE = {
    'ATA':'I', 'ATC':'I', 'ATT':'I', 'ATG':'M',
    'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACT':'T',
    'AAC':'N', 'AAT':'N', 'AAA':'K', 'AAG':'K',
    'AGC':'S', 'AGT':'S', 'AGA':'R', 'AGG':'R',
    'CTA':'L', 'CTC':'L', 'CTG':'L', 'CTT':'L',
    'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCT':'P',
    'CAC':'H', 'CAT':'H', 'CAA':'Q', 'CAG':'Q',
    'CGA':'R', 'CGC':'R', 'CGG':'R', 'CGT':'R',
    'GTA':'V', 'GTC':'V', 'GTG':'V', 'GTT':'V',
    'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCT':'A',
    'GAC':'D', 'GAT':'D', 'GAA':'E', 'GAG':'E',
    'GGA':'G', 'GGC':'G', 'GGG':'G', 'GGT':'G',
    'TCA':'S', 'TCC':'S', 'TCG':'S', 'TCT':'S',
    'TTC':'F', 'TTT':'F', 'TTA':'L', 'TTG':'L',
    'TAC':'Y', 'TAT':'Y', 'TAA':'_', 'TAG':'_',
    'TGC':'C', 'TGT':'C', 'TGA':'_', 'TGG':'W',
}

RESTRICTION_SITES = {
    'EcoRI': 'GAATTC',
    'BamHI': 'GGATCC',
    'HindIII': 'AAGCTT',
    'NotI': 'GCGGCCGC',
    'XhoI': 'CTCGAG'
}

def translate_dna(dna_sequence: str) -> str:
    """Translates a DNA sequence into an amino acid peptide using codon dictionary."""
    seq = dna_sequence.upper()
    peptide = []
    for i in range(0, len(seq) - 2, 3):
        codon = seq[i:i+3]
        amino_acid = CODON_TABLE.get(codon, '?')
        if amino_acid == '_':
            peptide.append('*')  # Stop codon
            break
        peptide.append(amino_acid)
    return "".join(peptide)

def find_unique_kmers(sequence: str, k: int = 4) -> set:
    """Extracts all unique k-mers from a sequence using a set."""
    kmers = set()
    seq = sequence.upper()
    for i in range(len(seq) - k + 1):
        kmers.add(seq[i:i+k])
    return kmers

def find_restriction_sites(sequence: str) -> list:
    """Scans sequence for standard restriction enzyme sites."""
    seq = sequence.upper()
    matches = []
    for enzyme, motif in RESTRICTION_SITES.items():
        start = 0
        while True:
            idx = seq.find(motif, start)
            if idx == -1:
                break
            matches.append((enzyme, motif, idx + 1))  # 1-based coordinate tuple
            start = idx + 1
    return matches

def main():
    parser = argparse.ArgumentParser(description="Chapter 2: Genomic data structures demo.")
    parser.add_argument("-s", "--seq", help="Input DNA sequence")
    args = parser.parse_args()

    seq = args.seq.strip() if args.seq else "ATGGTGCACCTGACTCCTGAGGAGAAGTCTGCCGAATTCGTTACTGCCGGATCCCTGTGGTAA"
    print(f"Analyzing Sequence ({len(seq)} bp): {seq}\n")

    # 1. Translation via Dictionary
    protein = translate_dna(seq)
    print(f"[1] Translation: {protein}")

    # 2. Unique k-mers via Set
    kmers = find_unique_kmers(seq, k=4)
    print(f"[2] Unique 4-mers (Set count: {len(kmers)}): {sorted(list(kmers))[:8]}...")

    # 3. Restriction sites as Tuples in a List
    sites = find_restriction_sites(seq)
    print(f"[3] Restriction Enzyme Matches (List of Tuples):")
    if sites:
        for enzyme, motif, pos in sites:
            print(f"    - {enzyme} ({motif}) at position {pos}")
    else:
        print("    - No restriction sites found.")

if __name__ == "__main__":
    main()