ARGenus 0.4.0

ARG detection and genus-level classification using flanking sequence analysis
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
//! Genus Classifier Module
//!
//! Classifies the source genus of detected ARGs using flanking sequence analysis.
//! Compares extracted flanking regions against a pre-built database of known
//! gene-genus associations.
//!
//! # Classification Method
//! 1. Extract upstream and downstream flanking sequences from contig
//! 2. Query the flanking database for the detected ARG
//! 3. Align query flanking sequences against reference flanking sequences
//! 4. Score genus candidates based on alignment identity and coverage
//! 5. Report top genus with confidence and specificity metrics
//!
//! # Key Metrics
//! - **Confidence**: Alignment identity score (0-100%)
//! - **Specificity**: Gene-genus association strength in the database (0-100%)

use anyhow::{Context, Result};
use rustc_hash::FxHashMap;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
use std::os::unix::fs::FileExt;
use std::path::Path;
use std::process::Command;

use crate::snp::{self, SnpStatus};

const FDB_MAGIC: &[u8; 8] = b"FLANKDB\0";

// ============================================================================
// Embedded phylogenetic-context tables (compiled into the binary)
// ============================================================================
// The GTDB genus distance/lineage tables and the conformal calibration are embedded
// (zstd-compressed) so a deployed binary is self-contained — only the flanking DB (FDB)
// ships separately. The kernel constants (lambda, coherence radius, absent distance) are
// calibrated to THIS table's patristic scale, so binding them into one binary prevents an
// old-table/new-constant version mismatch. An external file of the same name next to the
// FDB still overrides (for in-place updates); the loader logs which source it used.
const EMB_GENUS_DIST: &[u8] = include_bytes!("embedded/genus_dist.tsv.zst");
const EMB_GENUS_LINEAGE: &[u8] = include_bytes!("embedded/genus_lineage.tsv.zst");
const EMB_CONFORMAL: &[u8] = include_bytes!("embedded/conformal.tsv.zst");

/// Text of a phylogenetic-context table plus a label for its source. Prefers an external file
/// of `filename` in `db_dir` (lets a newer table override without recompiling); otherwise
/// decompresses the copy embedded in the binary at build time. The embedded copy always
/// exists, so these tables are never missing.
fn context_table_text(
    db_dir: Option<&Path>,
    filename: &str,
    embedded_zst: &[u8],
) -> Result<(String, String)> {
    if let Some(dir) = db_dir {
        let p = dir.join(filename);
        if p.exists() {
            let s = std::fs::read_to_string(&p).with_context(|| format!("read {:?}", p))?;
            return Ok((s, format!("{:?}", p)));
        }
    }
    let bytes =
        zstd::decode_all(embedded_zst).with_context(|| format!("decode embedded {}", filename))?;
    let s = String::from_utf8(bytes).with_context(|| format!("embedded {} utf8", filename))?;
    Ok((s, format!("embedded:{}", filename)))
}

// ============================================================================
// Data Structures
// ============================================================================

/// ARG hit with position information for flanking extraction.
///
/// Contains all information needed to extract and classify flanking sequences.
#[derive(Debug, Clone)]
pub struct ArgPosition {
    /// ARG gene name (e.g., "blaTEM-1").
    pub arg_name: String,
    /// Contig identifier where ARG was detected.
    pub contig_name: String,
    /// Full contig nucleotide sequence.
    pub contig_seq: String,
    /// Contig length in base pairs.
    pub contig_len: usize,
    /// ARG start position on contig (0-based).
    pub arg_start: usize,
    /// ARG end position on contig.
    pub arg_end: usize,
    /// Strand orientation ('+' or '-').
    pub strand: char,
    /// Redundant PanRes reference IDs tied at this locus (includes `arg_name`). Their
    /// flanking reference sets are unioned for classification. Never empty.
    pub members: Vec<String>,
}

/// Genus classification result for a single ARG.
/// Two genera are "tied" (indistinguishable) when their scores are within this many
/// identity points. Used to decide multi-genus reporting.
pub const GENUS_TIE_PCT: f64 = 1.0;

#[derive(Debug, Clone)]
pub struct GenusResult {
    /// ARG gene name.
    pub arg_name: String,
    /// Contig identifier.
    pub contig_name: String,
    /// Classified genus (None if unresolved).
    pub genus: Option<String>,
    /// Classification confidence (alignment identity, 0-100).
    pub confidence: f64,
    /// Gene-genus specificity in database (0-100).
    pub specificity: f64,
    /// Extracted upstream flanking length.
    pub upstream_len: usize,
    /// Extracted downstream flanking length.
    pub downstream_len: usize,
    /// Top genus matches with scores: [(genus, score), ...].
    pub top_matches: Vec<(String, f64)>,
    /// SNP verification status (for point mutation ARGs).
    pub snp_status: SnpStatus,
    /// Extracted upstream flanking sequence (retained for export/inspection of
    /// unresolved loci — e.g. flanking present but no flanking-DB match).
    pub upstream_seq: String,
    /// Extracted downstream flanking sequence (retained for export/inspection).
    pub downstream_seq: String,
    /// Replicon context of the matched reference flanking: "plasmid" (mobile, genus
    /// unreliable), "chromosome", "ambiguous", or "NA" (no plasmid list loaded).
    pub context: String,
    /// Number of genera within GENUS_TIE_PCT of the top score (from the full score
    /// list, not just top-5). ≥2 means the call cannot be pinned to a single genus.
    pub n_genera_tied: usize,
    /// Best species (binomial) when a species map is loaded and matches clear the
    /// (higher) species identity threshold; None if unavailable/undeterminable.
    pub species: Option<String>,
    /// Top species matches with scores, for multi-species reporting.
    pub species_top_matches: Vec<(String, f64)>,
    /// Number of species within GENUS_TIE_PCT of the top species score.
    pub n_species_tied: usize,
    /// Credible set: genera accumulating posterior mass to the conformal threshold, each
    /// with its posterior — the structured form of the answer (a set of taxa, not a single
    /// call). The TSV emits its grouped display (`credible_set_grouped`); this Vec is kept
    /// for programmatic consumers.
    #[allow(dead_code)]
    pub credible_set: Vec<(String, f64)>,
    /// Posterior mass covered by `credible_set` (≈0.9+). UNCALIBRATED until conformal
    /// calibration is fitted on a ground-truth set; treat as a raw confidence for now.
    pub support: f64,
    /// Mash radius of the credible set: max distance from its medoid to any member. 0 for
    /// a single-genus call; larger = the shared context spans a broader clade.
    pub resolution_distance: f64,
    /// Why resolution stopped: "query" (flank truncated by the contig edge), "biology"
    /// (full flank but context genuinely shared across genera), or "none" (single genus).
    pub limited_by: String,
    /// Ragged rank of the answer: the LCA rank of the credible set (species/genus/family/
    /// …/root). This is the taxonomic level at which the ARG's context is actually shared.
    pub resolution_rank: String,
    /// LCA taxon name at `resolution_rank` (e.g. "Enterobacteriaceae").
    pub resolution_taxon: String,
    /// The credible set rendered with its close-genera sub-clusters grouped (see
    /// [`group_credible_set`]): `g1(p1),g2(p2) | g3(p3)`. This is the display form of
    /// `credible_set`; empty when there is none.
    pub credible_set_grouped: String,
}

