#!/bin/bash

# Define root data directory and input VCF file
INPUT_VCF=$1
OUTPUT_VCF=$2
DATA_DIR=$3

# Create a dictionary to map labels to classes
declare -A label_2_class

# Construct bedtools intersect command
bedtools_cmd="bedtools intersect -loj -a $INPUT_VCF"
for region_class in $(ls $DATA_DIR); do
    region_class_path="$DATA_DIR/$region_class"
    for file in $(ls $region_class_path); do
        if [[ $file == *.bed.gz && $file != *notin* ]]; then
            file_path="$region_class_path/$file"
            label_region=$(basename $file .bed.gz | cut -d'_' -f2-)
            bedtools_cmd="$bedtools_cmd -b $file_path -names $label_region"
            label_2_class[$label_region]=$region_class
        fi
    done
done

# Add header flag to bedtools command
bedtools_cmd="$bedtools_cmd -header"

# Convert label_2_class to a string for passing to awk
label_2_class_str=""
for label in "${!label_2_class[@]}"; do
    label_2_class_str+="$label=${label_2_class[$label]};"
done

# Print the label_2_class dictionary for verification
echo "Label to Class mapping:"
for label in "${!label_2_class[@]}"; do
    echo "$label : ${label_2_class[$label]}"
done

# Run bedtools intersect and process with awk
eval $bedtools_cmd | \
awk -F'\t' -v OFS='\t' -v label_2_class_str="$label_2_class_str" '
BEGIN {
    split(label_2_class_str, pairs, ";");
    for (i in pairs) {
        split(pairs[i], kv, "=");
        label_2_class[kv[1]] = kv[2];
    }
    # Print new header lines for GA4GHC and GA4GHL
    print "##INFO=<ID=GA4GHC,Number=.,Type=String,Description=\"Stratification were developed with the Global Alliance for Genomic Health (GA4GH) Class\">";
    print "##INFO=<ID=GA4GHL,Number=.,Type=String,Description=\"Stratification were developed with the Global Alliance for Genomic Health (GA4GH) Label\">";
}
{
    if ($1 ~ /^##/ || $1 ~ /^#CHROM/) {
        print $0;
    } else {
        key = $1"_"$2"_"$3"_"$4"_"$5;
        info = $8;
        region_labels = $(NF-3);
        
        if (region_labels == ".") {
            variants_region_class[key] = ".";
            variants_region_label[key] = ".";
        } else {
            if (variants_region_label[key]) {
                variants_region_label[key] = variants_region_label[key] "," region_labels;
            } else {
                variants_region_label[key] = region_labels;
            }
            
            split(variants_region_label[key], labels, ",");
            unique_classes = "";
            unique_labels = "";
            delete seen_classes;
            delete seen_labels;
            
            for (i in labels) {
                label = labels[i];
                if (!(label in seen_labels) && label in label_2_class) {
                    seen_labels[label] = 1;
                    unique_labels = (unique_labels == "" ? label : unique_labels "," label);
                    class = label_2_class[label];
                    if (!(class in seen_classes)) {
                        seen_classes[class] = 1;
                        unique_classes = (unique_classes == "" ? class : unique_classes "," class);
                    }
                }
            }
            variants_region_class[key] = unique_classes;
            variants_region_label[key] = unique_labels;
        }
        
        if (variants_region_class[key] != "" && variants_region_label[key] != "") {
            info = info ";GA4GHC=" variants_region_class[key] ";GA4GHL=" variants_region_label[key];
        } else {
            info = info ";GA4GHC=.;GA4GHL=.";
        }
        
        variant_data[key] = $1"\t"$2"\t"$3"\t"$4"\t"$5"\t"$6"\t"$7"\t"info;
        for (i = 9; i <= (NF - 4); i++) {
            variant_data[key] = variant_data[key] "\t" $i;
        }
    }
} END {
    for (key in variant_data) {
        print variant_data[key];
    }
}' > $OUTPUT_VCF
