#!/usr/bin/env bash
# =============================================================================
# Male Data Processing Pipeline
# TaRGET II / Liver Adult / BPA ATAC-seq
# Directory : ~/TaRGET_II/Li_adt/BPA/
# Date      : 2026-04-03
# Filter    : Tissue=Li, Age=adt, Sex=M, Exposure=BPA10mg or BPA10ug, Lab=BA
# =============================================================================

set -euo pipefail

# --- Paths -------------------------------------------------------------------
BPA_DIR="/home/yan/TaRGET_II/Li_adt/BPA"
DAR_DIR="/home/yan/TaRGET_II/Li_adt/DAR"
SHARE_DIR="${BPA_DIR}/share"
OUT_DIR="${BPA_DIR}/BPA_male"
TSS="/home/yan/TaRGET_II/Li_adt/Target_expose/gencode_mm10_gene_TSS_pc_lncRNA.sorted.bed"

OCR_ALL="${DAR_DIR}/OCR_allTissue.txt"
DEG_ALL="${SHARE_DIR}/"

OCR_M="${BPA_DIR}/OCR_Li_adt_M_BPA_BA.txt"
DEG_M="${OUT_DIR}/DEG_Bartolomei_liver_20weeks_male.txt"

mkdir -p "${OUT_DIR}"

# =============================================================================
# STEP M1 — Extract male OCR subset
# =============================================================================
# Filter OCR_allTissue.txt: Tissue=Li, Age=adt, Sex=M, BPA10mg or BPA10ug, Lab=BA
# Add status column: MORE (logFC > 0) / LESS (logFC < 0)
# Output: OCR_Li_adt_M_BPA_BA.txt

echo "[M1] Extracting male OCR subset..."

awk 'NR==1 || ($7=="Li" && $8=="adt" && $9=="M" &&
     ($10=="BPA10mg" || $10=="BPA10ug") && $11=="BA")' \
    "${OCR_ALL}" > "${OCR_M}"

awk 'BEGIN{OFS="\t"} NR==1{print $0,"status"}
     NR>1{print $0, ($2>0?"MORE":"LESS")}' \
    "${OCR_M}" > "${OCR_M}.tmp" && mv "${OCR_M}.tmp" "${OCR_M}"

echo "  Done: ${OCR_M}"
echo "  Rows: $(( $(wc -l < "${OCR_M}") - 1 )) (excl. header)"

# =============================================================================
# STEP M2 — Build union DAR BED files (MORE and LESS)
# =============================================================================
# From OCR_Li_adt_M_BPA_BA.txt, keep DAR==TRUE peaks.
# Build the union of BPA10mg and BPA10ug per direction (MORE / LESS).
# Annotate each peak: shared / 10mg_unique / 10ug_unique.
# Outputs:
#   BPA_male/MORE_DAR_M_Li_adt_BPA_BA.bed
#   BPA_male/LESS_DAR_M_Li_adt_BPA_BA.bed

echo "[M2] Building union DAR BED files..."

python3 - "${OCR_M}" "${OUT_DIR}" <<'PYEOF'
import sys, csv

ocr_file = sys.argv[1]
out_dir  = sys.argv[2]

sets = {
    "MORE": {"BPA10mg": set(), "BPA10ug": set()},
    "LESS": {"BPA10mg": set(), "BPA10ug": set()},
}

with open(ocr_file) as f:
    for row in csv.DictReader(f, delimiter='\t'):
        if row['DAR'] == 'TRUE':
            sets[row['status']][row['Exposure']].add(row['OCR'])

def write_bed(set_mg, set_ug, outpath):
    union = set_mg | set_ug
    with open(outpath, 'w') as out:
        out.write("chr\tstart\tend\tannotation\n")
        for ocr in sorted(union):
            chrom, start, end = ocr.split(',')
            if ocr in set_mg and ocr in set_ug:
                ann = "shared"
            elif ocr in set_mg:
                ann = "10mg_unique"
            else:
                ann = "10ug_unique"
            out.write(f"{chrom}\t{start}\t{end}\t{ann}\n")
    print(f"  Written: {outpath}  ({len(union)} peaks)")
    shared  = sum(1 for o in union if o in set_mg and o in set_ug)
    mg_only = sum(1 for o in union if o in set_mg and o not in set_ug)
    ug_only = sum(1 for o in union if o not in set_mg and o in set_ug)
    print(f"    shared={shared}  10mg_unique={mg_only}  10ug_unique={ug_only}")

