#! /usr/bin/env python

import sys
import argparse
import math
import numpy as np
import re
from Bio import SeqIO

re1 = re.compile(r"[^acgtuACGTU]")

parser = argparse.ArgumentParser()
parser.add_argument(
    "infile", nargs="?", default=sys.stdin, help="filename (default: stdin)"
)
parser.add_argument(
    "--show-lens", action="store_true", help="show all sequence lengths"
)
args = parser.parse_args()

count = 0
total_len = 0
unkn_len = 0
total_len_1k = 0
lens = []

for r in SeqIO.parse(args.infile, "fasta"):
    count += 1
    stext = str(r.seq)
    if args.show_lens:
        # print(str(r.id), "{:_}".format(len(stext)))
        print(f"{str(r.id)}\t{len(stext)}")
    total_len += len(stext)
    if len(stext) > 999:
        total_len_1k += len(stext)
    lens.append(len(stext))
    # print(stext)
    ttext = re1.sub("", stext)
    # print(ttext)
    unkn_len += len(stext) - len(ttext)

if len(lens) == 0:
    exit(0)

# N50 calculation
n50 = 0
lens = np.sort(lens)
t = 0
mid = total_len / 2
for l in lens:
    if t <= mid and mid < t + l:
        n50 = l
        break
    t += l

# Reverse the sorted lengths array
rev_lens = ["{:_}".format(i) for i in lens[::-1]]


print("Number of sequences              :", "{:_}".format(count))
print("Total number of characters       :", "{:_}".format(total_len))
if total_len > 0:
    print("Percent non-ACGTU characters     :", round(100 * unkn_len / total_len, 1))
    print(
        "Percent bases in 1000+ base seqs :", round(100 * total_len_1k / total_len, 1)
    )
print("Min len                          :", "{:_}".format(np.amin(lens)))
print("Max len                          :", "{:_}".format(np.amax(lens)))
print("Longest                          :", end=" "),
for l in rev_lens[:5]:
    print(l, "", end="")
print("")
print(
    "Quartiles:                       :",
    "{:_}".format(np.percentile(lens, 25)),
    "{:_}".format(np.percentile(lens, 50)),
    "{:_}".format(np.percentile(lens, 75)),
)
if count > 4:
    print("N50                              :", "{:_}".format(n50))
