#!/usr/bin/env python3
"""
Chapter 5: Python for Next Gen Biologist
Script: ch05_fasta_csv_io.py
Description: Robust FASTA reader/writer handling multi-line sequences without high memory usage;
             CSV/TSV expression matrix parser calculating per-gene mean and variance.
Usage:
    python ch05_fasta_csv_io.py --fasta sample.fasta
    python ch05_fasta_csv_io.py --matrix counts.tsv
"""

import argparse
import csv
import io
import math

def read_fasta(filepath):
    """Generator yielding (header, sequence) records from multi-line FASTA files."""
    header = None
    seq_chunks = []
    with open(filepath, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            if line.startswith('>'):
                if header is not None:
                    yield header, "".join(seq_chunks)
                header = line[1:]
                seq_chunks = []
            else:
                seq_chunks.append(line)
        if header is not None:
            yield header, "".join(seq_chunks)

def write_fasta(records, output_filepath, line_width=60):
    """Writes list of (header, sequence) tuples to FASTA with wrapped lines."""
    with open(output_filepath, 'w') as f:
        for header, seq in records:
            f.write(f">{header}\n")
            for i in range(0, len(seq), line_width):
                f.write(f"{seq[i:i+line_width]}\n")

def parse_expression_tsv(filepath):
    """Parses a gene expression count matrix (GeneID \t S1 \t S2 ...) and computes mean/variance."""
    results = []
    with open(filepath, 'r') as f:
        reader = csv.reader(f, delimiter='\t')
        headers = next(reader)
        samples = headers[1:]
        
        for row in reader:
            if not row:
                continue
            gene_id = row[0]
            counts = [float(x) for x in row[1:]]
            n = len(counts)
            mean = sum(counts) / n if n > 0 else 0
            variance = sum((x - mean) ** 2 for x in counts) / (n - 1) if n > 1 else 0
            results.append({
                'gene_id': gene_id,
                'mean': round(mean, 2),
                'variance': round(variance, 2),
                'std_dev': round(math.sqrt(variance), 2)
            })
    return samples, results

def main():
    parser = argparse.ArgumentParser(description="Chapter 5: FASTA & CSV/TSV File I/O.")
    parser.add_argument("--fasta", help="Path to input FASTA")
    parser.add_argument("--matrix", help="Path to input TSV expression matrix")
    args = parser.parse_args()

    # Demo FASTA
    demo_fasta = "/tmp/demo_genes.fasta"
    with open(demo_fasta, "w") as f:
        f.write(">geneA Homo_sapiens chr1:100-300\nATGCGATCGATCGATCGATCGATCGATC\nGATCGATCGATCGATCGATCGATCGATCG\n")
        f.write(">geneB Mus_musculus chr2:500-750\nATGGCCAATTGGCCAAATTTGGCCCAAA\nTTTAAACCCGGGTTTAAACCCGGGTTTA\n")

    print("=== Demo 1: Reading Multi-line FASTA ===")
    for h, s in read_fasta(demo_fasta):
        print(f"Header: {h} | Length: {len(s)} bp | Preview: {s[:25]}...")

    # Demo TSV Matrix
    demo_tsv = "/tmp/demo_counts.tsv"
    with open(demo_tsv, "w") as f:
        f.write("GeneID\tControl_1\tControl_2\tTreated_1\tTreated_2\n")
        f.write("ENSG000001\t120\t115\t450\t480\n")
        f.write("ENSG000002\t2300\t2450\t2280\t2310\n")
        f.write("ENSG000003\t10\t12\t95\t105\n")

    print("\n=== Demo 2: Parsing TSV Count Matrix ===")
    samples, stats = parse_expression_tsv(demo_tsv)
    print(f"Samples ({len(samples)}): {', '.join(samples)}")
    for g in stats:
        print(f"Gene: {g['gene_id']:12} | Mean Count: {g['mean']:6} | StdDev: {g['std_dev']:6}")

if __name__ == "__main__":
    main()