#!/usr/bin/env python3
"""
Chapter 3: Python for Next Gen Biologist
Script: ch03_functions_orfs.py
Description: Defining modular functions, scanning all 6 reading frames for Open Reading Frames (ORFs),
             passing arguments, returning structured data, and recursion.
Usage:
    python ch03_functions_orfs.py --seq ATGAAACCCGGGTTTTAA
"""

import argparse

def reverse_complement(sequence: str) -> str:
    comp = str.maketrans("ACGTacgt", "TGCAtgca")
    return sequence.translate(comp)[::-1]

def find_orfs_in_frame(sequence: str, frame: int, min_len_aa: int = 15) -> list:
    """Find open reading frames in a specific reading frame (0, 1, or 2)."""
    seq = sequence.upper()
    start_codons = {"ATG"}
    stop_codons = {"TAA", "TAG", "TGA"}
    orfs = []

    start_idx = None
    for i in range(frame, len(seq) - 2, 3):
        codon = seq[i:i+3]
        if codon in start_codons and start_idx is None:
            start_idx = i
        elif codon in stop_codons and start_idx is not None:
            orf_dna = seq[start_idx:i+3]
            length_aa = len(orf_dna) // 3
            if length_aa >= min_len_aa:
                orfs.append({
                    'start': start_idx + 1,  # 1-based coordinate
                    'end': i + 3,
                    'length_nt': len(orf_dna),
                    'length_aa': length_aa,
                    'dna': orf_dna
                })
            start_idx = None
    return orfs

def scan_all_six_frames(sequence: str, min_len_aa: int = 10) -> dict:
    """Scan both forward (+) and reverse complement (-) strands across 3 frames each."""
    forward_seq = sequence.upper()
    revcomp_seq = reverse_complement(forward_seq)
    
    results = {'forward': {}, 'reverse': {}}
    for f in range(3):
        results['forward'][f"+F{f+1}"] = find_orfs_in_frame(forward_seq, f, min_len_aa)
        results['reverse'][f"-F{f+1}"] = find_orfs_in_frame(revcomp_seq, f, min_len_aa)
    return results

def recursive_nucleotide_distance(seq1: str, seq2: str) -> int:
    """Recursive Hamming edit distance calculation between two equal-length sequences."""
    if len(seq1) != len(seq2):
        raise ValueError("Sequences must be of equal length for Hamming distance")
    if len(seq1) == 0:
        return 0
    head_dist = 1 if seq1[0] != seq2[0] else 0
    return head_dist + recursive_nucleotide_distance(seq1[1:], seq2[1:])

def main():
    parser = argparse.ArgumentParser(description="Chapter 3: Functions and 6-frame ORF scanner.")
    parser.add_argument("-s", "--seq", help="Input DNA sequence")
    parser.add_argument("--min-aa", type=int, default=10, help="Minimum ORF length in amino acids")
    args = parser.parse_args()

    seq = args.seq.strip() if args.seq else (
        "GGGATGGCCACCGATGGCATGCCCAAGCTGTACGACTACGTGTACGAGCGCCTCGACGAGACCGAGGAG"
        "TTCGCCCGCTTCGAGGCCGCCAACATGCTCCGCTACCGCGCCCGCACCCGCTAGTTTTTT"
    )

    print(f"Scanning sequence ({len(seq)} bp) for ORFs (Min AA: {args.min_aa})...\n")
    six_frames = scan_all_six_frames(seq, min_len_aa=args.min_aa)
    
    for strand, frames in six_frames.items():
        for frame_name, orfs in frames.items():
            if orfs:
                print(f"Strand: {strand.upper()} | Frame: {frame_name} ({len(orfs)} ORFs)")
                for orf in orfs:
                    print(f"  -> {orf['start']}-{orf['end']} ({orf['length_nt']} bp, {orf['length_aa']} AA): {orf['dna'][:30]}...")

if __name__ == "__main__":
    main()