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

//! Data structure of the lookup table.

use crate::tree::Tree;
use crate::point::Point;
use super::wirelength_vector::*;
use super::position_sequence::PositionSequence;
use super::hanan_grid::HananGridWithCoords;
use num_traits::{Num, Zero, PrimInt};
use std::iter::Sum;
use std::ops::Add;
use super::HananCoord;
use super::marker_types::Canonical;

/// Lookup-table for nets of low degree.
#[derive(Clone, Debug)]
pub struct LookupTable {
    /// Lookup tables for each degree of net.
    pub(crate) sub_lut_by_degree: Vec<SubLookupTable>
}

impl LookupTable {
    /// Find a rectilinear steiner minimal tree which connects all pins.
    /// Number of de-duplicated pins cannot be larger than the maximum supported number of pins of the
    /// lookup table.
    ///
    /// Returns: (tree edges, tree weight).
    pub fn rsmt_low_degree<T>(&self, mut pins: Vec<Point<T>>) -> (Vec<(Point<T>, Point<T>)>, T)
        where T: Ord + Copy + Num + Zero + Sum + std::fmt::Debug {

        // Normalize pins.
        pins.sort();
        pins.dedup();

        let num_pins = pins.len();

        if num_pins <= 1 {
            return (vec![], Zero::zero());
        }

        assert!(pins.len() <= self.max_degree(), "number of pins exceeds the supported maximal net degree");

        let position_sequence = PositionSequence::from_points(&pins);
        let group_index = position_sequence.group_index();

        let sub_lut = self.get_sub_lut(num_pins);
        debug_assert_eq!(sub_lut.values.len(), (1..num_pins + 1).product());

        // Resolve eventual redirect to other group index.
        let (powvs, transform) = sub_lut.resolve_group_index(group_index);

        // Transform the points into the target group.
        pins.iter_mut()
            .for_each(|p| {
                *p = transform.transform_point(*p);
            });

        // Find coordinates of Hanan grid lines.
        let xs = {
            let mut xs: Vec<_> = pins.iter().map(|p| p.x).collect();
            xs.sort();
            xs
        };
        let ys = {
            let mut ys: Vec<_> = pins.iter().map(|p| p.y).collect();
            ys.sort();
            ys
        };

        // Compute distances between Hanan grid lines.
        let (edge_widths_horizontal, edge_widths_vertical) = {
            let mut w = xs.clone();
            let mut h = ys.clone();

            for i in 0..xs.len() - 1 {
                w[i] = w[i + 1] - w[i];
                h[i] = h[i + 1] - h[i];
            }
            w.pop();
            h.pop();
            (w, h)
        };

        // Evaluate POWVs, find the POWV which leads to shortest wirelength.
        let best_powv_and_post = {

            // Compute scalar product of wirelength vector and edge count vector.
            let compute_wire_length = |powv: &WirelenghtVector| -> T {
                let (h, v) = powv.hv_vectors();
                debug_assert_eq!(h.len(), edge_widths_horizontal.len());
                debug_assert_eq!(v.len(), edge_widths_vertical.len());

                let wl_horiz: T = h.into_iter()
                    .zip(&edge_widths_horizontal)
                    .map(|(count, w)| (0..*count).map(|_| *w).sum())
                    // .fold(Zero::zero(), |acc,x| acc+x); // sum
                    .sum();

                let wl_vert: T = v.into_iter()
                    .zip(&edge_widths_vertical)
                    .map(|(count, w)| (0..*count).map(|_| *w).sum())
                    .sum();

                wl_horiz + wl_vert
            };

            powvs.into_iter()
                .map(|powv| (powv, compute_wire_length(&powv.potential_optimal_wirelength_vector)))
                .min_by_key(|(_powv, wire_len)| *wire_len)

        };

        let (best_powv_and_post, wire_length) = best_powv_and_post.expect("no steiner tree found");

        // Transform a point from the hanan grid with unit distances to the original grid.
        let transform_point_to_original_grid = |p: Point<HananCoord>| -> Point<T> {
            Point::new(xs[p.x as usize], ys[p.y as usize])
        };

        let tree_edges = best_powv_and_post.potential_optimal_steiner_tree
            .all_edges()
            .map(|e| (e.start(), e.end())) // Remove zero-length edges.
            .map(|(a, b)| (transform_point_to_original_grid(a), transform_point_to_original_grid(b)));

        // Transform back to original group.
        let inverse_transform = transform.inverse();

        let tree = tree_edges
            .map(|(a, b)|
                     (inverse_transform.transform_point(a), inverse_transform.transform_point(b))
            )
            .filter(|(a, b)| a != b)
            .collect();

        (tree, wire_length)
    }

    /// Get maximum number of pins which is supported by the lookup-table.
    pub fn max_degree(&self) -> usize {
        self.sub_lut_by_degree.len() + 1
    }

    /// Get the lookup-table for a net of degree `num_pins`.
    pub fn get_sub_lut(&self, num_pins: usize) -> &SubLookupTable {
        assert!(num_pins >=2, "no LUT for less than 2 pins");
        &self.sub_lut_by_degree[num_pins-2]
    }

