// Copyright (c) 2020-2021 Thomas Kramer.
// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Compute the placement density or placement density maps of layouts.

use crate::db;
use ndarray::Array2;
use num_traits::{FromPrimitive, Num, ToPrimitive};

/// Pixelized representation of the cell locations.
/// Cells are drawn to the DensityMap like to a rasterized image.
/// The 2D array is the input to the FFT based electrostatic force computation.
///
/// Use `DensityMap::new()` and `DensityMap::from_data()` to create a DensityMap struct.
///
/// * `F`: Data type of coordinates.
/// * `Z`: Data type of values.
///
/// # Example
/// ```
/// use libreda_pnr::db;
/// use libreda_pnr::metrics::placement_density::DensityMap;
///
/// // Create a 'bin image' for accumulating densities.
/// let mut c = DensityMap::new(
///     db::Rect::new((0.0, 0.0), (10.0, 10.0)),
///     (10, 10)
/// );
/// // Draw a rectangle to the density image.
/// let r = db::Rect::new((1.0, 1.0), (2.0, 3.0));
/// c.draw_rect(&r, 1.0);
///
/// // Query the density.
/// assert_eq!(c.density_at((1.5, 1.5).into()), 1.);
///
/// // Directly access the density bins.
/// // The indices may not coincide with the coordinates!
/// // Here this is only the case because of the very specific dimension of the density map.
/// assert_eq!(c.data[[0, 0]], 0.0);
/// assert_eq!(c.data[[1, 1]], 1.0);
/// assert_eq!(c.data[[2, 3]], 0.0);
///
/// ```
#[derive(Clone)]
pub struct DensityMap<F, Z> {
    /// Offset and dimension of the drawable DensityMap.
    pub dimension: db::Rect<F>,
    /// Raster data of the DensityMap. Hold the sum of values, not densities.
    pub data: Array2<Z>,
}

impl<F, Z> DensityMap<F, Z>
where
    F: Copy,
{
    /// Consume this object and return the underlying data array.
    pub fn get_data(self) -> Array2<Z> {
        self.data
    }

    /// Get reference to underlying data array.
    pub fn get_data_ref(&self) -> &Array2<Z> {
        &self.data
    }

    /// Get mutable reference to underlying data array.
    pub fn get_data_ref_mut(&mut self) -> &mut Array2<Z> {
        &mut self.data
    }

    /// Get the area of the DensityMap as a rectangle.
    pub fn dimension(&self) -> db::Rect<F> {
        self.dimension
    }

    /// Get the number of bins in `x` and `y` direction.
    fn num_bins(&self) -> (usize, usize) {
        self.data.dim()
    }
}

impl<F, Z> DensityMap<F, Z>
where
    F: Copy + Num + PartialOrd + FromPrimitive + ToPrimitive,
    Z: Copy + std::ops::Div<F, Output = Z>,
{
    /// Read the density at a given coordinate `p`.
    /// The values are interpolated by the 'nearest neighbour' strategy.
    ///
    /// # Panics
    /// Panics when the point `p` is outside of the defined area of the density map.
    pub fn density_at(&self, p: db::Point<F>) -> Z {
        self.get_density_at(p)
            .expect("Point `p` is outside of the defined region.")
    }

    /// Read the density at a given coordinate `p`.
    /// The values are interpolated by the 'nearest neighbour' strategy.
    ///
    /// Returns `None` if `p` is outside of the defined region.
    pub fn get_density_at(&self, p: db::Point<F>) -> Option<Z> {
        self.get_value_at(p).map(|v| v / self.bin_area())
    }
}

impl<F, Z> DensityMap<F, Z>
where
    F: Copy + Num + FromPrimitive + ToPrimitive + PartialOrd,
    Z: Copy,
{
    /// Get real dimension (width, height) of a bin.
    pub fn bin_dimension(&self) -> (F, F) {
        let (w, h) = self.num_bins();
        let bin_width = self.dimension.width() / F::from_usize(w).unwrap();
        let bin_height = self.dimension.height() / F::from_usize(h).unwrap();
        (bin_width, bin_height)
    }

    /// Get the area of a bin.
    pub fn bin_area(&self) -> F {
        let (w, h) = self.bin_dimension();
        w * h
    }

    /// Convert a coordinate into array indices.
    pub fn coordinates_to_indices(&self, p: db::Point<F>) -> (usize, usize) {
        let r = self.dimension;
        assert!(r.contains_point(p), "Point is not inside boundary.");
        let (w, h) = self.data.dim();
        let x = (p.x - r.lower_left.x) * F::from_usize(w).unwrap() / r.width();
        let y = (p.y - r.lower_left.y) * F::from_usize(h).unwrap() / r.height();
        (x.to_usize().unwrap(), y.to_usize().unwrap())
    }

    /// Read the accumulated value at a given coordinate `p`.
    /// The values are interpolated by the 'nearest neighbour' strategy.
    ///
    /// # Panics
    /// Panics when the point `p` is outside of the defined area of the density map.
    /// `get_value_at()` returns an `Option` instead of panicking.
    pub fn value_at(&self, p: db::Point<F>) -> Z {
        self.get_value_at(p)
            .expect("Point `p` is out of the defined region.")
    }

    /// Read the accumulated value at a given coordinate `p`.
    /// The values are interpolated by the 'nearest neighbour' strategy.
    /// Returns `None` if `p` is outside of the map region.
    pub fn get_value_at(&self, p: db::Point<F>) -> Option<Z> {
        if self.dimension.contains_point(p) {
            let (x, y) = self.coordinates_to_indices(p);
            Some(self.data[[x, y]])
        } else {
            None
        }
    }
}

