#!/bin/bash
set -euo pipefail

# Function to display usage information
usage() {
  echo "Usage: $0 -s <sample_id> -i <input_bam_or_cram> -g <gene_class> -o <output_directory> [-r <reference>]"
  echo "  -s  Sample ID or gene ID (required)"
  echo "  -i  Input BAM or CRAM file mapped to hg38 (required)"
  echo "  -g  Gene class, one of: HLA, KIR, CYP, IG_TR (required)"
  echo "  -o  Output directory (required)"
  echo "  -r  Reference file (required if input is CRAM)"
  exit 1
}

help() {
	 sed -rn 's/^### ?//;T;p' "$0"
}

if [[ $# == 0 ]] || [[ "$1" == "-h" ]]; then
	usage
	exit 1
fi


# Parse command line arguments
sample_id=""
input_file=""
gene_class=""
output_directory=""
reference=""

while getopts ":s:i:g:o:r:" opt; do
  case ${opt} in
    s)
      sample_id=$OPTARG
      ;;
    i)
      input_file=$OPTARG
      ;;
    g)
      gene_class=$OPTARG
      ;;
    o)
      output_directory=$OPTARG
      ;;
    r)
      reference=$OPTARG
      ;;
    \?)
      echo "Invalid option: -$OPTARG" 1>&2
      usage
      ;;
    :)
      echo "Option -$OPTARG requires an argument." 1>&2
      usage
      ;;
  esac
done

# Check if all required parameters are provided
if [ -z "$sample_id" ] || [ -z "$input_file" ] || [ -z "$gene_class" ] || [ -z "$output_directory" ]; then
  echo "Error: Missing required parameter(s)." 1>&2
  usage
fi

# Check if the reference file is provided for CRAM input
if [[ "$input_file" == *.cram && -z "$reference" ]]; then
  echo "Error: Reference file is required for CRAM input." 1>&2
  usage
fi

# Tools
samtools=samtools
# if output_directory does not exist, create it
if [ ! -d "$output_directory" ]; then
  mkdir -p "$output_directory"
fi

# Output files
extract_bam="${output_directory}/${sample_id}.${gene_class}.bam"
out_fq="${output_directory}/${sample_id}.${gene_class}.fastq"

