// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use super::hanan_grid::UnitHananGrid;
use super::iterator_set_operations::{intersection, symmetric_difference};
use super::marker_types::*;
use super::rectangle::Rect;
use super::wirelength_vector::*;
use super::{HananCoord, Point};

use std::collections::HashSet;
use std::marker::PhantomData;

use bitvec::prelude::BitVec;
use itertools::Itertools;
use std::cmp::Ordering;
use std::hash::{Hash, Hasher};

#[derive(Clone, Debug)]
pub struct Tree<C: CanonicalMarker = NonCanonical> {
    edges: Vec<TreeEdge>,
    _is_canonical: PhantomData<C>,
}

impl PartialEq for Tree<Canonical> {
    fn eq(&self, other: &Self) -> bool {
        self.edges.eq(&other.edges)
    }
}

impl Eq for Tree<Canonical> {}

impl PartialOrd for Tree<Canonical> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.edges.partial_cmp(&other.edges)
    }
}

impl Ord for Tree<Canonical> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.edges.cmp(&other.edges)
    }
}

impl Hash for Tree<Canonical> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.edges.hash(state)
    }
}

impl Tree {
    /// Create an empty tree.
    pub fn empty_non_canonical() -> Self {
        Self {
            edges: Default::default(),
            _is_canonical: Default::default(),
        }
    }
}

impl Tree<Canonical> {
    /// Create an empty tree.
    pub fn empty_canonical() -> Self {
        Self {
            edges: Default::default(),
            _is_canonical: Default::default(),
        }
    }

    /// Create a three from a set of edges.
    /// Returns an error if the set of edges does not form a tree (i.e. if it has cycles or is not connected).
    pub fn from_edges(mut edges: Vec<TreeEdge>) -> Result<Self, ()> {
        edges.sort_unstable();

        let maybe_tree = Tree {
            edges,
            _is_canonical: Default::default(),
        };

        if maybe_tree.is_tree() {
            Ok(maybe_tree)
        } else {
            // Is not a valid tree.
            Err(())
        }
    }

    /// Compute the XOR of the tree edges.
    /// For two equal trees, the result will be an empty iterator.
    pub(crate) fn symmetric_difference<'a>(
        &'a self,
        other: &'a Self,
    ) -> impl Iterator<Item = TreeEdge> + 'a {
        symmetric_difference(self.all_edges(), other.all_edges())
    }
}

#[test]
fn test_tree_from_valid_edges() {
    let edges = vec![
        TreeEdge::new((0, 0).into(), EdgeDirection::Right),
        TreeEdge::new((1, 0).into(), EdgeDirection::Right),
        TreeEdge::new((1, 0).into(), EdgeDirection::Up),
        TreeEdge::new((2, 0).into(), EdgeDirection::Right),
    ];

    let tree = Tree::from_edges(edges);
    assert!(tree.is_ok())
}

#[test]
fn test_tree_from_invalid_edges() {
    // A cycle:
    let edges = vec![
        TreeEdge::new((0, 0).into(), EdgeDirection::Right),
        TreeEdge::new((0, 0).into(), EdgeDirection::Up),
        TreeEdge::new((1, 0).into(), EdgeDirection::Up),
        TreeEdge::new((0, 1).into(), EdgeDirection::Right),
    ];

    let tree = Tree::from_edges(edges);
    assert!(tree.is_err())
}

