#!/bin/bash

# ============================================================================
# SCRIPT 02: Coverage Preprocessing - Convert PAF files to bedGraph format
# ============================================================================
# PURPOSE: Convert PAF alignments to coverage tracks (bedGraph format)
# INPUT: PAF files from outputs/01_mapping/
# OUTPUT: bedGraph files in outputs/02_coverage/
# DEPENDENCIES: bedtools, awk, sort
# ============================================================================

set -e

# CONFIGURATION
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# INPUTS_DIR for this script is always outputs/01_mapping (PAF files)
INPUTS_DIR="${PROJECT_ROOT}/outputs/01_mapping"

# Use environment variables if set, otherwise use defaults
REFS_DIR="${REFS_DIR:-${PROJECT_ROOT}/references/chr21_only}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/outputs/02_coverage}"

# Determine chrom.sizes file name (could be chr21.chrom.sizes, hs1.chrom.sizes, etc.)
if [[ -f "${REFS_DIR}/chr21.chrom.sizes" ]]; then
  CHROM_SIZES="${REFS_DIR}/chr21.chrom.sizes"
elif [[ -f "${REFS_DIR}/hs1.chrom.sizes" ]]; then
  CHROM_SIZES="${REFS_DIR}/hs1.chrom.sizes"
elif [[ -f "${REFS_DIR}/hg38.chrom.sizes" ]]; then
  CHROM_SIZES="${REFS_DIR}/hg38.chrom.sizes"
elif [[ -f "${REFS_DIR}/genome.chrom.sizes" ]]; then
  CHROM_SIZES="${REFS_DIR}/genome.chrom.sizes"
else
  # Fallback: use first .chrom.sizes file found
  CHROM_SIZES=$(find "${REFS_DIR}" -maxdepth 1 -name "*.chrom.sizes" -type f | head -1)
  if [[ -z "$CHROM_SIZES" ]]; then
    echo "ERROR: No chrom.sizes file found in ${REFS_DIR}"
    exit 1
  fi
fi

# ============================================================================
# VALIDATION
# ============================================================================

echo "============================================================================"
echo "SCRIPT 02: Coverage Preprocessing"
echo "============================================================================"
echo "Input PAF directory: ${INPUTS_DIR}"
echo "Chromosome sizes: ${CHROM_SIZES}"
echo "Output directory: ${OUTPUT_DIR}"
echo ""

# CHECK CHROMOSOME SIZES FILE EXISTS
if [[ ! -f "${CHROM_SIZES}" ]]; then
  echo "ERROR: Chromosome sizes file not found: ${CHROM_SIZES}"
  exit 1
fi

# CREATE OUTPUT DIRECTORY
mkdir -p "${OUTPUT_DIR}"

# ============================================================================
# FUNCTION: Convert control sample (standard coverage)
# ============================================================================

process_control_sample() {
  local PAF_FILE=$1
  local SAMPLE_NAME=$2
  local OUTPUT_BASE=$3

  echo "Processing control sample: ${SAMPLE_NAME}"

  # STEP 1: Extract alignment information from PAF
  # PAF columns: query_name, query_length, query_start, query_end, strand, target_name, target_length, target_start, target_end, num_matches, alignment_length, mapping_quality
  # We extract: target_name, target_start, target_end

  awk 'BEGIN {OFS="\t"} {
    print $6, $8, $9, $1, $12, $5
    }' "${PAF_FILE}" >"${OUTPUT_BASE}.bed"

  # STEP 2: Sort by chromosome and position
  sort -k1,1 -k2,2n "${OUTPUT_BASE}.bed" >"${OUTPUT_BASE}.sorted.bed"

  # STEP 3: Generate coverage track using bedtools
  bedtools genomecov -i "${OUTPUT_BASE}.sorted.bed" -g "${CHROM_SIZES}" -bg >"${OUTPUT_BASE}.bedgraph"

  # CLEANUP INTERMEDIATE FILES
  rm -f "${OUTPUT_BASE}.bed" "${OUTPUT_BASE}.sorted.bed"

  LINE_COUNT=$(wc -l <"${OUTPUT_BASE}.bedgraph")
  echo "✓ Coverage track generated: ${LINE_COUNT} coverage windows"
  echo ""
}

# ============================================================================
# FUNCTION: Convert RCA sample (pass-weighted coverage for concatemers)
# ============================================================================