# Define regions for each gene class
declare -A regions
regions["HLA"]="chr6:29700000-33200000  chr6_GL000250v2_alt: chr6_GL000251v2_alt: chr6_GL000252v2_alt: chr6_GL000253v2_alt: chr6_GL000254v2_alt: chr6_GL000255v2_alt: chr6_GL000256v2_alt: \
      	HLA-A*01:01:01:01: HLA-A*01:01:01:02N: HLA-A*01:01:38L: HLA-A*01:02: HLA-A*01:03: HLA-A*01:04N: HLA-A*01:09: HLA-A*01:11N: HLA-A*01:14: HLA-A*01:16N: HLA-A*01:20: HLA-A*02:01:01:01: HLA-A*02:01:01:02L: HLA-A*02:01:01:03: HLA-A*02:01:01:04: HLA-A*02:02:01: HLA-A*02:03:01: HLA-A*02:03:03: HLA-A*02:05:01: HLA-A*02:06:01: HLA-A*02:07:01: HLA-A*02:10: HLA-A*02:251: HLA-A*02:259: HLA-A*02:264: HLA-A*02:265: HLA-A*02:266: HLA-A*02:269: HLA-A*02:279: HLA-A*02:32N: HLA-A*02:376: HLA-A*02:43N: HLA-A*02:455: HLA-A*02:48: HLA-A*02:51: HLA-A*02:533: HLA-A*02:53N: HLA-A*02:57: HLA-A*02:60:01: HLA-A*02:65: HLA-A*02:68: HLA-A*02:77: HLA-A*02:81: HLA-A*02:89: HLA-A*02:95: HLA-A*03:01:01:01: HLA-A*03:01:01:02N: HLA-A*03:01:01:03: HLA-A*03:02:01: HLA-A*03:11N: HLA-A*03:21N: HLA-A*03:36N: HLA-A*11:01:01: HLA-A*11:01:18: HLA-A*11:02:01: HLA-A*11:05: HLA-A*11:110: HLA-A*11:25: HLA-A*11:50Q: HLA-A*11:60: HLA-A*11:69N: HLA-A*11:74: HLA-A*11:75: HLA-A*11:77: HLA-A*23:01:01: HLA-A*23:09: HLA-A*23:38N: HLA-A*24:02:01:01: HLA-A*24:02:01:02L: HLA-A*24:02:01:03: HLA-A*24:02:03Q: HLA-A*24:02:10: HLA-A*24:03:01: HLA-A*24:07:01: HLA-A*24:08: HLA-A*24:09N: HLA-A*24:10:01: HLA-A*24:11N: HLA-A*24:152: HLA-A*24:20: HLA-A*24:215: HLA-A*24:61: HLA-A*24:86N: HLA-A*25:01:01: HLA-A*26:01:01: HLA-A*26:11N: HLA-A*26:15: HLA-A*26:50: HLA-A*29:01:01:01: HLA-A*29:01:01:02N: HLA-A*29:02:01:01: HLA-A*29:02:01:02: HLA-A*29:46: HLA-A*30:01:01: HLA-A*30:02:01:01: HLA-A*30:02:01:02: HLA-A*30:04:01: HLA-A*30:89: HLA-A*31:01:02: HLA-A*31:01:23: HLA-A*31:04: HLA-A*31:14N: HLA-A*31:46: HLA-A*32:01:01: HLA-A*32:06: HLA-A*33:01:01: HLA-A*33:03:01: HLA-A*33:07: HLA-A*34:01:01: HLA-A*34:02:01: HLA-A*36:01: HLA-A*43:01: HLA-A*66:01:01: HLA-A*66:17: HLA-A*68:01:01:01: HLA-A*68:01:01:02: HLA-A*68:01:02:01: HLA-A*68:01:02:02: HLA-A*68:02:01:01: HLA-A*68:02:01:02: HLA-A*68:02:01:03: HLA-A*68:02:02: HLA-A*68:03:01: HLA-A*68:08:01: HLA-A*68:113: HLA-A*68:17: HLA-A*68:18N: HLA-A*68:22: HLA-A*68:71: HLA-A*69:01: HLA-A*74:01: HLA-A*74:02:01:01: HLA-A*74:02:01:02: HLA-A*80:01:01:01: HLA-A*80:01:01:02: HLA-B*07:02:01: HLA-B*07:05:01: HLA-B*07:06: HLA-B*07:156: HLA-B*07:33:01: HLA-B*07:41: HLA-B*07:44: HLA-B*07:50: HLA-B*08:01:01: HLA-B*08:08N: HLA-B*08:132: HLA-B*08:134: HLA-B*08:19N: HLA-B*08:20: HLA-B*08:33: HLA-B*08:79: HLA-B*13:01:01: HLA-B*13:02:01: HLA-B*13:02:03: \
	HLA-B*13:02:09: HLA-B*13:08: HLA-B*13:15: HLA-B*13:25: HLA-B*14:01:01: HLA-B*14:02:01: HLA-B*14:07N: HLA-B*15:01:01:01: HLA-B*15:01:01:02N: HLA-B*15:01:01:03: HLA-B*15:02:01: HLA-B*15:03:01: HLA-B*15:04:01: HLA-B*15:07:01: HLA-B*15:108: HLA-B*15:10:01: HLA-B*15:11:01: HLA-B*15:13:01: HLA-B*15:16:01: HLA-B*15:17:01:01: HLA-B*15:17:01:02: HLA-B*15:18:01: HLA-B*15:220: HLA-B*15:25:01: HLA-B*15:27:01: HLA-B*15:32:01: HLA-B*15:42: HLA-B*15:58: HLA-B*15:66: HLA-B*15:77: HLA-B*15:83: HLA-B*18:01:01:01: HLA-B*18:01:01:02: HLA-B*18:02: HLA-B*18:03: HLA-B*18:17N: HLA-B*18:26: HLA-B*18:94N: HLA-B*27:04:01: HLA-B*27:05:02: HLA-B*27:05:18: HLA-B*27:06: HLA-B*27:07:01: HLA-B*27:131: HLA-B*27:24: HLA-B*27:25: HLA-B*27:32: HLA-B*35:01:01:01: HLA-B*35:01:01:02: HLA-B*35:01:22: HLA-B*35:02:01: HLA-B*35:03:01: HLA-B*35:05:01: HLA-B*35:08:01: HLA-B*35:14:02: HLA-B*35:241: HLA-B*35:41: HLA-B*37:01:01: HLA-B*37:01:05: HLA-B*38:01:01: HLA-B*38:02:01: HLA-B*38:14: HLA-B*39:01:01:01: HLA-B*39:01:01:02L: HLA-B*39:01:01:03: HLA-B*39:01:03: HLA-B*39:01:16: HLA-B*39:01:21: HLA-B*39:05:01: HLA-B*39:06:02: HLA-B*39:10:01: HLA-B*39:13:02: HLA-B*39:14: HLA-B*39:34: HLA-B*39:38Q: HLA-B*40:01:01: HLA-B*40:01:02: HLA-B*40:02:01: HLA-B*40:03: HLA-B*40:06:01:01: HLA-B*40:06:01:02: HLA-B*40:10:01: HLA-B*40:150: HLA-B*40:40: HLA-B*40:72:01: HLA-B*40:79: HLA-B*41:01:01: HLA-B*41:02:01: HLA-B*42:01:01: HLA-B*42:02: HLA-B*42:08: HLA-B*44:02:01:01: HLA-B*44:02:01:02S: HLA-B*44:02:01:03: HLA-B*44:02:17: HLA-B*44:02:27: HLA-B*44:03:01: HLA-B*44:03:02: HLA-B*44:04: HLA-B*44:09: HLA-B*44:138Q: HLA-B*44:150: HLA-B*44:23N: HLA-B*44:26: HLA-B*44:46: HLA-B*44:49: HLA-B*44:56N: HLA-B*45:01:01: HLA-B*45:04: HLA-B*46:01:01: HLA-B*46:01:05: HLA-B*47:01:01:01: HLA-B*47:01:01:02: HLA-B*48:01:01: HLA-B*48:03:01: HLA-B*48:04: HLA-B*48:08: HLA-B*49:01:01: HLA-B*49:32: HLA-B*50:01:01: HLA-B*51:01:01: HLA-B*51:01:02: HLA-B*51:02:01: HLA-B*51:07:01: HLA-B*51:42: HLA-B*52:01:01:01: HLA-B*52:01:01:02: HLA-B*52:01:01:03: HLA-B*52:01:02: HLA-B*53:01:01: HLA-B*53:11: HLA-B*54:01:01: HLA-B*54:18: HLA-B*55:01:01: HLA-B*55:01:03: HLA-B*55:02:01: HLA-B*55:12: HLA-B*55:24: HLA-B*55:48: HLA-B*56:01:01: HLA-B*56:03: HLA-B*56:04: HLA-B*57:01:01: HLA-B*57:03:01: HLA-B*57:06: HLA-B*57:11: HLA-B*57:29: HLA-B*58:01:01: HLA-B*58:31N: HLA-B*59:01:01:01: HLA-B*59:01:01:02: HLA-B*67:01:01: HLA-B*67:01:02: HLA-B*67:02: HLA-B*73:01: HLA-B*78:01:01: HLA-B*81:01: HLA-B*82:02:01: HLA-C*01:02:01: HLA-C*01:02:11: HLA-C*01:02:29: HLA-C*01:02:30: \
	HLA-C*01:03: HLA-C*01:06: HLA-C*01:08: HLA-C*01:14: HLA-C*01:21: HLA-C*01:30: HLA-C*01:40: HLA-C*02:02:02:01: HLA-C*02:02:02:02: HLA-C*02:10: HLA-C*02:11: HLA-C*02:16:02: HLA-C*02:69: HLA-C*02:85: HLA-C*02:86: HLA-C*02:87: HLA-C*03:02:01: HLA-C*03:02:02:01: HLA-C*03:02:02:02: HLA-C*03:02:02:03: HLA-C*03:03:01: HLA-C*03:04:01:01: HLA-C*03:04:01:02: HLA-C*03:04:02: HLA-C*03:04:04: HLA-C*03:05: HLA-C*03:06: HLA-C*03:100: HLA-C*03:13:01: HLA-C*03:20N: HLA-C*03:219: HLA-C*03:261: HLA-C*03:40:01: HLA-C*03:41:02: HLA-C*03:46: HLA-C*03:61: HLA-C*04:01:01:01: HLA-C*04:01:01:02: HLA-C*04:01:01:03: HLA-C*04:01:01:04: HLA-C*04:01:01:05: HLA-C*04:01:62: HLA-C*04:03:01: HLA-C*04:06: HLA-C*04:09N: HLA-C*04:128: HLA-C*04:161: HLA-C*04:177: HLA-C*04:70: HLA-C*04:71: HLA-C*05:01:01:01: HLA-C*05:01:01:02: HLA-C*05:08: HLA-C*05:09:01: HLA-C*05:93: HLA-C*06:02:01:01: HLA-C*06:02:01:02: HLA-C*06:02:01:03: HLA-C*06:23: HLA-C*06:24: HLA-C*06:46N: HLA-C*07:01:01:01: HLA-C*07:01:01:02: HLA-C*07:01:02: HLA-C*07:01:19: HLA-C*07:01:27: HLA-C*07:01:45: HLA-C*07:02:01:01: HLA-C*07:02:01:02: HLA-C*07:02:01:03: HLA-C*07:02:01:04: HLA-C*07:02:01:05: HLA-C*07:02:05: HLA-C*07:02:06: HLA-C*07:02:64: HLA-C*07:04:01: HLA-C*07:04:02: HLA-C*07:06: HLA-C*07:149: HLA-C*07:18: HLA-C*07:19: HLA-C*07:26: HLA-C*07:30: HLA-C*07:32N: HLA-C*07:384: HLA-C*07:385: HLA-C*07:386: HLA-C*07:391: HLA-C*07:392: HLA-C*07:49: HLA-C*07:56:02: HLA-C*07:66: HLA-C*07:67: HLA-C*08:01:01: HLA-C*08:01:03: HLA-C*08:02:01:01: HLA-C*08:02:01:02: HLA-C*08:03:01: HLA-C*08:04:01: HLA-C*08:112: HLA-C*08:20: HLA-C*08:21: HLA-C*08:22: HLA-C*08:24: HLA-C*08:27: HLA-C*08:36N: HLA-C*08:40: HLA-C*08:41: HLA-C*08:62: HLA-C*12:02:02: HLA-C*12:03:01:01: HLA-C*12:03:01:02: HLA-C*12:08: HLA-C*12:13: HLA-C*12:19: HLA-C*12:22: HLA-C*12:99: HLA-C*14:02:01: HLA-C*14:03: HLA-C*14:21N: HLA-C*14:23: HLA-C*15:02:01: HLA-C*15:05:01: HLA-C*15:05:02: HLA-C*15:13: HLA-C*15:16: HLA-C*15:17: HLA-C*15:96Q: HLA-C*16:01:01: HLA-C*16:02:01: HLA-C*16:04:01: HLA-C*17:01:01:01: HLA-C*17:01:01:02: HLA-C*17:01:01:03: HLA-C*17:03: HLA-C*18:01: HLA-DQA1*01:01:02: HLA-DQA1*01:02:01:01: HLA-DQA1*01:02:01:02: HLA-DQA1*01:02:01:03: HLA-DQA1*01:02:01:04: HLA-DQA1*01:03:01:01: HLA-DQA1*01:03:01:02: HLA-DQA1*01:04:01:01: HLA-DQA1*01:04:01:02: HLA-DQA1*01:05:01: HLA-DQA1*01:07: HLA-DQA1*01:10: HLA-DQA1*01:11: HLA-DQA1*02:01: HLA-DQA1*03:01:01: HLA-DQA1*03:02: HLA-DQA1*03:03:01: HLA-DQA1*04:01:02:01: HLA-DQA1*04:01:02:02: HLA-DQA1*04:02: HLA-DQA1*05:01:01:01: HLA-DQA1*05:01:01:02: \
	HLA-DQA1*05:03: HLA-DQA1*05:05:01:01: HLA-DQA1*05:05:01:02: HLA-DQA1*05:05:01:03: HLA-DQA1*05:11: HLA-DQA1*06:01:01: HLA-DQB1*02:01:01: HLA-DQB1*02:02:01: HLA-DQB1*03:01:01:01: HLA-DQB1*03:01:01:02: HLA-DQB1*03:01:01:03: HLA-DQB1*03:02:01: HLA-DQB1*03:03:02:01: HLA-DQB1*03:03:02:02: HLA-DQB1*03:03:02:03: HLA-DQB1*03:05:01: HLA-DQB1*05:01:01:01: HLA-DQB1*05:01:01:02: HLA-DQB1*05:03:01:01: HLA-DQB1*05:03:01:02: HLA-DQB1*06:01:01: HLA-DQB1*06:02:01: HLA-DQB1*06:03:01: HLA-DQB1*06:09:01: HLA-DRB1*01:01:01: HLA-DRB1*01:02:01: HLA-DRB1*03:01:01:01: HLA-DRB1*03:01:01:02: HLA-DRB1*04:03:01: HLA-DRB1*07:01:01:01: HLA-DRB1*07:01:01:02: HLA-DRB1*08:03:02: HLA-DRB1*09:21: HLA-DRB1*10:01:01: HLA-DRB1*11:01:01: HLA-DRB1*11:01:02: HLA-DRB1*11:04:01: HLA-DRB1*12:01:01: HLA-DRB1*12:17: HLA-DRB1*13:01:01: HLA-DRB1*13:02:01: HLA-DRB1*14:05:01: HLA-DRB1*14:54:01: HLA-DRB1*15:01:01:01: HLA-DRB1*15:01:01:02: HLA-DRB1*15:01:01:03: HLA-DRB1*15:01:01:04: HLA-DRB1*15:02:01: HLA-DRB1*15:03:01:01: HLA-DRB1*15:03:01:02: HLA-DRB1*16:02:01: "
