#!/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.5
#  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)
#    - Integrons (IntegronFinder)
#    - 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
usage() {
    cat <<'HELP'

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

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

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

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.01]

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]
  --integrons     Run IntegronFinder
  --eval          IntegronFinder e-value threshold [Default: 1e-9]  

Other options:
  --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.5
HELP
    exit 0
}

[[ "$1" == "--help" || "$1" == "-h" ]] && usage

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.5"
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=14
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
RUN_INTEGRONS=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)"
    "Integron Detection (IntegronFinder)"
    "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
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# 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 integron_section=""
    local ice_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")
    integron_finder_v=$(integron_finder --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()

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

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

    # 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>
            <p style='color:#7f8c8d; font-size:0.85rem;'>
                This table lists every antimicrobial resistance (AMR) and virulence gene identified by AMRFinder+ across the assembly.
            </p>
            <ul style='color:#7f8c8d; font-size:0.85rem; line-height:1.6;'>
                <li><strong>Gene:</strong> The AMRFinder+ reference gene symbol matched in the assembly.</li>
                <li><strong>Class:</strong> The antimicrobial drug class or resistance mechanism the gene is predicted to confer .</li>
                <li><strong>Identity:</strong> The percent sequence identity between the detected gene and the AMRFinder+ reference sequence.</li>
                <li><strong>Contig:</strong> Which assembled contig the gene was found on. Genes on smaller, non-chromosomal contigs may indicate plasmid-borne resistance.</li>
            </ul>
            <div style='overflow-x:auto;'>
            ${amr_html}
            </div>
        </div>"
    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 -E 's/^[0-9]+_//; 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;'>
        This circular map summarizes the genome assembly's key genomic features across four
        concentric tracks, read from the center outward:
    </p>
    <ul style='color:#7f8c8d; font-size:0.85rem; line-height:1.6;'>
        <li><strong>Innermost (grey):</strong> GC content deviation from the genome-wide average, calculated in sliding windows around the sequence.</li>
        <li><strong>Middle-inner (cyan):</strong> ICE / conjugation systems detected by MacSyFinder (CONJScan), labeled by system type and size in kb. Arrows indicate strand orientation.</li>
        <li><strong>Middle-outer (green):</strong> Prophage regions identified by geNomad, labeled with size in kb.</li>
        <li><strong>Outermost:</strong> AMR and virulence genes identified by AMRFinder+, color-coded by resistance class (see legend). On plasmid/extrachromosomal maps, all genes are shown in red regardless of class, reflecting the elevated horizontal transfer risk associated with plasmid-borne resistance determinants.</li>
    </ul>
    <p style='color:#7f8c8d; font-size:0.85rem;'>
        The main chromosome is shown separately from any additional contigs (plasmids or extrachromosomal elements), each rendered as its own circular map.
    </p>
    <p style='color:#7f8c8d; font-size:0.85rem;'>
        Note: Details about the detected ICE elements or AMR genes can be reviewed in their respective output folders. Mobile genetic elements identified by MacSyFinder with less than 50% system completeness are excluded from this visualization to reduce visual noise from low-confidence partial hits. All detected systems, including those below this threshold, are retained in <code>ice_detection/best_solution.tsv</code> for manual review.
    </p>
    <p style='color:#7f8c8d; font-size:0.85rem;'>
        Hint: Click any map to zoom. Move mouse to pan while zoomed.
    </p>
    <div class='map-grid'>$map_content</div>
    </div>"
        fi
    fi

        # Mobile resistome section (geNomad)
    mobile_section=""
    if [[ "${RUN_GENOMAD}" == "true" ]]; then
        local mob_csv="$OUTDIR/intermediate/mob_sum.csv"
        local mob_html="<tr><td colspan='5' style='text-align:center; color:#7f8c8d;'>No mobile resistance elements detected</td></tr>"
        
        if [[ -f "$mob_csv" ]]; then
            mob_html=$(python3 - "$mob_csv" "$MAIN_ID_VAL" <<'PYMOB'
import sys, pandas as pd
try:
    df = pd.read_csv(sys.argv[1])
    if df.empty: sys.exit(0)
    main_id = sys.argv[2].strip()
    rows = ""
    for _, r in df.iterrows():
        gene = str(r.iloc[5])
        contig = str(r.iloc[1])
        cls = str(r.iloc[10])
        score = float(r.iloc[-1])
        risk = "HIGH" if score >= 0.8 else "LOW"
        c_id = contig.split()[0].strip()
        color = '#2c3e50'
        
        if c_id != main_id: color = '#e74c3c'
        elif 'virulence' in cls.lower(): color = '#000000'
        elif any(x in cls.lower() for x in ['beta-lactam', 'penicillin', 'carbapenem']): color = '#3498db'
        elif 'tetracycline' in cls.lower(): color = '#9b59b6'
        elif 'aminoglycoside' in cls.lower(): color = '#e67e22'
        elif any(x in cls.lower() for x in ['macrolide', 'lincosamide', 'streptogramin']): color = '#f1c40f'
        elif 'quinolone' in cls.lower() or 'fluoroquinolone' in cls.lower(): color = '#27ae60'
        elif risk == 'HIGH': color = '#e74c3c'

        r_class = "class='risk-high'" if risk == 'HIGH' else ""
        rows += f"<tr><td><span style='border-left:4px solid {color};padding-left:10px;font-weight:600;'>{gene}</span></td><td>{contig}</td><td>{cls}</td><td>{score:.2f}</td><td><span {r_class}>{risk}</span></td></tr>"
    if rows: print(rows)
except Exception:
    pass
PYMOB
            )
            [[ -z "$mob_html" ]] && mob_html="<tr><td colspan='5' style='text-align:center; color:#7f8c8d;'>No mobile resistance elements detected</td></tr>"
        fi

        mobile_section="<section class='card'>
            <h2>Mobile Resistome &amp; Genomic Plasticity</h2>
            <p style='color:#7f8c8d; font-size:0.85rem; margin-bottom:1rem; line-height:1.5;'>
                This section tracks antimicrobial resistance (AMR) genes located on mobile genetic elements, such as plasmids. Unlike chromosomal mutations, plasmid-borne resistance can easily transfer horizontally between different bacteria, representing a much higher risk for the rapid spread of drug resistance.
            </p>
            <div class='alert-box alert-red'>
                <strong>Plasmid-borne AMR detected:</strong>
                <span class='big-num'>${HIGH_RISK_COUNT:-0}</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>${mob_html}</tbody>
            </table>
            <ul style='color:#7f8c8d; font-size:0.85rem; line-height:1.6; margin-top:1rem;'>
                <li><strong>Gene:</strong> The specific antimicrobial resistance gene identified by the pipeline.</li>
                <li><strong>Contig:</strong> The assembled sequence (e.g., a specific plasmid) where the gene is physically located.</li>
                <li><strong>Class:</strong> The category of antibiotics to which this gene confers resistance.</li>
                <li><strong>Score:</strong> The confidence metric or identity match percentage for the gene detection.</li>
                <li><strong>Risk:</strong> The potential clinical or epidemiological threat level associated with this specific resistance profile.</li>
            </ul>
        </section>"
    fi

    # Prophage section (geNomad)
    phage_section=""
    if [[ "${RUN_GENOMAD}" == "true" ]]; then
        local phage_tsv=$(find "$OUTDIR/genomad" -name "*_virus_summary.tsv" 2>/dev/null | head -n1 || true)
        local phage_html="<tr><td colspan='4' style='text-align:center; color:#7f8c8d;'>No prophages detected</td></tr>"

        if [[ -f "$phage_tsv" ]]; then
            phage_html=$(python3 - "$phage_tsv" <<'PYPHAGE'
import sys, pandas as pd
try:
    df = pd.read_csv(sys.argv[1], sep='\t')
    if df.empty: sys.exit(0)
    rows=""
    for _, r in df.iterrows():
        contig = str(r.get('seq_name', 'N/A'))
        length = f"{int(r.get('length', 0)):,} bp"
        score = float(r.get('virus_score', 0))
        ptype = "Intact" if score >= 0.8 else "Incomplete"
        rows += f"<tr><td><strong>{contig}</strong></td><td>{length}</td><td>{score:.2f}</td><td><span class='phage-tag'>{ptype}</span></td></tr>"
    if rows: print(rows)
except Exception:
    pass
PYPHAGE
            )
            [[ -z "$phage_html" ]] && phage_html="<tr><td colspan='4' style='text-align:center; color:#7f8c8d;'>No prophages detected</td></tr>"
        fi

        phage_section="<section class='card'>
            <h2>Prophage Tracker</h2>
            <p style='color:#7f8c8d; font-size:0.85rem; margin-bottom:1rem; line-height:1.5;'>
                Prophages are viral sequences (bacteriophages) that have integrated directly into the bacterial DNA. They act as major drivers of bacterial evolution and often carry accessory genes—such as virulence factors or toxins—that can enhance the organism's adaptability and pathogenicity.
            </p>
            <div class='alert-box alert-green'>
                <strong>Total Prophages Detected:</strong>
                <span class='big-num'>${PHAGE_COUNT:-0}</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>${phage_html}</tbody>
            </table>
            <ul style='color:#7f8c8d; font-size:0.85rem; line-height:1.6; margin-top:1rem;'>
                <li><strong>Contig ID:</strong> The specific genomic region where the viral DNA is integrated.</li>
                <li><strong>Length:</strong> The total size of the integrated viral sequence in base pairs.</li>
                <li><strong>Score:</strong> The probability or confidence metric indicating the likelihood that the sequence is a true prophage (higher is better).</li>
                <li><strong>Type:</strong> The classification of the viral sequence (e.g., intact, incomplete, or questionable).</li>
            </ul>
        </section>"
    fi

    # ICE / Conjugation Systems section (MacSyFinder / CONJScan)
    ice_section=""
    local ice_best=$(find "$OUTDIR/ice_detection" -name "best_solution.tsv" 2>/dev/null | head -n1 || true)
    
    if [[ -f "$ice_best" ]]; then
        local ice_html="<tr><td colspan='4' style='text-align:center; color:#7f8c8d;'>No conjugation systems detected</td></tr>"
        local ice_total="0"

        ice_raw=$(python3 - "$ice_best" <<'PYICE'
import sys, pandas as pd
try:
    df = pd.read_csv(sys.argv[1], sep='\t', comment='#')
    if df.empty: sys.exit(0)
    
    # Adapt to different MacSyFinder versions
    rep_col = 'replicon' if 'replicon' in df.columns else df.columns[0]
    sys_id_col = 'sys_id' if 'sys_id' in df.columns else 'system'
    gene_col = 'gene_name' if 'gene_name' in df.columns else 'hit_id'
    
    rows = ""
    total_systems = 0
    
    if rep_col in df.columns and sys_id_col in df.columns:
        # Group by system ID (a system ID belongs to one replicon)
        for (rep, sys_id), group in df.groupby([rep_col, sys_id_col]):
            total_systems += 1
            
            # Extract System Type
            sys_type = ""
            if 'model_fqn' in group.columns:
                sys_type = str(group['model_fqn'].iloc[0])
            else:
                sys_type = str(sys_id).rsplit('_', 1)[0]
                
            # Clean up sys_type (e.g. 'CONJScan/ICE/T4SS_typeF' becomes 'T4SS_typeF (ICE)')
            if '/' in sys_type:
                parts = sys_type.split('/')
                if len(parts) >= 3:
                    sys_type = f"{parts[-1]} ({parts[-2]})"
                else:
                    sys_type = parts[-1]
            
            # Extract Completeness
            wholeness = "N/A"
            if 'sys_wholeness' in group.columns:
                try:
                    w_val = float(group['sys_wholeness'].iloc[0])
                    wholeness = f"{w_val:.2f}"
                except:
                    pass
            
            # Extract Genes/Contents
            elements = []
            if gene_col in group.columns:
                for val in group[gene_col].dropna():
                    clean_val = str(val).split('-NCBIFAM')[0].replace('_like', '')
                    elements.append(f"<span style='background:#f1f2f6; padding:2px 6px; border-radius:4px; border:1px solid #dcdde1; margin-right:4px; display:inline-block; margin-bottom:4px;'>{clean_val}</span>")
            contents = "".join(elements)
            
            rows += (
                f"<tr>"
                f"<td><strong>{rep}</strong></td>"
                f"<td><strong style='color:#006080;'>{sys_type}</strong></td>"
                f"<td>{wholeness}</td>"
                f"<td style='font-size:0.85em; max-width: 300px; line-height: 1.6;'>{contents}</td>"
                f"</tr>"
            )
            
    if rows: 
        print(f"{total_systems}:::{rows}")
except Exception:
    pass
PYICE
        )
        
        # Split the python output into the total count and the HTML rows
        if [[ "$ice_raw" == *":::"* ]]; then
            ice_total=$(echo "$ice_raw" | awk -F':::' '{print $1}')
            ice_html=$(echo "$ice_raw" | awk -F':::' '{print $2}')
        fi

        ice_section="<section class='card'>
            <h2>Integrative &amp; Conjugative Elements (ICE)</h2>
            <p style='color:#7f8c8d; font-size:0.85rem; margin-bottom:1rem; line-height:1.5;'>
                Integrative and Conjugative Elements (ICEs) and other secretion systems are massive mobile genetic elements capable of self-transfer between bacteria via conjugation. They are primary vehicles for the horizontal spread of antibiotic resistance and virulence factors.
            </p>
            <div class='alert-box alert-cyan'>
                <strong>Total Conjugation Systems Detected:</strong>
                <span class='big-num'>${ice_total}</span>
            </div>
            <table class='resistome-table' id='iceTable'>
                <thead>
                    <tr>
                        <th>Replicon</th>
                        <th>System Type</th>
                        <th>Completeness</th>
                        <th>System Contents (Genes)</th>
                    </tr>
                </thead>
                <tbody>${ice_html}</tbody>
            </table>
            <ul style='color:#7f8c8d; font-size:0.85rem; line-height:1.6; margin-top:1rem;'>
                <li><strong>Replicon:</strong> The genomic contig or plasmid where the conjugative system was localized.</li>
                <li><strong>System Type:</strong> The classification of the secretion system (e.g., T4SS, ICE, IME).</li>
                <li><strong>Completeness:</strong> A wholeness score (where 1.0 = 100%) indicating how many core mandatory genes of the system's model were successfully identified.</li>
                <li><strong>System Contents:</strong> The structural and functional genes that construct the secretion apparatus and conjugation machinery.</li>
            </ul>
        </section>"
    fi

        # Integron section (IntegronFinder)
    integron_section=""
    if [[ "$RUN_INTEGRONS" == "true" ]]; then
        # Find both files dynamically based on their actual extensions
        local int_summ=$(find "$OUTDIR/integron_detection" -name "*.summary" 2>/dev/null | head -n1 || true)
        local int_details=$(find "$OUTDIR/integron_detection" -name "*.integrons" 2>/dev/null | head -n1 || true)
        
        local int_html="<tr><td colspan='6' style='text-align:center; color:#7f8c8d;'>No integrons detected</td></tr>"
        local int_total="0"

        if [[ -f "$int_summ" ]]; then
            int_html=$(python3 - "$int_summ" "$int_details" <<'PYINT'
import sys, pandas as pd
try:
    # 1. Parse the summary file
    df_sum = pd.read_csv(sys.argv[1], sep='\t', comment='#')
    if df_sum.empty: sys.exit(0)
    
    # 2. Parse the details file to get the genetic contents (if it exists)
    contents_map = {}
    details_file = sys.argv[2].strip()
    if details_file:
        try:
            df_int = pd.read_csv(details_file, sep='\t', comment='#')
            if not df_int.empty:
                # Group all genetic elements by the replicon they belong to
                for rep, group in df_int.groupby('ID_replicon'):
                    elements = []
                    for val in group['annotation'].dropna():
                        val_str = str(val)
                        # Clean up messy database names (e.g., 'trim_DfrA1_like-NCBIFAM' -> 'trim_DfrA1')
                        clean_val = val_str.split('-NCBIFAM')[0].replace('_like', '')
                        elements.append(f"<span style='background:#f1f2f6; padding:2px 6px; border-radius:4px; border:1px solid #dcdde1; margin-right:4px; display:inline-block; margin-bottom:4px;'>{clean_val}</span>")
                    contents_map[str(rep)] = "".join(elements)
        except:
            pass
            
    # 3. Build the HTML rows
    rows=""
    for _, r in df_sum.iterrows():
        comp = int(r.get('complete', 0))
        in0 = int(r.get('In0', 0))
        calin = int(r.get('CALIN', 0))
        tot = comp + in0 + calin
        
        # Hide contigs with 0 integrons
        if tot == 0:
            continue
            
        replicon = str(r.get('ID_replicon', r.get('replicon', 'N/A')))
        contents = contents_map.get(replicon, "<span style='color:#7f8c8d; font-style:italic;'>Unknown</span>")
        
        rows += (
            f"<tr>"
            f"<td><strong>{replicon}</strong></td>"
            f"<td>{comp}</td><td>{in0}</td><td>{calin}</td>"
            f"<td><strong>{tot}</strong></td>"
            f"<td style='font-size:0.85em; max-width: 300px; line-height: 1.6;'>{contents}</td>"
            f"</tr>"
        )
        
    if rows: print(rows)
except Exception:
    pass
PYINT
            )
            [[ -z "$int_html" ]] && int_html="<tr><td colspan='6' style='text-align:center; color:#7f8c8d;'>No integrons detected</td></tr>"
            
            # Recalculate total for the badge
            int_total=$(awk -F'\t' '!/^#/ && $1!~"replicon" {sum += $2 + $3 + $4} END {print sum+0}' "$int_summ" 2>/dev/null || echo "0")
        fi

        integron_section="<section class='card'>
            <h2>Integron Profile</h2>
            <p style='color:#7f8c8d; font-size:0.85rem; margin-bottom:1rem; line-height:1.5;'>
                Integrons are genetic elements capable of capturing and expressing mobile gene cassettes via site-specific recombination.
            </p>
            <div class='alert-box alert-pink'>
                <strong>Total Integrons Detected:</strong>
                <span class='big-num'>${int_total}</span>
            </div>
            <table class='resistome-table' id='integronTable'>
                <thead>
                    <tr>
                        <th>Replicon</th>
                        <th>Complete Integrons</th>
                        <th>In0 (Integrase only)</th>
                        <th>CALIN (attC clusters)</th>
                        <th>Total</th>
                        <th>Integron Contents</th>
                    </tr>
                </thead>
                <tbody>${int_html}</tbody>
            </table>
            <ul style='color:#7f8c8d; font-size:0.85rem; line-height:1.6; margin-top:1rem;'>
                <li><strong>Replicon:</strong> The genomic contig or plasmid where the integron structure was localized.</li>
                <li><strong>Complete Integrons:</strong> Structures containing both the integrase gene and an adjacent attC site cluster.</li>
                <li><strong>In0 (Integrase only):</strong> Isolated integrase genes lacking detectable nearby attC cassettes.</li>
                <li><strong>CALIN (attC clusters):</strong> Clusters of attC sites lacking an identifiable integrase gene.</li>
                <li><strong>Total:</strong> The combined count of all integron-related structures found on this replicon.</li>
                <li><strong>Integron Contents:</strong> The predicted structural elements (e.g., attC, intI) and captured functional genes (e.g., AMR or virulence cassettes) found within the integron.</li>
            </ul>
        </section>"
    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;
        --blue:     #e1f5fe;
    }
    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); }
    .alert-blue  { background: #e1f5fe; border-left: 5px solid var(--blue);  }
    .alert-cyan  { background: #e6ffff; border-left: 5px solid #00f2ff; color: #006080; }
    .alert-pink  { background: #fdf2f7; border-left: 5px solid #e84393; color: #881b4c; }
    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; position: relative; }
    .map-zoom-container img { 
        width: 100%; 
        display: block; 
        transition: transform 0.35s cubic-bezier(0.25, 0.46, 0.45, 0.94); 
        will-change: transform; /* Forces GPU hardware acceleration */
        transform-origin: center center;
        pointer-events: none; /* Prevents accidental dragging of the image */ 
    }
    .map-zoom-container.zoomed { cursor: zoom-out; }
    .map-zoom-container.zoomed img { transform: scale(3.5) translateZ(0); } /* translateZ pushes it to the 3D GPU layer */

    .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}
${ice_section}
${integron_section}
<div class="card">
    <h2>Colour Legend for circular genomes</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 class="legend-item"><div class="swatch" style="background:#e84393; border-color:#aaa;"></div>Integron</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>
    // Circular Map Zoom Logic
    document.addEventListener("DOMContentLoaded", function() {
        var containers = document.querySelectorAll('.map-zoom-container');
        
        containers.forEach(function(container) {
            var img = container.querySelector('img');
            if (!img) return;
            
            container.addEventListener('mousemove', function(e) {
                if (!container.classList.contains('zoomed')) return;
                var r = container.getBoundingClientRect();
                var x = ((e.clientX - r.left) / r.width) * 100;
                var y = ((e.clientY - r.top) / r.height) * 100;
                
                // Disable transition while moving to prevent lag
                img.style.transition = 'none';
                img.style.transformOrigin = x + "% " + y + "%";
            });

            container.addEventListener('click', function(e) {
                var r = container.getBoundingClientRect();
                var x = ((e.clientX - r.left) / r.width) * 100;
                var y = ((e.clientY - r.top) / r.height) * 100;

                if (!container.classList.contains('zoomed')) {
                    // Restore transition for the zoom-in animation
                    img.style.transition = 'transform 0.35s cubic-bezier(0.25, 0.46, 0.45, 0.94)';
                    img.style.transformOrigin = x + "% " + y + "%";
                    container.classList.add('zoomed');
                } else {
                    img.style.transition = 'transform 0.35s cubic-bezier(0.25, 0.46, 0.45, 0.94)';
                    container.classList.remove('zoomed');
                    setTimeout(function() {
                        if (!container.classList.contains('zoomed')) {
                            img.style.transformOrigin = 'center center';
                        }
                    }, 350); 
                }
            });

            // Auto-zoom out if the mouse leaves the container entirely
            container.addEventListener('mouseleave', function() {
                if (container.classList.contains('zoomed')) {
                    img.style.transition = 'transform 0.35s cubic-bezier(0.25, 0.46, 0.45, 0.94)';
                    container.classList.remove('zoomed');
                    setTimeout(function() {
                        if (!container.classList.contains('zoomed')) {
                            img.style.transformOrigin = 'center center';
                        }
                    }, 350);
                }
            });
        });
    });
</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 ;;
        --integrons|-y)     RUN_INTEGRONS=true;            shift 1 ;;
        --eval|-Y)          USER_INTEGRON_EVALUE="$2";     shift 2 ;;
        --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")
