use crate::ir::ir_variable::IrVrVariableDescriptor;
use std::collections::{HashMap, HashSet};

pub type InterferenceGraph = HashMap<IrVrVariableDescriptor, HashSet<IrVrVariableDescriptor>>;
pub type LivenessMap = HashMap<usize, HashSet<IrVrVariableDescriptor>>;

pub trait HasVrUsers {
    fn vr_users(&self) -> Vec<&dyn VrUser>;
    fn vr_users_mut(&mut self) -> Vec<&mut dyn VrUser>;

    fn live_in_live_out(&self) -> (LivenessMap, LivenessMap) {
        let mut live_in: LivenessMap = HashMap::new();
        let mut live_out: LivenessMap = HashMap::new();

        loop {
            let mut did_work = false;
            for (block_index, statement) in self.vr_users().iter().enumerate().rev() {
                // init if necessary
                if !live_in.contains_key(&block_index) {
                    live_in.insert(block_index, HashSet::new());
                }
                if !live_out.contains_key(&block_index) {
                    live_out.insert(block_index, HashSet::new());
                }

                // out (union of successors ins)
                // for now, a statement can only have one successor
                // this will need to be updated when we add jumps
                if let Some(successor_live_in) = live_in.get(&(block_index + 1)) {
                    let statement_live_out = live_out.get_mut(&block_index).unwrap();
                    for vr_variable in successor_live_in {
                        if statement_live_out.insert(vr_variable.clone()) {
                            did_work = true;
                        }
                    }
                }

                // in: use(s) U ( out(s) - def(s) )
                let mut new_ins = statement
                    .vr_uses()
                    .iter()
                    .map(|u| (*u).clone())
                    .collect::<HashSet<_>>();
                let statement_live_out = live_out.get(&block_index).unwrap();
                let defs = statement
                    .vr_definitions()
                    .iter()
                    .map(|d| (*d).clone())
                    .collect::<HashSet<_>>();
                let rhs = statement_live_out - &defs;
                new_ins.extend(rhs);

                let statement_live_in = live_in.get_mut(&block_index).unwrap();
                for new_in in new_ins {
                    if statement_live_in.insert(new_in) {
                        did_work = true;
                    }
                }
            }
            if !did_work {
                break;
            }
        }
        (live_in, live_out)
    }

    fn interference_graph(&self) -> InterferenceGraph {
        let mut all_vr_variables: HashSet<IrVrVariableDescriptor> = HashSet::new();

        for vr_user in self.vr_users() {
            all_vr_variables.extend(vr_user.vr_definitions());
            all_vr_variables.extend(vr_user.vr_uses());
        }

        let mut graph: InterferenceGraph = HashMap::new();
        for variable in all_vr_variables {
            graph.insert(variable, HashSet::new());
        }

        let (_, live_out) = self.live_in_live_out();

        for (index, vr_user) in self.vr_users().iter().enumerate() {
            let user_live_in = live_out.get(&index).unwrap();
            for definition_vr_variable in vr_user.vr_definitions() {
                for live_out_variable in user_live_in {
                    // we do the following check to avoid adding an edge to itself
                    if definition_vr_variable.ne(live_out_variable) {
                        graph
                            .get_mut(&definition_vr_variable)
                            .unwrap()
                            .insert(live_out_variable.clone());
                        graph
                            .get_mut(live_out_variable)
                            .unwrap()
                            .insert(definition_vr_variable.clone());
                    }
                }
            }
        }

        graph
    }

    fn assign_registers(
        &mut self,
        register_count: usize,
    ) -> (HashMap<IrVrVariableDescriptor, usize>, isize) {
        let mut spills: HashSet<IrVrVariableDescriptor> = HashSet::new();
        loop {
            let mut interference_graph = self.interference_graph();
            let (registers, new_spills) =
                registers_and_spills(&mut interference_graph, register_count);

            if spills != new_spills {
                spills = new_spills;
                // propagate all spills, since those won't be used for the next interference graph
                for vr_user in &mut self.vr_users_mut() {
                    vr_user.propagate_spills(&spills);
                }
            } else {
                // we've calculated final assignments, so propagate them
                for vr_user in self.vr_users_mut() {
                    vr_user.propagate_register_assignments(&registers);
                }
                // also set offsets
                let mut offset_counter = OffsetCounter::new();
                for vr_user in self.vr_users_mut() {
                    vr_user.propagate_stack_offsets(&mut offset_counter);
                }

                return (registers, offset_counter.get_count());
            }
        }
    }
}

pub trait VrUser {
    fn vr_definitions(&self) -> HashSet<IrVrVariableDescriptor>;
    fn vr_uses(&self) -> HashSet<IrVrVariableDescriptor>;
    fn propagate_spills(&mut self, spills: &HashSet<IrVrVariableDescriptor>);
    fn propagate_register_assignments(
        &mut self,
        assignments: &HashMap<IrVrVariableDescriptor, usize>,
    );
    fn propagate_stack_offsets(&mut self, counter: &mut OffsetCounter);
}

pub struct OffsetCounter {
    counter: isize,
}

