# Setup: load libraries, Spark session, and global configuration
import os, sys, time, re, json, random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from scipy.stats import fisher_exact, binomtest, percentileofscore
from statsmodels.stats.multitest import multipletests
import statsmodels.api as sm
import statsmodels.formula.api as smf
from sklearn.utils import resample
import warnings
warnings.filterwarnings('ignore')

# Spark
sys.path.append('/opt/conda/lib/python3.13/site-packages')
from berdl_notebook_utils.setup_spark_session import get_spark_session
spark = get_spark_session()
from pyspark.sql import functions as F

# Constants
T4SS_CORE_COGS = ["COG0742","COG0606","COG0842","COG3502","COG3503","COG3494","COG3495"]
T4SS_CORE_PFAMS = [
    "PF01464","PF04962","PF05101","PF00437","PF06701","PF04610",
    "PF04918","PF02534","PF03135","PF00462","PF01312","PF00867",
    "PF00271","PF07387","PF07803","PF11130","PF00496","PF06586",
    "PF04956","PF02578","PF07917","PF05308","PF06957","PF06834",
    "PF07916","PF12287","PF07998","PF09673","PF03743","PF07934",
    "PF01926","PF12167","PF14157"
]
GLOBAL_KEGG_PATHS = ["01100","01110","01120","01130"]
pd.set_option('display.max_colwidth', None)
random.seed(42)

# Load high‑quality environmental MAGs and split into discovery/validation
print("Loading high‑quality environmental MAGs...")
genomes = spark.table("kescience_mgnify.genome").filter(
    (F.col("completeness") >= 90.0) & (F.col("contamination") <= 5.0)
)
biomes = spark.table("kescience_mgnify.biome").select("biome_id","biome_name","biome_lineage")
env_mags = genomes.join(biomes, on="biome_id", how="inner") \
    .filter(~F.lower(F.col("biome_lineage")).like("%host-associated%")) \
    .select("genome_id","length","lineage","biome_id") \
    .cache()

total_mags = env_mags.count()
print(f"Total high‑quality environmental MAGs: {total_mags}")

(discovery_mags, validation_mags) = env_mags.randomSplit([0.7,0.3], seed=42)
discovery_mags.cache()
validation_mags.cache()
print(f"Discovery: {discovery_mags.count()} | Validation: {validation_mags.count()}")

Loading high‑quality environmental MAGs...
Total high‑quality environmental MAGs: 30497
Discovery: 21474 | Validation: 9023

# T4SS classification with sensitivity analysis (≥2, ≥3, ≥4 distinct domains)
def classify_t4ss(mag_set, min_domains=2):
    # Pfam
    pfam_df = spark.table("kescience_mgnify.gene_eggnog") \
        .join(mag_set.select("genome_id"), on="genome_id", how="inner") \
        .filter(F.col("pfams").isNotNull() & (F.col("pfams") != "-") & (F.col("pfams") != "")) \
        .withColumn("pfam_list", F.split("pfams", ",")) \
        .withColumn("pfam_id", F.explode("pfam_list")) \
        .filter(F.col("pfam_id").isin(T4SS_CORE_PFAMS)) \
        .select("genome_id","pfam_id").distinct()

    pfam_cnt = pfam_df.groupBy("genome_id").agg(
        F.countDistinct("pfam_id").alias("n_t4ss_pfams")
    )

    # COG
    cog_df = spark.table("kescience_mgnify.gene_eggnog") \
        .join(mag_set.select("genome_id"), on="genome_id", how="inner") \
        .filter(F.col("eggnog_ogs").rlike("|".join(T4SS_CORE_COGS))) \
        .select("genome_id","eggnog_ogs").distinct()

    cog_cnt = cog_df.groupBy("genome_id").agg(
        F.countDistinct("eggnog_ogs").alias("n_t4ss_cogs")
    )

    joined = cog_cnt.join(pfam_cnt, on="genome_id", how="full") \
        .fillna(0, subset=["n_t4ss_cogs","n_t4ss_pfams"]) \
        .withColumn("total_domains", F.col("n_t4ss_cogs") + F.col("n_t4ss_pfams"))

    return joined.withColumn("has_t4ss", F.when(F.col("total_domains") >= min_domains, 1).otherwise(0))

# Base definition (≥2 domains) for discovery
hgt_engine_disc = classify_t4ss(discovery_mags, 2).select("genome_id","has_t4ss","total_domains").cache()
print(f"T4SS+ (≥2 domains) discovery: {hgt_engine_disc.filter('has_t4ss=1').count()}")

T4SS+ (≥2 domains) discovery: 5537

# Sensitivity: define with ≥3 and ≥4 for robustness checks
hgt_engine_disc_strict3 = classify_t4ss(discovery_mags, 3).select("genome_id", F.col("has_t4ss").alias("has_t4ss_strict"))
print(f"T4SS+ (≥3 domains) discovery: {hgt_engine_disc_strict3.filter('has_t4ss=1').count()}")
hgt_engine_disc_strict4 = classify_t4ss(discovery_mags, 4).select("genome_id", F.col("has_t4ss").alias("has_t4ss_strict4"))
print(f"T4SS+ (≥4 domains) discovery: {hgt_engine_disc_strict4.filter('has_t4ss=1').count()}")

T4SS+ (≥3 domains) discovery: 4894
T4SS+ (≥4 domains) discovery: 4135

# Global Mobilome Load: Hub of Plasticity Test
# Calculate the total MGE density (total mobilome events per megabase) and compare T4SS+ vs. T4SS- MAGs.

import scipy.stats as stats

# Count all distinct MGEs per genome
global_mge = spark.table("kescience_mgnify.gene_mobilome") \
    .groupBy("genome_id") \
    .agg(F.count("contig_id").alias("total_mge_count"))

