use ahash::{random_state::RandomState, HashMap, HashSet};
use indicatif::{ProgressBar, ProgressStyle};
use partitions::partition_vec::PartitionVec;
use mset::MultiSet;
use std::{hash::Hash, cmp::Eq};

use crate::{MetricElem, TspAlg};

fn get_hashmap<K, V>(hash_seed: &Option<u64>, capacity: Option<usize>) -> HashMap<K, V> {
    let hasher = match hash_seed {
        Some(s) => RandomState::with_seeds(*s, 0, 0, 0),
        None => RandomState::new(),
    };

    match capacity {
        Some(sz) => HashMap::with_capacity_and_hasher(sz, hasher),
        None => HashMap::with_hasher(hasher),
    }
}

fn get_hashset<V>(hash_seed: &Option<u64>, capacity: Option<usize>) -> HashSet<V> {
    let hasher = match hash_seed {
        Some(s) => RandomState::with_seeds(*s, 0, 0, 0),
        None => RandomState::new(),
    };

    match capacity {
        Some(sz) => HashSet::with_capacity_and_hasher(sz, hasher),
        None => HashSet::with_hasher(hasher),
    }
}

fn get_multiset<V: Hash + Eq>(hash_seed: &Option<u64>, capacity: Option<usize>) -> MultiSet<V, RandomState> {
    let hasher = match hash_seed {
        Some(s) => RandomState::with_seeds(*s, 0, 0, 0),
        None => RandomState::new(),
    };

    match capacity {
        Some(sz) => MultiSet::with_capacity_and_hasher(sz, hasher),
        None => MultiSet::with_hasher(hasher),
    }
}

fn get_mst(
    dist_cache: &HashMap<(usize, usize), f64>,
    num_embeds: usize,
    hash_seed: &Option<u64>,
) -> HashMap<usize, Vec<usize>> {
    let mut possible_edges = Vec::with_capacity((num_embeds * num_embeds - num_embeds) / 2);
    let mut mst = get_hashmap(hash_seed, Some(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)
}

// this is a greedy, non-exact implementation. The polynomial time solution would be
// in O(n^3), which is too large for my taste, while this is O(n^2).
/// 'verts' must be an even number of vertices with odd degree
fn min_weight_matching(
    dist_cache: &HashMap<(usize, usize), f64>,
    verts: &[usize],
    hash_seed: &Option<u64>,
) -> 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 = get_hashmap(hash_seed, None);

    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, MultiSet<usize, RandomState>>,
    hash_seed: &Option<u64>,
) -> (usize, HashMap<usize, Vec<(usize, usize, usize)>>) {
    let mut r: HashMap<_, Vec<_>> = get_hashmap(hash_seed, None);
    let mut partially_explored = get_hashset(hash_seed, None);

    // TODO das hier brauch fixing. bitte nochmal algorithmus verstehen vorher.
    // 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>>,
    hash_seed: &Option<u64>,
) -> (Vec<usize>, f64) {
    let mut ext_mst: HashMap<usize, MultiSet<usize, RandomState>> = get_hashmap(hash_seed, Some(mst.len()));
    for (k, v) in mst.into_iter() {
        let mut as_mset = get_multiset(hash_seed, Some(v.len()));
        for l in v.into_iter() {
            as_mset.insert(l);
        }
        ext_mst.insert(k, as_mset);
    }

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

    // from here on, 'mst' would be 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, hash_seed) {
        ext_mst.get_mut(&a).unwrap().insert(b);
    }

    let mut r = Vec::new();
    let mut total_dist = 0.;
    let mut visited_verts = get_hashset(hash_seed, None);
    visited_verts.insert(usize::MAX);

    let mut pprev = usize::MAX;
    let mut prev = usize::MAX;
    let (mut cur, euler_tour) = euler_tour(ext_mst, hash_seed);
    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_2_opt(
    dist_cache: &HashMap<(usize, usize), f64>,
    tour: Vec<usize>,
    tour_len: f64,
) -> (Vec<usize>, f64) {
    if tour.len() < 4 {
        return (tour, tour_len);
    }

    // convert the tour into a doubly linked-list. instead of pointers, we use
    // an array of index pairs.
    let mut tour: Vec<_> = tour
        .into_iter()
        .enumerate()
        .map(|(i, v)| (i - 1, v, i + 1))
        .collect();
    let n = tour.len();
    // fix boundary
    tour[0].0 = n - 1;
    tour[n - 1].2 = 0;

    // this will be an implementation of 2-opt swapping where the swapping takes O(1).
    // to do this, the links of a node no are no longer "backward"/"forward"; instead,
    // the forward direction will be in the direction of the link you didn't come from.
    // This implicit direction indicator allows for reversing a sublist in O(1)
    let mut /* "left edge" */ le = (
        /* previous tour index */ n - 1,
        /* tour index */ 0,
    );

    let adv = |(pi, i): (usize, usize)| {
        if tour[i].0 == pi {
            (i, tour[i].2) // forward
        } else {
            (i, tour[i].0) // reverse
        }
    };

    while le.0 != n - 1 {
        // TODO so nehmen wir die boundary nicht mit.
        let mut re = adv(le);
        println!("{:?} <-> {:?}", le, re);

        le = adv(le);
    }

    todo!()
}

fn get_dist_cache<M>(embeds: &[M], hash_seed: &Option<u64>) -> HashMap<(usize, usize), f64>
where
    M: MetricElem,
{
    let n = embeds.len();
    let mut r = get_hashmap(hash_seed, Some((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, hash_seed: &Option<u64>) -> (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, hash_seed);

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

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

    bar.finish();

    r
}