//! Dominance and Dominance Frontier Analysis
//!
//! Computes dominator trees and dominance frontiers for CFGs.

use crate::algorithms::astar::CfgGraphNode;
use std::collections::{HashMap, HashSet};

#[derive(Debug, Clone)]
pub struct DominanceResult {
    pub entry: u64,
    pub dominators: HashMap<u64, HashSet<u64>>,
    pub idom: HashMap<u64, u64>,
    pub dom_tree: HashMap<u64, Vec<u64>>,
}

#[derive(Debug, Clone)]
pub struct DominanceFrontierResult {
    pub frontier: HashMap<u64, HashSet<u64>>,
}

pub fn compute_dominance(nodes: &[CfgGraphNode], entry_id: u64) -> DominanceResult {
    let _node_map: HashMap<u64, &CfgGraphNode> = nodes.iter().map(|n| (n.id, n)).collect();
    let node_ids: HashSet<u64> = nodes.iter().map(|n| n.id).collect();

    // Compute predecessors
    let mut predecessors: HashMap<u64, Vec<u64>> = HashMap::new();
    for node in nodes {
        for &succ in &node.successors {
            predecessors.entry(succ).or_default().push(node.id);
        }
    }

    // Initialize dominators
    let mut dominators: HashMap<u64, HashSet<u64>> = HashMap::new();
    for &id in &node_ids {
        if id == entry_id {
            let mut entry_dom = HashSet::new();
            entry_dom.insert(entry_id);
            dominators.insert(id, entry_dom);
        } else {
            dominators.insert(id, node_ids.clone());
        }
    }

    // Iterative dataflow
    let mut changed = true;
    while changed {
        changed = false;
        for &id in &node_ids {
            if id == entry_id {
                continue;
            }

            let preds = predecessors.get(&id).cloned().unwrap_or_default();
            if preds.is_empty() {
                continue;
            }

            let mut new_dom: HashSet<u64> = dominators.get(&preds[0]).cloned().unwrap_or_default();

            for &pred in &preds[1..] {
                let pred_dom = dominators.get(&pred).cloned().unwrap_or_default();
                new_dom = new_dom.intersection(&pred_dom).copied().collect();
            }
            new_dom.insert(id);

            if new_dom != *dominators.get(&id).unwrap() {
                dominators.insert(id, new_dom);
                changed = true;
            }
        }
    }

    // Compute immediate dominators (closest dominator = largest proper subset)
    let mut idom: HashMap<u64, u64> = HashMap::new();
    for &id in &node_ids {
        if id == entry_id {
            continue;
        }

        let doms = dominators.get(&id).cloned().unwrap_or_default();
        // Find closest dominator (not self) - the one with most dominators
        let mut closest: Option<u64> = None;
        let mut closest_size = 0;

        for &dom in &doms {
            if dom == id {
                continue;
            }
            let dom_size = dominators.get(&dom).map(|d| d.len()).unwrap_or(0);
            if dom_size > closest_size {
                closest_size = dom_size;
                closest = Some(dom);
            }
        }

        if let Some(c) = closest {
            idom.insert(id, c);
        }
    }

    // Build dominance tree
    let mut dom_tree: HashMap<u64, Vec<u64>> = HashMap::new();
    for (&child, &parent) in &idom {
        dom_tree.entry(parent).or_default().push(child);
    }

    DominanceResult {
        entry: entry_id,
        dominators,
        idom,
        dom_tree,
    }
}

pub fn compute_dominance_frontier(
    nodes: &[CfgGraphNode],
    entry_id: u64,
) -> DominanceFrontierResult {
    let dom_result = compute_dominance(nodes, entry_id);

    let mut predecessors: HashMap<u64, Vec<u64>> = HashMap::new();
    for node in nodes {
        for &succ in &node.successors {
            predecessors.entry(succ).or_default().push(node.id);
        }
    }

    let mut frontier: HashMap<u64, HashSet<u64>> = HashMap::new();

    for node in nodes {
        let preds = predecessors.get(&node.id).cloned().unwrap_or_default();
        if preds.len() >= 2 {
            for &pred in &preds {
                let mut runner = pred;
                while runner != entry_id
                    && dom_result.idom.get(&node.id).is_none_or(|&id| id != runner)
                {
                    frontier.entry(runner).or_default().insert(node.id);
                    runner = *dom_result.idom.get(&runner).unwrap_or(&entry_id);
                }
            }
        }
    }

    DominanceFrontierResult { frontier }
}

pub fn find_dominators_to(target: u64, dom_result: &DominanceResult) -> HashSet<u64> {
    dom_result
        .dominators
        .get(&target)
        .cloned()
        .unwrap_or_default()
}

pub fn dominates(a: u64, b: u64, dom_result: &DominanceResult) -> bool {
    dom_result
        .dominators
        .get(&b)
        .map(|d| d.contains(&a))
        .unwrap_or(false)
}

pub fn strictly_dominates(a: u64, b: u64, dom_result: &DominanceResult) -> bool {
    a != b && dominates(a, b, dom_result)
}

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

    #[test]
    fn test_compute_dominance_linear() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![2],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];

        let result = compute_dominance(&nodes, 0);

        assert!(result.dominators.get(&0).unwrap().contains(&0));
        assert!(result.dominators.get(&1).unwrap().contains(&0));
        assert!(result.dominators.get(&1).unwrap().contains(&1));
        assert!(result.dominators.get(&2).unwrap().contains(&0));
        assert_eq!(result.idom.get(&1), Some(&0));
        assert_eq!(result.idom.get(&2), Some(&1));
    }

    #[test]
    fn test_compute_dominance_diamond() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1, 2],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 3,
                x: 3.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];

        let result = compute_dominance(&nodes, 0);

        let dom_3 = result.dominators.get(&3).unwrap();
        assert!(dom_3.contains(&0));
        assert!(dom_3.contains(&3));
    }

    #[test]
    fn test_dominates() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];

        let result = compute_dominance(&nodes, 0);

        assert!(dominates(0, 0, &result));
        assert!(dominates(0, 1, &result));
        assert!(!dominates(1, 0, &result));
    }

    #[test]
    fn test_strictly_dominates() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];

        let result = compute_dominance(&nodes, 0);

        assert!(!strictly_dominates(0, 0, &result));
        assert!(strictly_dominates(0, 1, &result));
    }

    #[test]
    fn test_compute_dominance_frontier() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1, 2],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 3,
                x: 3.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];

        let result = compute_dominance_frontier(&nodes, 0);
        assert!(result.frontier.get(&1).is_some_and(|set| set.contains(&3)));
        assert!(result.frontier.get(&2).is_some_and(|set| set.contains(&3)));
    }

    #[test]
    fn test_find_dominators_to() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![2],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];

        let result = compute_dominance(&nodes, 0);

        let doms_2 = find_dominators_to(2, &result);
        assert!(doms_2.contains(&0));
        assert!(doms_2.contains(&1));
        assert!(doms_2.contains(&2));
    }
}