# Join with discovery MAGs and T4SS status
plasticity_df = discovery_mags.select("genome_id", "length") \
    .join(hgt_engine_disc.select("genome_id", "has_t4ss"), on="genome_id", how="left") \
    .join(global_mge, on="genome_id", how="left") \
    .fillna({"has_t4ss": 0, "total_mge_count": 0}) \
    .withColumn("mge_density", F.col("total_mge_count") / (F.col("length") / 1e6)) \
    .toPandas()

# Statistical test
t4ss_pos_density = plasticity_df[plasticity_df['has_t4ss'] == 1]['mge_density']
t4ss_neg_density = plasticity_df[plasticity_df['has_t4ss'] == 0]['mge_density']

u_stat, p_val = stats.mannwhitneyu(t4ss_pos_density, t4ss_neg_density, alternative='greater')

print(f"Median MGE Density (T4SS+): {t4ss_pos_density.median():.2f} elements/MB")
print(f"Median MGE Density (T4SS-): {t4ss_neg_density.median():.2f} elements/MB")
print(f"Mann-Whitney U P-value: {p_val:.4e}")

Median MGE Density (T4SS+): 0.00 elements/MB
Median MGE Density (T4SS-): 0.00 elements/MB
Mann-Whitney U P-value: 3.4832e-74

# Trans-acting Mobilization: Are CAZymes on plasmids?
# Check if GT2/CAZy genes are disproportionately located on contigs predicted as plasmids (using your geNomad data) in T4SS+ genomes.
gene_coords_all = spark.table("kescience_mgnify.gene") \
    .join(discovery_mags.select("genome_id"), on="genome_id", how="inner") \
    .select("genome_id","gene_id","contig_id",
            F.col("start").alias("start"), F.col("end").alias("end"))

# Identify plasmid/viral contigs from geNomad
plasmid_contigs = spark.table("kescience_mgnify.gene_mobilome") \
    .filter((F.col("tool") == "geNomad") & (F.col("mobile_element_type").like("%plasmid%"))) \
    .select("genome_id", "contig_id").distinct() \
    .withColumn("is_plasmid", F.lit(1))

# Get all CAZy genes and link to their contigs
cazy_localization = spark.table("kescience_mgnify.gene_eggnog") \
    .filter(F.col("cazy").isNotNull() & (F.col("cazy") != "-")) \
    .select("genome_id", "gene_id", "cazy") \
    .join(gene_coords_all.select("genome_id", "gene_id", "contig_id"), on=["genome_id", "gene_id"]) \
    .join(plasmid_contigs, on=["genome_id", "contig_id"], how="left") \
    .fillna({"is_plasmid": 0})

# Aggregate to see if T4SS+ genomes have more plasmid-borne CAZymes
cazy_plasmid_stats = cazy_localization \
    .join(hgt_engine_disc.select("genome_id", "has_t4ss"), on="genome_id", how="inner") \
    .groupBy("has_t4ss") \
    .agg(
        F.count("gene_id").alias("total_cazy_genes"),
        F.sum("is_plasmid").alias("plasmid_borne_cazy")
    ) \
    .withColumn("pct_on_plasmid", (F.col("plasmid_borne_cazy") / F.col("total_cazy_genes")) * 100)

display(cazy_plasmid_stats.toPandas())

# Ecological Co-selection: Biome-level patterns
# Look at the correlation between T4SS prevalence and CAZy density across different environmental biomes.

cazy_div = spark.table("kescience_mgnify.genome_cazy").groupBy("genome_id").agg(
    F.countDistinct("cazy_family").alias("feat_carbon_versatility")
)

# Calculate average T4SS prevalence and CAZy diversity per biome
biome_stats = discovery_mags \
    .join(hgt_engine_disc.select("genome_id", "has_t4ss"), on="genome_id", how="left") \
    .join(cazy_div, on="genome_id", how="left") \
    .fillna({"has_t4ss": 0, "feat_carbon_versatility": 0}) \
    .groupBy("biome_id", "lineage") \
    .agg(
        F.count("genome_id").alias("n_mags"),
        F.avg("has_t4ss").alias("t4ss_prevalence"),
        F.avg("feat_carbon_versatility").alias("avg_cazy_families")
    ) \
    .filter(F.col("n_mags") >= 50) \
    .toPandas()

# Plot the correlation
plt.figure(figsize=(8, 6))
plt.scatter(biome_stats['t4ss_prevalence'], biome_stats['avg_cazy_families'], alpha=0.7)
plt.title('Ecological Co-selection: T4SS vs. CAZy Repertoire by Biome')
plt.xlabel('Proportion of MAGs with T4SS')
plt.ylabel('Average distinct CAZy families per MAG')

# Add trendline
z = np.polyfit(biome_stats['t4ss_prevalence'], biome_stats['avg_cazy_families'], 1)
p = np.poly1d(z)
plt.plot(biome_stats['t4ss_prevalence'], p(biome_stats['t4ss_prevalence']), "r--")
plt.show()

print(f"Spearman correlation across biomes: {stats.spearmanr(biome_stats['t4ss_prevalence'], biome_stats['avg_cazy_families'])[0]:.3f}")

Spearman correlation across biomes: -0.165

	has_t4ss	total_cazy_genes	plasmid_borne_cazy	pct_on_plasmid
0	1	236328	770	0.325818
1	0	6364	0	0.000000

04 Mobilome Scan

The "Hub of Plasticity" Test (Global Mobilome Load)¶

Trans-acting Mobilization (Plasmid & Viral Carriage)¶

Niche Co-selection (Biome-level Analysis)¶