#!/usr/bin/env bash
set -euo pipefail
trap 'echo "[ERROR] Pipeline failed at line $LINENO"; exit 1' ERR

# ==============================================================================
#  MOSTAR - Microbial Oxford Sequencing: Typing, Assembly & Resistome
#  Author  : Nermin Zecic <nermze@gmail.com>
#  Version : 1.0.0
#  License : MIT
# ==============================================================================
#
#  A comprehensive bioinformatics pipeline for processing ONT reads.
#
#  Features:
#    - Hybrid assembly (ONT + Illumina) or ONT-only assembly
#    - ONT polishing (Medaka) + Hybrid polishing (Polypolish)
#    - Functional annotation (Bakta)
#    - AMR profiling (AMRFinder+)
#    - Taxonomic profiling (Kraken2)
#    - Genomic plasticity & mobile elements (geNomad)
#    - ICE / conjugative element detection (MacSyFinder / CONJScan)
#    - Circular genome maps with GC, ICE, and AMR tracks
#    - Interactive HTML report
#
# ==============================================================================


#  =============================================================================
#  Environment Validation & Path Sanitization 
# ==============================================================================

# Enforce Conda environment
if [[ -z "${CONDA_PREFIX:-}" ]]; then
    echo "[ERROR] No Conda environment detected!"
    echo "Please run: conda activate mostar_env"
    exit 1
fi

# Sanitize PATH to prioritize Conda and lock out local user bin conflicts
export PATH="$CONDA_PREFIX/bin:$PATH"

# 3. Verify critical tools are strictly loading from Conda
REQUIRED_TOOLS=(
    "flye" "medaka_consensus" "samtools" "bwa" "amrfinder"
)
MISSING_TOOLS=0
for tool in "${REQUIRED_TOOLS[@]}"; do
    TOOL_PATH=$(command -v "$tool" || true)
    if [[ -z "$TOOL_PATH" ]]; then
        echo "[ERROR] Required tool not found: $tool"
        MISSING_TOOLS=1
    elif [[ "$TOOL_PATH" != "$CONDA_PREFIX"* ]]; then
        echo "[ERROR] Tool hijack detected! '$tool' is loading from outside the environment:"
        echo "        Found at: $TOOL_PATH"
        echo "        Expected: $CONDA_PREFIX/bin/$tool"
        MISSING_TOOLS=1
    fi
done

if [[ "$MISSING_TOOLS" -eq 1 ]]; then
    echo "[ERROR] Environment validation failed. Please check your Conda installation."
    exit 1
fi


# ------------------------------------------------------------------------------
# Constants
# ------------------------------------------------------------------------------
readonly VERSION="1.0.0"
readonly START_TIME=$(date +%s)

# ANSI colours
readonly BOLD='\033[1m'
readonly BLUE='\033[1;34m'
readonly GREEN='\033[0;32m'
readonly NC='\033[0m'

# Progress tracking
TOTAL_STEPS=15
CURRENT_STEP=0

# ------------------------------------------------------------------------------
# Default parameter values — overridden by argument parsing below
# ------------------------------------------------------------------------------
THREADS=10
MEDAKA_MODEL="r1041_e82_400bps_sup_v5.2.0"
FILTLONG_COV=100
ONT=""
R1=""
R2=""
OUTDIR=""
GSIZE=""
ORGANISM=""
BAKTA_DB=""
BAKTA_COMPLETE=""
BAKTA_REGIONS=""
KRAKEN2_DB=""
KRAKEN_CONFIDENCE="0.01"
GENOMAD_DB=""
USE_META=""
CLEANUP=false
RUN_KRAKEN2=false
RUN_ICE=false
RUN_GENOMAD=false
SKIP_BAKTA=false

# Runtime metrics — populated during the run
FINAL_ASM=""
ASM_SPECIES_FULL="Unknown"
POS_CHANGED=0
MEAN_DEPTH="N/A"
BWA_MAPPING_RATE="N/A"
MEDAKA_LEN=0
MEDAKA_CNT=0
MEDAKA_N50=0
FINAL_LEN=0
FINAL_CNT=0
FINAL_N50=0
PHAGE_COUNT=0
HIGH_RISK_COUNT=0
RESISTOME_JS_DATA="[]"
PHAGE_JS_DATA="[]"
TAX_CARDS=""
HYBRID_CARDS=""
METRICS_BLOCK=""

# ------------------------------------------------------------------------------
# OS detection — used for format_bp compatibility
# ------------------------------------------------------------------------------
OS_TYPE="$(uname -s)"
case "$OS_TYPE" in
    Linux*)  PLATFORM="linux"  ;;
    Darwin*) PLATFORM="macos"  ;;
    *) echo "[ERROR] Unsupported OS: $OS_TYPE"; exit 1 ;;
esac

# ------------------------------------------------------------------------------
# format_bp <number>
#   Formats an integer with thousands separators.
#   Uses pure shell/sed to avoid GNU printf "%'d" which breaks on macOS.
# ------------------------------------------------------------------------------
format_bp() {
    printf "%d" "${1:-0}" | rev | sed 's/[0-9]\{3\}/&,/g' | rev | sed 's/^,//'
}