impl<C: CanonicalMarker> Tree<C> {
    /// Create an empty tree.
    pub fn empty() -> Self {
        Self {
            edges: vec![],
            _is_canonical: Default::default(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.edges.is_empty()
    }

    pub(crate) fn num_edges(&self) -> usize {
        self.edges.len()
    }

    /// Check if the set of edges forms a valid tree, i.e. a single connected and non-cyclic graph.
    /// An empty graph is also a tree.
    pub(crate) fn is_tree(&self) -> bool {
        if self.edges.len() <= 1 {
            true
        } else {
            let mut visited_edges = vec![false; self.edges.len()];
            let mut front = vec![self.edges[0].start()];

            while let Some(current_node) = front.pop() {
                // Find adjacent edges of the current node.
                let mut num_second_visits = 0;
                for (idx, adjacent_edge) in self.adjacent_edges_with_index(current_node) {
                    if visited_edges[idx] {
                        num_second_visits += 1;
                        if num_second_visits > 1 {
                            // Visiting the edge the second time. Cycle detected.
                            // This is not tree.
                            return false;
                        }
                    } else {
                        // Edge was not yet visited.
                        visited_edges[idx] = true;

                        // Grow the front.
                        if adjacent_edge.start() != current_node {
                            front.push(adjacent_edge.start());
                        }
                        if adjacent_edge.end() != current_node {
                            front.push(adjacent_edge.end());
                        }
                    }
                }
            }

            true
        }
    }

    /// Iterate over all nodes of the tree. A node of degree `d` appears `d` times.
    fn all_nodes(&self) -> impl Iterator<Item = Point<HananCoord>> + '_ {
        self.edges
            .iter()
            .map(|e| e.start())
            .chain(self.edges.iter().map(|e| e.end()))
    }

    pub fn all_edges(&self) -> impl Iterator<Item = TreeEdge> + '_ {
        self.edges.iter().copied()
    }

