eidos-kernel 0.1.0

use std::collections::HashSet;

use serde::Serialize;

use crate::graph_index::GraphIndex;
use crate::retrieval::ground_subgraph;
use crate::schema::{Edge, Graph};

use super::{ContextEdgeExpectation, ContextOmissionExpectation, GoldenCase};

/// Per-case context outcome.
#[derive(Serialize, Clone, Debug)]
pub struct ContextCaseResult {
    pub query: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expected_focus_route: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub expected_focus_routes: Vec<String>,
    pub focus_route: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expected_compound_primary: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub compound_primary: Option<String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub vocabulary_gap_terms: Vec<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expected_no_vocabulary_gap: Option<bool>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub editable_vocab_targets: Vec<String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub non_editable_vocab_targets: Vec<String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub omitted_context: Vec<ContextOmissionExpectation>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub omitted_compound_anchors: Vec<ContextOmissionExpectation>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub vocabulary_gap_ok: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub no_vocabulary_gap_ok: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub editable_vocab_targets_ok: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub non_editable_vocab_targets_ok: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub omitted_context_ok: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub omitted_compound_anchors_ok: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expected_partition: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub focus_route_partition: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub partition_ok: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub focus_route_ok: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub compound_primary_ok: Option<bool>,
    /// Fraction of `context_must` ids present in the activated context (None if none judged).
    pub recall: Option<f64>,
    /// Fraction of `context_must_not` ids present (noise — lower is better; None if none judged).
    pub noise: Option<f64>,
    /// Fraction of context items that carry a source receipt (source file or span).
    pub receipt_coverage: f64,
    /// Number of node ids in the context pack.
    pub pack_size: usize,
    /// Fraction of `context_edges_must` present in the focused edge set.
    pub edge_recall: Option<f64>,
    pub missing: Vec<String>,
    pub leaked: Vec<String>,
    pub missing_edges: Vec<ContextEdgeExpectation>,
    pub missing_omitted_context: Vec<ContextOmissionExpectation>,
    pub missing_omitted_compound_anchors: Vec<ContextOmissionExpectation>,
}

/// Aggregate context-quality metrics over cases carrying `context_must` / `context_must_not`.
#[derive(Serialize, Clone, Debug)]
pub struct ContextReport {
    pub context_cases: usize,
    /// Mean recall over cases that declared `context_must`.
    pub context_recall: f64,
    /// Mean noise over cases that declared `context_must_not` (fraction that wrongly appeared).
    pub context_noise: f64,
    /// Mean receipt coverage over all judged context cases.
    pub receipt_coverage: f64,
    /// Mean context pack size over all judged context cases.
    pub mean_pack_size: f64,
    /// Mean edge recall over cases that declared `context_edges_must`.
    pub edge_recall: f64,
    /// Match rate over cases with `expected_focus_route`; `None` means no route judgment.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub focus_route_match_rate: Option<f64>,
    /// Match rate over cases with `expected_compound_primary`; `None` means no compound judgment.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub compound_primary_match_rate: Option<f64>,
    /// Match rate over cases with `expected_vocabulary_gap_terms`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub vocabulary_gap_match_rate: Option<f64>,
    /// Match rate over cases with `expected_no_vocabulary_gap`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub no_vocabulary_gap_match_rate: Option<f64>,
    /// Match rate over cases with `expected_vocab_editable_targets`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub editable_vocab_target_match_rate: Option<f64>,
    /// Match rate over cases with `expected_vocab_non_editable_targets`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub non_editable_vocab_target_match_rate: Option<f64>,
    /// Match rate over cases with `expected_omitted_context`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub omitted_context_match_rate: Option<f64>,
    /// Match rate over cases with `expected_omitted_compound_anchors`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub omitted_compound_anchor_match_rate: Option<f64>,
    /// Match rate over cases with `expected_partition`; `None` means no partition judgment.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub partition_match_rate: Option<f64>,
    pub cases: Vec<ContextCaseResult>,
}