impl Default for GenusResult {
    fn default() -> Self {
        Self {
            arg_name: String::new(),
            contig_name: String::new(),
            genus: None,
            confidence: 0.0,
            specificity: 0.0,
            upstream_len: 0,
            downstream_len: 0,
            top_matches: vec![],
            snp_status: SnpStatus::NotApplicable,
            upstream_seq: String::new(),
            downstream_seq: String::new(),
            context: "NA".to_string(),
            n_genera_tied: 0,
            species: None,
            species_top_matches: vec![],
            n_species_tied: 0,
            credible_set: vec![],
            support: 0.0,
            resolution_distance: 0.0,
            limited_by: "none".to_string(),
            resolution_rank: "NA".to_string(),
            resolution_taxon: "NA".to_string(),
            credible_set_grouped: String::new(),
        }
    }
}

/// Flanking database record from FDB file.
///
/// Represents a single flanking sequence entry in the database.
#[derive(Debug, Clone)]
pub struct FlankingRecord {
    /// Source contig identifier.
    pub contig: String,
    /// Source organism genus.
    pub genus: String,
    /// Upstream flanking sequence.
    pub upstream: String,
    /// Downstream flanking sequence.
    pub downstream: String,
}

// ============================================================================
// FDB Index Entry
// ============================================================================

/// Index entry for compressed gene blocks in FDB format.
#[derive(Debug, Clone)]
struct FdbIndexEntry {
    offset: u64,
    compressed_len: u32,
    record_count: u32,
}

// ============================================================================
// Flanking Database Reader
// ============================================================================

/// Reader for compressed flanking database (.fdb) files.
///
/// The FDB format stores flanking sequences grouped by gene,
/// with zstd compression for each gene block.
///
/// # File Format
/// ```text
/// [Header: 8 bytes magic + 4 bytes version + 4 bytes gene_count + 8 bytes index_offset]
/// [Gene blocks: zstd compressed TSV data]
/// [Index: gene name -> (offset, compressed_len, record_count)]
/// ```
pub struct FlankingDatabase {
    file: File,
    index: FxHashMap<String, FdbIndexEntry>,
    /// Maps gene name (first field before '|') to full key
    /// Enables lookup by just gene name when FDB uses full header format (e.g., "mexQ|DRUG|CLASS|CODE")
    gene_name_to_key: FxHashMap<String, String>,
}

impl FlankingDatabase {
    /// Opens a flanking database file.
    ///
    /// Reads and validates the header, then loads the gene index into memory.
    ///
    /// # Arguments
    /// * `path` - Path to the .fdb file
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let mut file = File::open(path.as_ref())
            .with_context(|| format!("Failed to open fdb: {}", path.as_ref().display()))?;

        // Read and verify header
        let mut magic = [0u8; 8];
        file.read_exact(&mut magic)?;
        if &magic != FDB_MAGIC {
            anyhow::bail!("Invalid fdb magic");
        }

        let mut buf4 = [0u8; 4];
        let mut buf8 = [0u8; 8];

        file.read_exact(&mut buf4)?;
        let _version = u32::from_le_bytes(buf4);

        file.read_exact(&mut buf4)?;
        let gene_count = u32::from_le_bytes(buf4);

        file.read_exact(&mut buf8)?;
        let index_offset = u64::from_le_bytes(buf8);

        // Read index from end of file
        file.seek(SeekFrom::Start(index_offset))?;
        let mut index = FxHashMap::default();

        for _ in 0..gene_count {
            let mut buf2 = [0u8; 2];
            file.read_exact(&mut buf2)?;
            let name_len = u16::from_le_bytes(buf2) as usize;

            let mut name_buf = vec![0u8; name_len];
            file.read_exact(&mut name_buf)?;
            let gene = String::from_utf8(name_buf)?;

            file.read_exact(&mut buf8)?;
            let offset = u64::from_le_bytes(buf8);

            file.read_exact(&mut buf4)?;
            let compressed_len = u32::from_le_bytes(buf4);

            file.read_exact(&mut buf4)?;
            let record_count = u32::from_le_bytes(buf4);

            index.insert(gene, FdbIndexEntry {
                offset,
                compressed_len,
                record_count,
            });
        }

        // Build gene_name -> full_key mapping for flexible lookup
        // This handles cases where FDB keys are "gene|class|class|code" but lookup uses just "gene"
        let mut gene_name_to_key = FxHashMap::default();
        for full_key in index.keys() {
            // Extract gene name (first field before '|')
            let gene_name = full_key.split('|').next().unwrap_or(full_key);
            // Only add if not already present (first match wins)
            if !gene_name_to_key.contains_key(gene_name) {
                gene_name_to_key.insert(gene_name.to_string(), full_key.clone());
            }
        }

        Ok(Self { file, index, gene_name_to_key })
    }

    /// Checks if a gene exists in the database.
    /// First tries direct key lookup, then falls back to gene name mapping.
    pub fn has_gene(&self, gene: &str) -> bool {
        // Try direct lookup first
        if self.index.contains_key(gene) {
            return true;
        }
        // Fall back to gene name mapping (for "gene" -> "gene|class|class|code" lookup)
        self.gene_name_to_key.contains_key(gene)
    }

    /// Retrieves all flanking records for a specific gene.
    ///
    /// Decompresses the gene block on demand.
    /// Supports both direct key lookup and gene name mapping (e.g., "mexQ" -> "mexQ|DRUG|CLASS|CODE").
    pub fn get_gene_records(&self, gene: &str) -> Result<Vec<FlankingRecord>> {
        // Try direct lookup, then gene name mapping
        let lookup_key = if self.index.contains_key(gene) {
            gene.to_string()
        } else if let Some(full_key) = self.gene_name_to_key.get(gene) {
            full_key.clone()
        } else {
            anyhow::bail!("Gene not found: {}", gene);
        };

        let entry = self.index.get(&lookup_key)
            .ok_or_else(|| anyhow::anyhow!("Gene not found in index: {}", lookup_key))?;

        // Positioned read at the block offset — does not move a shared file cursor, so
        // this is &self and safe to call concurrently from many threads.
        let mut compressed = vec![0u8; entry.compressed_len as usize];
        self.file.read_exact_at(&mut compressed, entry.offset)?;

        // Decompress with zstd
        let decompressed = zstd::decode_all(&compressed[..])?;
        let content = String::from_utf8(decompressed)?;

        // Parse TSV content
        let mut records = Vec::with_capacity(entry.record_count as usize);
        let mut lines = content.lines();

        // Skip header line
        let _header = lines.next();

        for line in lines {
            if line.is_empty() {
                continue;
            }
            let fields: Vec<&str> = line.split('\t').collect();
            // Format: Gene | Contig | Genus | Start | End | Upstream | Downstream
            if fields.len() < 7 {
                continue;
            }

            records.push(FlankingRecord {
                contig: fields[1].to_string(),
                genus: fields[2].to_string(),
                upstream: fields[5].to_string(),
                downstream: fields[6].to_string(),
            });
        }

        Ok(records)
    }

}

// ============================================================================
// Genus Classifier
// ============================================================================

/// Mash distance between two genus representatives from a neighbor map; 0 if identical,
/// 1.0 (≈ w→0) if the pair is absent (pruned as too distant, or a novel label).
fn ou_lookup(neighbors: &FxHashMap<String, FxHashMap<String, f64>>, a: &str, b: &str) -> f64 {
    if a == b {
        return 0.0;
    }
    if let Some(m) = neighbors.get(a) {
        if let Some(&d) = m.get(b) {
            return d;
        }
    }
    if let Some(m) = neighbors.get(b) {
        if let Some(&d) = m.get(a) {
            return d;
        }
    }
    ABSENT_PATRISTIC
}

/// Distance assigned when a genus pair is absent from the GTDB patristic table. Set beyond the
/// inter-domain median (~2.76 subs/site) so an absent relative contributes ≈0 kernel weight
/// (exp(-3/λ) with λ=0.3 ≈ 5e-5): "unknown neighbor" is treated as "maximally far", never as a
/// close borrow. (Was 1.0 on the old mash scale where distances saturated near 0.34.)
const ABSENT_PATRISTIC: f64 = 3.0;