# ------------------------------------------------------------------------------
# run_step <label> <run:true|false> <command>
#   Animates a progress bar for the current step, then runs the command if
#   the condition is "true", otherwise sleeps briefly to advance the bar.
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# run_step <label> <run:true|false> <command>
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Pipeline Tools Initialization (Defines the upfront list)
# ------------------------------------------------------------------------------
PIPELINE_STEPS=(
    "Short-read QC (Fastp)"
    "Long-read QC (Filtlong)"
    "De-novo Assembly (Flye)"
    "ONT Polishing (Medaka)"
    "Hybrid Polishing (Polypolish)"
    "Fix Assembly Origin (Circlator)"
    "Taxonomy (Kraken2)"
    "Annotation (Bakta)"
    "AMR Profiling (AMRFinder+)"
    "ICE Detection (MacSyFinder)"
    "Plasticity (geNomad)"
    "Generating Circular Maps"
    "Generating HTML Report"
    "Cleaning up intermediate files"
)
TOTAL_STEPS=${#PIPELINE_STEPS[@]}
declare -a STEP_STATES
for ((i=0; i<TOTAL_STEPS; i++)); do STEP_STATES[i]="waiting"; done
CURRENT_STEP=0
DASHBOARD_DRAWN=""

# ------------------------------------------------------------------------------
# run_step (Dashboard Version)
# ------------------------------------------------------------------------------
run_step() {
    local name="$1"
    local run="$2"
    local cmd="$3"

    local idx=$CURRENT_STEP
    local pct_start=$(( (idx * 100) / TOTAL_STEPS ))
    local pct_end=$(( ((idx + 1) * 100) / TOTAL_STEPS ))

    local G='\033[0;32m' B='\033[1;34m' Y='\033[1;33m' K='\033[1;30m' N='\033[0m'
    
    # 1. Mark current tool as running or skipped
    if [[ "$run" == "true" ]]; then
        STEP_STATES[$idx]="running"
    else
        STEP_STATES[$idx]="skipped"
    fi

    draw_ui() {
        local current_pct=$1
        local elapsed=$2
        local bar=$3
        local spaces=$4

        # Move cursor back up to the top of the list (skip the very first time)
        if [[ -n "$DASHBOARD_DRAWN" ]]; then
            printf "\033[%dA" $((TOTAL_STEPS + 2))
        fi
        DASHBOARD_DRAWN="yes"

        # Redraw the list of tools
        for (( j=0; j<TOTAL_STEPS; j++ )); do
            local s_name="${PIPELINE_STEPS[j]}"
            local s_state="${STEP_STATES[j]}"
            
            if [[ "$s_state" == "complete" ]]; then
                printf "\033[K${G}%-35s complete${N}\n" "$s_name"
            elif [[ "$s_state" == "skipped" ]]; then
                printf "\033[K${K}%-35s skipped${N}\n" "$s_name"
            elif [[ "$s_state" == "running" ]]; then
                printf "\033[K${Y}%-35s running${N}\n" "$s_name"
            else
                printf "\033[K%-35s waiting\n" "$s_name"
            fi
        done
        
        # Redraw the empty line and the bottom progress bar
        echo -e "\033[K"
        printf "\033[K${B}[MOSTAR]${N} ${G}[%s%s]${N} %d%% : %-30s ${B}[%s]${N}\n" \
               "$bar" "$spaces" "$current_pct" "$name" "$elapsed"
    }

    # Animate progress bar for the phase
    local i now diff elapsed filled empty bar spaces
    for (( i = pct_start; i <= pct_end; i++ )); do
        now=$(date +%s)
        diff=$(( now - START_TIME ))
        elapsed=$(printf "%02d:%02d:%02d" $((diff/3600)) $((diff%3600/60)) $((diff%60)))
        filled=$(( i / 4 ))
        empty=$(( 25 - filled ))
        bar=$(printf "%${filled}s" | tr ' ' '#')
        spaces=$(printf "%${empty}s" | tr ' ' '-')
        
        draw_ui "$i" "$elapsed" "$bar" "$spaces"
        
        if [[ "$run" == "true" ]]; then
            sleep 0.01
        else
            sleep 0.002 # Speed through skipped steps
        fi
    done

    # Execute the tool command
    if [[ "$run" == "true" ]]; then
        eval "$cmd"
        STEP_STATES[$idx]="complete"
    fi

    # Final redraw for this step to instantly show it as 'complete'
    now=$(date +%s)
    diff=$(( now - START_TIME ))
    elapsed=$(printf "%02d:%02d:%02d" $((diff/3600)) $((diff%3600/60)) $((diff%60)))
    draw_ui "$pct_end" "$elapsed" "$bar" "$spaces"

    ((CURRENT_STEP++)) || true

    # 5. Success message at the very end
    if [[ "$pct_end" -eq 100 ]]; then
        echo -e "\n${G}✔ MOSTAR finished successfully.${N}"
    fi
}

# ------------------------------------------------------------------------------
# calc_assembly_stats <fasta>
#   Prints: <total_bp> <contig_count> <N50> <L50>
#   Always prints four integers; returns "0 0 0 0" on empty input.
# ------------------------------------------------------------------------------
calc_assembly_stats() {
    local fasta="$1"
    local stats

    stats=$(awk '/^>/ { if (l) print l; l=0; next } { l += length($0) } END { if (l) print l }' \
            "$fasta" | sort -rn)

    if [[ -z "$stats" ]]; then
        echo "0 0 0 0"
        return
    fi

    echo "$stats" | awk '
        { a[i++] = $1; s += $1 }
        END {
            if (s == 0 || i == 0) { print "0 0 0 0"; exit }
            for (j = 0; j < i; j++) {
                t += a[j]
                if (t >= s/2) { print s, i, a[j], j+1; exit }
            }
        }'
}

# ------------------------------------------------------------------------------
# usage — print help and exit
# ------------------------------------------------------------------------------
usage() {
    cat <<'HELP'

A comprehensive bioinformatics pipeline by Nermin Zecic <nermze@gmail.com>

      __      __   __    ____  _______   __     _____
      | |\  /| |  /  \  / ___| |_   _|  /  \   | __  \
      | | \/ | | | || | \___ \   | |   / /\ \  | |_) /
      |_|\__/|_|  \__/  |____/   |_|  /_/  \_\ |_\ \_\

Assembly, Polishing, Annotation, ICE detection & AMR profiling (v1.0.0)

Usage:
  Hybrid mode:   mostar --ont reads.fq.gz -g 2.1m -o outdir --r1 R1.fq --r2 R2.fq
  ONT-only mode: mostar --ont reads.fq.gz -g 2.1m -o outdir

Required:
  --ont           ONT long reads [.fastq.gz]
  --genome-size   Estimated genome size [e.g. 2.1m, 5.0m, 4g]
  --output        Output directory

Assembly & polishing:
  --model         Medaka model [Default: r1041_e82_400bps_sup_v5.2.0]
  --r1 / --r2     Illumina paired reads — enables hybrid mode
  --meta          Enable uneven-coverage mode in Flye [Default: off]

Annotation:
  --bakta-db      Path to Bakta database — enables annotation
  --bakta-ref     Bakta reference GFF for locus-tag transfer
  --complete      Tell Bakta all sequences are complete/circular

Classification:
  --kraken2-db     Path to Kraken2 database
  --confidence    Kraken2 confidence threshold [Default: 0.1]

Mobile element detection:
  --ice           Run MacSyFinder CONJScan (requires --bakta-db)
  --plasticity    Run geNomad plasmid/provirus detection
  --genomad-db    Path to geNomad database [Default: ~/genomad_db]

  --organism      Organism name for AMRFinder+ [e.g. Escherichia]
  --threads       CPU threads [Default: 10]
  --cleanup       Remove intermediate files after completion
  --help          Show this message and exit

MOSTAR v1.0.0
HELP
    exit 1
}