/// Score the CONTEXT payload (the activated subgraph's `context_order`), not just the route.
/// Only cases declaring `context_must` and/or `context_must_not` participate.
pub fn evaluate_context(
    graph: &Graph,
    cases: &[GoldenCase],
    limit: usize,
    depth: usize,
    width: usize,
) -> ContextReport {
    let mut results = Vec::new();
    let (mut recall_sum, mut recall_n) = (0.0f64, 0usize);
    let (mut noise_sum, mut noise_n) = (0.0f64, 0usize);
    let (mut edge_recall_sum, mut edge_recall_n) = (0.0f64, 0usize);
    let (mut partition_ok_sum, mut partition_ok_n) = (0usize, 0usize);
    let mut receipt_sum = 0.0f64;
    let mut pack_size_sum = 0usize;
    let (mut compound_primary_ok_sum, mut compound_primary_ok_n) = (0usize, 0usize);
    let (mut vocabulary_gap_ok_sum, mut vocabulary_gap_ok_n) = (0usize, 0usize);
    let (mut no_vocabulary_gap_ok_sum, mut no_vocabulary_gap_ok_n) = (0usize, 0usize);
    let (mut editable_vocab_target_ok_sum, mut editable_vocab_target_ok_n) = (0usize, 0usize);
    let (mut non_editable_vocab_target_ok_sum, mut non_editable_vocab_target_ok_n) =
        (0usize, 0usize);
    let (mut omitted_context_ok_sum, mut omitted_context_ok_n) = (0usize, 0usize);
    let (mut omitted_compound_anchor_ok_sum, mut omitted_compound_anchor_ok_n) = (0usize, 0usize);
    let index = GraphIndex::build(graph);

    for c in cases {
        if c.context_must.is_empty()
            && c.context_must_not.is_empty()
            && c.context_edges_must.is_empty()
            && c.expected_focus_route.is_none()
            && c.expected_compound_primary.is_none()
            && c.expected_vocabulary_gap_terms.is_empty()
            && !c.expected_no_vocabulary_gap
            && c.expected_vocab_editable_targets.is_empty()
            && c.expected_vocab_non_editable_targets.is_empty()
            && c.expected_omitted_context.is_empty()
            && c.expected_omitted_compound_anchors.is_empty()
        {
            continue;
        }
        let sg = ground_subgraph(graph, &c.query, limit, depth, width);
        let ctx: HashSet<&str> = sg.context_order.iter().map(String::as_str).collect();
        let focus_route_partition = index
            .node(&sg.route)
            .and_then(|node| node.partition.clone());
        let partition_ok = c
            .expected_partition
            .as_ref()
            .map(|expected| focus_route_partition.as_deref() == Some(expected.as_str()));
        if let Some(ok) = partition_ok {
            partition_ok_n += 1;
            partition_ok_sum += ok as usize;
        }
        let compound_primary = sg
            .compound
            .as_ref()
            .map(|compound| compound.primary.clone());
        let compound_primary_ok = c
            .expected_compound_primary
            .as_ref()
            .map(|expected| compound_primary.as_deref() == Some(expected.as_str()));
        if let Some(ok) = compound_primary_ok {
            compound_primary_ok_n += 1;
            compound_primary_ok_sum += ok as usize;
        }
        let vocabulary_gap_terms = sg
            .compound
            .as_ref()
            .map(|compound| compound.missing_terms.clone())
            .unwrap_or_default();
        let (editable_vocab_targets, non_editable_vocab_targets) =
            vocabulary_target_candidates(&index, sg.compound.as_ref());
        let vocabulary_gap_ok = (!c.expected_vocabulary_gap_terms.is_empty())
            .then(|| contains_all(&vocabulary_gap_terms, &c.expected_vocabulary_gap_terms));
        if let Some(ok) = vocabulary_gap_ok {
            vocabulary_gap_ok_n += 1;
            vocabulary_gap_ok_sum += ok as usize;
        }
        let no_vocabulary_gap_ok = c
            .expected_no_vocabulary_gap
            .then_some(vocabulary_gap_terms.is_empty());
        if let Some(ok) = no_vocabulary_gap_ok {
            no_vocabulary_gap_ok_n += 1;
            no_vocabulary_gap_ok_sum += ok as usize;
        }
        let editable_vocab_targets_ok = (!c.expected_vocab_editable_targets.is_empty())
            .then(|| contains_all(&editable_vocab_targets, &c.expected_vocab_editable_targets));
        if let Some(ok) = editable_vocab_targets_ok {
            editable_vocab_target_ok_n += 1;
            editable_vocab_target_ok_sum += ok as usize;
        }
        let non_editable_vocab_targets_ok = (!c.expected_vocab_non_editable_targets.is_empty())
            .then(|| {
                contains_all(
                    &non_editable_vocab_targets,
                    &c.expected_vocab_non_editable_targets,
                )
            });
        if let Some(ok) = non_editable_vocab_targets_ok {
            non_editable_vocab_target_ok_n += 1;
            non_editable_vocab_target_ok_sum += ok as usize;
        }
        let omitted_context = sg
            .omitted_context_candidates
            .iter()
            .map(|candidate| ContextOmissionExpectation {
                node_id: candidate.node_id.clone(),
                reason: candidate.reason.clone(),
            })
            .collect::<Vec<_>>();
        let omitted_compound_anchors = sg
            .compound
            .as_ref()
            .map(|compound| {
                compound
                    .omitted_anchors
                    .iter()
                    .map(|anchor| ContextOmissionExpectation {
                        node_id: anchor.node_id.clone(),
                        reason: anchor.reason.clone(),
                    })
                    .collect::<Vec<_>>()
            })
            .unwrap_or_default();
        let missing_omitted_context =
            missing_omissions(&omitted_context, &c.expected_omitted_context);
        let omitted_context_ok =
            (!c.expected_omitted_context.is_empty()).then_some(missing_omitted_context.is_empty());
        if let Some(ok) = omitted_context_ok {
            omitted_context_ok_n += 1;
            omitted_context_ok_sum += ok as usize;
        }
        let missing_omitted_compound_anchors = missing_omissions(
            &omitted_compound_anchors,
            &c.expected_omitted_compound_anchors,
        );
        let omitted_compound_anchors_ok = (!c.expected_omitted_compound_anchors.is_empty())
            .then_some(missing_omitted_compound_anchors.is_empty());
        if let Some(ok) = omitted_compound_anchors_ok {
            omitted_compound_anchor_ok_n += 1;
            omitted_compound_anchor_ok_sum += ok as usize;
        }

        let recall = if c.context_must.is_empty() {
            None
        } else {
            let hit = c
                .context_must
                .iter()
                .filter(|m| ctx.contains(m.as_str()))
                .count();
            recall_sum += hit as f64 / c.context_must.len() as f64;
            recall_n += 1;
            Some(hit as f64 / c.context_must.len() as f64)
        };
        let noise = if c.context_must_not.is_empty() {
            None
        } else {
            let leaked = c
                .context_must_not
                .iter()
                .filter(|m| ctx.contains(m.as_str()))
                .count();
            noise_sum += leaked as f64 / c.context_must_not.len() as f64;
            noise_n += 1;
            Some(leaked as f64 / c.context_must_not.len() as f64)
        };
        let missing = c
            .context_must
            .iter()
            .filter(|m| !ctx.contains(m.as_str()))
            .cloned()
            .collect();
        let leaked = c
            .context_must_not
            .iter()
            .filter(|m| ctx.contains(m.as_str()))
            .cloned()
            .collect();
        let missing_edges = c
            .context_edges_must
            .iter()
            .filter(|expected| !edge_present(expected, &sg.edges))
            .cloned()
            .collect::<Vec<_>>();
        let edge_recall = if c.context_edges_must.is_empty() {
            None
        } else {
            let recall = (c.context_edges_must.len() - missing_edges.len()) as f64
                / c.context_edges_must.len() as f64;
            edge_recall_sum += recall;
            edge_recall_n += 1;
            Some(recall)
        };
        let pack_size = sg.context_order.len();
        let receipt_coverage = super::receipt_coverage(&index, &sg.context_order);
        receipt_sum += receipt_coverage;
        pack_size_sum += pack_size;
        results.push(ContextCaseResult {
            query: c.query.clone(),
            expected_focus_route: c.expected_focus_route.clone(),
            expected_focus_routes: c.expected_focus_routes.clone(),
            focus_route: sg.route.clone(),
            expected_compound_primary: c.expected_compound_primary.clone(),
            compound_primary,
            vocabulary_gap_terms,
            expected_no_vocabulary_gap: c.expected_no_vocabulary_gap.then_some(true),
            editable_vocab_targets,
            non_editable_vocab_targets,
            omitted_context,
            omitted_compound_anchors,
            vocabulary_gap_ok,
            no_vocabulary_gap_ok,
            editable_vocab_targets_ok,
            non_editable_vocab_targets_ok,
            omitted_context_ok,
            omitted_compound_anchors_ok,
            expected_partition: c.expected_partition.clone(),
            focus_route_partition,
            partition_ok,
            focus_route_ok: acceptable_focus_routes(c).map(|expected| expected.contains(&sg.route)),
            compound_primary_ok,
            recall,
            noise,
            receipt_coverage,
            pack_size,
            edge_recall,
            missing,
            leaked,
            missing_edges,
            missing_omitted_context,
            missing_omitted_compound_anchors,
        });
    }

    let context_cases = results.len();

    ContextReport {
        context_cases,
        context_recall: if recall_n == 0 {
            0.0
        } else {
            recall_sum / recall_n as f64
        },
        context_noise: if noise_n == 0 {
            0.0
        } else {
            noise_sum / noise_n as f64
        },
        receipt_coverage: if context_cases == 0 {
            0.0
        } else {
            receipt_sum / context_cases as f64
        },
        mean_pack_size: if context_cases == 0 {
            0.0
        } else {
            pack_size_sum as f64 / context_cases as f64
        },
        edge_recall: if edge_recall_n == 0 {
            0.0
        } else {
            edge_recall_sum / edge_recall_n as f64
        },
        focus_route_match_rate: focus_route_match_rate(&results),
        compound_primary_match_rate: (compound_primary_ok_n > 0)
            .then_some(compound_primary_ok_sum as f64 / compound_primary_ok_n as f64),
        vocabulary_gap_match_rate: (vocabulary_gap_ok_n > 0)
            .then_some(vocabulary_gap_ok_sum as f64 / vocabulary_gap_ok_n as f64),
        no_vocabulary_gap_match_rate: (no_vocabulary_gap_ok_n > 0)
            .then_some(no_vocabulary_gap_ok_sum as f64 / no_vocabulary_gap_ok_n as f64),
        editable_vocab_target_match_rate: (editable_vocab_target_ok_n > 0)
            .then_some(editable_vocab_target_ok_sum as f64 / editable_vocab_target_ok_n as f64),
        non_editable_vocab_target_match_rate: (non_editable_vocab_target_ok_n > 0).then_some(
            non_editable_vocab_target_ok_sum as f64 / non_editable_vocab_target_ok_n as f64,
        ),
        omitted_context_match_rate: (omitted_context_ok_n > 0)
            .then_some(omitted_context_ok_sum as f64 / omitted_context_ok_n as f64),
        omitted_compound_anchor_match_rate: (omitted_compound_anchor_ok_n > 0)
            .then_some(omitted_compound_anchor_ok_sum as f64 / omitted_compound_anchor_ok_n as f64),
        partition_match_rate: (partition_ok_n > 0)
            .then_some(partition_ok_sum as f64 / partition_ok_n as f64),
        cases: results,
    }
}