impl<F, Z> DensityMap<F, Z>
where
    F: Copy + Num + PartialOrd + FromPrimitive + ToPrimitive,
    Z: Num + std::ops::AddAssign + Copy + Clone + std::ops::Mul<F, Output = Z>,
{
    /// Create an all-zero `w`x`h` array.
    /// `r` defines the spanned region in the euclidean plane.
    pub fn new(r: db::Rect<F>, (w, h): (usize, usize)) -> Self {
        Self {
            dimension: r,
            data: Array2::zeros((w, h)),
        }
    }

    /// Create a DensityMap from existing data.
    /// `r` defines the spanned region in the euclidean plane.
    pub fn from_data(r: db::Rect<F>, data: Array2<Z>) -> Self {
        Self { dimension: r, data }
    }

    /// Set all values to zero.
    pub fn clear(&mut self) {
        self.data.iter_mut().for_each(|x| *x = Z::zero());
    }

    /// Get the location of the lower left corner of bin with index `[x, y]`.
    fn bin_lower_left_corner(&self, (x, y): (usize, usize)) -> db::Point<F> {
        let (bin_width, bin_height) = self.bin_dimension();
        // Lower left corner of the bin.
        let (x, y) = (F::from_usize(x).unwrap(), F::from_usize(y).unwrap());
        self.dimension.lower_left() + db::Point::new(x * bin_width, y * bin_height)
    }

    /// Get the location of the center of bin with index `[x, y]`.
    pub fn bin_center(&self, (x, y): (usize, usize)) -> db::Point<F> {
        let (bin_width, bin_height) = self.bin_dimension();
        let _2 = F::one() + F::one();
        let bin_center = db::Point::new(bin_width / _2, bin_height / _2);
        self.bin_lower_left_corner((x, y)) + bin_center
    }

    /// Get the rectangle shape of the bin at index `(i, j)`.
    pub fn get_bin_shape(&self, (x, y): (usize, usize)) -> db::Rect<F> {
        // Lower left corner of the bin
        let start = self.bin_lower_left_corner((x, y));
        // Upper right corner of the bin.
        let (bin_width, bin_height) = self.bin_dimension();
        let end = start + db::Point::new(bin_width, bin_height);

        db::Rect::new(start, end)
    }

    /// Draw the rectangle `r` to the DensityMap by adding the `value`
    /// to all bins that interact with `r`. If a bin overlaps only partially
    /// with `r` then `a*value` is added to it where `a` is the fraction of the overlap.
    pub fn draw_rect(&mut self, r: &db::Rect<F>, value: Z) {
        // Crop rectangle to dimension of density bins.
        if let Some(r) = r.intersection(&self.dimension) {
            // Find indices of bins that interact with the corners of the rectangle `r`.
            let (xstart, ystart) = self.coordinates_to_indices(r.lower_left);
            let (xend, yend) = self.coordinates_to_indices(r.upper_right);

            let xend = xend.min(self.data.dim().0 - 1);
            let yend = yend.min(self.data.dim().1 - 1);

            let bin_area = self.bin_area();
            // Loop over all bins that interact with the rectangle `r`.
            for x in xstart..xend + 1 {
                for y in ystart..yend + 1 {
                    if x == xstart || x == xend || y == ystart || y == yend {
                        // Corner cases.
                        // Compute increment by the overlap of the bin shape and the rectangle.
                        let bin_shape = self.get_bin_shape((x, y));
                        let overlap_area = bin_shape
                            .intersection(&r)
                            .map(|r| r.width() * r.height())
                            .unwrap_or(F::zero());
                        self.data[[x, y]] += value * overlap_area;
                    } else {
                        self.data[[x, y]] += value * bin_area;
                    }
                }
            }
        }
    }
}

