#!/usr/bin/env bash
# Step 40's 23andMe-layout file read back by scripts/chip-to-vcf.sh, chrM
# included. Step 40 writes dbSNP's rCRS positions for the MT sites; a vendor
# 23andMe file numbers them on hg19's chrM, which chip-to-vcf.sh moves to the
# rCRS. It knows step 40's file by its first line ("# This data file was
# generated by OwnGenome") and keeps those positions.
#   1. step 40 on a copy of case 21's gVCF, with the CI slice of the site table
#      (tests/fixtures/consumer_export/sites.tsv.gz, 285 chrM sites);
#   2. chip-to-vcf.sh on the 23andMe file says it kept the MT positions, and
#      every called chrM site lands on the table's rCRS position with the
#      reference base there and the allele of case 21's call (hom-ref where
#      case 21 called nothing);
#   2b. the same file with a UTF-8 byte order mark in front, as an editor
#      that adds one saves it, is still known as step 40's file;
#   3. control: the same file without its first line, as a vendor file has
#      no such line, is moved through the chain's chrM block, so most called
#      chrM sites no longer land on their rCRS position with that call.
#      Before the fix, check 2 failed the same way.
. "$(dirname "$0")/lib.sh"

G="$GENOME_DIR"
GVCF="${G}/${SAMPLE}/vcf/${SAMPLE}.g.vcf.gz"
check "case 21 wrote a gVCF" vcf_ok "${SAMPLE}/vcf/${SAMPLE}.g.vcf.gz"
SLICE="${REPO}/tests/fixtures/consumer_export/sites.tsv.gz"
SITES="${G}/reference/consumer/23andme_v5_sites_dbsnp${CONSUMER_DBSNP_BUILD}.tsv.gz"
OWN_SITES=0
if [ ! -e "$SITES" ]; then
  mkdir -p "$(dirname "$SITES")"
  cp "$SLICE" "$SITES"
  OWN_SITES=1
fi
CHAIN_URL=https://hgdownload.soe.ucsc.edu/goldenPath/hg19/liftOver/hg19ToHg38.over.chain.gz
CHAIN_SHA256=5c0598e500ceb5a78c73086929e8ef993aec309bcafb595139b53d440b125a1d
E=HG002CMT

# --- 1. step 40 on a copy of the gVCF ------------------------------------------
rm -rf "${G:?}/${E}"
mkdir -p "${G}/${E}/vcf"
cp "$GVCF" "${G}/${E}/vcf/${E}.g.vcf.gz"
cp "${GVCF}.tbi" "${G}/${E}/vcf/${E}.g.vcf.gz.tbi"
run_step 40-export-consumer.sh "$E" male
check_step_exit 40-export-consumer.sh
python3 -c 'import sys, zipfile; z = zipfile.ZipFile(sys.argv[1]); open(sys.argv[2], "wb").write(z.read(z.namelist()[0]))' \
  "${G}/${E}/consumer_export/${E}_23andme_v5.zip" "${CASE_TMP}/export.txt"
check "the 23andMe file starts with step 40's first line" \
  has '^# This data file was generated by OwnGenome' "$(head -n 1 "${CASE_TMP}/export.txt")"
N_MT=$(awk -F'\t' '!/^#/ && $2 == "MT"' "${CASE_TMP}/export.txt" | wc -l | tr -d ' ')
check_ge "chrM sites in the 23andMe file" "$N_MT" 200

# --- the GRCh37 side chip-to-vcf.sh reads ----------------------------------------
# A FASTA with the table's GRCh37 base at each site (N elsewhere), as in
# consumer-export-1, and UCSC's chain.
mkdir -p "${G}/reference_hg19" "${G}/liftover"
python3 - "$SITES" "${G}/reference_hg19/human_g1k_v37.fasta" <<'PY'
import gzip, sys
sites, fa = sys.argv[1:3]
order = [str(i) for i in range(1, 23)] + ["X", "Y", "MT"]
bases = {}
for l in gzip.open(sites, "rt"):
    if l.startswith("#"):
        continue
    c = l.rstrip("\n").split("\t")
    bases.setdefault(c[1], {})[int(c[2])] = c[3]