for status in ("MORE", "LESS"):
    write_bed(
        sets[status]["BPA10mg"],
        sets[status]["BPA10ug"],
        f"{out_dir}/{status}_DAR_M_Li_adt_BPA_BA.bed"
    )
PYEOF

# =============================================================================
# STEP M3 — Venn diagram (R)
# =============================================================================
# Two-panel Venn: MORE DAR (red) and LESS DAR (blue), BPA10mg vs BPA10ug.
# Output: BPA_male/plot_venn_DAR_M.pdf

echo "[M3] Drawing Venn diagram..."

Rscript - "${OUT_DIR}" <<'REOF'
args    <- commandArgs(trailingOnly=TRUE)
out_dir <- args[1]

library(VennDiagram)
library(grid)

make_venn <- function(bed_file, status) {
  df <- read.table(bed_file, header=TRUE, sep="\t", stringsAsFactors=FALSE)
  df$ocr <- paste(df$chr, df$start, df$end, sep=",")

  set_mg <- df$ocr[df$annotation %in% c("shared", "10mg_unique")]
  set_ug <- df$ocr[df$annotation %in% c("shared", "10ug_unique")]

  col <- if (status == "MORE") c("#FF6666","#FF9999") else c("#6699FF","#99BBFF")

  venn.diagram(
    x        = list(BPA10mg=set_mg, BPA10ug=set_ug),
    filename = NULL,
    fill     = col,
    alpha    = 0.5,
    main     = paste(status, "DAR — Male Li adt BA"),
    main.cex = 1.2,
    cex      = 1.2,
    cat.cex  = 1.1
  )
}

pdf_out <- file.path(out_dir, "plot_venn_DAR_M.pdf")
pdf(pdf_out, width=10, height=5)
grid.newpage()

# Left panel: MORE
pushViewport(viewport(layout=grid.layout(1, 2)))
pushViewport(viewport(layout.pos.col=1))
grid.draw(make_venn(
  file.path(out_dir, "MORE_DAR_M_Li_adt_BPA_BA.bed"), "MORE"))
popViewport()

# Right panel: LESS
pushViewport(viewport(layout.pos.col=2))
grid.draw(make_venn(
  file.path(out_dir, "LESS_DAR_M_Li_adt_BPA_BA.bed"), "LESS"))
popViewport()

dev.off()
cat("  Written:", pdf_out, "\n")
REOF

# =============================================================================
# STEP M4 — Nearest-gene annotation
# =============================================================================
# Assign each DAR peak to its single nearest TSS genome-wide (bedtools closest).
# TSS reference: GENCODE vM10 mm10 protein-coding + lncRNA.
# Outputs:
#   BPA_male/MORE_DAR_M_Li_adt_BPA_BA_nearestGene.txt
#   BPA_male/LESS_DAR_M_Li_adt_BPA_BA_nearestGene.txt

echo "[M4] Annotating nearest genes..."

for status in MORE LESS; do
    bed="${OUT_DIR}/${status}_DAR_M_Li_adt_BPA_BA.bed"
    out="${OUT_DIR}/${status}_DAR_M_Li_adt_BPA_BA_nearestGene.txt"

    tail -n +2 "${bed}" \
      | bedtools sort -i - \
      | bedtools closest -a - -b "${TSS}" -D a -k 1 -sorted \
      | awk 'BEGIN{OFS="\t";
               print "chr","start","end","annotation",\
                     "gene_id","gene_name","gene_type","strand","dist_to_TSS"}
             {print $1,$2,$3,$4,$8,$9,$10,$11,$NF}' \
      > "${out}"

    echo "  ${status}: ${out}  ($(( $(wc -l < "${out}") - 1 )) peaks)"