impl OffsetCounter {
    pub fn new() -> Self {
        Self { counter: 0 }
    }

    pub fn next(&mut self) -> isize {
        let offset = self.counter;
        self.counter += 1;
        offset
    }

    pub fn get_count(&self) -> isize {
        self.counter
    }
}

#[derive(Debug)]
struct WorkItem {
    vr: IrVrVariableDescriptor,
    edges: HashSet<IrVrVariableDescriptor>,
    color: bool,
}

pub fn registers_and_spills(
    interference_graph: &mut InterferenceGraph,
    k: usize,
) -> (
    HashMap<IrVrVariableDescriptor, usize>,
    HashSet<IrVrVariableDescriptor>,
) {
    let mut work_stack: Vec<WorkItem> = vec![];

    while !interference_graph.is_empty() {
        work_stack.push(next_work_item(interference_graph, k));
    }

    // 3. assign colors to registers
    let mut rebuilt_graph: InterferenceGraph = HashMap::new();
    let mut register_assignments: HashMap<IrVrVariableDescriptor, usize> = HashMap::new();
    let mut spills: HashSet<IrVrVariableDescriptor> = HashSet::new();

    while let Some(work_item) = work_stack.pop() {
        if work_item.color {
            assign_register(&work_item, &mut rebuilt_graph, k, &mut register_assignments);
        } else if can_optimistically_color(&work_item, &mut register_assignments, k) {
            assign_register(&work_item, &mut rebuilt_graph, k, &mut register_assignments);
        } else {
            // spill
            spills.insert(work_item.vr.clone());
        }
    }

    (register_assignments, spills)
}

fn assign_register(
    work_item: &WorkItem,
    graph: &mut InterferenceGraph,
    k: usize,
    register_assignments: &mut HashMap<IrVrVariableDescriptor, usize>,
) {
    rebuild_vr_and_edges(graph, work_item);

    // find a register which is not yet shared by all outgoing edges' vertices
    'outer: for i in 0..k {
        for edge in graph.get_mut(&work_item.vr).unwrap().iter() {
            if register_assignments.contains_key(edge) {
                let assignment = register_assignments.get(edge).unwrap();
                if *assignment == i {
                    continue 'outer;
                }
            }
        }
        register_assignments.insert(work_item.vr.clone(), i);
        break;
    }
}

fn find_vr_lt_k(
    interference_graph: &InterferenceGraph,
    k: usize,
) -> Option<&IrVrVariableDescriptor> {
    interference_graph.iter().find_map(
        |(vr, neighbors)| {
            if neighbors.len() < k { Some(vr) } else { None }
        },
    )
}

fn remove_vr_and_edges(
    interference_graph: &mut InterferenceGraph,
    vr: &IrVrVariableDescriptor,
) -> HashSet<IrVrVariableDescriptor> {
    // first, outgoing
    let outgoing_edges = interference_graph.remove(vr).unwrap();

    // second, incoming
    for neighbor in &outgoing_edges {
        let neighbor_edges = interference_graph.get_mut(neighbor).unwrap();
        neighbor_edges.remove(vr);
    }

    outgoing_edges
}

fn next_work_item(interference_graph: &mut InterferenceGraph, k: usize) -> WorkItem {
    // try to find a node (virtual register) with less than k outgoing edges, and mark as color
    // for step 3.
    // if not, pick any, and mark as spill for step 3.
    let register_lt_k = find_vr_lt_k(interference_graph, k);
    if let Some(vr) = register_lt_k {
        let vr = vr.clone();

        // remove edges; save outgoing to work_item
        let edges = remove_vr_and_edges(interference_graph, &vr);

        // push to work stack
        WorkItem {
            vr,
            edges,
            color: true,
        }
    } else {
        // pick any
        let vr = interference_graph.iter().last().unwrap().0.clone();

        // remove edges
        let edges = remove_vr_and_edges(interference_graph, &vr);

        WorkItem {
            vr,
            edges,
            color: false, // spill
        }
    }
}

fn rebuild_vr_and_edges(graph: &mut InterferenceGraph, work_item: &WorkItem) {
    // init the vertex
    graph.insert(work_item.vr.clone(), HashSet::new());

    // outgoing
    for neighbor in &work_item.edges {
        // check if neighbor exists in the graph first; if it was marked spill earlier and could
        // not optimistically color, it won't be in the graph
        if graph.contains_key(neighbor) {
            // get outgoing set and insert neighbor
            graph
                .get_mut(&work_item.vr)
                .unwrap()
                .insert(neighbor.clone());
        }
    }

    // incoming
    for neighbor in &work_item.edges {
        // like above, neighbor may not have been added because of failure to optimistically
        // color
        if graph.contains_key(neighbor) {
            graph
                .get_mut(neighbor)
                .unwrap()
                .insert(work_item.vr.clone());
        }
    }
}