impl<F, Z> DensityMap<F, Z>
where
    F: Copy + Num + PartialOrd + FromPrimitive + ToPrimitive,
    Z: Num + std::ops::AddAssign + Copy + Clone + std::ops::Mul<F, Output = Z> + FromPrimitive,
{
    /// Create a density map with lower resolution.
    ///
    /// Down-sampling is done by creating `n*n` bins. Therefore the `reduction_factor` must
    /// divide the number of bins in both x and y direction.
    pub fn downsample(&self, reduction_factor: usize) -> Self {
        assert!(reduction_factor >= 1);
        let (w, h) = (self.data.shape()[0], self.data.shape()[1]);
        assert_eq!(
            w % reduction_factor,
            0,
            "Dimension must be divisible by the reduction factor."
        );
        assert_eq!(
            h % reduction_factor,
            0,
            "Dimension must be divisible by the reduction factor."
        );

        let w_new = w / reduction_factor;
        let h_new = h / reduction_factor;

        let mut new_data = Array2::zeros((w_new, h_new));

        for x in 0..w {
            let x_new = x / reduction_factor;
            for y in 0..h {
                let y_new = y / reduction_factor;

                new_data[[x_new, y_new]] += self.data[[x, y]];
            }
        }

        // // Normalize
        // let f = Z::from_usize(reduction_factor * reduction_factor).unwrap();
        // new_data.iter_mut()
        //     .for_each(|d| *d = *d / f);

        Self::from_data(self.dimension, new_data)
    }
}

#[test]
fn test_coordinates_to_indices() {
    let c: DensityMap<_, f64> = DensityMap::new(db::Rect::new((0.0, 0.0), (10.0, 20.0)), (10, 20));
    assert_eq!(c.coordinates_to_indices(db::Point::new(0.0, 0.0)), (0, 0));
    assert_eq!(
        c.coordinates_to_indices(db::Point::new(10.0, 20.0)),
        (10, 20)
    );
    assert_eq!(c.coordinates_to_indices(db::Point::new(0.5, 0.5)), (0, 0));
}

#[test]
fn test_bin() {
    let c: DensityMap<_, f64> = DensityMap::new(db::Rect::new((0.0, 0.0), (10.0, 20.0)), (10, 10));
    assert_eq!(c.num_bins(), (10, 10));
    assert_eq!(
        c.get_bin_shape((0, 0)),
        db::Rect::new((0.0, 0.0), (1.0, 2.0))
    );
}

#[test]
fn test_draw_rect() {
    use db::traits::DoubledOrientedArea;

    let mut c = DensityMap::new(db::Rect::new((0.0, 0.0), (10.0, 10.0)), (10, 10));
    let r = db::Rect::new((1.0, 1.0), (2.0, 3.0));
    c.draw_rect(&r, 1.0);
    assert_eq!(c.data[[0, 0]], 0.0);
    assert_eq!(c.data[[1, 1]], 1.0);
    assert_eq!(c.data[[2, 3]], 0.0);

    let sum: f64 = c.data.iter().sum();
    assert_eq!(2.0 * sum, r.area_doubled_oriented());
}

#[test]
fn test_draw_rect_with_partial_bins() {
    use db::traits::DoubledOrientedArea;

    let mut c = DensityMap::new(db::Rect::new((0.0, 0.0), (10.0, 10.0)), (10, 10));
    let r = db::Rect::new((1.5, 1.5), (5.25, 6.0));
    c.draw_rect(&r, 1.0);
    assert_eq!(c.data[[0, 0]], 0.0);
    assert_eq!(c.data[[1, 1]], 0.25);
    assert_eq!(c.data[[1, 2]], 0.5);
    assert_eq!(c.data[[2, 2]], 1.0);

    // Total sum of DensityMap must correspond to the total value of the drawn rectangle.
    let sum: f64 = c.data.iter().sum();
    assert_eq!(2.0 * sum, r.area_doubled_oriented());
}

#[test]
fn test_draw_oversize_rect() {
    let mut c = DensityMap::new(db::Rect::new((0.0, 0.0), (10.0, 10.0)), (10, 10));
    let r = db::Rect::new((-1.0, -1.0), (11.0, 11.0));
    c.draw_rect(&r, 1.0);
    assert_eq!(c.data[[0, 0]], 1.0);
    assert_eq!(c.data[[9, 9]], 1.0);
}

#[test]
fn test_draw_oversize_rect_1x1() {
    let mut c = DensityMap::new(db::Rect::new((0.0, 0.0), (10.0, 10.0)), (1, 1));
    let r = db::Rect::new((-1.0, -1.0), (11.0, 11.0));
    c.draw_rect(&r, 1.0);
    assert_eq!(c.density_at((0.0, 0.0).into()), 1.0);
    assert_eq!(c.data[[0, 0]], 100.0);
}