fn acceptable_focus_routes(c: &GoldenCase) -> Option<Vec<String>> {
    let mut expected = c.expected_focus_routes.clone();
    if let Some(route) = &c.expected_focus_route {
        expected.push(route.clone());
    }
    expected.sort();
    expected.dedup();
    (!expected.is_empty()).then_some(expected)
}

fn vocabulary_target_candidates(
    index: &GraphIndex<'_>,
    compound: Option<&crate::retrieval::CompoundSubgraph>,
) -> (Vec<String>, Vec<String>) {
    let Some(compound) = compound else {
        return (Vec::new(), Vec::new());
    };
    let mut candidates = Vec::new();
    candidates.push(compound.primary.clone());
    candidates.extend(compound.anchors.iter().map(|anchor| anchor.node_id.clone()));
    candidates.extend(
        compound
            .omitted_anchors
            .iter()
            .map(|anchor| anchor.node_id.clone()),
    );
    candidates.sort();
    candidates.dedup();

    let mut editable = Vec::new();
    let mut non_editable = Vec::new();
    for node_id in candidates {
        let target = if index.node(&node_id).is_some_and(vocab_target_is_editable) {
            &mut editable
        } else {
            &mut non_editable
        };
        target.push(node_id);
    }
    (editable, non_editable)
}

fn vocab_target_is_editable(node: &crate::schema::Node) -> bool {
    node.source_files
        .first()
        .is_some_and(|source| source.ends_with(".md"))
}

fn contains_all(actual: &[String], expected: &[String]) -> bool {
    expected.iter().all(|item| actual.contains(item))
}

fn missing_omissions(
    actual: &[ContextOmissionExpectation],
    expected: &[ContextOmissionExpectation],
) -> Vec<ContextOmissionExpectation> {
    expected
        .iter()
        .filter(|expected| !actual.contains(expected))
        .cloned()
        .collect()
}

fn focus_route_match_rate(results: &[ContextCaseResult]) -> Option<f64> {
    let judged = results
        .iter()
        .filter_map(|case| case.focus_route_ok)
        .collect::<Vec<_>>();
    if judged.is_empty() {
        None
    } else {
        Some(judged.iter().filter(|ok| **ok).count() as f64 / judged.len() as f64)
    }
}

fn edge_present(expected: &ContextEdgeExpectation, edges: &[Edge]) -> bool {
    edges.iter().any(|edge| {
        edge.from == expected.from && edge.to == expected.to && edge.relation == expected.relation
    })
}