/// Phylogenetic OU-kernel posterior over genera from per-genus likelihoods.
///
/// post[T] = Σ_{T'} w(T,T')·L[T'] / Σ_{T'} w(T,T'),  w(d) = exp(-d/λ), then normalized to
/// sum 1. Sparse genera borrow evidence from close relatives (Escherichia↔Shigella
/// d≈0.03), so a novel query's mass climbs to the correct clade instead of collapsing
/// onto one reference — this is what kills the argmax/DB-depth bias. Returns (genus,
/// posterior) sorted descending. With an empty neighbor map every off-diagonal weight is
/// exp(-ABSENT_PATRISTIC/λ)≈0, so the result degenerates to the normalized likelihoods (no
/// smoothing).
fn ou_posterior(
    likelihoods: &FxHashMap<String, f64>,
    neighbors: &FxHashMap<String, FxHashMap<String, f64>>,
    lambda: f64,
) -> Vec<(String, f64)> {
    let genera: Vec<&String> = likelihoods.keys().collect();
    let mut post: Vec<(String, f64)> = Vec::with_capacity(genera.len());
    for &ta in &genera {
        let (mut num, mut den) = (0.0f64, 0.0f64);
        for &tb in &genera {
            let w = (-ou_lookup(neighbors, ta, tb) / lambda).exp();
            num += w * likelihoods[tb];
            den += w;
        }
        post.push((ta.clone(), if den > 0.0 { num / den } else { 0.0 }));
    }
    let z: f64 = post.iter().map(|(_, p)| *p).sum();
    if z > 0.0 {
        for p in post.iter_mut() {
            p.1 /= z;
        }
    }
    post.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    post
}

/// A credible set whose patristic radius is within this is "phylogenetically coherent" — a
/// tight cluster of close genera — and is reported as the genus list rather than rolled up.
/// The criterion is phylogenetic COHERENCE (the radius), NOT the number of genera: two far
/// genera and ten close ones are different answers even at the same family. Scaled to the GTDB
/// patristic table (genus-pair median ≈0.19, family-pair median ≈0.60), so 0.5 sits just below
/// the family diameter: within-family sisters list individually; a set spilling past one family
/// rolls up to its LCA. (Was 0.12 on the old mash scale, which saturated near 0.30 at family.)
const GENUS_COHERENCE_RADIUS: f64 = 0.5;

/// Renders the credible set's resolution notation. The credible set itself is the answer
/// (in `Credible_Set`); this chooses the human label for `Resolution_Rank`/`_Taxon` from the
/// set's phylogenetic COHERENCE (its mash `radius`), not its size:
///   - one genus                    → ("genus", that genus)
///   - several genera, radius tight  → ("genus", "Escherichia|Shigella") — a coherent
///     cluster of close genera, named individually (strictly more specific than the family)
///   - several genera, radius spread → (LCA rank, LCA taxon) — the genera fan out across a
///     family/order/…, so the clade name is the honest summary
/// {Escherichia, Shigella} (radius ≈0.03) lists both genera; a set fanning across
/// Enterobacteriaceae (radius ≈0.17) reports "Enterobacteriaceae" even if it's only 2 genera.
fn render_resolution(members: &[&str], lca_rank: &str, lca_taxon: &str, radius: f64)
    -> (String, String) {
    if members.len() == 1 {
        ("genus".to_string(), members[0].to_string())
    } else if radius <= GENUS_COHERENCE_RADIUS {
        ("genus".to_string(), members.join("|"))
    } else {
        (lca_rank.to_string(), lca_taxon.to_string())
    }
}

