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

use crate::{MetricElem, TspBaseAlg};

struct DistCache(HashMap<(usize, usize), f64>);
impl DistCache {
    fn dist(&self, a: usize, b: usize) -> f64 {
        self.0[&(a.min(b), a.max(b))]
    }
}

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: &DistCache,
    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.dist(a, 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);
        }
    }

    mst
}

fn tsp_from_mst(mst: HashMap<usize, Vec<usize>>) -> Vec<usize> {
    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);

    tsp_path
}

// 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: &DistCache,
    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.dist(x, 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
}

#[allow(clippy::type_complexity)]
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: &DistCache,
    mst: HashMap<usize, Vec<usize>>,
    hash_seed: &Option<u64>,
) -> Vec<usize> {
    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 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);
                }
                pprev = prev;
                prev = cur;
                cur = next;
            }
        }
    }

    r
}

fn refine_2_opt(dist_cache: &DistCache, tour: Vec<usize>) -> (bool, Vec<usize>) {
    if tour.len() < 4 {
        return (false, tour);
    }

    // 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.wrapping_sub(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), tour: &Vec<(usize, usize, usize)>| {
        if tour[i].0 == pi {
            (i, tour[i].2) // forward
        } else {
            (i, tour[i].0) // reverse
        }
    };
    let replace_matching = |(prev, v, next): (usize, usize, usize), old: usize, new: usize| {
        if prev == old {
            (new, v, next)
        } else {
            assert!(next == old);
            (prev, v, new)
        }
    };

    let mut improved = false;

    // for each combination of edges ...
    while le.0 != n - 2 {
        let mut re = adv(le, &tour);
        while re.0 != n - 1 {
            // ... check whether a 2-opt swap improves the tour

            let le_elems = (tour[le.0].1, tour[le.1].1);
            let re_elems = (tour[re.0].1, tour[re.1].1);
            let cur_cost =
                dist_cache.dist(le_elems.0, le_elems.1) + dist_cache.dist(re_elems.0, re_elems.1);
            let swap_cost =
                dist_cache.dist(le_elems.0, re_elems.0) + dist_cache.dist(le_elems.1, re_elems.1);

            if cur_cost <= swap_cost {
                re = adv(re, &tour);
                continue;
            }

            // logically, to swap, we want: swap(le.1, re.0)
            // however, changing the link in the `tour` is much more cumbersome
            // as the tour may contain forward/backward nodes, so we use a helper.
            tour[le.0] = replace_matching(tour[le.0], le.1, re.0);
            tour[le.1] = replace_matching(tour[le.1], le.0, re.1);
            tour[re.0] = replace_matching(tour[re.0], re.1, le.0);
            tour[re.1] = replace_matching(tour[re.1], re.0, le.1);

            swap(&mut le.1, &mut re.0);

            improved = true;

            re = adv(re, &tour);
        }

        le = adv(le, &tour);
    }

    // calculate tour as vector from linked-list
    let mut tour_flat = Vec::with_capacity(n);
    let mut cur = (n - 1, 0);

    tour_flat.push(tour[0].1);
    cur = adv(cur, &tour);
    while cur.1 != 0 {
        tour_flat.push(tour[cur.1].1);
        cur = adv(cur, &tour);
    }

    (improved, tour_flat)
}

fn rotate_towards_optimum(dist_cache: &DistCache, tour: &mut [usize]) {
    let mut max_dist = dist_cache.dist(tour[0], tour[tour.len() - 1]);
    let mut max_dist_at = tour.len() - 1;

    for i in 0..tour.len() - 2 {
        let next_dist = dist_cache.dist(tour[i], tour[i + 1]);
        if next_dist > max_dist {
            max_dist = next_dist;
            max_dist_at = i;
        }
    }

    if max_dist_at != tour.len() - 1 {
        tour.rotate_left(max_dist_at + 1);
    }
}

fn get_dist_cache<M>(embeds: &[M], hash_seed: &Option<u64>) -> DistCache
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]));
        }
    }

    DistCache(r)
}

pub(crate) fn tsp<M>(
    embeds: &[M],
    alg: &TspBaseAlg,
    refinements: usize,
    hash_seed: &Option<u64>,
) -> (Vec<usize>, f64)
where
    M: MetricElem,
{
    match embeds.len() {
        0 => return (vec![], 0.0),
        1 => return (vec![0], 0.0),
        _ => (),
    }

    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 mut tour = match alg {
        TspBaseAlg::MstDfs => tsp_from_mst(mst),
        TspBaseAlg::Christofides => christofides(&dc, mst, hash_seed),
    };

    for _ in 0..refinements {
        rotate_towards_optimum(&dc, &mut tour);
        let res = refine_2_opt(&dc, tour);
        tour = res.1;
        if !res.0 {
            break; // stop early at convergence
        }
    }

    rotate_towards_optimum(&dc, &mut tour);

    let mut total_dist = 0.;
    for i in 0..tour.len() - 1 {
        let (a, b) = (tour[i], tour[i + 1]);
        total_dist += dc.dist(a, b);
    }

    bar.finish();

    (tour, total_dist)
}