use indicatif::{ProgressBar, ProgressStyle};
use multiset::HashMultiSet;
use partitions::partition_vec::PartitionVec;
use std::collections::{HashMap, HashSet};

use crate::{MetricElem, TspAlg};

fn get_mst(dist_cache: &HashMap<(usize, usize), f64>, num_embeds: usize) -> HashMap<usize, Vec<usize>> {
    let mut possible_edges = Vec::with_capacity((num_embeds * num_embeds - num_embeds) / 2);
    let mut mst = HashMap::with_capacity(num_embeds);

    // insert all edges we could ever use
    mst.insert(0, Vec::new());
    for a in 0..num_embeds {
        for b in a + 1..num_embeds {
            possible_edges.push((dist_cache[&(a.min(b), a.max(b))], a, b));
        }
    }

    let mut subtrees: PartitionVec<usize> = (0..num_embeds).collect();

    possible_edges.sort_unstable_by(|(da, _, _), (db, _, _)| da.partial_cmp(db).unwrap());
    for (_, a, b) in possible_edges.into_iter() {
        if !subtrees.same_set(a, b) {
            mst.entry(a).or_default().push(b);
            mst.entry(b).or_default().push(a);

            subtrees.union(a, b);
        }
        if mst.len() >= num_embeds {
            break;
        }
    }

    mst
}

fn tsp_from_mst(dist_cache: &HashMap<(usize, usize), f64>, mst: HashMap<usize, Vec<usize>>) -> (Vec<usize>, f64) {
    fn dfs(cur: usize, prev: usize, t: &HashMap<usize, Vec<usize>>, into: &mut Vec<usize>) {
        into.push(cur);
        t.get(&cur).unwrap().iter().for_each(|&c| {
            if c != prev {
                dfs(c, cur, t, into)
            }
        });
    }
    let mut tsp_path = Vec::with_capacity(mst.len());
    dfs(0, usize::MAX, &mst, &mut tsp_path);

    let mut total_dist = 0.;
    for i in 0..tsp_path.len() - 1 {
        let (a, b) = (tsp_path[i], tsp_path[i + 1]);
        total_dist += dist_cache[&(a.min(b), a.max(b))];
    }

    (tsp_path, total_dist)
}

// TODO this is a non-ideal, greedy algorithm. there are better algorithms for this,
// and i should probably implement one.
/// 'verts' must be an even number of vertices with odd degree
fn min_weight_matching(dist_cache: &HashMap<(usize, usize), f64>, verts: &[usize]) -> HashMap<usize, usize> {
    let num_odd = verts.len();
    assert!(num_odd % 2 == 0);

    let mut possible_edges = Vec::with_capacity((num_odd * num_odd - num_odd) / 2);
    for &x in verts {
        for &y in verts {
            if x != y {
                possible_edges.push((dist_cache[&(x.min(y), x.max(y))], x, y));
            }
        }
    }
    possible_edges.sort_unstable_by(|(da, _, _), (db, _, _)| da.partial_cmp(db).unwrap());

    let mut res = HashMap::new();

    for (_, a, b) in possible_edges.into_iter() {
        if res.len() >= num_odd {
            break;
        }

        if !res.contains_key(&a) && !res.contains_key(&b) {
            res.insert(a, b);
            res.insert(b, a);
        }
    }

    res
}