[[ "$RUN_INTEGRONS" == "true" ]] && REQUIRED_TOOLS+=("integron_finder")

# 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 "[WARNING] Kraken2 database not found — taxonomy will be skipped."
    RUN_KRAKEN2=false
fi
if [[ -z "$GENOMAD_DB" || ! -d "$GENOMAD_DB" ]]; then
    echo "[WARNING] 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 "Fix Assembly 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-dir /opt/macsy-models \
                --models CONJScan/Chromosome all \
                -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." >> "$OUTDIR/logs/pipeline_warnings.log"
        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
from statistics import median

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)]

MIN_WHOLENESS = 0.5
best = best[best['sys_wholeness'] >= MIN_WHOLENESS]
if best.empty:
    exit(0)

# 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])

sys_hits = {}
sys_score = {}
for sys_id, group in best.groupby('sys_id'):
    sys_hits[sys_id] = set(group['hit_id'])
    sys_score[sys_id] = group['sys_wholeness'].iloc[0]

to_drop = set()
sys_ids = list(sys_hits)
for i, a in enumerate(sys_ids):
    for b in sys_ids[i+1:]:
        overlap = len(sys_hits[a] & sys_hits[b])
        smaller = min(len(sys_hits[a]), len(sys_hits[b]))
        if smaller and overlap / smaller > 0.5:
            loser = a if sys_score[a] < sys_score[b] else b
            to_drop.add(loser)