    /// Get all edges adjacent to `p`.
    pub fn adjacent_edges(&self, p: Point<HananCoord>) -> impl Iterator<Item = TreeEdge> + '_ {
        self.adjacent_edges_with_index(p).map(|(_, e)| e)
    }

    /// Get all edges adjacent to `p` and their position in the list of edges.
    fn adjacent_edges_with_index(
        &self,
        p: Point<HananCoord>,
    ) -> impl Iterator<Item = (usize, TreeEdge)> + '_ {
        // TODO: There's a faster way than brute-force search when the edges are sorted.
        self.edges
            .iter()
            .copied()
            .enumerate()
            .filter(move |(_, e)| e.start() == p || e.end() == p)
    }

    /// Check if the tree uses this node.
    pub fn contains_node(&self, p: Point<HananCoord>) -> bool {
        self.adjacent_edges(p).next().is_some()
    }

    /// Check if the tree contains this edge.
    pub fn contains_edge(&self, edge: &TreeEdge) -> bool {
        if C::is_canonical() {
            self.edges.binary_search(edge).is_ok()
        } else {
            self.edges.iter().find(|e| *e == edge).is_some()
        }
    }

    pub fn add_edge(&mut self, edge: TreeEdge) -> &mut Self {
        {
            // Check that the tree remains a connected tree after insertion of the edge.
            // 1) No cycle should be created.
            // 2) New edge must connect to the tree.

            // Simpler but probably slower:
            // let has_existing_edge_at_a = self.adjacent_edges(a).next().is_some();
            // let has_existing_edge_at_b = self.adjacent_edges(b).next().is_some();
            //
            // assert!( !(has_existing_edge_at_a && has_existing_edge_at_b), "insertion of this edge creates a cycle");
            // assert!( (has_existing_edge_at_a || has_existing_edge_at_b) || self.is_empty(), "insertion of this edge creates an unconnected component");

            let a = edge.start();
            let b = edge.end();

            let mut existing_points = self.all_nodes();

            let contains_a_or_b = existing_points.find(|&p| p == a || p == b);
            if let Some(found_point) = contains_a_or_b {
                let other = if found_point == a { b } else { a };
                let contains_other = existing_points.find(|&p| p == other).is_some();
                assert!(!contains_other, "insertion of this edge creates a cycle");
            }

            assert!(
                contains_a_or_b.is_some() || self.is_empty(),
                "insertion of this edge creates an unconnected component"
            );
        }

        self.add_edge_unchecked(edge)
    }

    /// Add edge without guarantee that the tree stays a tree.
    pub(crate) fn add_edge_unchecked(&mut self, edge: TreeEdge) -> &mut Self {
        if C::is_canonical() {
            // Insert edge such that the edges remain sorted.
            match self.edges.binary_search(&edge) {
                Ok(_) => {
                    panic!("edge already exists");
                }
                Err(pos) => self.edges.insert(pos, edge),
            }

            debug_assert!(
                self.edges
                    .iter()
                    .zip(self.edges.iter().skip(1))
                    .all(|(a, b)| a <= b),
                "edges must remain sorted"
            );
        } else {
            self.edges.push(edge)
        }

        self
    }

    pub fn remove_edge(&mut self, edge: TreeEdge) -> &mut Self {
        let a = edge.start();
        let b = edge.end();

        dbg!(a, b);
        let num_adj_a = self.adjacent_edges(a).take(2).count();
        let num_adj_b = self.adjacent_edges(b).take(2).count();
        dbg!(num_adj_a, num_adj_b);

        let will_disconnect_tree = num_adj_a > 1 && num_adj_b > 1;
        assert!(
            !will_disconnect_tree,
            "removing this edge disconnects the tree"
        );

        // Remove the edge.
        if C::is_canonical() {
            // Remove edge while preserving the ordering of the other edges.
            self.edges.retain(|e| e != &edge);
        } else {
            // Faster removal of the edge but destroy the ordering.
            if let Some((index, _)) = self.edges.iter().enumerate().find(|(_, e)| e == &&edge) {
                self.edges.swap_remove(index);
            }
        }

        self
    }

    pub fn into_canonical_form(mut self) -> Tree<Canonical> {
        if !C::is_canonical() {
            self.edges.sort_unstable();
        }
        Tree {
            edges: self.edges,
            _is_canonical: Default::default(),
        }
    }

    pub fn into_non_canonical_form(self) -> Tree<NonCanonical> {
        Tree {
            edges: self.edges,
            _is_canonical: Default::default(),
        }
    }

    /// Shift the whole tree by the vector `(dx, dy)`.
    pub fn translate(&mut self, (dx, dy): (HananCoord, HananCoord)) {
        self.edges
            .iter_mut()
            .for_each(|e| *e = e.translate((dx, dy)))
    }

    /// Rotate the tree around the origin by 90 degrees counter-clock-wise.
    pub fn rotate_90ccw(mut self) -> Tree<NonCanonical> {
        self.edges.iter_mut().for_each(|e| *e = e.rotate_ccw90());

        self.into_non_canonical_form()
    }

    /// Mirror the tree at the y axis.
    pub fn mirror_at_y_axis(mut self) -> Tree<NonCanonical> {
        self.edges
            .iter_mut()
            .for_each(|e| *e = e.mirror_at_y_axis());

        self.into_non_canonical_form()
    }

    /// Merge the other tree into this tree.
    /// The trees must either intersect in exactly one node or at least one of the trees must be empty.
    /// Otherwise the result will not be a connected tree, hence an error is returned.
    pub fn merge(&mut self, other: &Self) -> Result<(), ()> {
        if self.is_empty() {
            self.edges = other.edges.clone();
            return Ok(());
        }

        if other.is_empty() {
            return Ok(());
        }

        // Trees can be merged if their intersection is a non-empty tree (a single node is ok).

        // TODO: Compute intersection with less allocations.
        let nodes_a: HashSet<_> = self.all_nodes().collect();
        let num_node_intersections = other.all_nodes().filter(|n| nodes_a.contains(n)).count();

        let intersection_edges: Vec<_> = if C::is_canonical() {
            // Edges are sorted.
            intersection(self.all_edges(), other.all_edges()).collect()
        } else {
            self.all_edges()
                .filter(|e| other.contains_edge(e))
                .collect()
        };

        let intersection_tree = Tree::from_edges(intersection_edges)?; // Return Err, if intersection is not a tree.

        let additional_edges = other
            .all_edges()
            .filter(|e| !intersection_tree.contains_edge(e));

        if num_node_intersections > 0 {
            if C::is_canonical() {
                // Join edges while keeping the ordering.

                let new_edges = self.all_edges().merge(additional_edges).collect();
                self.edges = new_edges;

                Ok(())
            } else {
                // Don't maintain ordering of edges.
                self.edges.extend(additional_edges);
                Ok(())
            }
        } else {
            dbg!(num_node_intersections);
            Err(())
        }
    }

    /// Compute amount of used edges for each column and row.
    ///
    pub(crate) fn compute_wirelength_vector(&self, w: &mut WirelenghtVector) {
        let (horizontal_edge_count, vertical_edge_count) = w.hv_vectors_mut();

        // let ll = self.bounding_box()
        //     .unwrap_or(Rect::new((0, 0).into(), (0, 0).into()))
        //     .lower_left();

        // Check that the nodes in the tree are contained in the upper right quadrant and do not exceed the coordinates covered by the
        // supplied slices.
        debug_assert!(self
            .all_edges()
            .all(|e| e.start().x >= 0 && e.start().x as usize <= horizontal_edge_count.len()));
        debug_assert!(self
            .all_edges()
            .all(|e| e.start().y >= 0 && e.start().y as usize <= vertical_edge_count.len()));

        self.all_edges().for_each(|e| match e.direction {
            EdgeDirection::Right => {
                horizontal_edge_count[e.start().x as usize] += 1;
            }
            EdgeDirection::Up => {
                vertical_edge_count[e.start().y as usize] += 1;
            }
        });
    }

    /// Compute the corners of the minimal bounding box of all tree nodes.
    /// Returns none if the tree is empty.
    pub fn bounding_box(&self) -> Option<Rect<HananCoord>> {
        let mut nodes = self.all_nodes();
        nodes.next().map(|first_node| {
            let mut lower_left = first_node;
            let mut upper_right = first_node;

            for n in nodes {
                lower_left.x = lower_left.x.min(n.x);
                lower_left.y = lower_left.y.min(n.y);
                upper_right.x = upper_right.x.max(n.x);
                upper_right.y = upper_right.y.max(n.y);
            }

            Rect::new(lower_left, upper_right)
        })
    }

    /// Translate the tree such that the lower left corner of the bounding box is at `(0, 0)`;
    pub fn translate_to_origin(mut self) -> Self {
        if let Some(r) = self.bounding_box() {
            self.translate((-r.lower_left().x, -r.lower_left().y))
        }
        self
    }

    /// Encode the tree as a vector of bits. Each bit is associated with a possible tree edge
    /// and stores if the edge is present or not.
    /// This is used for fast computation of the Hamming distance between trees.
    pub fn to_bitvec(&self) -> BitVec {
        let bbox = self
            .bounding_box()
            .unwrap_or(Rect::new(Point::new(0, 0), Point::new(0, 0)));
        let ll = bbox.lower_left();
        let ur = bbox.upper_right();
        assert!(
            ll.x >= 0 && ll.y >= 0,
            "tree must be in upper right quadrant"
        );

        let bbox = Rect::new(Point::new(0, 0), ur);
        let grid = UnitHananGrid::new(bbox);

        let mut bits = BitVec::new();

        for p in grid.all_points_spiral() {
            let horizontal_edge = TreeEdge::new(p, EdgeDirection::Right);
            let vertical_edge = TreeEdge::new(p, EdgeDirection::Up);
            bits.push(self.contains_edge(&horizontal_edge));
            bits.push(self.contains_edge(&vertical_edge));
        }

        bits
    }
}