    /// Print number of POWVs and POSTs.
    pub fn print_statistics(&self) {
        for (degree, sub_lut) in self.sub_lut_by_degree.iter().enumerate() {
            let degree = degree + 2;
            let max_num_powvs = sub_lut.values.iter()
                .filter_map(|v| match v {
                    ValuesOrRedirect::Values(v) => Some(v.len()),
                    ValuesOrRedirect::Redirect { .. } => None
                })
                .max()
                .unwrap_or(0);

            let total_num_powvs: usize = sub_lut.values.iter()
                .map(|v| match v {
                    ValuesOrRedirect::Values(v) => v.len(),
                    ValuesOrRedirect::Redirect { group_index, .. } => {
                        match &sub_lut.values[*group_index] {
                            ValuesOrRedirect::Values(v) => v.len(),
                            _ => panic!("invalid double-redirect")
                        }
                    }
                })
                .sum();
            let average_num_pows = (total_num_powvs as f64) / (sub_lut.values.len() as f64);

            println!("degree={}, max. number of POWVs = {}, avg. num. POWVs = {:.3}", degree, max_num_powvs, average_num_pows);
        }
    }
}

#[test]
fn test_rsmt_low_degree() {
    let lut = super::gen_lut::gen_full_lut(5);

    // 3 points
    let points = vec![(5,0).into(), (0, 10).into(), (10, 10).into()];
    let (rsmt, wl) = lut.rsmt_low_degree(points);
    assert_eq!(rsmt.len(), 3);
    assert_eq!(wl, 20);

    // 4 points
    let points = vec![(0,0).into(), (2, 1).into(), (2, 4).into(), (4, 2).into()];
    let (rsmt, wl) = lut.rsmt_low_degree(points);
    assert_eq!(rsmt.len(), 5);
    assert_eq!(wl, 8);

    // 5 points (cross with a center)
    let points = vec![(10,5).into(), (5, 0).into(), (5, 5).into(), (0, 5).into(), (5, 10).into()];
    let (rsmt, wl) = lut.rsmt_low_degree(points);
    assert_eq!(rsmt.len(), 4);
    assert_eq!(wl, 20);
}

/// Lookup-table for a single degree of nets.
#[derive(Clone, Debug)]
pub struct SubLookupTable {
    /// Potential optimal wire-length vectors and trees, indexed by their 'group index'.
    pub(crate) values: Vec<ValuesOrRedirect>
}

impl SubLookupTable {
    /// Get the POWVs and POSTs for a group index. Follow symlinks to other groups.
    ///
    /// Returns the vector of POWVs and POSTs plus the transform which must be applied to convert
    /// the pins.
    fn resolve_group_index(&self, group_index: usize) -> (&Vec<LutValue>, Transform){
        let lut_value = &self.values[group_index];

        // Resolve eventual redirect to other group index.
        match lut_value {
            ValuesOrRedirect::Values(values) => (values, Transform::identity()),
            &ValuesOrRedirect::Redirect { group_index, transform } => {
                // Resolve redirection.
                let values = match &self.values[group_index] {
                    ValuesOrRedirect::Values(values) => values,
                    _ => panic!("lookup-table contains nested redirects")
                };
                (values, transform)
            }
        }
    }
}

/// A pair of a wirelength vector with a tree.
#[derive(Clone, Debug)]
pub struct LutValue {
    /// 'POWV'
    pub(crate) potential_optimal_wirelength_vector: WirelenghtVector,
    /// 'POST'
    pub(crate) potential_optimal_steiner_tree: Tree<Canonical>,
}


/// Contains either the POWVs for a group or redirects to another group which is equivalent
/// under rotations and/or mirroring transforms.
#[derive(Debug, Clone)]
pub enum ValuesOrRedirect {
    /// Wirelength vectors and trees for this group.
    Values(Vec<LutValue>),
    /// Redirection to another equivalent group (under the given transformation).
    Redirect {
        /// Index of the redirect.
        group_index: usize,
        /// Transformation that needs to be applied to get from the current group to the redirect.
        transform: Transform,
    },
}

/// Geometrical transformation of pins and trees.
/// Note: After applying the transformation, the points and trees must be shifted
/// such that the lower left corner of the bounding box is on `(0, 0)`.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct Transform {
    /// Mirror at y-axis.
    mirror_x: bool,
    /// Rotations by multiple of 90 degrees.
    rotate: u8,
}

impl Transform {
    /// Create a new transformation.
    pub fn new(rotate: u8, mirror_x: bool) -> Self {
        Self {
            rotate: rotate % 4,
            mirror_x,
        }
    }

    /// Identity transform.
    pub fn identity() -> Self {
        Self::new(0, false)
    }

    /// Get the inverse transform.
    pub fn inverse(&self) -> Self {
        let r = if self.mirror_x {
            self.rotate
        } else {
            4 - self.rotate
        };
        Self::new(r, self.mirror_x)
    }

    /// Apply the transform to a point.
    pub fn transform_point<T>(&self, mut p: Point<T>) -> Point<T>
        where T: Copy + Num + Zero {

        for _ in 0..self.rotate {
            p.rotate_ccw90();
        }

        if self.mirror_x {
            p.x = T::zero()-p.x;
        }

        p
    }
}