off = 0
with open(fa, "w") as f, open(fa + ".fai", "w") as fai:
    for chrom in [c for c in order if c in bases]:
        n = max(bases[chrom]) + 1
        head = f">{chrom}\n"
        f.write(head)
        off += len(head)
        fai.write(f"{chrom}\t{n}\t{off}\t60\t61\n")
        seq = bytearray(b"N" * n)
        for p, b in bases[chrom].items():
            seq[p - 1] = ord(b)
        for i in range(0, n, 60):
            f.write(seq[i:i + 60].decode() + "\n")
        off += n + (n + 59) // 60
PY
if curl -fsSL --retry 3 -o "${G}/liftover/hg19ToHg38.over.chain.gz" "$CHAIN_URL"; then
  check_eq "UCSC's hg19ToHg38 chain has the pinned sha256" \
    "$(sha256sum "${G}/liftover/hg19ToHg38.over.chain.gz" | cut -d' ' -f1)" "$CHAIN_SHA256"
else
  fail "could not download ${CHAIN_URL}"
fi
bcf query -f '%CHROM\t%POS\t%REF\t%ALT\t%FILTER[\t%GT]\n' "${SAMPLE}/vcf/${SAMPLE}.vcf.gz" 2>/dev/null | awk -F'\t' '$1 == "chrM"' > "${CASE_TMP}/calls.tsv"

# mt_roundtrip NAME TXT: chip-to-vcf.sh on TXT as a 23andMe file of sample
# NAME, its log in CASE_TMP/NAME.log; prints "called got on_place bad": the
# called chrM sites of TXT, the chrM records with a call, those on the table's
# rCRS position with its reference base, and those of them whose allele is
# not case 21's call there.
mt_roundtrip() {
  local n=$1
  rm -rf "${G:?}/${n}"
  mkdir -p "${G}/${n}/raw"
  cp "$2" "${G}/${n}/raw/${n}_raw.txt"
  local status=0
  "${REPO}/scripts/chip-to-vcf.sh" "$n" 23andme > "${CASE_TMP}/${n}.log" 2>&1 || status=$?
  echo "$status" > "${CASE_TMP}/${n}.status"
  echo "chip-to-vcf.sh ${n}: exit ${status}" >&2
  tail -n 25 "${CASE_TMP}/${n}.log" >&2
  bcf query -f '%ID\t%CHROM\t%POS\t%REF\t%ALT[\t%GT]\n' "${n}/vcf/${n}.vcf.gz" 2>/dev/null | awk -F'\t' '$2 == "chrM"' > "${CASE_TMP}/${n}_chrM.tsv"
  python3 - "${CASE_TMP}/${n}_chrM.tsv" "${CASE_TMP}/calls.tsv" "$SITES" "$2" <<'PY'
import gzip, re, sys
lifted, calls, sites, txt = sys.argv[1:5]
table = {}
for l in gzip.open(sites, "rt"):
    if not l.startswith("#"):
        c = l.rstrip("\n").split("\t")
        if c[1] == "MT":
            table[c[0]] = c
dv = {}
for l in open(calls):
    chrom, pos, ref, alt, flt, gt = l.rstrip("\n").split("\t")
    if flt == "PASS":
        al = [ref] + alt.split(",")
        dv[int(pos)] = {al[int(i)][0] for i in re.split(r"[/|]", gt) if i not in (".", "")}
called = sum(1 for l in open(txt)
             if not l.startswith("#") and l.split("\t")[0] in table and l.rstrip("\n").split("\t")[3] != "--")
got = on_place = bad = 0
for l in open(lifted):
    rid, chrom, pos, ref, alt, gt = l.rstrip("\n").split("\t")
    t = table.get(rid)
    idx = [i for i in re.split(r"[/|]", gt) if i not in (".", "")]
    if t is None or not idx:
        continue
    got += 1
    if int(pos) != int(t[4]) or ref != t[5]:
        if got - on_place <= 5:
            print(f"off place {rid}: rCRS {t[4]} {t[5]}, read back at chrM:{pos} {ref}", file=sys.stderr)
        continue
    on_place += 1
    al = [ref] + alt.split(",")
    have = {al[int(i)][0] for i in idx}
    want = dv.get(int(pos), {t[5]})
    if have != want:
        bad += 1
        if bad <= 5:
            print(f"mismatch {rid} chrM:{pos} read back {sorted(have)}, case 21 {sorted(want)}", file=sys.stderr)
print(called, got, on_place, bad)
PY
}