/// An edge of the tree.
/// Edges are defined by a point and a direction into positive x or y direction ('right' or 'up').
/// An edge always has unit length, therefore the end point can be easily computed.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub struct TreeEdge {
    start: Point<HananCoord>,
    direction: EdgeDirection,
}

impl TreeEdge {
    pub fn new(start: Point<HananCoord>, direction: EdgeDirection) -> Self {
        Self { start, direction }
    }

    /// Returns `Err` when the points are not adjacent on the grid.
    pub fn from_points(a: Point<HananCoord>, b: Point<HananCoord>) -> Result<Self, ()> {
        let dx = b.x - a.x;
        let dy = b.y - a.y;

        if dx.abs() + dy.abs() != 1 {
            Err(())
        } else {
            let edge = match (dx, dy) {
                (1, _) => TreeEdge::new(a, EdgeDirection::Right),
                (-1, _) => TreeEdge::new(b, EdgeDirection::Right),
                (_, 1) => TreeEdge::new(a, EdgeDirection::Up),
                (_, -1) => TreeEdge::new(b, EdgeDirection::Up),
                _ => unreachable!(),
            };

            debug_assert!(edge.start() == a || edge.start() == b);
            debug_assert!(edge.end() == a || edge.end() == b);

            Ok(edge)
        }
    }