out_rows = []
MAX_CLUSTER_SPREAD = 100_000  

for sys_id, group in best.groupby('sys_id'):
    if sys_id in to_drop:
        continue

    raw_hits = [coord_map[hid] for hid in group['hit_id'] if hid in coord_map]
    if not raw_hits:
        continue

    # Restrict to the dominant contig for this system
    contigs = [h[0] for h in raw_hits]
    main_contig = max(set(contigs), key=contigs.count)
    same_contig = [h for h in raw_hits if h[0] == main_contig]

    # Drop hits far from the main cluster's median midpoint
    midpoints = [(h[1] + h[2]) / 2 for h in same_contig]
    med = median(midpoints)
    hits = [h for h, mid in zip(same_contig, midpoints)
            if abs(mid - med) <= MAX_CLUSTER_SPREAD] or same_contig

    model_name = group['model_fqn'].iloc[0].split('/')[-1]
    out_rows.append([hits[0][0], model_name, 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 — Integron Detection (IntegronFinder)  [optional]
# ==============================================================================
INTEGRON_EVALUE="${USER_INTEGRON_EVALUE:-1e-9}"   # Adjust after checking real output distributions

if [[ "$RUN_INTEGRONS" == "true" ]]; then
    run_step "Integron Detection (IntegronFinder)" "true" \
        "rm -rf '$OUTDIR/integron_detection' && integron_finder \
            --local-max \
            --gbk \
            --func-annot \
            --cpu $THREADS \
            --outdir '$OUTDIR/integron_detection' \
            '$FINAL_ASM' \
            >'$OUTDIR/logs/integron_finder.log' 2>&1"

    INTEGRON_TSV=$(find "$OUTDIR/integron_detection" -name "*.integrons" 2>/dev/null | head -n1 || true)
    INTEGRON_SUMMARY=$(find "$OUTDIR/integron_detection" -name "*.summary" 2>/dev/null | head -n1 || true)

        if [[ -f "$INTEGRON_TSV" ]]; then
        python3 - <<PYINTEGRON
import pandas as pd
import traceback
import os

try:
    # Read the file
    try:
        df = pd.read_csv('$INTEGRON_TSV', sep='\t', comment='#')
    except pd.errors.EmptyDataError:
        df = pd.DataFrame()  # Create an empty dataframe if no data exists

    out_rows = []
    
    # Only process if dataframe is not empty
    if not df.empty:
        # Ensure evalue is numeric (coerces weird strings to NaN), then fill NaN with 0
        df['evalue'] = pd.to_numeric(df['evalue'], errors='coerce').fillna(0)
        
        # Filter based on threshold
        df = df[df['evalue'] <= float('$INTEGRON_EVALUE')]

        # GROUP BY Integron ID to create one solid block per integron!
        if 'ID_integron' in df.columns:
            for integron_id, group in df.groupby('ID_integron'):
                replicon = group['ID_replicon'].iloc[0]
                start = group['pos_beg'].min()
                end = group['pos_end'].max()
                strand = group['strand'].iloc[0]
                
                # Name it "Integron" so the map script colors it Deep Pink
                out_rows.append([replicon, "Integron", start, end, strand])

    if out_rows:
        # Append to ICE_FILE so they get plotted on the circular map!
        pd.DataFrame(out_rows).to_csv('$ICE_FILE', sep='\t', index=False, header=False, mode='a')

    # --- REPORT GENERATION ---
    summary_file = '$INTEGRON_SUMMARY'
    report_out = '$OUTDIR/reports/integron_summary.md'
    
    # Ensure report directory exists
    os.makedirs(os.path.dirname(report_out), exist_ok=True)
    
    if os.path.exists(summary_file) and os.path.getsize(summary_file) > 0:
        try:
            # Added comment='#' here as well to prevent pandas from crashing on the header!
            df_sum = pd.read_csv(summary_file, sep='\t', comment='#')
        except pd.errors.EmptyDataError:
            df_sum = pd.DataFrame()
            
        if not df_sum.empty:
            # Get sums of different integron types (default to 0 if column is somehow missing)
            tot_calin = df_sum['CALIN'].sum() if 'CALIN' in df_sum.columns else 0
            tot_in0 = df_sum['In0'].sum() if 'In0' in df_sum.columns else 0
            tot_complete = df_sum['complete'].sum() if 'complete' in df_sum.columns else 0
            
            with open(report_out, 'w') as f:
                f.write("### Integron Profile\n\n")
                f.write("| Element Type | Count |\n")
                f.write("|--------------|-------|\n")
                f.write(f"| Complete Integrons | {tot_complete} |\n")
                f.write(f"| In0 (Integrase only) | {tot_in0} |\n")
                f.write(f"| CALIN (attC clusters) | {tot_calin} |\n\n")
                
                f.write("**Breakdown by Replicon:**\n\n")
                
                # Generate a manual Markdown table to avoid needing the 'tabulate' pip package
                cols = df_sum.columns.tolist()
                f.write("| " + " | ".join(cols) + " |\n")
                f.write("|" + "|".join(["---"] * len(cols)) + "|\n")
                for _, row in df_sum.iterrows():
                    row_vals = [str(row[c]) for c in cols]
                    f.write("| " + " | ".join(row_vals) + " |\n")
                f.write("\n")
        else:
            with open(report_out, 'w') as f:
                f.write("### Integron Profile\n\nNo integrons found in this assembly.\n")
    else:
        with open(report_out, 'w') as f:
            f.write("### Integron Profile\n\nNo integrons found in this assembly.\n")

except Exception:
    with open('$OUTDIR/logs/integron_parse_error.log', 'w') as fh:
        fh.write(traceback.format_exc())
PYINTEGRON
    else
        echo "[WARNING] No .integrons file found — skipping integron plotting integration."
    fi
else
    run_step "Integron Detection (IntegronFinder)" "false" ""
fi

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

if [[ "$RUN_GENOMAD" == "true" && -d "$GENOMAD_DB" ]]; 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 13 — 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 is 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

            # Check if this element is an Integron or an ICE
            symbol_name = str(row['Symbol'])
            is_integron = 'Integron' in symbol_name
            
            # Use Deep Pink for Integrons, Cyan for ICE
            arc_color = '#e84393' if is_integron else '#00f2ff'
            txt_color = '#b71540' if is_integron else '#006080'

            # ICE arc
            ax.plot(np.linspace(s, e, 100), [0.85] * 100,
                    color=arc_color, 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_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\nGC: {rec_gc:.1f}%\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=""
if [[ "$RUN_KRAKEN2" == "true" ]]; then
    TAX_CARDS="<div class='qc-card'>
        <span class='qc-label'>Taxonomy report</span>
        <span class='qc-val' style='color:#2e7d32; font-weight:600;'>${ASM_SPECIES_FULL:-Unknown}</span>
        <span class='qc-sub' style='color:#7f8c8d; font-size:0.8rem; display:block; margin-top:0.3rem;'>
            The most likely match at the specified confidence score: ${KRAKEN_CONFIDENCE}. 
            Note that if no species-level match is found, the report will show the closest higher-level taxon. 
        </span>
</div>"
fi
    
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>
  <p style='color:#7f8c8d; font-size:0.85rem; margin-top:1rem;'>
      This table summarizes key assembly quality metrics before and after final polishing.
  </p>
  <ul style='color:#7f8c8d; font-size:0.85rem; line-height:1.6;'>
      <li><strong>Total Length:</strong> The combined length of all assembled contigs, in base pairs. This should approximate the expected genome size for the organism.</li>
      <li><strong>Contig Count:</strong> The number of separate, non-overlapping sequences produced by assembly. A single circular chromosome ideally assembles as 1 contig; additional contigs may represent plasmids, extrachromosomal elements, or — in a less complete assembly — unresolved chromosomal fragments.</li>
      <li><strong>N50:</strong> A standard measure of assembly contiguity: the length of the shortest contig needed, when contigs are ordered largest to smallest, to cover at least 50% of the total assembly length. A higher N50 relative to contig count generally indicates a more contiguous, higher-quality assembly.</li>
      <li><strong>After Medaka (ONT polish):</strong> Metrics following long-read consensus polishing with Medaka, before any further correction. </li>  
      <li><strong>Final (polished):</strong> The final polishing/correction metric for either Hybrid or ONT-only mode. Note that in ONT-only mode, no additional polishing step is applied after Medaka.</li>
  </ul>
</div>"

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

# =================================================================r=============
#  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 ""
