use crate::ir::ir_block::IrBlock;
use crate::ir::ir_variable::IrVariableId;
use crate::ir::variable_locations::VariableLocations;
use crate::offset_counter::OffsetCounter;
use dvm_lib::instruction::{Register, StackFrameOffset};
use std::collections::{HashMap, HashSet};

pub type RegisterAssignment = Register;
pub type InterferenceGraph = HashMap<IrVariableId, HashSet<IrVariableId>>;
pub type LivenessSets = Vec<HashSet<IrVariableId>>;

pub fn block_live_in_live_out(ir_block: &IrBlock) -> (LivenessSets, LivenessSets) {
    // init
    let n_statements = ir_block.statements().len();
    let mut live_in: LivenessSets = vec![HashSet::new(); n_statements];
    let mut live_out: LivenessSets = vec![HashSet::new(); n_statements];

    loop {
        let mut did_work = false;
        // Go backwards for efficiency
        for (statement_index, ir_statement) in ir_block.statements().iter().enumerate().rev() {
            // out (union of successors ins)
            // for now, a statement can only have one successor
            // this will need to be updated when we add jumps
            let statement_live_out = &mut live_out[statement_index];
            let successor_live_in = &live_in[statement_index + 1];
            for ir_variable_id in successor_live_in {
                if statement_live_out.insert(*ir_variable_id) {
                    did_work = true;
                }
            }

            // in: use(s) U ( out(s) - def(s) )
            let use_s = ir_statement
                .vr_uses()
                .iter()
                .map(|u| *u)
                .collect::<HashSet<_>>();
            let out_s = &live_out[statement_index];
            let def_s = ir_statement
                .vr_definitions()
                .iter()
                .map(|d| *d)
                .collect::<HashSet<_>>();
            let rhs = out_s - &def_s;
            let new_ins = use_s.union(&rhs).map(|v| *v).collect::<HashSet<_>>();

            // add new ins to statement's live in
            let statement_live_in = &mut live_in[statement_index];
            for new_in in new_ins {
                if statement_live_in.insert(new_in) {
                    // if we added a previously unadded variable, we did work
                    did_work = true;
                }
            }
        }
        // break only if we did nothing!
        if !did_work {
            break;
        }
    }
    (live_in, live_out)
}

pub fn block_interference_graph(
    ir_block: &IrBlock,
    spilled: &HashSet<IrVariableId>,
) -> InterferenceGraph {
    // create a set of all variables used in the block
    let mut all_vr_variables: HashSet<IrVariableId> = HashSet::new();
    for statement in ir_block.statements() {
        all_vr_variables.extend(statement.vr_definitions());
        all_vr_variables.extend(statement.vr_uses());
    }

    // create graph and init all variables' outgoing-edge sets
    let mut graph: InterferenceGraph = HashMap::new();
    for variable in all_vr_variables {
        graph.insert(variable, HashSet::new());
    }

    let (_, live_out) = block_live_in_live_out(ir_block);

    for (statement_index, ir_statement) in ir_block.statements().iter().enumerate() {
        let statement_live_out = &live_out[statement_index];
        for definition_vr_variable in ir_statement.vr_definitions() {
            for live_out_variable in statement_live_out {
                // we do the following check to avoid adding an edge to itself
                if definition_vr_variable != *live_out_variable {
                    // add edges to both sets (two-way)
                    let definition_edges = graph.get_mut(&definition_vr_variable).unwrap();
                    definition_edges.insert(*live_out_variable);
                    let live_out_variable_edges = graph.get_mut(live_out_variable).unwrap();
                    live_out_variable_edges.insert(definition_vr_variable);
                }
            }
        }
    }

    graph
}

pub fn block_assign_registers(ir_block: &IrBlock, register_count: usize) -> VariableLocations {
    let mut spilled: HashSet<IrVariableId> = HashSet::new();
    loop {
        let mut interference_graph = block_interference_graph(ir_block, &spilled);
        let (registers, new_spills) = registers_and_spills(&mut interference_graph, register_count);

        if spilled != new_spills {
            spilled = new_spills; // todo: figure out if this works algorithmically
        } else {
            let mut spills_to_stack_offsets: HashMap<IrVariableId, StackFrameOffset> =
                HashMap::new();
            let mut offset_counter = 0isize;
            for spill in spilled {
                spills_to_stack_offsets.insert(spill, offset_counter);
                offset_counter += 1;
            }
            return VariableLocations::new(registers, spills_to_stack_offsets);
        }
    }
}

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

pub trait VrUser {
    fn vr_definitions(&self) -> HashSet<IrVariableId> {
        HashSet::new()
    }

    #[deprecated]
    fn vr_uses(&self) -> HashSet<IrVariableId>;

    #[deprecated]
    fn propagate_spills(&mut self, _spills: &HashSet<IrVariableId>) {
        // default no-op
    }

    #[deprecated]
    fn propagate_register_assignments(
        &mut self,
        _assignments: &HashMap<IrVariableId, RegisterAssignment>,
    ) {
        // default no-op
    }

    #[deprecated]
    fn propagate_stack_offsets(&mut self, _counter: &mut OffsetCounter) {
        // default no-op
    }
}

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

pub fn registers_and_spills(
    interference_graph: &mut InterferenceGraph,
    k: usize,
) -> (
    HashMap<IrVariableId, RegisterAssignment>,
    HashSet<IrVariableId>,
) {
    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<IrVariableId, RegisterAssignment> = HashMap::new();
    let mut spills: HashSet<IrVariableId> = 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<IrVariableId, RegisterAssignment>,
) {
    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<IrVariableId> {
    interference_graph.iter().find_map(
        |(vr, neighbors)| {
            if neighbors.len() < k { Some(*vr) } else { None }
        },
    )
}

/// Returns the (removed) outgoing edges for the given vr
fn remove_vr_and_edges(
    interference_graph: &mut InterferenceGraph,
    vr: &IrVariableId,
) -> HashSet<IrVariableId> {
    // 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<IrVariableId, 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();

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

    fn triangle_graph() -> InterferenceGraph {
        let mut graph: InterferenceGraph = HashMap::new();

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

    fn get_vrs() -> Vec<IrVariableId> {
        vec![0, 1, 2]
    }

    #[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() == 1 || found.unwrap() == 2);
    }

    #[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();

        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();
        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);
    }
}