    pub fn is_vertical(&self) -> bool {
        self.direction == EdgeDirection::Up
    }

    pub fn is_horizontal(&self) -> bool {
        self.direction == EdgeDirection::Right
    }

    /// Get the start point of the edge.
    pub fn start(&self) -> Point<HananCoord> {
        self.start
    }

    /// Get the end point of the edge.
    pub fn end(&self) -> Point<HananCoord> {
        match self.direction {
            EdgeDirection::Right => Point::new(self.start.x + 1, self.start.y),
            EdgeDirection::Up => Point::new(self.start.x, self.start.y + 1),
        }
    }

    /// Shift the edge by a vector `(dx, dy)`.
    fn translate(mut self, (dx, dy): (HananCoord, HananCoord)) -> Self {
        self.start.x += dx;
        self.start.y += dy;
        self
    }

    /// Mirror the edge at the y-axis.
    fn mirror_at_y_axis(mut self) -> Self {
        match self.direction {
            EdgeDirection::Right => {
                self.start.x = -self.end().x;
            }
            EdgeDirection::Up => {
                self.start.x = -self.start.x;
            }
        }
        self
    }

    /// Rotate around the origin by 90 degrees counter-clock-wise.
    fn rotate_ccw90(&self) -> Self {
        match self.direction {
            EdgeDirection::Right => Self {
                start: self.start().rotate_ccw90(),
                direction: EdgeDirection::Up,
            },
            EdgeDirection::Up => Self {
                start: self.end().rotate_ccw90(),
                direction: EdgeDirection::Right,
            },
        }
    }
}

/// Direction into positive x or y direction.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub enum EdgeDirection {
    /// Horizontal edge, directed in positive x direction.
    Right,
    /// Vertical edge, directed in positive y direction.
    Up,
}

#[test]
fn test_mirror_edge() {
    let e = TreeEdge::new((1, 2).into(), EdgeDirection::Right);
    let mirrored = e.mirror_at_y_axis();

    let expected = TreeEdge::new((-2, 2).into(), EdgeDirection::Right);
    assert_eq!(mirrored, expected);
}

#[test]
fn test_rotate_edge() {
    let e = TreeEdge::new((1, 2).into(), EdgeDirection::Right);
    let rotated = e.rotate_ccw90();

    let expected = TreeEdge::new((-2, 1).into(), EdgeDirection::Up);
    assert_eq!(rotated, expected);
}

#[test]
fn test_rotate_edge_4_times() {
    let e = TreeEdge::new((1, 2).into(), EdgeDirection::Right);
    let rotated = e
        .rotate_ccw90()
        .rotate_ccw90()
        .rotate_ccw90()
        .rotate_ccw90();

    assert_eq!(rotated, e);
}

#[test]
fn test_add_edges() {
    let mut tree = Tree::empty_non_canonical();

    tree.add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Right))
        .add_edge(TreeEdge::new((1, 0).into(), EdgeDirection::Right));
}

#[test]
fn test_bounding_box() {
    let mut tree = Tree::empty_non_canonical();

    tree.add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Right))
        .add_edge(TreeEdge::new((1, 0).into(), EdgeDirection::Up))
        .add_edge(TreeEdge::new((1, 1).into(), EdgeDirection::Up))
        .add_edge(TreeEdge::new((-1, 0).into(), EdgeDirection::Right));

    assert_eq!(
        tree.bounding_box(),
        Some(Rect::new((-1, 0).into(), (1, 2).into()))
    );
}