# --- 2. step 40's file keeps its chrM positions --------------------------------------
read -r called got on_place bad <<< "$(mt_roundtrip HG002CMR "${CASE_TMP}/export.txt")"
echo "step 40's file: ${called:-?} called chrM sites, ${got:-?} read back, ${on_place:-?} on their rCRS position, ${bad:-?} with another allele than case 21's call"
check_eq "chip-to-vcf.sh exits 0 on step 40's file" "$(cat "${CASE_TMP}/HG002CMR.status")" 0
check "chip-to-vcf.sh says it kept the MT positions of step 40's file" \
  has "written by OwnGenome step 40, ${N_MT} site\(s\) already on rCRS positions, kept" "$(cat "${CASE_TMP}/HG002CMR.log")"
check "chip-to-vcf.sh moved no MT site of step 40's file" lacks 'moved from hg19 chrM' "$(cat "${CASE_TMP}/HG002CMR.log")"
check_ge "called chrM sites in step 40's file" "${called:-0}" 100
check_eq "called chrM sites read back" "${got:-x}" "${called:-y}"
check_eq "called chrM sites read back on their rCRS position with its reference base" "${on_place:-x}" "${called:-y}"
check_eq "chrM sites read back with another allele than case 21's call" "${bad:-x}" 0

# --- 2b. the same file with a UTF-8 byte order mark ---------------------------------
{ printf '\357\273\277'; cat "${CASE_TMP}/export.txt"; } > "${CASE_TMP}/bom.txt"
read -r b_called b_got b_on b_bad <<< "$(mt_roundtrip HG002CMB "${CASE_TMP}/bom.txt")"
echo "byte order mark: ${b_called:-?} called chrM sites, ${b_got:-?} read back, ${b_on:-?} on their rCRS position, ${b_bad:-?} with another allele"
check "chip-to-vcf.sh keeps the MT positions of step 40's file with a byte order mark" \
  has "written by OwnGenome step 40, ${N_MT} site\(s\) already on rCRS positions, kept" "$(cat "${CASE_TMP}/HG002CMB.log")"
check_eq "byte order mark: called chrM sites read back on their rCRS position" "${b_on:-x}" "${called:-y}"
check_eq "byte order mark: chrM sites with another allele than case 21's call" "${b_bad:-x}" 0

# --- 3. control: the same file without step 40's first line ----------------------------
tail -n +2 "${CASE_TMP}/export.txt" > "${CASE_TMP}/vendor.txt"
read -r v_called v_got v_on v_bad <<< "$(mt_roundtrip HG002CMV "${CASE_TMP}/vendor.txt")"
echo "control, no first line: ${v_called:-?} called chrM sites, ${v_got:-?} read back, ${v_on:-?} on their rCRS position, ${v_bad:-?} with another allele"
check_eq "control: chip-to-vcf.sh exits 0" "$(cat "${CASE_TMP}/HG002CMV.status")" 0
check_ge "control: called chrM records read back" "${v_got:-0}" 1
check "control: chip-to-vcf.sh moves the MT sites of a file without step 40's first line" \
  has 'site\(s\) moved from hg19 chrM to rCRS positions' "$(cat "${CASE_TMP}/HG002CMV.log")"
check "control: fewer than half of the called chrM sites stay on their rCRS position with case 21's call (${v_on:-?} of ${v_called:-?})" \
  test "$(( (${v_on:-0} - ${v_bad:-0}) * 2 ))" -lt "${v_called:-0}"
printf '#### chrM round trip of step 40\n\n%s called chrM sites, %s read back on their rCRS position, %s mismatches; without the first line %s stay\n\n' \
  "${called:-?}" "${on_place:-?}" "${bad:-?}" "${v_on:-?}" >> "$E2E_NOTES"

rm -rf "${G:?}/${E}" "${G:?}/HG002CMR" "${G:?}/HG002CMB" "${G:?}/HG002CMV" "${G:?}/reference_hg19" "${G:?}/liftover"
[ "$OWN_SITES" = 0 ] || rm -f "$SITES"
finish