fn can_optimistically_color(
    work_item: &WorkItem,
    register_assignments: &HashMap<IrVrVariableDescriptor, usize>,
    k: usize,
) -> bool {
    // see if we can optimistically color
    // find how many assignments have been made for the outgoing edges
    // if it's less than k, we can do it
    let mut number_of_assigned_edges = 0;
    for edge in &work_item.edges {
        if register_assignments.contains_key(edge) {
            number_of_assigned_edges += 1;
        }
    }
    number_of_assigned_edges < k
}

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

    fn line_graph() -> InterferenceGraph {
        let mut graph: InterferenceGraph = HashMap::new();
        let v0 = IrVrVariableDescriptor::new("v0".into(), 0);
        let v1 = IrVrVariableDescriptor::new("v1".into(), 0);
        let v2 = IrVrVariableDescriptor::new("v2".into(), 0);

        // v1 -- v0 -- v2
        graph.insert(v0.clone(), HashSet::from([v1.clone(), v2.clone()]));
        graph.insert(v1.clone(), HashSet::from([v0.clone()]));
        graph.insert(v2.clone(), HashSet::from([v0.clone()]));
        graph
    }

    fn triangle_graph() -> InterferenceGraph {
        let mut graph: InterferenceGraph = HashMap::new();
        let v0 = IrVrVariableDescriptor::new("v0".into(), 0);
        let v1 = IrVrVariableDescriptor::new("v1".into(), 0);
        let v2 = IrVrVariableDescriptor::new("v2".into(), 0);

        // triangle: each has two edges
        // v0
        // |  \
        // v1--v2
        graph.insert(v0.clone(), HashSet::from([v1.clone(), v2.clone()]));
        graph.insert(v1.clone(), HashSet::from([v0.clone(), v2.clone()]));
        graph.insert(v2.clone(), HashSet::from([v0.clone(), v1.clone()]));
        graph
    }

    fn get_vrs(graph: &InterferenceGraph) -> Vec<IrVrVariableDescriptor> {
        let v0 = graph.keys().find(|k| k.name() == "v0").unwrap();
        let v1 = graph.keys().find(|k| k.name() == "v1").unwrap();
        let v2 = graph.keys().find(|k| k.name() == "v2").unwrap();
        vec![v0.clone(), v1.clone(), v2.clone()]
    }

    #[test]
    fn find_vr_lt_k_when_k_2() {
        let graph = line_graph();
        let found = find_vr_lt_k(&graph, 2);
        assert!(found.is_some());
        assert!(found.unwrap().name() == "v1" || found.unwrap().name() == "v2");
    }

    #[test]
    fn find_vr_lt_k_when_k_1() {
        let graph = line_graph();
        let found = find_vr_lt_k(&graph, 1);
        assert!(found.is_none());
    }

    #[test]
    fn remove_edges_v0() {
        let mut graph = line_graph();
        let vrs = get_vrs(&graph);

        let v0_outgoing = remove_vr_and_edges(&mut graph, &vrs[0]);
        assert!(v0_outgoing.contains(&vrs[1]));
        assert!(v0_outgoing.contains(&vrs[2]));

        // check that incoming edges were removed
        let v1_outgoing = graph.get(&vrs[1]).unwrap();
        assert!(v1_outgoing.is_empty());
        let v2_outgoing = graph.get(&vrs[2]).unwrap();
        assert!(v2_outgoing.is_empty());
    }

    fn triangle_work_stack_k_2() -> Vec<WorkItem> {
        let k = 2;
        let mut graph = triangle_graph();

        let mut work_stack = vec![];

        // run three times, once for each register
        work_stack.push(next_work_item(&mut graph, k));
        work_stack.push(next_work_item(&mut graph, k));
        work_stack.push(next_work_item(&mut graph, k));

        work_stack
    }

    #[test]
    fn next_work_item_k_2() {
        let work_stack = triangle_work_stack_k_2();

        // the actual edges may be different, depending on the underlying order in the sets
        // (HashSet seems to use randomness in order)
        // however, the bottommost item must be a spill, and the edge counts must be (from the
        // bottom of the stack) 2-1-0
        assert!(!work_stack[0].color);
        assert_eq!(work_stack[0].edges.len(), 2);
        assert_eq!(work_stack[1].edges.len(), 1);
        assert_eq!(work_stack[2].edges.len(), 0);
    }

    #[test]
    fn rebuild_graph_triangle_k_2() {
        let mut work_stack = triangle_work_stack_k_2();
        let mut rebuilt_graph: InterferenceGraph = HashMap::new();

        // it should be possible to rebuild the graph from the stack, without yet worrying
        // about spilling/etc.
        while let Some(work_item) = work_stack.pop() {
            rebuild_vr_and_edges(&mut rebuilt_graph, &work_item);
        }

        // we should have a triangle graph again
        let vrs = get_vrs(&rebuilt_graph);
        for vr in &vrs {
            assert!(rebuilt_graph.contains_key(vr));
            assert_eq!(rebuilt_graph.get(vr).unwrap().len(), 2);
        }
    }

    #[test]
    fn registers_and_spills_triangle_k_2() {
        let mut graph = triangle_graph();
        let (registers, spills) = registers_and_spills(&mut graph, 2);
        // there should be one spill when k is 2
        assert_eq!(registers.len(), 2);
        assert_eq!(spills.len(), 1);
    }
}