#[test]
fn test_adjacent_edges() {
    let mut tree = Tree::empty_non_canonical();

    let a = TreeEdge::new((0, 0).into(), EdgeDirection::Right);
    let b = TreeEdge::new((1, 0).into(), EdgeDirection::Right);

    assert_eq!(tree.adjacent_edges(a.start()).count(), 0);
    assert_eq!(tree.adjacent_edges(a.end()).count(), 0);

    tree.add_edge(a);
    tree.add_edge(b);

    assert_eq!(tree.adjacent_edges(a.start()).next(), Some(a));
    assert_eq!(tree.adjacent_edges(a.end()).count(), 2);
    assert_eq!(tree.adjacent_edges(b.start()).count(), 2);
    assert_eq!(tree.adjacent_edges(b.end()).next(), Some(b));
}

#[test]
fn test_remove_edges() {
    let mut tree = Tree::empty_non_canonical();
    let a = TreeEdge::new((0, 0).into(), EdgeDirection::Right);
    let b = TreeEdge::new((1, 0).into(), EdgeDirection::Right);
    let c = TreeEdge::new((2, 0).into(), EdgeDirection::Right);

    tree.add_edge(a)
        .add_edge(b)
        .add_edge(c)
        .remove_edge(a)
        .remove_edge(c)
        .remove_edge(b);
}

#[test]
#[should_panic]
fn test_remove_disconnecting_edge() {
    let mut tree = Tree::empty_non_canonical();

    let a = TreeEdge::new((0, 0).into(), EdgeDirection::Right);
    let b = TreeEdge::new((1, 0).into(), EdgeDirection::Right);
    let c = TreeEdge::new((2, 0).into(), EdgeDirection::Right);

    tree.add_edge(a)
        .add_edge(b)
        .add_edge(c)
        // Remove the middle edge.
        .remove_edge(b);
}

#[test]
fn test_canonical_tree_equivalence() {
    let mut tree1 = Tree::empty_canonical();
    let mut tree2 = Tree::empty_canonical();

    assert_eq!(tree1, tree2);

    tree1.add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Right));
    tree1.add_edge(TreeEdge::new((1, 0).into(), EdgeDirection::Right));

    // Add edges in reverse order.
    tree2.add_edge(TreeEdge::new((1, 0).into(), EdgeDirection::Right));
    tree2.add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Right));

    assert_eq!(tree1, tree2);
}

#[test]
#[should_panic]
fn test_create_unconnected_tree() {
    let mut tree = Tree::empty_non_canonical();
    tree.add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Right));
    tree.add_edge(TreeEdge::new((2, 0).into(), EdgeDirection::Right));
}

#[test]
#[should_panic]
fn test_create_tree_with_cycle() {
    let mut tree = Tree::empty_non_canonical();
    tree.add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Right));
    tree.add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Up));
    tree.add_edge(TreeEdge::new((1, 0).into(), EdgeDirection::Up));
    // Close the cycle.
    tree.add_edge(TreeEdge::new((0, 1).into(), EdgeDirection::Right));
}

#[test]
fn test_wirelength_vector() {
    let mut tree = Tree::empty_non_canonical();

    tree.add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Right))
        .add_edge(TreeEdge::new((1, 0).into(), EdgeDirection::Up))
        .add_edge(TreeEdge::new((1, 1).into(), EdgeDirection::Up))
        .add_edge(TreeEdge::new((1, 2).into(), EdgeDirection::Up))
        .add_edge(TreeEdge::new((0, 0).into(), EdgeDirection::Up));

    let mut w = WirelenghtVector::zero(2, 3);

    tree.compute_wirelength_vector(&mut w);

    assert_eq!(w.hv_vectors().0, vec![1, 0]);
    assert_eq!(w.hv_vectors().1, vec![2, 1, 1]);
}