/// Renders the credible set with its internal phylogenetic structure exposed: genera are
/// single-linkage clustered by mash distance (two genera join if within `threshold`), so a
/// tight sub-cluster (Escherichia+Pseudescherichia) is shown as one group and the genera it
/// fans out to are shown separately. Clusters are sorted by total posterior mass; within a
/// cluster, members by posterior. Format: `g1(p1),g2(p2) | g3(p3) | g4(p4)` — commas inside
/// a tight cluster, ` | ` between clusters. A flat list hides that the mass is really on one
/// close group plus a few outliers; this shows it.
fn group_credible_set(
    members: &[(String, f64)],
    neighbors: &FxHashMap<String, FxHashMap<String, f64>>,
    threshold: f64,
) -> String {
    let n = members.len();
    if n == 0 {
        return "NA".to_string();
    }
    if n == 1 {
        return format!("{}({:.2})", members[0].0, members[0].1);
    }
    // Single-linkage union-find over pairs within `threshold`.
    let mut parent: Vec<usize> = (0..n).collect();
    fn find(parent: &mut [usize], mut x: usize) -> usize {
        while parent[x] != x {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        x
    }
    for i in 0..n {
        for j in (i + 1)..n {
            if ou_lookup(neighbors, &members[i].0, &members[j].0) <= threshold {
                let (ri, rj) = (find(&mut parent, i), find(&mut parent, j));
                if ri != rj {
                    parent[ri] = rj;
                }
            }
        }
    }
    let mut clusters: FxHashMap<usize, Vec<usize>> = FxHashMap::default();
    for i in 0..n {
        let r = find(&mut parent, i);
        clusters.entry(r).or_default().push(i);
    }
    let mut rendered: Vec<(f64, String)> = clusters
        .values()
        .map(|idxs| {
            let mut mem = idxs.clone();
            mem.sort_by(|&a, &b| {
                members[b].1.partial_cmp(&members[a].1).unwrap_or(std::cmp::Ordering::Equal)
            });
            let mass: f64 = mem.iter().map(|&k| members[k].1).sum();
            let s = mem.iter()
                .map(|&k| format!("{}({:.2})", members[k].0, members[k].1))
                .collect::<Vec<_>>()
                .join(",");
            (mass, s)
        })
        .collect();
    rendered.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
    rendered.into_iter().map(|(_, s)| s).collect::<Vec<_>>().join(" | ")
}

/// Lowest common ancestor of a set of genera over their lineages: the finest standard
/// rank at which every member shares one non-empty taxon. Returns (rank, taxon). A single
/// genus resolves to ("genus", that genus); genera that agree only at the root give
/// ("root", "unclassified"). This renders the credible set as one ragged-rank label —
/// multi-genus {Escherichia, Salmonella} → ("family", "Enterobacteriaceae").
fn lca_of(genera: &[&str], lineage: &FxHashMap<String, [String; 7]>) -> (String, String) {
    let lins: Vec<&[String; 7]> = genera.iter().filter_map(|g| lineage.get(*g)).collect();
    if lins.is_empty() {
        // No lineage info: fall back to the genus label itself if singleton.
        return if genera.len() == 1 {
            ("genus".to_string(), genera[0].to_string())
        } else {
            ("root".to_string(), "unclassified".to_string())
        };
    }
    // Finest → coarsest, capped at genus (idx 5): the credible set is a set of GENERA, so
    // the species column (a representative genome's species) is not a valid resolution —
    // species-level calls come from the separate species tally, not this LCA.
    for r in (0..6).rev() {
        let first = &lins[0][r];
        if !first.is_empty() && lins.iter().all(|l| &l[r] == first) {
            return (LINEAGE_RANKS[r].to_string(), first.clone());
        }
    }
    ("root".to_string(), "unclassified".to_string())
}

/// Minimap2-based genus classifier using flanking sequence alignment.
///
/// Classifies source genus by aligning extracted flanking sequences
/// against reference flanking sequences in the database.
pub struct GenusClassifier {
    db: FlankingDatabase,
    minimap2_path: String,
    min_identity: f64,
    min_align_len: usize,
    max_flanking: usize,
    /// Source contig accessions known to be plasmids (from PLSDB-derived flanking).
    /// Empty if no plasmid list was provided → Context reported as "NA".
    plasmid_contigs: rustc_hash::FxHashSet<String>,
    /// Higher identity threshold (0-1) required to call species. 0 disables species.
    species_identity: f64,
    /// Source contig accession → species (binomial). Empty disables species calls.
    species_map: FxHashMap<String, String>,
    /// Plasmid-fraction thresholds for the Context call: frac >= `plasmid_hi`
    /// → "plasmid", frac <= `plasmid_lo` → "chromosome", else "ambiguous".
    plasmid_hi: f64,
    plasmid_lo: f64,
    /// Optional cap on reference flanks aligned per locus (0 = unlimited). A few genes
    /// carry tens of thousands of flanks; capping trades some accuracy for speed.
    max_ref_flanks: usize,
    /// Phylogenetic OU kernel: genus -> [(neighbor_genus, mash_distance)]. Loaded from a
    /// sibling `genus_dist.tsv` next to the FDB (genus-representative mash distances,
    /// k=21 s=100000). Empty => no smoothing (posterior = normalized per-genus likelihood).
    genus_neighbors: FxHashMap<String, FxHashMap<String, f64>>,
    /// OU relaxation length (= 1/alpha). ln2*lambda = distance at which borrowed
    /// evidence halves. NOT Pagel's lambda.
    kernel_lambda: f64,
    /// Likelihood sharpness: L = exp(-(1 - ani)/tau) per alignment, ani = identity/100.
    kernel_tau: f64,
    /// genus -> its standard 7-rank lineage [superkingdom, phylum, class, order, family,
    /// genus, species]. Loaded from sibling `genus_lineage.tsv`. Enables reporting the
    /// credible set's LCA as a ragged rank (e.g. multi-genus Enterobacteriaceae → family).
    genus_lineage: FxHashMap<String, [String; 7]>,
    /// Conformal credible-set mass threshold θ for the requested target coverage, read
    /// from sibling `conformal.tsv`. The set grows until its posterior mass reaches θ,
    /// which guarantees the true host's clade is covered at ≥ the target rate. Falls back
    /// to `DEFAULT_CREDIBLE_MASS` (uncalibrated) when the table is absent.
    conformal_theta: f64,
}

/// Uncalibrated fallback credible-set mass when no `conformal.tsv` is available.
const DEFAULT_CREDIBLE_MASS: f64 = 0.9;

/// Standard ranks, coarsest → finest, matching the columns of `genus_lineage.tsv`.
const LINEAGE_RANKS: [&str; 7] =
    ["superkingdom", "phylum", "class", "order", "family", "genus", "species"];

impl GenusClassifier {
    /// Creates a new genus classifier.
    ///
    /// # Arguments
    /// * `db_path` - Path to flanking database (.fdb)
    /// * `minimap2_path` - Path to minimap2 executable
    /// * `min_identity` - Minimum alignment identity (0-1)
    /// * `min_align_len` - Minimum alignment length in bp
    /// * `max_flanking` - Maximum flanking length to extract
    pub fn new<P: AsRef<Path>>(
        db_path: P,
        minimap2_path: &str,
        min_identity: f64,
        min_align_len: usize,
        max_flanking: usize,
        plasmid_contigs_path: Option<&Path>,
        species_identity: f64,
        species_map_path: Option<&Path>,
        plasmid_hi: f64,
        plasmid_lo: f64,
        max_ref_flanks: usize,
        // Requested coverage for the conformal credible set (e.g. 0.9 = the true host's
        // clade is inside the reported set >=90% of the time). Selects θ from conformal.tsv.
        target_coverage: f64,
    ) -> Result<Self> {
        let db_dir = db_path.as_ref().parent().map(|p| p.to_path_buf());
        let db = FlankingDatabase::open(db_path)?;
        let mut plasmid_contigs = rustc_hash::FxHashSet::default();
        if let Some(p) = plasmid_contigs_path {
            let f = File::open(p).with_context(|| format!("open plasmid list {:?}", p))?;
            for line in BufReader::new(f).lines() {
                let acc = line?.trim().to_string();
                if !acc.is_empty() { plasmid_contigs.insert(acc); }
            }
        }
        let mut species_map = FxHashMap::default();
        if let Some(p) = species_map_path {
            let f = File::open(p).with_context(|| format!("open species map {:?}", p))?;
            for line in BufReader::new(f).lines() {
                let line = line?;
                if let Some((c, sp)) = line.split_once('\t') {
                    let (c, sp) = (c.trim(), sp.trim());
                    if !c.is_empty() && !sp.is_empty() { species_map.insert(c.to_string(), sp.to_string()); }
                }
            }
        }
        // Phylogenetic OU kernel distances: sibling `genus_dist.tsv` next to the FDB.
        // Format: genus_a \t genus_b \t mash_distance (both directions present). Absent
        // => empty map => posterior degenerates to normalized per-genus likelihood.
        let mut genus_neighbors: FxHashMap<String, FxHashMap<String, f64>> = FxHashMap::default();
        {
            let (gd_text, gd_src) =
                context_table_text(db_dir.as_deref(), "genus_dist.tsv", EMB_GENUS_DIST)?;
            let mut npair = 0usize;
            for line in gd_text.lines() {
                let mut it = line.split('\t');
                if let (Some(a), Some(b), Some(d)) = (it.next(), it.next(), it.next()) {
                    if let Ok(d) = d.trim().parse::<f64>() {
                        genus_neighbors.entry(a.to_string()).or_default()
                            .insert(b.to_string(), d);
                        npair += 1;
                    }
                }
            }
            eprintln!("[kernel] loaded {} genus-distance pairs from {}", npair, gd_src);
        }
        // Ragged-rank lineage: sibling `genus_lineage.tsv`, header
        // `genus<TAB>superkingdom..species`. Absent => LCA reporting disabled (genus only).
        let mut genus_lineage: FxHashMap<String, [String; 7]> = FxHashMap::default();
        {
            let (gl_text, gl_src) =
                context_table_text(db_dir.as_deref(), "genus_lineage.tsv", EMB_GENUS_LINEAGE)?;
            for (i, line) in gl_text.lines().enumerate() {
                if i == 0 {
                    continue; // header
                }
                let cols: Vec<&str> = line.split('\t').collect();
                if cols.len() >= 8 {
                    let mut lin: [String; 7] = Default::default();
                    for r in 0..7 {
                        lin[r] = cols[r + 1].to_string();
                    }
                    genus_lineage.insert(cols[0].to_string(), lin);
                }
            }
            eprintln!("[kernel] loaded {} genus lineages from {}", genus_lineage.len(), gl_src);
        }
        // Conformal θ: sibling `conformal.tsv`, rows `target_coverage<TAB>theta<TAB>...`.
        // Pick the row whose target is closest to the requested coverage.
        let mut conformal_theta = DEFAULT_CREDIBLE_MASS;
        let mut conformal_target = 0.0;
        {
            let (cf_text, cf_src) =
                context_table_text(db_dir.as_deref(), "conformal.tsv", EMB_CONFORMAL)?;
            let mut best = f64::MAX;
            for line in cf_text.lines() {
                if line.starts_with('#') || line.trim().is_empty() {
                    continue;
                }
                let c: Vec<&str> = line.split('\t').collect();
                if c.len() >= 2 {
                    if let (Ok(t), Ok(th)) = (c[0].parse::<f64>(), c[1].parse::<f64>()) {
                        let d = (t - target_coverage).abs();
                        if d < best {
                            best = d;
                            conformal_theta = th;
                            conformal_target = t;
                        }
                    }
                }
            }
            if conformal_target > 0.0 {
                eprintln!("[kernel] conformal θ={:.3} for target coverage {:.2} (from {})",
                          conformal_theta, conformal_target, cf_src);
            }
        }
        Ok(Self {
            db,
            minimap2_path: minimap2_path.to_string(),
            min_identity,
            min_align_len,
            max_flanking,
            plasmid_contigs,
            species_identity,
            species_map,
            plasmid_hi,
            plasmid_lo,
            max_ref_flanks,
            genus_neighbors,
            kernel_lambda: 0.3,
            kernel_tau: 0.01,
            genus_lineage,
            conformal_theta,
        })
    }

    /// Phylogenetic OU-kernel posterior over genera from per-genus likelihoods, using
    /// this classifier's loaded genus distances and lambda. Thin wrapper over the free
    /// [`ou_posterior`] (extracted so the kernel math is unit-testable without an FDB).
    fn kernel_posterior(&self, likelihoods: &FxHashMap<String, f64>) -> Vec<(String, f64)> {
        ou_posterior(likelihoods, &self.genus_neighbors, self.kernel_lambda)
    }

    /// Classifies genus for multiple ARG positions.
    ///
    /// Two phases: (1) a serial pass reads every gene's flanking records and genus
    /// distribution from the .fdb — the only work needing `&mut self` (the reader
    /// seeks the file) — deduplicated by gene; (2) the flanking alignment + scoring
    /// per locus, which is independent and CPU/subprocess-bound, runs in parallel.
    /// Each locus writes its own indexed temp files, so there are no conflicts.
    pub fn classify_batch(&mut self, positions: &[ArgPosition], threads: usize) -> Result<Vec<GenusResult>> {
        use rayon::prelude::*;
        use std::collections::HashSet;

        let this: &GenusClassifier = self;
        let pool = rayon::ThreadPoolBuilder::new().num_threads(threads.max(1)).build()?;

        // Unique genes present in the DB — dedup so a gene shared by many loci is read
        // and decompressed only once.
        let mut seen: HashSet<&str> = HashSet::new();
        let unique_genes: Vec<&str> = positions.iter()
            .flat_map(|p| p.members.iter().map(|s| s.as_str()))
            .filter(|&g| this.db.has_gene(g) && seen.insert(g))
            .collect();

        pool.install(|| -> Result<Vec<GenusResult>> {
            // Phase 1 (parallel): positioned reads (read_at) decompress each gene block
            // concurrently; the genus distribution is derived from the same records so
            // the block is decompressed only once.
            let loaded: Vec<(String, Vec<FlankingRecord>, FxHashMap<String, usize>)> = unique_genes
                .par_iter()
                .map(|&g| {
                    let recs = this.db.get_gene_records(g)?;
                    let mut dist: FxHashMap<String, usize> = FxHashMap::default();
                    for rec in &recs { *dist.entry(rec.genus.clone()).or_default() += 1; }
                    Ok((g.to_string(), recs, dist))
                })
                .collect::<Result<Vec<_>>>()?;

            let mut recs_by_gene: FxHashMap<String, Vec<FlankingRecord>> = FxHashMap::default();
            let mut dist_by_gene: FxHashMap<String, FxHashMap<String, usize>> = FxHashMap::default();
            for (g, recs, dist) in loaded {
                recs_by_gene.insert(g.clone(), recs);
                dist_by_gene.insert(g, dist);
            }

            // Phase 2 (parallel): flanking alignment + scoring, one locus per task.
            positions
                .par_iter()
                .enumerate()
                .map(|(idx, pos)| this.classify_prepared(pos, idx, &recs_by_gene, &dist_by_gene))
                .collect()
        })
    }

    /// Classifies genus for a single ARG position using preloaded per-gene data
    /// (so it needs only `&self` and is safe to run in parallel across loci).
    ///
    /// # Algorithm
    /// 1. Extract flanking sequences from contig
    /// 2. Write query and reference FASTA files (indexed by `idx` to avoid conflicts)
    /// 3. Run minimap2 alignment (single-threaded; parallelism is across loci)
    /// 4. Parse PAF and score genus candidates
    /// 5. Return top genus with confidence metrics
    fn classify_prepared(
        &self,
        pos: &ArgPosition,
        idx: usize,
        recs_by_gene: &FxHashMap<String, Vec<FlankingRecord>>,
        dist_by_gene: &FxHashMap<String, FxHashMap<String, usize>>,
    ) -> Result<GenusResult> {
        // Extract flanking sequences
        let (upstream, downstream) = self.extract_flanking_regions(pos);

        let upstream_len = upstream.len();
        let downstream_len = downstream.len();

        // Verify SNP for point mutation genes
        let snp_status = snp::verify_snp(
            &pos.contig_seq,
            &pos.arg_name,
            0,
            pos.arg_end - pos.arg_start,
            pos.arg_start,
            pos.arg_end,
            pos.strand,
        );

        // Require minimum flanking for classification
        if upstream_len < 50 && downstream_len < 50 {
            return Ok(GenusResult {
                arg_name: pos.arg_name.clone(),
                contig_name: pos.contig_name.clone(),
                genus: None,
                confidence: 0.0,
                specificity: 0.0,
                upstream_len,
                downstream_len,
                top_matches: vec![],
                snp_status,
                upstream_seq: upstream.clone(),
                downstream_seq: downstream.clone(),
                context: "NA".to_string(),
                n_genera_tied: 0,
                species: None,
                species_top_matches: vec![],
                n_species_tied: 0,
                credible_set: vec![],
                support: 0.0,
                resolution_distance: 0.0,
                limited_by: "none".to_string(),
                resolution_rank: "NA".to_string(),
                resolution_taxon: "NA".to_string(),
                credible_set_grouped: String::new(),
            });
        }

        // Check if gene exists in database (any tied member present == has_gene)
        if !pos.members.iter().any(|m| recs_by_gene.contains_key(m)) {
            return Ok(GenusResult {
                arg_name: pos.arg_name.clone(),
                contig_name: pos.contig_name.clone(),
                genus: None,
                confidence: 0.0,
                specificity: 0.0,
                upstream_len,
                downstream_len,
                top_matches: vec![("gene_not_in_db".to_string(), 0.0)],
                snp_status,
                upstream_seq: upstream.clone(),
                downstream_seq: downstream.clone(),
                context: "NA".to_string(),
                n_genera_tied: 0,
                species: None,
                species_top_matches: vec![],
                n_species_tied: 0,
                credible_set: vec![],
                support: 0.0,
                resolution_distance: 0.0,
                limited_by: "none".to_string(),
                resolution_rank: "NA".to_string(),
                resolution_taxon: "NA".to_string(),
                credible_set_grouped: String::new(),
            });
        }

        // Union the flanking references of all tied redundant refs at this locus, so the
        // same physical gene's evidence from every source DB is used together.
        let mut ref_records: Vec<&FlankingRecord> = pos.members.iter()
            .filter_map(|m| recs_by_gene.get(m))
            .flatten()
            .collect();
        // Optional speed cap (self.max_ref_flanks, 0 = unlimited/default): some genes
        // carry tens of thousands of flanks (up to ~128k), which dominate wall-clock
        // (writing + indexing + aligning). An evenly-strided sample keeps most of the
        // genus signal but DOES shift the call on those heavy loci, so it is off by
        // default; the specificity denominator always uses the full counts.
        if self.max_ref_flanks > 0 && ref_records.len() > self.max_ref_flanks {
            let step = ref_records.len() / self.max_ref_flanks;
            ref_records = ref_records.iter().step_by(step).copied().take(self.max_ref_flanks).collect();
        }
        if ref_records.is_empty() {
            return Ok(GenusResult {
                arg_name: pos.arg_name.clone(),
                contig_name: pos.contig_name.clone(),
                genus: None,
                confidence: 0.0,
                specificity: 0.0,
                upstream_len,
                downstream_len,
                top_matches: vec![("no_ref_records".to_string(), 0.0)],
                snp_status,
                upstream_seq: upstream.clone(),
                downstream_seq: downstream.clone(),
                context: "NA".to_string(),
                n_genera_tied: 0,
                species: None,
                species_top_matches: vec![],
                n_species_tied: 0,
                credible_set: vec![],
                support: 0.0,
                resolution_distance: 0.0,
                limited_by: "none".to_string(),
                resolution_rank: "NA".to_string(),
                resolution_taxon: "NA".to_string(),
                credible_set_grouped: String::new(),
            });
        }

        // Create temporary files for alignment. Prefer tmpfs (/dev/shm) to avoid disk
        // I/O, and key names by pid+idx so parallel loci never collide.
        let temp_dir = if Path::new("/dev/shm").is_dir() {
            std::path::PathBuf::from("/dev/shm")
        } else {
            std::env::temp_dir()
        };
        let pid = std::process::id();
        let query_path = temp_dir.join(format!("argenus_query_{}_{}.fas", pid, idx));
        let ref_path = temp_dir.join(format!("argenus_ref_{}_{}.fas", pid, idx));
        let paf_path = temp_dir.join(format!("argenus_align_{}_{}.paf", pid, idx));

        // Write query FASTA
        {
            let mut query_file = BufWriter::new(File::create(&query_path)?);
            if !upstream.is_empty() {
                writeln!(query_file, ">upstream")?;
                writeln!(query_file, "{}", upstream)?;
            }
            if !downstream.is_empty() {
                writeln!(query_file, ">downstream")?;
                writeln!(query_file, "{}", downstream)?;
            }
        }

        // Write reference FASTA (grouped by genus)
        {
            let mut ref_file = BufWriter::new(File::create(&ref_path)?);
            for (i, rec) in ref_records.iter().enumerate() {
                if !rec.upstream.is_empty() {
                    writeln!(ref_file, ">{}|{}|up_{}", rec.genus, rec.contig, i)?;
                    writeln!(ref_file, "{}", rec.upstream)?;
                }
                if !rec.downstream.is_empty() {
                    writeln!(ref_file, ">{}|{}|down_{}", rec.genus, rec.contig, i)?;
                    writeln!(ref_file, "{}", rec.downstream)?;
                }
            }
        }

        // Run minimap2 with sr preset for short queries
        let output = Command::new(&self.minimap2_path)
            .args(["-x", "sr", "-t", "1", "-c", "--secondary=yes", "-N", "100", "-k", "15", "-w", "5"])
            .arg(&ref_path)
            .arg(&query_path)
            .arg("-o").arg(&paf_path)
            .stderr(std::process::Stdio::null())
            .output()
            .context("Failed to run minimap2")?;

        if !output.status.success() {
            // Cleanup and return error result
            let _ = std::fs::remove_file(&query_path);
            let _ = std::fs::remove_file(&ref_path);
            let _ = std::fs::remove_file(&paf_path);

            return Ok(GenusResult {
                arg_name: pos.arg_name.clone(),
                contig_name: pos.contig_name.clone(),
                genus: None,
                confidence: 0.0,
                specificity: 0.0,
                upstream_len,
                downstream_len,
                top_matches: vec![("minimap2_failed".to_string(), 0.0)],
                snp_status,
                upstream_seq: upstream.clone(),
                downstream_seq: downstream.clone(),
                context: "NA".to_string(),
                n_genera_tied: 0,
                species: None,
                species_top_matches: vec![],
                n_species_tied: 0,
                credible_set: vec![],
                support: 0.0,
                resolution_distance: 0.0,
                limited_by: "none".to_string(),
                resolution_rank: "NA".to_string(),
                resolution_taxon: "NA".to_string(),
                credible_set_grouped: String::new(),
            });
        }

        // Parse PAF and calculate genus scores + plasmid provenance + species scores.
        let (genus_likelihood, genus_confidence, plasmid_frac, species_scores) =
            self.calculate_genus_scores(&paf_path)?;

        // Cleanup temporary files
        let _ = std::fs::remove_file(&query_path);
        let _ = std::fs::remove_file(&ref_path);
        let _ = std::fs::remove_file(&paf_path);

        // Context from provenance of the matched reference flanking (mobility proxy).
        // "plasmid" = the flanking mostly matched PLSDB-derived plasmid references
        // (genus is then unreliable — the ARG is on a mobile element). NA when no
        // plasmid list is loaded OR the flanking matched nothing (no basis to judge).
        let context = if self.plasmid_contigs.is_empty() || genus_likelihood.is_empty() {
            "NA".to_string()
        } else if plasmid_frac >= self.plasmid_hi {
            "plasmid".to_string()
        } else if plasmid_frac <= self.plasmid_lo {
            "chromosome".to_string()
        } else {
            "ambiguous".to_string()
        };

        // Calculate genus specificity from the DB, combined over all tied members.
        let mut genus_dist: FxHashMap<String, usize> = FxHashMap::default();
        for m in &pos.members {
            if let Some(d) = dist_by_gene.get(m) {
                for (k, v) in d { *genus_dist.entry(k.clone()).or_default() += *v; }
            }
        }
        let total_in_db: usize = genus_dist.values().sum();

        // Determine best genus via the phylogenetic OU-kernel posterior (replaces the
        // old count-weighted identity score). `sorted_scores` now holds (genus, posterior)
        // in descending posterior order; DB depth no longer biases the ranking.
        let sorted_scores: Vec<(String, f64)> = self.kernel_posterior(&genus_likelihood);

        let (genus, confidence, specificity) = if let Some((best_genus, _post)) = sorted_scores.first() {
            let genus_count = genus_dist.get(best_genus).copied().unwrap_or(0);
            let specificity = if total_in_db > 0 {
                (genus_count as f64 / total_in_db as f64) * 100.0
            } else {
                0.0
            };
            // Confidence stays on the 0-100 identity scale (mean alignment identity of the
            // called genus), independent of the posterior used for ranking.
            let conf = genus_confidence.get(best_genus).copied().unwrap_or(0.0);
            (Some(best_genus.clone()), conf, specificity)
        } else {
            (None, 0.0, 0.0)
        };

        // Credible set: the smallest prefix of the posterior-sorted genera whose mass
        // reaches the conformal threshold θ. θ is calibrated so that the true host's clade
        // falls inside the set's LCA at ≥ the target coverage — this is what makes Support
        // mean something. Uncalibrated (θ = DEFAULT_CREDIBLE_MASS) when conformal.tsv is
        // absent. Single genus = resolved; several = shared/mobile context.
        let mut credible_set: Vec<(String, f64)> = Vec::new();
        let mut support = 0.0f64;
        for (g, p) in &sorted_scores {
            credible_set.push((g.clone(), *p));
            support += *p;
            if support >= self.conformal_theta {
                break;
            }
        }
        let n_genera_tied = credible_set.len();

        // Resolution distance = mash radius of the credible set around its posterior-
        // weighted medoid (the member minimizing Σ p·d to the others). 0 for one genus.
        let resolution_distance = if credible_set.len() < 2 {
            0.0
        } else {
            let members: Vec<&str> = credible_set.iter().map(|(g, _)| g.as_str()).collect();
            let medoid = members.iter().min_by(|&&a, &&b| {
                let sa: f64 = credible_set.iter()
                    .map(|(g, p)| p * ou_lookup(&self.genus_neighbors, a, g)).sum();
                let sb: f64 = credible_set.iter()
                    .map(|(g, p)| p * ou_lookup(&self.genus_neighbors, b, g)).sum();
                sa.partial_cmp(&sb).unwrap_or(std::cmp::Ordering::Equal)
            }).copied().unwrap_or(members[0]);
            members.iter()
                .map(|&g| ou_lookup(&self.genus_neighbors, medoid, g))
                .fold(0.0f64, f64::max)
        };

        // Limited-by: was resolution capped by the query (contig edge cut the flank) or by
        // biology (full flank, but the context is genuinely shared across genera)?
        let reach_up = pos.arg_start.min(self.max_flanking);
        let reach_dn = pos.contig_len.saturating_sub(pos.arg_end).min(self.max_flanking);
        let contig_limited = reach_up < self.max_flanking || reach_dn < self.max_flanking;
        let limited_by = if credible_set.len() < 2 {
            "none".to_string()
        } else if contig_limited {
            "query".to_string()
        } else {
            "biology".to_string()
        };

        // Ragged rank: the LCA of the credible set is the taxonomic level at which this
        // ARG's context is actually shared (single genus → genus; several → their family
        // or higher). This is the honest answer, not a forced single genus.
        let (resolution_rank, resolution_taxon) = if credible_set.is_empty() {
            ("NA".to_string(), "NA".to_string())
        } else {
            let members: Vec<&str> = credible_set.iter().map(|(g, _)| g.as_str()).collect();
            let (lca_rank, lca_taxon) = lca_of(&members, &self.genus_lineage);
            render_resolution(&members, &lca_rank, &lca_taxon, resolution_distance)
        };
        let credible_set_grouped = if credible_set.is_empty() {
            String::new()
        } else {
            group_credible_set(&credible_set, &self.genus_neighbors, GENUS_COHERENCE_RADIUS)
        };

        let top_matches: Vec<(String, f64)> = sorted_scores.into_iter().take(5).collect();

        // Species resolution (parallel to genus, from the higher-threshold tally).
        let mut sorted_species: Vec<(String, f64)> = species_scores.into_iter().collect();
        sorted_species.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        let species = sorted_species.first().map(|(s, _)| s.clone());
        let n_species_tied = match sorted_species.first() {
            Some((_, best)) => sorted_species.iter()
                .filter(|(s, sc)| !s.is_empty() && *sc >= best - GENUS_TIE_PCT)
                .count(),
            None => 0,
        };
        let species_top_matches: Vec<(String, f64)> = sorted_species.into_iter().take(5).collect();

        Ok(GenusResult {
            arg_name: pos.arg_name.clone(),
            contig_name: pos.contig_name.clone(),
            genus,
            confidence,
            specificity,
            upstream_len,
            downstream_len,
            top_matches,
            snp_status,
            upstream_seq: upstream.clone(),
            downstream_seq: downstream.clone(),
            context,
            n_genera_tied,
            species,
            species_top_matches,
            n_species_tied,
            credible_set,
            support,
            resolution_distance,
            limited_by,
            resolution_rank,
            resolution_taxon,
            credible_set_grouped,
        })
    }

    /// Extracts flanking sequences from a contig.
    ///
    /// Handles strand orientation automatically.
    fn extract_flanking_regions(&self, pos: &ArgPosition) -> (String, String) {
        let seq = &pos.contig_seq;

        // Extract upstream (before ARG)
        let upstream_end = pos.arg_start;
        let upstream_start = upstream_end.saturating_sub(self.max_flanking);
        let upstream = if upstream_end > upstream_start {
            seq[upstream_start..upstream_end].to_string()
        } else {
            String::new()
        };

        // Extract downstream (after ARG)
        let downstream_start = pos.arg_end;
        let downstream_end = (downstream_start + self.max_flanking).min(seq.len());
        let downstream = if downstream_end > downstream_start {
            seq[downstream_start..downstream_end].to_string()
        } else {
            String::new()
        };

        // Handle reverse strand
        if pos.strand == '-' {
            (reverse_complement(&downstream), reverse_complement(&upstream))
        } else {
            (upstream, downstream)
        }
    }

    /// Parses PAF alignment file and calculates genus scores.
    /// Parses the PAF and returns per-genus (mean likelihood, mean identity), the plasmid
    /// fraction, and per-species mean identity. The kernel posterior consumes the
    /// likelihoods; the identities feed confidence/context (kept on the 0-100 scale).
    fn calculate_genus_scores(&self, paf_path: &Path)
        -> Result<(FxHashMap<String, f64>, FxHashMap<String, f64>, f64, FxHashMap<String, f64>)> {
        let file = File::open(paf_path)?;
        let reader = BufReader::new(file);

        let mut genus_matches: FxHashMap<String, Vec<f64>> = FxHashMap::default();
        let min_identity_pct = self.min_identity * 100.0;
        // Track plasmid provenance of the matched reference flanking records.
        let (mut plasmid_hits, mut total_hits) = (0usize, 0usize);
        // Species tally (only alignments clearing the higher species threshold whose
        // source contig has a species label).
        let mut species_matches: FxHashMap<String, Vec<f64>> = FxHashMap::default();
        let species_pct = self.species_identity * 100.0;
        let do_species = self.species_identity > 0.0 && !self.species_map.is_empty();

        for line in reader.lines() {
            let line = line?;
            let fields: Vec<&str> = line.split('\t').collect();
            if fields.len() < 12 {
                continue;
            }

            let block_len: usize = fields[10].parse().unwrap_or(0);
            let matches: usize = fields[9].parse().unwrap_or(0);

            if block_len < self.min_align_len {
                continue;
            }

            let identity = if block_len > 0 {
                (matches as f64 / block_len as f64) * 100.0
            } else {
                0.0
            };

            if identity < min_identity_pct {
                continue;
            }

            // Target name is "genus|contig|direction_idx".
            let target_name = fields[5];
            let mut toks = target_name.split('|');
            let genus = toks.next().unwrap_or("");
            let contig = toks.next().unwrap_or("");
            // Skip records with an empty genus label (DB label gaps) — they must not
            // win the classification (fix for the empty-"" winner bug).
            if genus.is_empty() {
                continue;
            }
            total_hits += 1;
            if !contig.is_empty() && self.plasmid_contigs.contains(contig) {
                plasmid_hits += 1;
            }
            genus_matches.entry(genus.to_string()).or_default().push(identity);

            // Species tally: stricter identity + a species label for this contig.
            if do_species && identity >= species_pct {
                if let Some(sp) = self.species_map.get(contig) {
                    species_matches.entry(sp.clone()).or_default().push(identity);
                }
            }
        }

        // Per-genus mean likelihood (kernel input) and mean identity (confidence). The
        // likelihood is averaged WITHIN genus so DB depth (many flanks of one genus) can
        // no longer inflate the score — the old count bonus did exactly that.
        let tau = self.kernel_tau;
        let mut genus_likelihood: FxHashMap<String, f64> = FxHashMap::default();
        let mut genus_confidence: FxHashMap<String, f64> = FxHashMap::default();
        for (genus, scores) in genus_matches {
            if scores.is_empty() {
                continue;
            }
            let n = scores.len() as f64;
            let mean_id = scores.iter().sum::<f64>() / n;
            let mean_lik = scores.iter()
                .map(|id| (-(1.0 - id / 100.0) / tau).exp())
                .sum::<f64>() / n;
            genus_likelihood.insert(genus.clone(), mean_lik);
            genus_confidence.insert(genus, mean_id);
        }

        let mut species_scores: FxHashMap<String, f64> = FxHashMap::default();
        for (sp, scores) in species_matches {
            if scores.is_empty() { continue; }
            let avg = scores.iter().sum::<f64>() / scores.len() as f64;
            species_scores.insert(sp, avg);
        }

        let plasmid_frac = if total_hits > 0 { plasmid_hits as f64 / total_hits as f64 } else { 0.0 };
        Ok((genus_likelihood, genus_confidence, plasmid_frac, species_scores))
    }

}

// ============================================================================
// Utility Functions
// ============================================================================

/// Computes the reverse complement of a DNA sequence.
fn reverse_complement(seq: &str) -> String {
    seq.chars()
        .rev()
        .map(|c| match c.to_ascii_uppercase() {
            'A' => 'T',
            'T' => 'A',
            'G' => 'C',
            'C' => 'G',
            _ => 'N',
        })
        .collect()
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_reverse_complement() {
        assert_eq!(reverse_complement("ATGC"), "GCAT");
        assert_eq!(reverse_complement("AAAA"), "TTTT");
        assert_eq!(reverse_complement(""), "");
    }

    #[test]
    fn test_genus_result_default() {
        let result = GenusResult::default();
        assert!(result.genus.is_none());
        assert_eq!(result.confidence, 0.0);
    }

    fn lik(pairs: &[(&str, f64)]) -> FxHashMap<String, f64> {
        pairs.iter().map(|(g, l)| (g.to_string(), *l)).collect()
    }

    #[test]
    fn test_ou_posterior_no_neighbors_is_normalized_likelihood() {
        // Empty distance map => every off-diagonal weight ≈ 0 => posterior is just the
        // likelihoods renormalized (no phylogenetic smoothing).
        let neighbors = FxHashMap::default();
        let l = lik(&[("Escherichia", 3.0), ("Salmonella", 1.0)]);
        let post = ou_posterior(&l, &neighbors, 0.1);
        assert_eq!(post[0].0, "Escherichia");
        // Not exactly 0.75/0.25: the absent pair still carries w=exp(-ABSENT_PATRISTIC/λ)=
        // exp(-30)≈9e-14 of residual smoothing, which is the intended "≈0" degenerate limit.
        assert!((post[0].1 - 0.75).abs() < 1e-3, "got {}", post[0].1);
        assert!((post[1].1 - 0.25).abs() < 1e-3, "got {}", post[1].1);
    }

    #[test]
    fn test_ou_posterior_borrows_from_close_relative() {
        // Escherichia has all the likelihood; Salmonella has none but sits d=0.0406 away
        // (same family, GTDB patristic), while Bacillus is far (absent => d=ABSENT_PATRISTIC).
        // The kernel must lift Salmonella's posterior well above Bacillus's — borrowed by
        // proximity.
        let mut neighbors: FxHashMap<String, FxHashMap<String, f64>> = FxHashMap::default();
        neighbors.entry("Escherichia".into()).or_default().insert("Salmonella".into(), 0.0406);
        neighbors.entry("Salmonella".into()).or_default().insert("Escherichia".into(), 0.0406);
        let l = lik(&[("Escherichia", 1.0), ("Salmonella", 0.0), ("Bacillus", 0.0)]);
        let post = ou_posterior(&l, &neighbors, 0.3);
        let p: FxHashMap<_, _> = post.into_iter().collect();
        assert!(p["Salmonella"] > p["Bacillus"], "close relative must outrank far one");
        assert!(p["Salmonella"] > 0.05, "Salmonella should borrow real mass, got {}", p["Salmonella"]);
        assert!(p["Escherichia"] > p["Salmonella"], "the observed genus still leads");
    }

    fn ent(g: &str, fam: &str, gen: &str) -> (String, [String; 7]) {
        // [superkingdom, phylum, class, order, family, genus, species]
        (g.to_string(), [
            "Bacteria".into(), "Pseudomonadota".into(), "Gammaproteobacteria".into(),
            "Enterobacterales".into(), fam.into(), gen.into(), format!("{} sp.", gen),
        ])
    }

    #[test]
    fn test_lca_ragged_rank() {
        let mut lin: FxHashMap<String, [String; 7]> = FxHashMap::default();
        for (k, v) in [ent("Escherichia", "Enterobacteriaceae", "Escherichia"),
                       ent("Salmonella", "Enterobacteriaceae", "Salmonella")] {
            lin.insert(k, v);
        }
        lin.insert("Bacillus".into(), [
            "Bacteria".into(), "Bacillota".into(), "Bacilli".into(), "Bacillales".into(),
            "Bacillaceae".into(), "Bacillus".into(), "Bacillus subtilis".into(),
        ]);
        // Single genus → resolves to genus.
        assert_eq!(lca_of(&["Escherichia"], &lin), ("genus".into(), "Escherichia".into()));
        // Two Enterobacteriaceae genera → family LCA (the ragged-rank answer).
        assert_eq!(lca_of(&["Escherichia", "Salmonella"], &lin),
                   ("family".into(), "Enterobacteriaceae".into()));
        // Across phyla → superkingdom.
        assert_eq!(lca_of(&["Escherichia", "Bacillus"], &lin),
                   ("superkingdom".into(), "Bacteria".into()));
    }

    #[test]
    fn test_render_resolution_notation() {
        // Single genus → that genus (radius irrelevant).
        assert_eq!(render_resolution(&["Escherichia"], "genus", "Escherichia", 0.0),
                   ("genus".into(), "Escherichia".into()));
        // Tight cluster (small radius) → list the genera, NOT the family — coherence, not count.
        assert_eq!(render_resolution(&["Escherichia", "Shigella"], "family", "Enterobacteriaceae", 0.03),
                   ("genus".into(), "Escherichia|Shigella".into()));
        // Many genera but still tight → still listed (count does NOT trigger rollup).
        let tight_many = ["A", "B", "C", "D", "E", "F"];
        assert_eq!(render_resolution(&tight_many, "family", "Enterobacteriaceae", 0.10),
                   ("genus".into(), "A|B|C|D|E|F".into()));
        // Two genera spread past the family diameter (radius > 0.5) → the family name.
        assert_eq!(render_resolution(&["Escherichia", "Klebsiella"], "family", "Enterobacteriaceae", 0.6),
                   ("family".into(), "Enterobacteriaceae".into()));
        // Genera spanning families (large patristic radius, coarse LCA) → the LCA clade.
        assert_eq!(render_resolution(&["Escherichia", "Bacillus"], "superkingdom", "Bacteria", 1.5),
                   ("superkingdom".into(), "Bacteria".into()));
    }

    #[test]
    fn test_group_credible_set_clusters_by_distance() {
        // Escherichia+Pseudescherichia are tight (0.08 patristic); Cronobacter and Klebsiella
        // sit >0.5 from them and from each other → three clusters. The tight pair (mass 0.55)
        // leads; within it the higher-posterior genus is first.
        let mut nb: FxHashMap<String, FxHashMap<String, f64>> = FxHashMap::default();
        let mut set = |a: &str, b: &str, d: f64| {
            nb.entry(a.into()).or_default().insert(b.into(), d);
            nb.entry(b.into()).or_default().insert(a.into(), d);
        };
        set("Escherichia", "Pseudescherichia", 0.08);
        set("Escherichia", "Cronobacter", 0.90);
        set("Escherichia", "Klebsiella", 0.70);
        set("Pseudescherichia", "Cronobacter", 0.90);
        set("Pseudescherichia", "Klebsiella", 0.70);
        set("Cronobacter", "Klebsiella", 0.80);
        let members = vec![
            ("Escherichia".to_string(), 0.30),
            ("Cronobacter".to_string(), 0.25),
            ("Pseudescherichia".to_string(), 0.25),
            ("Klebsiella".to_string(), 0.20),
        ];
        let out = group_credible_set(&members, &nb, GENUS_COHERENCE_RADIUS);
        assert_eq!(out, "Escherichia(0.30),Pseudescherichia(0.25) | Cronobacter(0.25) | Klebsiella(0.20)");
    }
}