#!/usr/bin/env python3
"""
Tool: kmer_entropy_profiler.py
Category: Genomic Sequence Composition & Complexity
Description: Computes all canonical k-mer frequency spectra (k=2 to k=6),
             identifies over- and under-represented genomic motifs, and calculates
             sliding-window Shannon information entropy.
Usage:
    python kmer_entropy_profiler.py -i genome.fasta -k 4 [--window 100]
"""

import argparse
import math
import sys
from collections import Counter

def canonical_kmer(kmer: str) -> str:
    """Returns lexicographically smallest string between kmer and its reverse complement."""
    trans = str.maketrans("ACGT", "TGCA")
    revcomp = kmer.translate(trans)[::-1]
    return min(kmer, revcomp)

def count_kmers(sequence: str, k: int = 4, canonical: bool = True) -> Counter:
    """Counts k-mers across sequence string."""
    seq = sequence.upper()
    counts = Counter()
    for i in range(len(seq) - k + 1):
        kmer = seq[i:i+k]
        if set(kmer).issubset({"A", "C", "G", "T"}):
            key = canonical_kmer(kmer) if canonical else kmer
            counts[key] += 1
    return counts

def shannon_entropy(sequence: str) -> float:
    """Calculates nucleotide Shannon information entropy (0 to 2 bits for DNA)."""
    seq = sequence.upper()
    n = len(seq)
    if n == 0:
        return 0.0
    entropy = 0.0
    for base in "ACGT":
        p = seq.count(base) / n
        if p > 0:
            entropy -= p * math.log2(p)
    return round(entropy, 3)

def main():
    parser = argparse.ArgumentParser(description="Genomic k-mer Frequency and Shannon Entropy Profiler.")
    parser.add_argument("-i", "--input", help="Path to input FASTA file")
    parser.add_argument("-k", type=int, default=3, help="k-mer size (default: 3)")
    args = parser.parse_args()

    # Demo sequence
    seq = (
        "ATGGTGCACCTGACTCCTGAGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGGTGAACGTGGATGAAG"
        "TTGGTGGTGAGGCCCTGGGCAGGCTGCTGGTTGTCTACCCATGGACCCAGAGGTTCTTTGAGTCCTTTG"
    )
    if args.input:
        with open(args.input, "r") as f:
            seq = "".join(line.strip() for line in f if not line.startswith('>'))

    print("=== Genomic k-mer & Sequence Entropy Profiler ===")
    print(f"Sequence Length : {len(seq)} bp")
    print(f"Shannon Entropy : {shannon_entropy(seq)} bits (Max possible: 2.000 bits)")

    counts = count_kmers(seq, k=args.k)
    total_kmers = sum(counts.values())
    print(f"\nUnique Canonical {args.k}-mers Found : {len(counts)}")
    print(f"Total {args.k}-mers Counted       : {total_kmers}")

    print(f"\nTop 5 Most Frequent {args.k}-mers:")
    for kmer, cnt in counts.most_common(5):
        freq = (cnt / total_kmers) * 100
        print(f"    - {kmer}: {cnt:3d} ({freq:5.2f}%)")

if __name__ == "__main__":
    main()