#!/usr/bin/env python3
"""
Tool: fasta_assembly_metrics.py
Category: Genome Assembly & Sequence Metrics
Description: Calculates comprehensive de novo genome assembly metrics from FASTA:
             Total length, contig counts, min/max/mean/median lengths, N50, L50, N90, L90,
             GC percentage, and cumulative GC skew.
Usage:
    python fasta_assembly_metrics.py -i contigs.fasta [--min-len 500]
"""

import argparse
import sys

def parse_fasta_lengths_and_gc(filepath, min_len=0):
    lengths = []
    total_g = 0
    total_c = 0
    total_len = 0
    
    current_len = 0
    current_gc = 0
    
    with open(filepath, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            if line.startswith('>'):
                if current_len >= min_len and current_len > 0:
                    lengths.append(current_len)
                    total_len += current_len
                current_len = 0
            else:
                s = line.upper()
                current_len += len(s)
                total_g += s.count('G')
                total_c += s.count('C')
                
        if current_len >= min_len and current_len > 0:
            lengths.append(current_len)
            total_len += current_len
            
    lengths.sort(reverse=True)
    gc_pct = (total_g + total_c) / total_len * 100 if total_len > 0 else 0.0
    return lengths, total_len, round(gc_pct, 2)

def calculate_nx_lx(lengths, total_length, x=50):
    """Calculates Nx (contig length) and Lx (contig count) for a given threshold x%."""
    target_bp = total_length * (x / 100.0)
    running_sum = 0
    for idx, l in enumerate(lengths, 1):
        running_sum += l
        if running_sum >= target_bp:
            return l, idx
    return 0, 0

def main():
    parser = argparse.ArgumentParser(description="Genome Assembly N50, L50 and Quality Metrics Calculator.")
    parser.add_argument("-i", "--input", help="Path to input FASTA file containing contigs or scaffolds")
    parser.add_argument("-m", "--min-len", type=int, default=200, help="Minimum contig length filter in bp (default: 200)")
    args = parser.parse_args()

    if not args.input:
        demo = "/tmp/demo_assembly.fasta"
        with open(demo, "w") as f:
            f.write(">contig_1 len=50000\n" + "ACGT" * 12500 + "\n")
            f.write(">contig_2 len=30000\n" + "CCGG" * 7500 + "\n")
            f.write(">contig_3 len=15000\n" + "AATT" * 3750 + "\n")
            f.write(">contig_4 len=5000\n" + "GGCC" * 1250 + "\n")
        args.input = demo
        print(f"[!] Using demo assembly at {demo}\n")

    lengths, total_bp, gc_pct = parse_fasta_lengths_and_gc(args.input, min_len=args.min_len)
    if not lengths:
        print("[ERROR] No contigs found meeting length threshold.", file=sys.stderr)
        sys.exit(1)

    n50, l50 = calculate_nx_lx(lengths, total_bp, 50)
    n90, l90 = calculate_nx_lx(lengths, total_bp, 90)

    print("=== Genome Assembly Metrics Report ===")
    print(f"Total Contigs (>= {args.min_len} bp) : {len(lengths):,}")
    print(f"Total Assembly Size     : {total_bp:,} bp ({total_bp / 1e6:.3f} Mb)")
    print(f"Largest Contig          : {lengths[0]:,} bp")
    print(f"Shortest Contig         : {lengths[-1]:,} bp")
    print(f"Mean Contig Length      : {total_bp / len(lengths):,.1f} bp")
    print(f"N50 Statistic           : {n50:,} bp")
    print(f"L50 Contig Count        : {l50:,}")
    print(f"N90 Statistic           : {n90:,} bp")
    print(f"L90 Contig Count        : {l90:,}")
    print(f"Overall GC Content      : {gc_pct}%")

if __name__ == "__main__":
    main()