fn euler_tour(
    mut graph: HashMap<usize, HashMultiSet<usize>>,
) -> (usize, HashMap<usize, Vec<(usize, usize, usize)>>) {
    let mut r: HashMap<_, Vec<_>> = HashMap::new();
    let mut partially_explored = HashSet::new();

    // initial setup: pretend that we have only the path 'INF -> root -> INF'
    // for some arbitrary root, and set cur to some node next to root.
    // This mimicks the state we're in just after a new phase (because it is)
    let &root = graph.keys().next().unwrap();
    r.insert(root, vec![(usize::MAX, usize::MAX, usize::MAX)]);
    let e = graph.get_mut(&root).unwrap();
    let &next = e.iter().next().unwrap();
    e.remove(&next);
    graph.get_mut(&next).unwrap().remove(&root);

    let mut cur = next;
    let mut prev = root;
    let mut pprev = usize::MAX;
    let mut circ_start_edge = cur;

    loop {
        let e = graph.get_mut(&cur).unwrap();
        if e.len() <= 1 {
            partially_explored.remove(&cur);
        } else {
            partially_explored.insert(cur);
        }

        match e.iter().next() {
            Some(&next) => {
                e.remove(&next);
                graph.get_mut(&next).unwrap().remove(&cur);

                r.entry(cur).or_default().push((pprev, prev, next));
                pprev = prev;
                prev = cur;
                cur = next;
            }
            None => {
                // we got stuck, which means we returned to the starting vertex of
                // the current phase. now, we need join the 2 formed trips

                // pick an arbitrary existing edge-pair going through cur
                let cur_vec = r.get_mut(&cur).unwrap();
                let (pp, p, n) = cur_vec[0];
                // reroute
                cur_vec[0] = (pp, p, circ_start_edge);
                cur_vec.push((pprev, prev, n));

                // after rerouting, the pprev value of the next node will be wrong
                match r.get_mut(&n) {
                    // should only happen with n == usize::MAX. no need to reroute the
                    // following node in that case, as there's no such node
                    None => (),
                    Some(rerouted_vec) => {
                        let p = rerouted_vec
                            .iter()
                            .position(|&(_, other_p, _)| other_p == cur)
                            .unwrap();
                        rerouted_vec[p] = (prev, cur, rerouted_vec[p].2);
                    }
                }

                // are there any partially explored vertices left?
                match partially_explored.iter().next() {
                    None => break, // graph fully explored :)
                    Some(&new_cur) => {
                        // reset our active point (note that we don't have to delete
                        // new_cur from partially_explored; that's done the next time we
                        // get there)
                        let e = graph.get_mut(&new_cur).unwrap();
                        let &next = e.iter().next().unwrap();
                        e.remove(&next);
                        graph.get_mut(&next).unwrap().remove(&new_cur);

                        circ_start_edge = next;
                        pprev = r[&new_cur][0].1;
                        prev = new_cur;
                        cur = next;
                    }
                }
            }
        }
    }

    (root, r)
}

fn christofides(dist_cache: &HashMap<(usize, usize), f64>, mst: HashMap<usize, Vec<usize>>) -> (Vec<usize>, f64) {
    let mut mst: HashMap<_, HashMultiSet<_>> = mst
        .into_iter()
        .map(|(k, v)| (k, v.into_iter().collect()))
        .collect();

    let odd_verts: Vec<_> = mst
        .iter()
        .filter_map(|(&i, s)| if s.len() % 2 == 0 { None } else { Some(i) })
        .collect();

    // from here on, 'mst' is a bit of a misnomer, as we're adding more edges such
    // that all vertices have even degree
    for (a, b) in min_weight_matching(dist_cache, &odd_verts) {
        mst.get_mut(&a).unwrap().insert(b);
    }

    let mut r = Vec::new();
    let mut total_dist = 0.;
    let mut visited_verts = HashSet::new();
    visited_verts.insert(usize::MAX);

    let mut pprev = usize::MAX;
    let mut prev = usize::MAX;
    let (mut cur, euler_tour) = euler_tour(mst);
    loop {
        match euler_tour.get(&cur) {
            None => break, // tour complete: this should happen iff 'cur == usize::MAX'
            Some(v) => {
                let &(_, _, next) = v
                    .iter()
                    .find(|(pp, p, _)| *pp == pprev && *p == prev)
                    .unwrap();

                if visited_verts.insert(next) {
                    // haven't visited 'next' yet
                    r.push(next);
                    total_dist += dist_cache[&(cur.min(next), cur.max(next))];
                }
                pprev = prev;
                prev = cur;
                cur = next;
            }
        }
    }

    (r, total_dist)
}

fn refine<M>(_: &[M], _: Vec<usize>, _: f64) -> (Vec<usize>, f64)
where
    M: MetricElem,
{
    // convert the tour into a linked-list representation. instead of pointers, we use
    // an array of indices
    //let tour_ll = Vec::new();
    todo!()
}

fn get_dist_cache<M>(embeds: &[M]) -> HashMap<(usize, usize), f64>
where
    M: MetricElem,
{
    let n = embeds.len();
    let mut r = HashMap::with_capacity((n * n - n) / 2);
    for a in 0..n {
        for b in a + 1..n {
            r.insert((a, b), embeds[a].dist(&embeds[b]));
        }
    }

    r
}

pub(crate) fn tsp<M>(embeds: &[M], alg: &TspAlg) -> (Vec<usize>, f64)
where
    M: MetricElem,
{
    let bar = ProgressBar::new_spinner();
    bar.set_style(ProgressStyle::with_template("{spinner} {msg}").unwrap());
    bar.enable_steady_tick(std::time::Duration::from_millis(100));

    bar.set_message("Calculating distances...");
    let dc = get_dist_cache(embeds);

    bar.set_message("Finding mst...");
    let mst = get_mst(&dc, embeds.len());

    bar.set_message("Finding path...");
    let r = match alg {
        TspAlg::MstDfs => tsp_from_mst(&dc, mst),
        TspAlg::Christofides => christofides(&dc, mst),
        TspAlg::ChristofidesRefined => {
            let (p, l) = christofides(&dc, mst);
            bar.set_message("Refining path...");
            refine(embeds, p, l)
        }
    };

    bar.finish();

    r
}