# ------------------------------------------------------------------------------
# generate_html_report
#   Reads all global variables produced during the run and writes the final
#   self-contained HTML report to $OUTDIR/MOSTAR_Final_Report.html.
# ------------------------------------------------------------------------------
generate_html_report() {
    local report_sample report_date duration
    local flye_v medaka_v amrfinder_v bakta_v kraken_v
    local fastp_v bwa_v filtlong_v polypolish_v
    local full_amr_section="" 
    local mobile_section="" 
    local phage_section=""
    local maps_section="" 
    local map_content="" 
    local img b64 label

    report_sample=$(basename "${ONT:-sample}")
    report_date=$(date '+%Y-%m-%d %H:%M:%S')
    duration=$(( ($(date +%s) - START_TIME) / 60 ))

    # Tool versions
    flye_v=$(flye --version 2>/dev/null || echo "N/A")
    medaka_v=$(medaka --version 2>/dev/null | head -n1 || echo "N/A")
    amrfinder_v=$(amrfinder --version 2>/dev/null || echo "N/A")
    bakta_v=$(bakta --version 2>/dev/null || echo "N/A")
    kraken_v=$(kraken2 --version 2>/dev/null | head -n1 | awk '{print $NF}' || echo "N/A")
    fastp_v=$(fastp --version 2>&1 | awk '{print $NF}' || echo "N/A")
    bwa_v=$(bwa 2>&1 | awk '/Version:/{print $NF}' || echo "N/A")
    filtlong_v=$(filtlong --version 2>/dev/null || echo "N/A")
    polypolish_v=$(polypolish --version 2>/dev/null || echo "N/A")

    local amr_tsv="$OUTDIR/amr_results/AMR_Report.tsv"
    local map_dir="$OUTDIR/amr_results/maps"
    local html_out="$OUTDIR/MOSTAR_Final_Report.html"

    local MAIN_ID_VAL
    MAIN_ID_VAL=$(awk '/^>/ {if(l) print l, id; id=$1; l=0; next} {l+=length($0)} END {print l, id}' "$FINAL_ASM" 2>/dev/null | sort -nr | head -n1 | awk '{print $2}' | sed 's/>//')
    [[ -z "$MAIN_ID_VAL" ]] && MAIN_ID_VAL="unknown"

    # AMR table (Using Python directly via stdin)
    if [[ -f "$amr_tsv" ]]; then
        # Pass MAIN_ID_VAL as the second argument to Python
        amr_html=$(python3 - "$amr_tsv" "$MAIN_ID_VAL" <<'PYAMR'
import sys, pandas as pd

# Clean the main ID to prevent whitespace matching errors
main_id = str(sys.argv[2]).strip()

def row_color(cls, contig):
    c = str(cls).lower()
    # Clean the contig ID
    c_id = str(contig).split()[0].strip()

    # 1. ULTIMATE PRIORITY: Is it on a plasmid? 
    # If yes, it is ALWAYS RED (Even if it is virulence)
    if c_id != main_id:
        return '#e74c3c'

    # 2. SECOND PRIORITY: Is it Virulence on the Main Chromosome?
    # If yes, it is BLACK
    if 'virulence' in c:
        return '#000000'

    # 3. STANDARD COLORS: For AMR genes on the Main Chromosome
    if any(x in c for x in ['beta-lactam', 'penicillin', 'carbapenem']): return '#3498db'
    if 'tetracycline'    in c: return '#9b59b6'
    if 'aminoglycoside'  in c: return '#e67e22'
    if any(x in c for x in ['macrolide', 'lincosamide', 'streptogramin']): return '#f1c40f'
    if 'quinolone' in c or 'fluoroquinolone' in c: return '#27ae60'
    if 'ice' in c or 'mobilome' in c: return '#00f2ff'
    
    return '#2c3e50'

df = pd.read_csv(sys.argv[1], sep='\t')
df['Class'] = df['Class'].fillna('Virulence')
df = df[['Element symbol', 'Class', '% Identity to reference', 'Contig id']]

rows = ""
for _, r in df.iterrows():
    # Pass both Class and Contig to the color logic
    color = row_color(r['Class'], r['Contig id'])
    rows += (
        f"<tr>"
        f"<td><span style='border-left:4px solid {color}; padding-left:10px; font-weight:600;'>{r['Element symbol']}</span></td>"
        f"<td style='color:{color}; font-weight:600;'>{r['Class']}</td>"
        f"<td>{r['% Identity to reference']:.1f}%</td>"
        f"<td style='color:#7f8c8d; font-size:0.85em;'>{r['Contig id']}</td>"
        f"</tr>"
    )

print(
    "<table class='resistome-table'>"
    "<thead><tr><th>Gene</th><th>Class</th><th>Identity</th><th>Contig</th></tr></thead>"
    f"<tbody>{rows}</tbody>"
    "</table>"
)
PYAMR
)
        full_amr_section="<div class='card'>
            <h2>Complete Resistome Profile (AMRFinder+)</h2>
            <p>All identified resistance and virulence elements.</p>
            <div style='overflow-x:auto;'>
            ${amr_html}
            </div>
        </div>"
    fi

    # Mobile resistome section (geNomad)
    mobile_section=""
    if [[ "${RUN_GENOMAD}" == "true" ]]; then
        mobile_section="<section class='card'>
            <h2>Mobile Resistome &amp; Genomic Plasticity</h2>
            <div class='alert-box alert-red'>
                <strong>Plasmid-borne AMR detected:</strong>
                <span class='big-num'>${HIGH_RISK_COUNT}</span>
            </div>
            <table class='resistome-table' id='resTable'>
                <thead><tr><th>Gene</th><th>Contig</th><th>Class</th><th>Score</th><th>Risk</th></tr></thead>
                <tbody></tbody>
            </table>
        </section>"
    fi

    # Prophage section (geNomad)
    phage_section=""
    if [[ "${RUN_GENOMAD}" == "true" ]]; then
        phage_section="<section class='card'>
            <h2>Prophage Tracker</h2>
            <div class='alert-box alert-green'>
                <strong>Total Prophages Detected:</strong>
                <span class='big-num'>${PHAGE_COUNT}</span>
            </div>
            <table class='resistome-table' id='pTable'>
                <thead><tr><th>Contig ID</th><th>Length</th><th>Score</th><th>Type</th></tr></thead>
                <tbody></tbody>
            </table>
        </section>"
    fi

    # Circular map gallery — embed PNGs as base64
    maps_section=""
    map_content=""
    if [[ -d "$map_dir" ]]; then
        for img in "$map_dir"/*.png; do
            [[ -f "$img" ]] || continue
            b64=$(base64 < "$img" | tr -d '\n')
            label=$(basename "$img" .png | sed 's/_/ /g')
            map_content+="<div class='map-card'>
                <h4>$label</h4>
                <div class='map-zoom-container'>
                    <img src='data:image/png;base64,$b64'>
                </div>
            </div>"
        done
        if [[ -n "$map_content" ]]; then
            maps_section="<div class='card'>
                <h2>Circular Genome Maps</h2>
                <p style='color:#7f8c8d; font-size:0.85rem;'>Click any map to zoom. Move mouse to pan while zoomed.</p>
                <div class='map-grid'>$map_content</div>
            </div>"
        fi
    fi

    cat <<EOF > "$html_out"
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>MOSTAR Report — $report_sample</title>
<style>
    :root {
        --primary:  #1a2a3a;
        --bg:       #f8f9fa;
        --border:   #edf2f7;
        --accent:   #3498db;
        --red:      #c62828;
        --green:    #2e7d32;
    }
    body {
        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
        background: var(--bg);
        color: var(--primary);
        padding: 30px;
        line-height: 1.6;
        max-width: 1400px;
        margin: 0 auto;
    }
    h2 { margin-top: 0; }
    .card {
        background: white;
        padding: 25px;
        border-radius: 12px;
        border: 1px solid var(--border);
        margin-bottom: 25px;
    }
    .qc-row {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(180px, 1fr));
        gap: 15px;
        margin-bottom: 25px;
    }
    .qc-card {
        background: white;
        padding: 15px;
        border-radius: 12px;
        border: 1px solid var(--border);
        text-align: center;
    }
    .qc-label {
        display: block;
        font-size: 0.7rem;
        color: #7f8c8d;
        text-transform: uppercase;
        letter-spacing: 0.05em;
        margin-bottom: 5px;
        font-weight: 600;
    }
    .qc-val  { font-weight: 700; font-size: 1.1rem; }
    .big-num { font-size: 1.2em; font-weight: 700; }
    .alert-box {
        margin-bottom: 20px;
        padding: 15px;
        border-radius: 8px;
    }
    .alert-red   { background: #fff5f5; border-left: 5px solid var(--red);   }
    .alert-green { background: #f0fff4; border-left: 5px solid var(--green); }
    table { width: 100%; border-collapse: collapse; }
    th { padding: 12px 8px; border-bottom: 2px solid var(--border); text-align: left; font-size: 0.85rem; }
    td { padding: 10px 8px; border-bottom: 1px solid var(--border); font-size: 0.9rem; }
    .val-new { font-weight: bold; color: var(--accent); }
    .map-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(480px, 1fr)); gap: 30px; }
    .map-card { text-align: center; background: white; border: 1px solid var(--border); padding: 15px; border-radius: 12px; }
    .map-card h4 { margin: 0 0 10px; font-size: 0.9rem; color: #7f8c8d; }
    .map-zoom-container { overflow: hidden; border-radius: 8px; cursor: zoom-in; }
    .map-zoom-container img { width: 100%; transition: transform 0.35s ease; display: block; }
    .map-zoom-container.zoomed { cursor: zoom-out; }
    .map-zoom-container.zoomed img { transform: scale(3.0); }
    .legend-grid {
        display: grid;
        grid-template-columns: repeat(auto-fill, minmax(200px, 1fr));
        gap: 10px;
        margin-top: 15px;
    }
    .legend-item { display: flex; align-items: center; gap: 8px; font-size: 0.85rem; font-weight: 600; }
    .swatch { width: 14px; height: 14px; border-radius: 50%; flex-shrink: 0; border: 1px solid #ccc; }
    .resistome-table { width: 100%; border-collapse: collapse; margin-top: 15px; }
    .resistome-table th,
    .resistome-table td { border: 1px solid #eee; padding: 9px 12px; text-align: left; }
    .resistome-table thead { background: var(--bg); }
    .risk-high { background: #ffebee; color: var(--red); font-weight: bold; padding: 2px 8px; border-radius: 4px; }
    .phage-tag { background: #e8f5e9; color: var(--green); font-weight: bold; padding: 2px 8px; border-radius: 12px; font-size: 0.75rem; }
    .version-grid {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(240px, 1fr));
        gap: 20px;
        font-size: 0.85rem;
    }
    .version-grid td:first-child { color: #7f8c8d; }
    .version-grid td:last-child  { font-weight: 600; }
</style>
</head>
<body>

<div class="card" style="background:var(--primary); color:white; padding:30px;">
    <h1 style="margin:0; font-size:1.6rem;">MOSTAR Run Report</h1>
    <p style="margin:8px 0 0; opacity:0.8;">
        <b>Sample:</b> $report_sample &nbsp;|&nbsp;
        <b>Date:</b> $report_date &nbsp;|&nbsp;
        <b>Duration:</b> ${duration} min
    </p>
</div>

<div class="qc-row">
    ${TAX_CARDS}
    ${HYBRID_CARDS}
</div>

${METRICS_BLOCK}
${mobile_section}
${phage_section}
<div class="card">
    <h2>Colour Legend</h2>
    <div class="legend-grid">
        <div class="legend-item"><div class="swatch" style="background:#3498db;"></div>Beta-lactams</div>
        <div class="legend-item"><div class="swatch" style="background:#9b59b6;"></div>Tetracyclines</div>
        <div class="legend-item"><div class="swatch" style="background:#e67e22;"></div>Aminoglycosides</div>
        <div class="legend-item"><div class="swatch" style="background:#f1c40f;"></div>Macrolides</div>
        <div class="legend-item"><div class="swatch" style="background:#27ae60;"></div>Quinolones</div>
        <div class="legend-item"><div class="swatch" style="background:#000000;"></div>Virulence</div>
        <div class="legend-item"><div class="swatch" style="background:#2e7d32;"></div>Prophage</div>
        <div class="legend-item"><div class="swatch" style="background:#e74c3c;"></div>Plasmid-borne</div>
        <div class="legend-item"><div class="swatch" style="background:#00f2ff; border-color:#aaa;"></div>ICE</div>
    </div>
</div>
${maps_section}
${full_amr_section}

<div class="card">
    <h2>Software Versions &amp; Traceability</h2>
    <div class="version-grid">
        <table>
            <thead><tr><th colspan="2">Assembly &amp; Polishing</th></tr></thead>
            <tr><td>Flye</td><td>$flye_v</td></tr>
            <tr><td>Medaka</td><td>$medaka_v</td></tr>
            <tr><td>Polypolish</td><td>$polypolish_v</td></tr>
        </table>
        <table>
            <thead><tr><th colspan="2">Processing</th></tr></thead>
            <tr><td>Fastp</td><td>$fastp_v</td></tr>
            <tr><td>Filtlong</td><td>$filtlong_v</td></tr>
            <tr><td>BWA-MEM</td><td>$bwa_v</td></tr>
        </table>
        <table>
            <thead><tr><th colspan="2">Annotation &amp; Profiling</th></tr></thead>
            <tr><td>AMRFinder+</td><td>$amrfinder_v</td></tr>
            <tr><td>Bakta</td><td>$bakta_v</td></tr>
            <tr><td>Kraken2</td><td>$kraken_v</td></tr>
        </table>
    </div>
</div>

<script>
    const MAIN_ID = "${MAIN_ID_VAL}";
    const pData = ${PHAGE_JS_DATA:-[]};
    const rData = ${RESISTOME_JS_DATA:-[]};

    // Populate mobile resistome table
    const resBody = document.querySelector('#resTable tbody');
    if (resBody && rData.length > 0) {
        rData.forEach(r => {
            let cls = r.class || 'Other';
            let type = (r.type || "").toLowerCase();
            let color = '#2c3e50'; // Default: Dark Grey

            // 1. ULTIMATE PRIORITY: Is it on a plasmid?
            // If the contig name doesn't match the main one, it's ALWAYS RED.
            if (r.contig && r.contig !== MAIN_ID) {
                color = '#e74c3c'; 
            } 
            // 2. SECOND PRIORITY: Is it Virulence?
            // If it's on the main chromosome and it's virulence, it's ALWAYS BLACK.
            else if (cls.toLowerCase().includes('virulence') || type === 'virulence') {
                color = '#000000';
            }
            // 3. THIRD PRIORITY: Standard colors for Main Chromosome
            else {
                const c = cls.toLowerCase();
                if (c.includes('beta-lactam') || c.includes('penicillin')) color = '#3498db';
                else if (c.includes('tetracycline'))   color = '#9b59b6';
                else if (c.includes('aminoglycoside')) color = '#e67e22'; 
                else if (c.includes('macrolide'))      color = '#f1c40f';
                else if (c.includes('quinolone'))      color = '#27ae60';
                else if (c.includes('ice') || c.includes('mobilome')) color = '#00f2ff';
                
                // Only if it's NOT virulence and ON the main chromosome do we use 
                // Red for "HIGH risk" flags from AMRFinder
                else if (r.risk === 'HIGH') color = '#e74c3c';
            }

            const tr = document.createElement('tr');
            tr.innerHTML =
                \`<td><span style="border-left:4px solid \${color};padding-left:10px;font-weight:600;">\${r.gene}</span></td>\` +
                \`<td>\${r.contig}</td><td>\${cls}</td><td>\${r.score || 'N/A'}</td>\` +
                \`<td><span class="\${r.risk === 'HIGH' ? 'risk-high' : ''}">\${r.risk}</span></td>\`;
            resBody.appendChild(tr);
        });
    }

    // Populate prophage table
    const phageBody = document.querySelector('#pTable tbody');
    if (phageBody && pData.length > 0) {
        pData.forEach(p => {
            const tr = document.createElement('tr');
            tr.innerHTML =
                \`<td>\${p.contig}</td><td>\${p.size}</td><td>\${p.score}</td>\` +
                \`<td><span class="phage-tag">Prophage</span></td>\`;
            phageBody.appendChild(tr);
        });
    }

    // Circular map zoom — follows cursor position
    document.querySelectorAll('.map-zoom-container').forEach(container => {
        const img = container.querySelector('img');
        container.addEventListener('mousemove', e => {
            if (!container.classList.contains('zoomed')) return;
            const r = container.getBoundingClientRect();
            img.style.transformOrigin =
                \`\${((e.clientX - r.left) / r.width) * 100}% \` +
                \`\${((e.clientY - r.top)  / r.height) * 100}%\`;
        });
        container.addEventListener('click', () => container.classList.toggle('zoomed'));
    });
</script>
</body>
</html>
EOF
}

# ==============================================================================
#  Argument parsing
# ==============================================================================
while [[ $# -gt 0 ]]; do
    case "$1" in
        --ont|-n)           ONT="$2";                      shift 2 ;;
        --r1|-1)            R1="$2";                       shift 2 ;;
        --r2|-2)            R2="$2";                       shift 2 ;;
        --genome-size|-g)   GSIZE="$2";                    shift 2 ;;
        --output|-o)        OUTDIR="$2";                   shift 2 ;;
        --model|-m)         MEDAKA_MODEL="$2";             shift 2 ;;
        --threads|-t)       THREADS="$2";                  shift 2 ;;
        --organism|-p)      ORGANISM="$2";                 shift 2 ;;
        --bakta-db|-b)      BAKTA_DB="$2";                 shift 2 ;;
        --bakta-ref|-a)     BAKTA_REGIONS="$2";            shift 2 ;;
        --complete|-i)      BAKTA_COMPLETE="--complete";   shift 1 ;;
        --kraken2-db|-k)     KRAKEN2_DB="$2"; RUN_KRAKEN2=true; shift 2 ;;
        --confidence|-c)    KRAKEN_CONFIDENCE="$2";        shift 2 ;;
        --ice|-I)           RUN_ICE=true;                  shift 1 ;;
        --plasticity|-s)    RUN_GENOMAD=true;              shift 1 ;;
        --genomad-db|-G)    GENOMAD_DB="$2";               shift 2 ;;
        --meta|-M)          USE_META="--meta";             shift 1 ;;
        --cleanup|-x)       CLEANUP=true;                  shift 1 ;;
        --help|-h)          usage ;;
        *) echo "[ERROR] Unknown argument: $1"; usage ;;
    esac
done

# ==============================================================================
#  Pre-flight validation
# ==============================================================================
MISSING=0

# ONT reads
if [[ -z "$ONT" ]]; then
    echo "[ERROR] ONT reads not specified (--ont)."; MISSING=1
elif [[ ! -f "$ONT" ]]; then
    echo "[ERROR] ONT file not found: $ONT"; MISSING=1
elif [[ ! -s "$ONT" ]]; then
    echo "[ERROR] ONT file is empty: $ONT"; MISSING=1
fi

# Genome size — required, exit immediately if missing
if [[ -z "$GSIZE" ]]; then
    echo "[ERROR] Genome size required (--genome-size, e.g. 2.1m)."; exit 1
fi
if ! [[ "$GSIZE" =~ ^[0-9]+(\.[0-9]+)?[kKmMgG]?$ ]]; then
    echo "[ERROR] Invalid genome size format: $GSIZE"; exit 1
fi

# ==============================================================================
#  Environment Validation & Strict Tool Auditing
# ==============================================================================

MISSING=0

# Enforce that the script is running inside a Conda environment
if [[ -z "${CONDA_PREFIX:-}" ]]; then
    echo "[ERROR] No Conda environment detected!"
    echo "Please run: conda activate mostar_env"
    exit 1
fi

# Sanitize the PATH to prioritize Conda and ignore local user bins (~/.local/bin)
export PATH="$CONDA_PREFIX/bin:$PATH"


# Base tools required for EVERY run
REQUIRED_TOOLS=(
    "flye" "medaka_consensus" "samtools" "amrfinder" "circlator" "filtlong"
)

# Conditionally require tools based on user flags
if [[ -n "$R1" && -n "$R2" ]]; then
    REQUIRED_TOOLS+=("fastp" "bwa" "polypolish")
fi
[[ "$SKIP_BAKTA" == "false" ]] && REQUIRED_TOOLS+=("bakta")
[[ "$RUN_ICE" == "true" ]]     && REQUIRED_TOOLS+=("macsyfinder")
[[ "$RUN_GENOMAD" == "true" ]] && REQUIRED_TOOLS+=("genomad")
[[ "$RUN_KRAKEN2" == "true" ]] && REQUIRED_TOOLS+=("kraken2")

# Verify tools exist AND live inside the active Conda environment
for tool in "${REQUIRED_TOOLS[@]}"; do
    TOOL_PATH=$(command -v "$tool" || true)
    
    if [[ -z "$TOOL_PATH" ]]; then
        echo "[ERROR] Required tool not found: $tool"
        MISSING=1
    elif [[ "$TOOL_PATH" != "$CONDA_PREFIX"* ]]; then
        echo "[ERROR] Tool hijack detected! '$tool' is loading from outside the environment:"
        echo "        Found at: $TOOL_PATH"
        echo "        Expected: $CONDA_PREFIX/bin/$tool"
        MISSING=1
    fi
done

# Output directory & Database Checks
if [[ -z "$OUTDIR" ]]; then
    OUTDIR="MOSTAR_$(date +%Y%m%d_%H%M%S)"
    echo "[WARNING] No output directory specified — using: $OUTDIR"
fi
mkdir -p "$OUTDIR"/{logs,intermediate,amr_results}

if [[ -z "$BAKTA_DB" || ! -d "$BAKTA_DB" ]]; then
    echo "[WARNING] Bakta database not found — annotation will be skipped."
    SKIP_BAKTA=true
fi
if [[ -z "$KRAKEN2_DB" || ! -d "$KRAKEN2_DB" ]]; then
    echo "[INFO] Kraken2 database not found — taxonomy will be skipped."
    RUN_KRAKEN2=false
fi
if [[ -z "$GENOMAD_DB" || ! -d "$GENOMAD_DB" ]]; then
    echo "[INFO] geNomad database not found — plasticity analysis will be skipped."
    RUN_GENOMAD=false
fi

# Python dependencies
if ! python3 -c "import pandas, matplotlib, Bio" &>/dev/null; then
    echo "[ERROR] Missing Python packages. Run: pip install pandas matplotlib biopython"
    MISSING=1
fi

if [[ "$MISSING" -eq 1 ]]; then
    echo "------------------------------------------------------------------------"
    echo "One or more required arguments or tools are missing — cannot continue."
    echo "Run: mostar --help"
    exit 1
fi

# ==============================================================================
#  Workspace initialisation
# ==============================================================================
REAL_OUT=$(cd "$OUTDIR" && pwd)
FINAL_ASM="$REAL_OUT/MOSTAR_Assembly.fasta"
MODE_LABEL=$([[ -n "$R1" && -n "$R2" ]] && echo "Hybrid Mode" || echo "ONT-only Mode")
KRAKEN_REPORT="$OUTDIR/taxonomy/assembly_kraken_report.txt"

# Filtlong target bases: genome_size_number × coverage_multiplier
NUM=$(echo  "$GSIZE" | tr -d '[:alpha:][:space:]')
UNIT=$(echo "$GSIZE" | tr -d '0-9.'    | tr '[:lower:]' '[:upper:]' | tr -d '[:space:]')
TARGET_BASES=$(awk -v n="$NUM" -v c="$FILTLONG_COV" 'BEGIN { printf "%.0f", n * c }')$UNIT

echo ""
echo "              __      __   __    ____  _______   __     _____             "
echo "              | |\  /| |  /  \  / ___| |_   _|  /  \   | __  \            "
echo "              | | \/ | | | || | \___ \   | |   / /\ \  | |_) /            "
echo "              |_|\__/|_|  \__/  |____/   |_|  /_/  \_\ |_\ \_\            "
echo ""
echo -e "\n${BOLD}${BLUE}>>> MOSTAR v${VERSION} — ${MODE_LABEL}${NC}\n"

# ==============================================================================
#  Step 1 — Short-read QC (Fastp)
# ==============================================================================
if [[ -n "$R1" && -n "$R2" ]]; then
    run_step "Short-read QC (Fastp)" "true" \
        "fastp \
            --in1 '$R1' --in2 '$R2' \
            --out1 '$OUTDIR/intermediate/sh_R1.fq.gz' \
            --out2 '$OUTDIR/intermediate/sh_R2.fq.gz' \
            --thread $THREADS \
            --json '$OUTDIR/logs/fastp.json' \
            --html '$OUTDIR/logs/fastp.html' \
            >/dev/null 2>&1"
else
    run_step "Short-read QC (Fastp)" "false" ""
fi

# ==============================================================================
#  Step 2 — Long-read QC (Filtlong)
# ==============================================================================
FILT_REF=""
[[ -n "$R1" && -n "$R2" ]] && FILT_REF="-1 '$R1' -2 '$R2'"

run_step "Long-read QC (Filtlong)" "true" \
    "filtlong --min_length 1000 --keep_percent 95 -t $TARGET_BASES $FILT_REF '$ONT' \
        2>'$OUTDIR/logs/filtlong.log' \
        | gzip > '$OUTDIR/intermediate/ont_filt.fq.gz'"

# ==============================================================================
#  Step 3 — De-novo Assembly (Flye)
# ==============================================================================
run_step "De-novo Assembly (Flye)" "true" \
    "flye \
        --nano-hq '$OUTDIR/intermediate/ont_filt.fq.gz' \
        -g $GSIZE -t $THREADS ${USE_META:-} \
        -o '$OUTDIR/flye' \
        >'$OUTDIR/logs/flye.log' 2>&1"

# ==============================================================================
#  Step 4 — ONT Polishing (Medaka)
# ==============================================================================
run_step "ONT Polishing (Medaka)" "true" \
    "medaka_consensus \
        -i '$OUTDIR/intermediate/ont_filt.fq.gz' \
        -d '$OUTDIR/flye/assembly.fasta' \
        -m $MEDAKA_MODEL -t $THREADS \
        -b 32 \
        -o '$OUTDIR/medaka' \
        >'$OUTDIR/logs/medaka.log' 2>&1"

# ==============================================================================
#  Step 5 — Hybrid Polishing (Polypolish)  [hybrid mode only]
# ==============================================================================
if [[ -n "$R1" && -n "$R2" ]]; then
    POLY_CMD="\
        bwa index '$OUTDIR/medaka/consensus.fasta' >/dev/null 2>&1 && \
        bwa mem -t $THREADS '$OUTDIR/medaka/consensus.fasta' \
            '$OUTDIR/intermediate/sh_R1.fq.gz' \
            >'$OUTDIR/intermediate/1.sam' 2>/dev/null && \
        bwa mem -t $THREADS '$OUTDIR/medaka/consensus.fasta' \
            '$OUTDIR/intermediate/sh_R2.fq.gz' \
            >'$OUTDIR/intermediate/2.sam' 2>/dev/null && \
        polypolish filter \
            --in1 '$OUTDIR/intermediate/1.sam' \
            --in2 '$OUTDIR/intermediate/2.sam' \
            --out1 '$OUTDIR/intermediate/f1.sam' \
            --out2 '$OUTDIR/intermediate/f2.sam' \
            >/dev/null 2>&1 && \
        polypolish polish \
            '$OUTDIR/medaka/consensus.fasta' \
            '$OUTDIR/intermediate/f1.sam' \
            '$OUTDIR/intermediate/f2.sam' \
            >'$FINAL_ASM' \
            2>'$OUTDIR/logs/polypolish.log'"

    run_step "Hybrid Polishing (Polypolish)" "true" "$POLY_CMD"

    POS_CHANGED=$(grep "positions changed" "$OUTDIR/logs/polypolish.log" \
                  | awk '{print $1}' | tr -d ',' || echo "0")
    MEAN_DEPTH=$(grep "mean read depth" "$OUTDIR/logs/polypolish.log" \
                 | awk '{print $4}' | tr -d 'x' || echo "0")
    BWA_MAPPING_RATE=$(samtools flagstat "$OUTDIR/intermediate/1.sam" \
                       | awk '/^[0-9]+ \+ 0 mapped/ {
                             match($0, /\(([0-9.]+)%/, a); print a[1] }' \
                       || echo "0")
else
    cp "$OUTDIR/medaka/consensus.fasta" "$FINAL_ASM"
    run_step "Hybrid Polishing (Polypolish)" "false" ""
fi

# ==============================================================================
#  Step 6 — Standardise Assembly Origin (Circlator)
# ==============================================================================
if command -v circlator &>/dev/null; then
    run_step "Fix Assembly Origin (Circlator)" "true" \
        "circlator fixstart '$FINAL_ASM' '$OUTDIR/intermediate/fixed_origin' \
        >'$OUTDIR/logs/circlator.log' 2>&1 || true"

    if [[ -f "$OUTDIR/intermediate/fixed_origin.fasta" ]]; then
        mv "$OUTDIR/intermediate/fixed_origin.fasta" "$FINAL_ASM"
    else
        echo "[INFO] Circlator did not find dnaA. Assembly left un-rotated." >> "$OUTDIR/logs/circlator.log"
    fi
else
    run_step "Standardising Origin (Circlator)" "false" ""
fi


# ==============================================================================
#  Step 7 — Whole-genome Taxonomy (Kraken2)  [optional]
# ==============================================================================
if [[ "$RUN_KRAKEN2" == "true" ]]; then
    mkdir -p "$OUTDIR/taxonomy"
    run_step "Taxonomy (Kraken2)" "true" \
        "kraken2 --db '$KRAKEN2_DB' --threads $THREADS --confidence '$KRAKEN_CONFIDENCE' --report '$KRAKEN_REPORT' '$FINAL_ASM' >/dev/null 2>&1"

    if [[ -s "$KRAKEN_REPORT" ]]; then
        # Prefer Subspecies/Strain (S1, S2...) if available
        TOP_HIT=$(awk -F'\t' '$4 == "S" || $4 == "S1"' "$KRAKEN_REPORT" | sort -k1,1nr | head -n1 || true)
        
        # Fallback to general Species (S)
        [[ -z "$TOP_HIT" ]] && \
            TOP_HIT=$(awk -F'\t' '$4 == "G"' "$KRAKEN_REPORT" | sort -k1,1nr | head -n1 || true)
        
        [[ -n "$TOP_HIT" ]] && \
            ASM_SPECIES_FULL=$(echo "$TOP_HIT" | cut -f6- | sed 's/^[[:space:]]*//')
    fi
else
    run_step "Taxonomy (Kraken2)" "false" ""
fi

# ==============================================================================
#  Step 8 — Functional Annotation (Bakta)  [optional]
# ==============================================================================
if [[ "$SKIP_BAKTA" == "false" && -n "$BAKTA_DB" ]]; then
    run_step "Annotation (Bakta)" "true" \
        "bakta \
            --output '$OUTDIR/annotation' \
            --db '$BAKTA_DB' \
            --threads $THREADS \
            --force \
            ${BAKTA_COMPLETE:-} \
            ${BAKTA_REGIONS:+--regions '$BAKTA_REGIONS'} \
            '$FINAL_ASM' \
            >'$OUTDIR/logs/bakta.log' 2>&1"
    
else
    run_step "Annotation (Bakta)" "false" ""
fi

# ==============================================================================
#  Step 9 — AMR Profiling (AMRFinder+)
# ==============================================================================
run_step "AMR Profiling (AMRFinder+)" "true" \
    "amrfinder --threads $THREADS --plus \
        ${ORGANISM:+--organism '$ORGANISM'} \
        -n '$FINAL_ASM' \
        >'$OUTDIR/amr_results/AMR_Report.tsv' \
        2>'$OUTDIR/logs/amrfinder.log'"

# ==============================================================================
#  Step 10 — ICE Detection (MacSyFinder / CONJScan)  [optional]
# ==============================================================================
ICE_FILE="$OUTDIR/amr_results/ICE_locations.tmp"

if [[ "$RUN_ICE" == "true" ]]; then
    P_FILE=$(ls "$OUTDIR/annotation/"*.faa 2>/dev/null | head -n1 || true)
    if [[ -f "$P_FILE" ]]; then
        run_step "ICE Detection (MacSyFinder)" "true" \
            "rm -rf '$OUTDIR/ice_detection' && macsyfinder \
                --db-type ordered_replicon \
                --models CONJScan \
                -o '$OUTDIR/ice_detection' \
                --sequence-db '$P_FILE' \
                --worker $THREADS \
                >'$OUTDIR/logs/macsyfinder.log' 2>&1"
    else
        echo "[WARNING] No .faa file in $OUTDIR/annotation/ — skipping ICE detection."
        run_step "ICE Detection (MacSyFinder)" "false" ""
    fi
else
    run_step "ICE Detection (MacSyFinder)" "false" ""
fi

# Parse MacSyFinder output → ICE_locations.tmp (used by the circular map)
if [[ "$RUN_ICE" == "true" && -d "$OUTDIR/ice_detection" ]]; then
    BEST_SOL=$(find "$OUTDIR/ice_detection" -name "best_solution.tsv"         2>/dev/null | head -n1 || true)
    BEST_SUM=$(find "$OUTDIR/ice_detection" -name "best_solution_summary.tsv" 2>/dev/null | head -n1 || true)

    if [[ -f "$BEST_SUM" && -f "$BEST_SOL" ]]; then

python3 - <<PYICE
import os, pandas as pd, re
from io import StringIO

def read_msf(path):
    with open(path) as f:
        lines = [l for l in f if not l.startswith('#') and l.strip()]
    return pd.read_csv(StringIO(''.join(lines)), sep='\t')

summary = read_msf('$BEST_SUM')
active_systems = [c for c in summary.columns if c != 'replicon' and summary[c].sum() > 0]
if not active_systems: exit(0)

best = read_msf('$BEST_SOL')
best = best[best['model_fqn'].isin(active_systems)]

# Get the source of truth from GFF (maps locus_tag to real contig ID)
gff_path = next(
    (os.path.join(root, fname)
     for root, dirs, files in os.walk('$OUTDIR/annotation')
     for fname in files
     if fname.endswith('.gff') or fname.endswith('.gff3')),
    None
)

coord_map = {}
if gff_path:
    with open(gff_path) as f:
        for line in f:
            if line.startswith('#') or '\t' not in line: continue
            p = line.strip().split('\t')
            lt = re.search(r'locus_tag=([^;]+)', p[8])
            if lt: coord_map[lt.group(1)] = (p[0], int(p[3]), int(p[4]), p[6])

# Build regions using the REAL contig name from the GFF (e.g., contig_1)
out_rows = []
for sys_id, group in best.groupby('model_fqn'):
    hits = [coord_map[hid] for hid in group['hit_id'] if hid in coord_map]
    if hits:
        out_rows.append([hits[0][0], sys_id.split('/')[-1], min(h[1] for h in hits), max(h[2] for h in hits), hits[0][3]])

if out_rows:
    pd.DataFrame(out_rows).to_csv('$ICE_FILE', sep='\t', index=False, header=False)
PYICE
    else
        echo "[WARNING] MacSyFinder output files are incomplete — cannot parse ICE locations."
    fi
fi    


# ==============================================================================
#  Step 11 — Genomic Plasticity (geNomad)  [optional]
# ==============================================================================
PHAGE_FILE="$OUTDIR/amr_results/prophage_locations.tmp"
export PHAGE_FILE
: > "$PHAGE_FILE"

if [[ "$RUN_GENOMAD" == "true" ]]; then
    run_step "Plasticity (geNomad)" "true" \
        "genomad end-to-end \
            --cleanup --splits 8 --threads $THREADS \
            '$FINAL_ASM' '$OUTDIR/genomad' '$GENOMAD_DB' \
            >'$OUTDIR/logs/genomad.log' 2>&1"

    P_SUMM=$(find "$OUTDIR/genomad" -name "*_plasmid_summary.tsv" 2>/dev/null | head -n1 || true)
    V_SUMM=$(find "$OUTDIR/genomad" -name "*_virus_summary.tsv"   2>/dev/null | head -n1 || true)

    # Mobile resistome — merge AMR hits with plasmid calls
    if [[ -f "$P_SUMM" ]]; then
        AMR_R="$OUTDIR/amr_results/AMR_Report.tsv"
        M_OUT="$OUTDIR/intermediate/mob_sum.csv"
        export AMR_R P_SUMM M_OUT

        python3 - <<'PYMERGE'
import pandas as pd, os
amr  = pd.read_csv(os.environ['AMR_R'],  sep='\t')
plas = pd.read_csv(os.environ['P_SUMM'], sep='\t')
amr['m_id']  = amr.iloc[:,  1].astype(str).str.split().str[0]
plas['m_id'] = plas.iloc[:, 0].astype(str).str.split('|').str[0].str.split().str[0]
merged = amr.merge(plas[['m_id', 'plasmid_score']], on='m_id', how='inner')
merged.to_csv(os.environ['M_OUT'], index=False)
PYMERGE

        if [[ -f "$M_OUT" ]]; then
            RESISTOME_JS_DATA=$(python3 - <<PYJSON
import pandas as pd, json
df = pd.read_csv('$M_OUT')
out = df.apply(lambda r: {
    'gene':   r.iloc[5],
    'contig': r.iloc[1],
    'class':  r.iloc[10],
    'type':   r.iloc[8],
    'score':  f"{r.iloc[-1]:.2f}",
    'risk':   'HIGH' if r.iloc[-1] >= 0.8 else 'LOW'
}, axis=1).tolist()
print(json.dumps(out))
PYJSON
)
            HIGH_RISK_COUNT=$(python3 - <<PYCOUNT
import pandas as pd
df = pd.read_csv('$M_OUT')
print((df.iloc[:, -1] >= 0.8).sum())
PYCOUNT
)
        fi
    fi

    # Prophage stats & map coordinate extraction
    if [[ -f "$V_SUMM" ]]; then
        PHAGE_COUNT=$(awk 'NR>1' "$V_SUMM" | wc -l | xargs)
        
        # Parse coordinates for the circular map
        python3 - <<PYPHAGE_MAP
import pandas as pd, os
df = pd.read_csv('$V_SUMM', sep='\t')
out_rows = []
for _, r in df.iterrows():
    # Split geNomad "start-end" string
    coords = str(r['coordinates']).split('-')
    start = coords[0] if len(coords) > 1 else 1
    end = coords[1] if len(coords) > 1 else r['length']
    contig = str(r['seq_name']).split('|')[0].split()[0]
    out_rows.append([contig, start, end, r.get('virus_score', 0), "+"])
if out_rows:
    pd.DataFrame(out_rows).to_csv(os.environ['PHAGE_FILE'], sep='\t', index=False, header=False)
PYPHAGE_MAP

        # JS data for the HTML report
        PHAGE_JS_DATA=$(python3 - <<PYPHAGE
import pandas as pd, json
df = pd.read_csv('$V_SUMM', sep='\t')
out = df.apply(lambda r: {
    'contig': r['seq_name'],
    'size':   f"{r['length']:,} bp",
    'score':  f"{r['virus_score']:.2f}"
}, axis=1).tolist()
print(json.dumps(out))
PYPHAGE
)
    fi
else
    run_step "Plasticity (geNomad)" "false" ""
fi

# ==============================================================================
#  Step 12 — Circular Genome Maps (Python/Matplotlib)
# ==============================================================================
export OUTDIR FINAL_ASM ASM_SPECIES_FULL \
       AMR_FILE="$OUTDIR/amr_results/AMR_Report.tsv" \
       MAP_DIR="$OUTDIR/amr_results/maps" \
       ICE_FILE="$OUTDIR/amr_results/ICE_locations.tmp" \
       PHAGE_FILE="$OUTDIR/amr_results/prophage_locations.tmp" \
       MOB_FILE="$OUTDIR/intermediate/mob_sum.csv"

VIZ_SCRIPT="$OUTDIR/intermediate/generate_maps.py"

cat <<'PYEOF' > "$VIZ_SCRIPT"
import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
from matplotlib.patches import Polygon
from Bio import SeqIO

# Robust GC calculator that completely avoids BioPython version compatibility issues
def calc_gc(seq):
    s = str(seq).upper()
    return (s.count('G') + s.count('C')) / max(1, len(s))

OUTDIR   = os.environ['OUTDIR']
FA       = os.environ['FINAL_ASM']
AMR_FILE = os.environ['AMR_FILE']
MAP_DIR  = os.environ['MAP_DIR']
ICE_FILE = os.environ['ICE_FILE']
PHAGE_FILE = os.environ['PHAGE_FILE']
MOB_FILE = os.environ.get('MOB_FILE', '')

os.makedirs(MAP_DIR, exist_ok=True)

# 1. Load high-risk mobile genes
high_risk_genes = set()
if os.path.exists(MOB_FILE) and os.path.getsize(MOB_FILE) > 0:
    try:
        m_df = pd.read_csv(MOB_FILE)
        # Score is last column (iloc[:, -1]) and gene name is iloc[:, 5]
        high_risk_genes = set(m_df[m_df.iloc[:, -1] >= 0.8].iloc[:, 5].astype(str))
    except Exception:
        pass


def amr_color(amr_class, amr_type="AMR", gene_name="", is_main=True):
    # ABSOLUTE PRIORITY: If it's on a plasmid, everything is ALWAYS Red
    if not is_main:
        return '#e74c3c'
        
    c = str(amr_class).lower()
    t = str(amr_type).lower()

    # Main Chromosome: Virulence is Black
    if 'virulence' in c or t == 'virulence': 
        return '#000000'
    
    # Main Chromosome: High-risk genes are Red
    try:
        if gene_name in high_risk_genes:
            return '#e74c3c'
    except NameError:
        pass

    # Main Chromosome: Standard AMR Colors
    if any(x in c for x in ['beta-lactam', 'penicillin', 'carbapenem']): return '#3498db'
    if 'tetracycline'   in c: return '#9b59b6'
    if 'aminoglycoside' in c: return '#e67e22'
    if any(x in c for x in ['macrolide', 'lincosamide', 'streptogramin']): return '#f1c40f'
    if 'quinolone' in c or 'fluoroquinolone' in c: return '#27ae60'
    if 'ice' in c or 'conjugation' in c:    return '#00f2ff'
    
    return '#2c3e50'


def draw_map(record, filename, centre_label, ice_df, phage_df, amr_df, is_main=False):
    gl        = len(record.seq)
    cid_clean = str(record.id).split()[0]
    two_pi    = 2 * np.pi

    fig = plt.figure(figsize=(10, 10))
    ax  = fig.add_subplot(111, projection='polar')
    ax.set_theta_zero_location('N')
    ax.set_theta_direction(-1)  # Clockwise direction
    ax.spines['polar'].set_visible(False)
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.grid(False)
    ax.set_ylim(0, 1.45)

    # Radial guide lines + bp position labels (8 divisions)
    for rad in np.linspace(0, two_pi, 8, endpoint=False):
        ax.plot([rad, rad], [0.3, 1.1], color='#d1d8e0', lw=0.8, ls='--', zorder=0)
        bp = int((rad / two_pi) * gl)
        if   bp == 0:         lbl = "0"
        elif bp >= 1_000_000: lbl = f"{bp/1_000_000:.1f} Mb"
        elif bp >= 1_000:     lbl = f"{bp/1_000:.0f} kb"
        else:                 lbl = f"{bp} bp"
        ax.text(rad, 1.25, lbl, ha='center', va='center',
                fontsize=6.5, color='#95a5a6', fontweight='500')

    # GC content (innermost ring, radius ~0.70)
    step     = max(1, gl // 500)
    gc_vals  = np.array([calc_gc(record.seq[j:j+step]) for j in range(0, gl, step)])
    theta_gc = np.linspace(0, two_pi, len(gc_vals))
    avg_gc   = np.mean(gc_vals)
    ax.fill_between(theta_gc,
                    0.70,
                    0.70 + (gc_vals - avg_gc) * 0.60,
                    color='#bdc3c7', alpha=0.6, zorder=2)

    # ICE / mobilome track (middle ring, radius 0.85)
    ax.plot(np.linspace(0, two_pi, 200), [0.85] * 200,
            color='#edf2f7', lw=10, alpha=0.3, zorder=1)

    if not ice_df.empty:
        contig_ice = ice_df[ice_df['Contig'].astype(str).str.split().str[0] == cid_clean]
        if contig_ice.empty and record == records[0]:
            contig_ice = ice_df

        for _, row in contig_ice.iterrows():
            s = (float(row['Start']) / gl) * two_pi
            e = (float(row['Stop'])  / gl) * two_pi
            mid = (s + e) / 2
            strand = str(row['Strand']).strip()
            size_kb = (float(row['Stop']) - float(row['Start'])) / 1000

            # ICE arc
            ax.plot(np.linspace(s, e, 100), [0.85] * 100,
                    color='#00f2ff', lw=8, alpha=0.9,
                    solid_capstyle='butt', zorder=9)

            # Arrow indicating direction
            arrow_offset = 0.025 if strand == '+' else -0.025 
            ax.annotate('',
                        xy=(mid + arrow_offset, 0.85),
                        xytext=(mid, 0.85),
                        arrowprops=dict(arrowstyle='-|>', color='#006080', lw=1.2, mutation_scale=16),
                        zorder=10)

            # ICE label
            label_text = f"{str(row['Symbol']).replace('_', ' ')}\n{size_kb:.1f} kb"
            ax.text(mid, 0.94, label_text,
                    ha='center', va='bottom', fontsize=6.5,
                    color='#006080', fontweight='bold',
                    path_effects=[pe.withStroke(linewidth=2, foreground='white')],
                    zorder=11)

    # Prophage integration 
    if not phage_df.empty:
        contig_phage = phage_df[phage_df['Contig'].astype(str).str.split().str[0] == cid_clean]

        if contig_phage.empty and is_main:
            contig_phage = phage_df

        for _, row in contig_phage.iterrows():
            s = (float(row['Start']) / gl) * two_pi
            e = (float(row['Stop'])  / gl) * two_pi
            mid = (s + e) / 2
            size_kb = (float(row['Stop']) - float(row['Start'])) / 1000

            color = '#2e7d32'

            ax.plot(np.linspace(s, e, 100), [0.865] * 100,
                    color=color, lw=8, alpha=0.85,
                    solid_capstyle='butt', zorder=8)

            label_text = f"Prophage\n{size_kb:.1f} kb"
            ax.text(mid, 0.78, label_text,
                    ha='center', va='top', fontsize=6.0,
                    color=color,
                    path_effects=[pe.withStroke(linewidth=2, foreground='white')],
                    zorder=9)

    # AMR track (outermost ring, radius 1.05)
    ax.plot(np.linspace(0, two_pi, 200), [1.05] * 200,
            color='#edf2f7', lw=16, alpha=0.4, zorder=1)

    if not amr_df.empty:
        # Filter for the current contig
        hits = amr_df[amr_df.iloc[:, 1].astype(str).str.split().str[0] == cid_clean]

        arcs = []
        for _, row in hits.iterrows():
            gene_name = str(row.iloc[5]) 
            color = amr_color(row.iloc[10], row.iloc[8], gene_name, is_main)
            s_rad  = (float(row.iloc[2]) / gl) * two_pi
            e_rad  = (float(row.iloc[3]) / gl) * two_pi
            # Safe strand parsing (column 4 is Strand in AMRFinder+)
            strand = str(row.iloc[4]).strip() if pd.notna(row.iloc[4]) else '+'
            arcs.append((s_rad, e_rad, (s_rad + e_rad) / 2, color, gene_name, strand))

        # DRAW BLOCKS + EXTENDING ARROWS
        for s_rad, e_rad, mid, color, name, strand in arcs:
            angular_len = abs(e_rad - s_rad)

            # Draw the thick colored block
            ax.plot(np.linspace(s_rad, e_rad, 100), [1.05] * 100,
                    color=color, lw=12, solid_capstyle='butt', zorder=10)

            # Dynamic arrowhead placement based on visual length
            if angular_len > 0.02:
                if strand == '+':
                    start_theta = mid - 0.005
                    end_theta   = mid + 0.005
                else:
                    start_theta = mid + 0.005
                    end_theta   = mid - 0.005

                ax.annotate('',
                            xy=(end_theta, 1.05),
                            xytext=(start_theta, 1.05),
                            arrowprops=dict(arrowstyle='-|>', color='black', lw=1.5, mutation_scale=12),
                            zorder=11)
            else:
                # Short gene: Extend block with an arrow of the same color
                if strand == '+':
                    start_theta = e_rad
                    end_theta   = e_rad + 0.030
                else:
                    start_theta = s_rad
                    end_theta   = s_rad - 0.030

                ax.annotate('',
                            xy=(end_theta, 1.05),
                            xytext=(start_theta, 1.05),
                            arrowprops=dict(arrowstyle='-|>', color=color, lw=1.2, mutation_scale=18),
                            zorder=11)

        arcs.sort(key=lambda x: x[2])
        MIN_SEP = 0.18
        placed  = []
        for s_rad, e_rad, mid, color, name, strand in arcs:
            level = 0
            for p_ang, p_lvl in placed[-6:]:
                dist = min(abs(mid - p_ang), two_pi - abs(mid - p_ang))
                if dist < MIN_SEP and p_lvl == level:
                    level += 1
            l_rad = 1.16 + level * 0.07
            placed.append((mid, level))
            if level > 0:
                ax.plot([mid, mid], [1.10, l_rad - 0.02],
                        color=color, lw=0.6, alpha=0.5, zorder=9)
            ax.text(mid, l_rad, name,
                    color=color, fontsize=6.5, fontweight='bold',
                    ha='center', va='bottom',
                    path_effects=[pe.withStroke(linewidth=1.8, foreground='white')],
                    zorder=11)

    ax.text(0, 0, centre_label,
            ha='center', va='center', fontsize=11, fontweight='bold',
            color='#1a2a3a')

    plt.savefig(os.path.join(MAP_DIR, filename),
                dpi=300, bbox_inches='tight', transparent=True)
    plt.close(fig)


try:
    ice_df = pd.read_csv(ICE_FILE, sep='\t',
                         names=['Contig', 'Symbol', 'Start', 'Stop', 'Strand']) \
             if os.path.exists(ICE_FILE) and os.path.getsize(ICE_FILE) > 0 \
             else pd.DataFrame()

    phage_df = pd.read_csv(PHAGE_FILE, sep='\t',
                       names=['Contig', 'Start', 'Stop', 'Score', 'Strand']) \
           if os.path.exists(PHAGE_FILE) and os.path.getsize(PHAGE_FILE) > 0 \
           else pd.DataFrame()

    amr_df = pd.read_csv(AMR_FILE, sep='\t').fillna('nan') \
             if os.path.exists(AMR_FILE) \
             else pd.DataFrame()

    # Sort all records by length so the largest is always at index 0
    records = sorted(list(SeqIO.parse(FA, 'fasta')),
                     key=lambda r: len(r.seq), reverse=True)

    # Identify largest one
    main_id = str(records[0].id).split()[0]
    
    species = os.environ.get('ASM_SPECIES_FULL', 'Assembly')
    
    # Draw the Main Chromosome 
    chrom = records[0]
    chrom_gc = calc_gc(chrom.seq) * 100 
    
    draw_map(chrom,
            filename="01_Main_Chromosome.png",
            centre_label=f"{len(chrom.seq):,} bp\nGC: {chrom_gc:.1f}%\n{species}", 
            ice_df=ice_df,
            phage_df=phage_df,
            amr_df=amr_df,
            is_main=True) # Tells the color function to use the rainbow scale

    # Draw everything else as Plasmids (Red)
    for i, rec in enumerate(records[1:], start=1):
        rec_gc = calc_gc(rec.seq) * 100 
        draw_map(rec,
                filename=f"{i+1:02d}_Extrachromosomal_{i}.png",
                centre_label=f"Plasmid\n{len(rec.seq)/1000:.1f} kb\n{str(rec.id).split()[0]}",
                ice_df=ice_df,
                phage_df=phage_df,
                amr_df=amr_df,
                is_main=False) 

except Exception as exc:
    import traceback, sys
    with open(os.path.join(OUTDIR, 'logs', 'viz_error.log'), 'w') as fh:
        fh.write(traceback.format_exc())
    sys.exit(1)
PYEOF

run_step "Generating Circular Maps" "true" "python3 '$VIZ_SCRIPT'"

# ==============================================================================
#  Assembly statistics
# ==============================================================================
M_FILE="$OUTDIR/medaka/consensus.fasta"
[[ ! -f "$M_FILE" ]] && M_FILE="$OUTDIR/flye/assembly.fasta"

if [[ -f "$M_FILE" ]]; then
    read -r MEDAKA_LEN MEDAKA_CNT MEDAKA_N50 _M_L50 \
        <<< "$(calc_assembly_stats "$M_FILE")"
    MEDAKA_LEN=${MEDAKA_LEN:-0}
    MEDAKA_CNT=${MEDAKA_CNT:-0}
    MEDAKA_N50=${MEDAKA_N50:-0}
else
    MEDAKA_LEN=0; MEDAKA_CNT=0; MEDAKA_N50=0
fi

if [[ -f "$FINAL_ASM" && -s "$FINAL_ASM" ]]; then
    read -r FINAL_LEN FINAL_CNT FINAL_N50 _F_L50 \
        <<< "$(calc_assembly_stats "$FINAL_ASM")"
    FINAL_LEN=${FINAL_LEN:-0}
    FINAL_CNT=${FINAL_CNT:-0}
    FINAL_N50=${FINAL_N50:-0}
else
    FINAL_LEN=0; FINAL_CNT=0; FINAL_N50=0
fi

# ==============================================================================
#  Build HTML component strings
# ==============================================================================
TAX_CARDS="<div class='qc-card'>
    <span class='qc-label'>Identified Species</span>
    <span class='qc-val'>${ASM_SPECIES_FULL}</span>
</div>"

if [[ -n "$R1" ]]; then
    HYBRID_CARDS="<div class='qc-card'>
        <span class='qc-label'>Hybrid Polishing</span>
        <span class='qc-val'>${POS_CHANGED} Changes</span>
        <div style='font-size:0.8rem; color:#7f8c8d; margin-top:4px;'>
            Depth: ${MEAN_DEPTH}x &nbsp;|&nbsp; Mapping: ${BWA_MAPPING_RATE}%
        </div>
    </div>"
else
    HYBRID_CARDS=""
fi

METRICS_BLOCK="<div class='card'>
  <h2>Assembly Statistics</h2>
  <table>
    <thead>
      <tr><th>Metric</th><th>After Medaka (ONT polish)</th><th>Final (polished)</th></tr>
    </thead>
    <tbody>
      <tr>
        <td>Total Length</td>
        <td>$(format_bp "$MEDAKA_LEN") bp</td>
        <td><span class='val-new'>$(format_bp "$FINAL_LEN") bp</span></td>
      </tr>
      <tr>
        <td>Contig Count</td>
        <td>$MEDAKA_CNT</td>
        <td><span class='val-new'>$FINAL_CNT</span></td>
      </tr>
      <tr>
        <td>N50</td>
        <td>$(format_bp "$MEDAKA_N50") bp</td>
        <td><span class='val-new'>$(format_bp "$FINAL_N50") bp</span></td>
      </tr>
    </tbody>
  </table>
</div>"

# ==============================================================================
#  Step 13 — HTML Report
# ==============================================================================
run_step "Generating HTML Report" "true" "generate_html_report"

# ==============================================================================
#  Step 14 — Cleanup Intermediate Files (Optional)
# ==============================================================================
if [[ "$CLEANUP" == "true" ]]; then
    run_step "Cleaning up intermediate files" "true" \
        "rm -rf '$OUTDIR/intermediate' '$OUTDIR/amr_results/ICE_locations.tmp' '$OUTDIR/amr_results/prophage_locations.tmp'"
else
    run_step "Cleaning up intermediate files" "false" ""
fi

# ==============================================================================
#  Terminal summary
# ==============================================================================
printf '%-18s %-24s %-24s\n' 'Metric'        'After Medaka'                  'Final'
printf '%-18s %-24s %-24s\n' '──────────────' '────────────────────────'      '────────────────────────'
printf '%-18s %-24s %-24s\n' 'Total Length'  "$(format_bp "$MEDAKA_LEN") bp" "$(format_bp "$FINAL_LEN") bp"
printf '%-18s %-24s %-24s\n' 'Contig Count'  "$MEDAKA_CNT"                   "$FINAL_CNT"
printf '%-18s %-24s %-24s\n' 'N50'           "$(format_bp "$MEDAKA_N50") bp" "$(format_bp "$FINAL_N50") bp"
echo ""
echo "  Mode    : $MODE_LABEL"
echo "  Report  : $OUTDIR/MOSTAR_Final_Report.html"
echo "  Version : MOSTAR v$VERSION"
echo ""
echo "  If you found this pipeline useful, please cite!"
echo "  Nermin Zecic, 2026"
echo ""