process_rca_sample() {
  local PAF_FILE=$1
  local SAMPLE_NAME=$2
  local OUTPUT_BASE=$3

  echo "Processing RCA sample: ${SAMPLE_NAME} (pass-weighted coverage)"

  # STEP 1: Extract alignments, counting passes per region
  # For RCA, we weight coverage by number of passes through same region
  # This accounts for concatemeric structure

  awk 'BEGIN {OFS="\t"} {
        read_id=$1
        ref_chr=$6
        ref_start=$8
        ref_end=$9
        region_key=ref_chr"_"ref_start"_"ref_end
        
        # Count how many times this read maps to this region
        if (!(read_id in pass_count) || !(region_key in pass_count[read_id])) {
            pass_count[read_id][region_key]=0
        }
        pass_count[read_id][region_key]++
        
        # Store first occurrence info for output
        if (!(read_id in first_occurrence) || !(region_key in first_occurrence[read_id])) {
            first_occurrence[read_id][region_key]=sprintf("%s\t%d\t%d", ref_chr, ref_start, ref_end)
        }
    }
    END {
        for (read_id in first_occurrence) {
            for (region_key in first_occurrence[read_id]) {
                # Output each region-read pair with pass count as coverage value
                # This effectively weights alignments by their multiplicity
                split(first_occurrence[read_id][region_key], fields, "\t")
                print fields[1], fields[2], fields[3], pass_count[read_id][region_key]
            }
        }
    }' "${PAF_FILE}" >"${OUTPUT_BASE}.pass_counts.bed"

  # STEP 2: Sort pass-count file
  sort -k1,1 -k2,2n "${OUTPUT_BASE}.pass_counts.bed" >"${OUTPUT_BASE}.pass_counts.sorted.bed"

  # STEP 3: Generate coverage track
  bedtools genomecov -i "${OUTPUT_BASE}.pass_counts.sorted.bed" -g "${CHROM_SIZES}" -bg >"${OUTPUT_BASE}.bedgraph"

  # CLEANUP INTERMEDIATE FILES
  rm -f "${OUTPUT_BASE}.pass_counts.bed" "${OUTPUT_BASE}.pass_counts.sorted.bed"

  LINE_COUNT=$(wc -l <"${OUTPUT_BASE}.bedgraph")
  echo "✓ Pass-weighted coverage track generated: ${LINE_COUNT} coverage windows"
  echo ""
}

# ============================================================================
# MAIN: Process all samples
# ============================================================================

# PROCESS CONTROL SAMPLES
for CONTROL_SAMPLE in "gDNA-NT" "gDNA-UV"; do
  INPUT_PAF="${INPUTS_DIR}/${CONTROL_SAMPLE}.paf"

  if [[ ! -f "${INPUT_PAF}" ]]; then
    echo "WARNING: PAF not found: ${INPUT_PAF} - SKIPPING"
    continue
  fi

  process_control_sample "${INPUT_PAF}" "${CONTROL_SAMPLE}" "${OUTPUT_DIR}/${CONTROL_SAMPLE}"
done

# PROCESS RCA SAMPLE (with pass-weighting)
RCA_PAF="${INPUTS_DIR}/RCA.paf"
if [[ -f "${RCA_PAF}" ]]; then
  process_rca_sample "${RCA_PAF}" "RCA" "${OUTPUT_DIR}/RCA"
else
  echo "WARNING: RCA PAF not found: ${RCA_PAF} - SKIPPING"
fi

# ============================================================================
# STANDARDIZE OUTPUT FILENAMES FOR DOWNSTREAM SCRIPTS
# ============================================================================
# Downstream scripts expect specific filenames for inputs
# Rename to match expected pattern

if [[ -f "${OUTPUT_DIR}/RCA.bedgraph" ]]; then
  mv "${OUTPUT_DIR}/RCA.bedgraph" "${OUTPUT_DIR}/rca_sorted.bedgraph"
fi

if [[ -f "${OUTPUT_DIR}/gDNA-NT.bedgraph" ]]; then
  mv "${OUTPUT_DIR}/gDNA-NT.bedgraph" "${OUTPUT_DIR}/control1_sorted.bedgraph"
fi

if [[ -f "${OUTPUT_DIR}/gDNA-UV.bedgraph" ]]; then
  mv "${OUTPUT_DIR}/gDNA-UV.bedgraph" "${OUTPUT_DIR}/control2_sorted.bedgraph"
fi

echo "============================================================================"
echo "SCRIPT 02 COMPLETE: All coverage tracks generated"
echo "Output files in: ${OUTPUT_DIR}"
echo ""
echo "Generated files:"
ls -lh "${OUTPUT_DIR}"/*.bedgraph 2>/dev/null || echo "No bedGraph files found"
echo "============================================================================"

# EXPORT FOR DOWNSTREAM SCRIPTS
export COVERAGE_DIR="${OUTPUT_DIR}"
export RCA_COVERAGE="${OUTPUT_DIR}/rca_sorted.bedgraph"
export CTRL1_COVERAGE="${OUTPUT_DIR}/control1_sorted.bedgraph"
export CTRL2_COVERAGE="${OUTPUT_DIR}/control2_sorted.bedgraph"