#!/usr/bin/env bash
# Usage: filter_one_count.sh SPECIES
# Filters global_index_count to keep only kmers specific to SPECIES,
# then selects the SPECIES column in-place.
# Outputs:
#   specific_index_count/SPECIES/index.done  (written by obikmer select)
#   stats/specific_kmer_count/SPECIES.stats  (one CSV data row, no header)
set -euo pipefail

SPECIES="$1"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BINARY="${SCRIPT_DIR}/../src/target/release/obikmer"

SOURCE="${SCRIPT_DIR}/global_index_count"
OUTPUT="${SCRIPT_DIR}/specific_index_count/${SPECIES}"
STATS_DIR="${SCRIPT_DIR}/stats/specific_kmer_count"
STATS_FILE="${STATS_DIR}/${SPECIES}.stats"

mkdir -p "${STATS_DIR}"

echo "[${SPECIES}] filter (count) → ${OUTPUT}"

LOG_FILTER=$(mktemp)
LOG_SELECT=$(mktemp)
trap 'rm -f "${LOG_FILTER}" "${LOG_SELECT}"' EXIT

"${BINARY}" filter \
    --output "${OUTPUT}" \
    --force \
    --ingroup "species=${SPECIES}" \
    --outgroup all \
    --min-frac 0.5 \
    --max-frac 1.0 \
    --max-outgroup-count 0 \
    "${SOURCE}" \
    2>"${LOG_FILTER}"

cat "${LOG_FILTER}" >&2

"${BINARY}" select \
    --in-place \
    --group "${SPECIES}:species=${SPECIES}" \
    --group-op "${SPECIES}:any" \
    --select "${SPECIES}" \
    "${OUTPUT}" \
    2>"${LOG_SELECT}"

cat "${LOG_SELECT}" >&2

python3 - "${SPECIES}" "${LOG_FILTER}" "${LOG_SELECT}" <<'PYEOF' >"${STATS_FILE}"
import sys, re

species, log_filter, log_select = sys.argv[1], sys.argv[2], sys.argv[3]

def strip_ansi(s):
    return re.sub(r'\x1b\[[\x30-\x3f]*[\x20-\x2f]*[\x40-\x7e]', '', s)

def parse_wall(s):
    s = s.strip()
    if s.endswith('ms'): return float(s[:-2]) / 1000.0
    if s.endswith('s'):  return float(s[:-1])
    return 0.0

def parse_rss(s):
    m = re.match(r'([\d.]+)\s*(GB|MB|KB|B)', s.strip())
    if not m: return 0
    return int(float(m.group(1)) * {'GB': 1<<30, 'MB': 1<<20, 'KB': 1024, 'B': 1}[m.group(2)])

def is_sep(s):
    return bool(s) and not re.search(r'[A-Za-z0-9]', s)

def parse_reporter(logfile):
    stats = {}
    state = 'scan'
    with open(logfile, errors='replace') as fh:
        for raw in fh:
            line = strip_ansi(raw.rstrip('\n'))
            s    = line.strip()
            if state == 'scan':
                if re.search(r'\bstage\b.*\bwall\b', line):
                    state = 'in_header'
            elif state == 'in_header':
                if is_sep(s): state = 'rows'
            elif state == 'rows':
                if is_sep(s): state = 'total'
                elif s:
                    parts = re.split(r'  +', s)
                    if len(parts) >= 4:
                        stats[parts[0]] = (parse_wall(parts[1]), parse_rss(parts[3]))
            elif state == 'total':
                if s:
                    parts = re.split(r'  +', s)
                    if len(parts) >= 3:
                        stats['TOTAL'] = (parse_wall(parts[1]),
                                          parse_rss(parts[3]) if len(parts) > 3 else 0)
                break
    return stats

f = parse_reporter(log_filter)
s = parse_reporter(log_select)

row = [species]
for stage, d in [('rebuild', f), ('pack', f), ('filter_total', f), ('select', s), ('select_total', s)]:
    key = 'TOTAL' if stage.endswith('_total') else stage
    w, r = d.get(key, ('', ''))
    row += [f'{w:.3f}' if isinstance(w, float) else '', str(r)]
print(','.join(row))
PYEOF