regions["KIR"]="chr19:53600000-56300000 chr19_KI270868v1_alt chr19_KI270865v1_alt chr19_GL383573v1_alt chr19_GL383575v2_alt chr19_GL383576v1_alt chr19_GL383574v1_alt chr19_KI270866v1_alt chr19_KI270867v1_alt chr19_GL949746v1_alt chr19_KI270882v1_alt chr19_KI270883v1_alt chr19_KI270884v1_alt chr19_KI270885v1_alt chr19_KI270886v1_alt chr19_KI270887v1_alt chr19_KI270888v1_alt chr19_KI270889v1_alt chr19_KI270890v1_alt chr19_KI270891v1_alt chr19_GL949747v2_alt chr19_KI270914v1_alt chr19_KI270915v1_alt chr19_KI270916v1_alt chr19_KI270917v1_alt chr19_KI270918v1_alt chr19_KI270919v1_alt chr19_KI270920v1_alt chr19_KI270921v1_alt chr19_KI270922v1_alt chr19_KI270923v1_alt chr19_GL949748v2_alt chr19_KI270929v1_alt chr19_KI270930v1_alt chr19_KI270931v1_alt chr19_KI270932v1_alt chr19_KI270933v1_alt chr19_GL000209v2_alt chr19_GL949749v2_alt chr19_GL949750v2_alt chr19_GL949751v2_alt chr19_GL949752v1_alt chr19_GL949753v2_alt chr19_KI270938v1_alt"
regions["CYP"]="chr13:48037726-48052755 chr8:18386301-18401218 chr12:21131194-21239796 chr22_KB663609v1_alt chr22_KI270928v1_alt chr22_GL383582v2_alt chr19:40800000-41220000 chr10:93000000-95100000 chr11:14837440-14952231 chr15:51178057-51368601 chr15:74689542-74755536 chr19:15858023-15918077 chr1:46727838-46849413 chr1:47107435-47179735 chr1:59873308-59946773 chr22:42106499-42150865 chr2:38046973-38129902 chr7:963181-1009640 chr7:99600000-99976102 chr7:99235817-99287619 chr20:49483874-49588137"
regions["IG_TR"]="chr14:104000000-107349540 chr22:22000000-25340000 chr22:16900000-17000000 chr14:21000000-23000000 chr7:142200000-143000000 chr2:88800000-92100000 chr2:97000000-97400000 chr16:28100000-34600000 chr7:38200000-38400000 chr15:19000000-23000000 chr9:33600000-33800000 chr2:94265000-94266000 chr2:94953000-94954000 chr21:10649000-10650000 chr22:10685500-10686000 chr8:47202000-47203000 chrY:11473700-11474100 chr7_KI270803v1_alt:260000-840000 chr14_KI270726v1_random:4000-42000 chr16_KI270728v1_random:1109000-1314000"

# Check if the gene class is valid
if [ -z "${regions[$gene_class]+_}" ]; then
  echo "Error: Invalid gene class: $gene_class" 1>&2
  usage
fi

# Extract the regions based on the input file type
if [[ "$input_file" == *.cram ]]; then
  echo "Processing CRAM file with reference..."
  $samtools view -b -T "$reference" -o "$extract_bam" "$input_file" ${regions[$gene_class]}
else
  echo "Processing BAM file..."
  $samtools view -b -o "$extract_bam" "$input_file" ${regions[$gene_class]}
fi

# Convert BAM to FASTQ
$samtools fastq -0 "$out_fq" "$extract_bam"
gzip -f "$out_fq"

echo "Extracted reads are in ${out_fq}.gz"