done

# =============================================================================
# STEP M5 — Extract male DEGs (Bartolomei lab)
# =============================================================================
# Filter: lab=Bartolomei, tissue=liver, stage=20weeks, sex=Male
# Exposures: BPA10mg_Male_adult, BPA10ug_Male_adult
# Output: BPA_male/DEG_Bartolomei_liver_20weeks_male.txt

echo "[M5] Extracting male DEGs..."

awk -F'\t' 'NR==1 || ($10=="Bartolomei" && $15=="liver" &&
            $16=="20weeks" && $17=="Male")' \
    "${DEG_ALL}" > "${DEG_M}"

echo "  Done: ${DEG_M}"
echo "  Rows: $(( $(wc -l < "${DEG_M}") - 1 )) (excl. header)"

# =============================================================================
# STEP M6 — Join DAR nearest-gene with male DEGs
# =============================================================================
# Inner join: nearestGene.gene_name == DEG.gene
# Keep only DAR peaks whose nearest gene appears in the male DEG list.
# Each peak expands once per matching DEG row (peak x DEG).
# Outputs:
#   BPA_male/MORE_DAR_M_Li_adt_BPA_BA_nearestGene_DEG.txt
#   BPA_male/LESS_DAR_M_Li_adt_BPA_BA_nearestGene_DEG.txt

echo "[M6] Joining nearest-gene with male DEGs..."

python3 - "${OUT_DIR}" "${DEG_M}" <<'PYEOF'
import sys, csv

out_dir  = sys.argv[1]
deg_file = sys.argv[2]

# Load DEG table keyed by gene name
deg_rows = {}
with open(deg_file) as f:
    for row in csv.DictReader(f, delimiter='\t'):
        deg_rows.setdefault(row['gene'], []).append(row)

for status in ("MORE", "LESS"):
    ng_file  = f"{out_dir}/{status}_DAR_M_Li_adt_BPA_BA_nearestGene.txt"
    out_file = f"{out_dir}/{status}_DAR_M_Li_adt_BPA_BA_nearestGene_DEG.txt"

    with open(ng_file) as f:
        ng_reader = csv.DictReader(f, delimiter='\t')
        ng_fields = ng_reader.fieldnames

        with open(deg_file) as d:
            deg_fields = csv.DictReader(d, delimiter='\t').fieldnames

        all_fields = ng_fields + deg_fields

        with open(out_file, 'w', newline='') as fout:
            writer = csv.DictWriter(fout, fieldnames=all_fields,
                                    delimiter='\t', extrasaction='ignore')
            writer.writeheader()

            n_rows = 0
            genes  = set()
            for ng in ng_reader:
                if ng['gene_name'] in deg_rows:
                    for deg in deg_rows[ng['gene_name']]:
                        writer.writerow({**ng, **deg})
                        n_rows += 1
                        genes.add(ng['gene_name'])

    print(f"  {status}: {out_file}")
    print(f"    {n_rows} rows  |  {len(genes)} unique genes matched")

    sig_rows = []
    with open(out_file) as f:
        for row in csv.DictReader(f, delimiter='\t'):
            if row.get('sig', '') == 'TRUE':
                sig_rows.append(row)
    up   = sum(1 for r in sig_rows if r.get('direction','') == 'up')
    down = sum(1 for r in sig_rows if r.get('direction','') == 'down')
    print(f"    sig=TRUE: {len(sig_rows)} rows  ({down} down, {up} up)")
PYEOF

# =============================================================================
echo ""
echo "=== Male pipeline complete ==="
echo "All outputs written to: ${OUT_DIR}"
echo ""
echo "Output files:"
for f in "${OUT_DIR}"/*; do
    rows=""
    if [[ "${f}" == *.txt ]] || [[ "${f}" == *.bed ]]; then
        rows="  ($(( $(wc -l < "${f}") - 1 )) data rows)"
    fi
    echo "  $(basename "${f}")${rows}"
done
PYEOF