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

//! Simple buffer tree generator.

use itertools::Itertools;
use libreda_pnr::db;
use libreda_pnr::db::{LayoutEdit, NetlistEdit, NetlistEditUtil, NetlistUtil, TerminalId};
use libreda_pnr::rebuffer::buffer_insertion::SimpleBufferInsertion;
use num_traits::{NumCast, PrimInt};

/// Simple buffer tree generator.
/// This generator creates buffer trees consisting of a single type of inverter cell.
///
/// The algorithm works as follows:
/// * Create clusters of the sink terminals based on their location. Each cluster must have maximum size `max_fanout`. Cluster sizes should be balanced.
/// * Insert a buffer to drive each cluster. Place the buffer at the 'center of mass' of the terminals, or such that the quadratic distance to the sinks is minimized.
/// * Repeat the same with the newly inserted buffers until the fan out limit is satisfied.
/// * Make sure that the signal is not being inverted by adding another buffer when necessary.
///
/// # Caveats
/// This algorithm assumes that the sinks are small standard-cells. This allows to easily find
/// the rough location of the actual pin shape by assuming it is equal to the position of the cell.
/// For large macro blocks this will not yield good results since the pins are substantially far away
/// from the center of the cell.
#[derive(Debug, Clone)]
pub struct SimpleBufferInsertionEngine {
    /// The name of the inverting buffer cell to be used in the buffer tree.
    pub inverting_buffer_cell: String,
    /// Maximal number of cells that should be driven by a buffer cell.
    /// Must be larger than `1`.
    pub max_fanout: u32,
}

impl SimpleBufferInsertionEngine {
    /// All terminals must be in the same parent cell.
    ///
    /// * `need_inversion`: Force the insertion of an inverter even if no buffer needs to be inserted.
    ///
    /// # Panics
    /// Panics when
    /// * The buffer cell cannot be found.
    /// * The source terminal is not connected to a net.
    /// * The sinks are not all connected to the same net as the source.
    fn insert_buffers_recursive<LN: NetlistEdit + LayoutEdit>(
        &self,
        chip: &mut LN,
        signal_source: TerminalId<LN>,
        signal_sinks: &Vec<TerminalId<LN>>,
        need_inversion: bool,
    ) -> Result<(Vec<LN::CellInstId>, Vec<LN::NetId>), ()>
    where
        LN::Coord: PrimInt,
    {
        if signal_sinks.len() <= self.max_fanout as usize {
            if need_inversion {
                // Insert a buffer anyway to make sure the signal has the right polarity.
            } else {
                log::debug!("Low fan-out. No need for buffer insertion.");
                return Ok((vec![], vec![]));
            }
        }

        log::debug!("Create buffer tree with {} sinks.", signal_sinks.len());

        log::debug!("Find buffer cell: {}", &self.inverting_buffer_cell);
        // Find buffer cell to be used.
        let buffer_name = &self.inverting_buffer_cell;
        log::debug!("Use buffer cell '{}'.", &buffer_name);
        let buffer_cell = chip
            .cell_by_name(buffer_name)
            .expect("Buffer cell not found.");

        let parent_cell = match &signal_source {
            TerminalId::PinId(p) => chip.parent_cell_of_pin(p),
            TerminalId::PinInstId(p) => chip.parent_cell(&chip.parent_of_pin_instance(p)),
        };

        // Sanity check: The source and all sinks must be connected to the same net.
        let source_net = chip.net_of_terminal(&signal_source);
        if source_net.is_none() {
            log::error!("Source terminal is not connected to any net.");
            // TODO: Pass a `Result::Err` instead o panicking?
            panic!("Source terminal is not connected to any net.");
        }

        {
            let is_connected_to_correct_nets = signal_sinks
                .iter()
                .all(|sink| chip.net_of_terminal(sink) == source_net);
            if !is_connected_to_correct_nets {
                log::error!("Sinks are not all connected to the same net as the source.");
                panic!("Sinks are not all connected to the same net as the source.");
            }
        }

        let source_net = source_net.unwrap();
        if chip.is_constant_net(&source_net) {
            log::warn!("Tie-cells should be used on a constant net, not buffers.")
        }

        // Detect the input and output of the buffer cell.
        let inputs: Vec<_> = chip
            .each_pin(&buffer_cell)
            .filter(|pin| chip.pin_direction(&pin).is_input())
            .collect();
        assert_eq!(inputs.len(), 1, "Buffer must have exactly one input pin.");

        let outputs: Vec<_> = chip
            .each_pin(&buffer_cell)
            .filter(|pin| chip.pin_direction(&pin).is_output())
            .collect();

        assert_eq!(outputs.len(), 1, "Buffer must have exactly one output pin.");

        let buffer_input = inputs[0].clone();
        let buffer_output = outputs[0].clone();

        // Find locations of terminals.
        let terminal_location = |t: &TerminalId<LN>| {
            match t {
                TerminalId::PinId(_t) => {
                    // TODO: Find the location of the pin shape.
                    db::Point::zero()
                }
                TerminalId::PinInstId(t) => {
                    let inst = chip.parent_of_pin_instance(t);
                    let tf = chip.get_transform(&inst);
                    // Assume that the physical pin is close to the origin of the cell.
                    // This is good enough only for small standard-cells, not for big macro blocks.
                    let location = tf.transform_point(db::Point::zero());
                    location
                }
            }
        };

        // Build clusters from the sinks.
        let clusters = {
            // Find sink locations.
            let sink_locations: Vec<_> = signal_sinks.iter().map(terminal_location).collect();

            // Find clusters.
            let (cluster_ids, num_clusters) = find_clusters(&sink_locations, self.max_fanout);

            // Create a nested list of terminals for each cluster.
            let mut clusters = vec![vec![]; num_clusters as usize];
            for (i, cluster_id) in cluster_ids.iter().enumerate() {
                clusters[*cluster_id as usize].push((signal_sinks[i].clone(), sink_locations[i]));
            }
            clusters
        };

        log::debug!("Number of clusters {:?}", clusters.len());

        // Create a buffer for each cluster.
        let mut buffer_inputs = vec![]; // Input pins to the buffer instances.
        let mut name_counter = 0; // Counter for unique names of buffer instances.
        let mut new_buffer_instances = vec![]; // Remember generated buffer instances.
        let mut new_nets = vec![]; // Remember generated nets.
        for cluster in clusters {
            let (sinks, positions): (Vec<_>, Vec<_>) = cluster.into_iter().unzip();
            let center = center_of_mass(&positions);

            // Create name for the buffer instance.
            let buffer_name = loop {
                let name = format!("_buffer_tree_{}", name_counter);
                name_counter += 1;
                if chip.cell_instance_by_name(&parent_cell, &name).is_none() {
                    break name;
                }
            };

            // Create name for the output net of the buffer.
            let buffer_net_name = {
                let source_net_name: String = chip
                    .net_name(&source_net)
                    .map(|n| n.to_string())
                    .unwrap_or_else(|| "".into());
                let mut ctr = 0;
                loop {
                    let name = if ctr == 0 {
                        format!("_{}_buf", source_net_name)
                    } else {
                        format!("_{}_buf.{}", source_net_name, ctr)
                    };

                    if chip.net_by_name(&parent_cell, &name).is_none() {
                        break name;
                    }
                    ctr += 1;
                }
            };

            let buffer_inst =
                chip.create_cell_instance(&parent_cell, &buffer_cell, Some(buffer_name.into()));
            // Create buffer output net.
            let buffer_net = chip.create_net(&parent_cell, Some(buffer_net_name.into()));

            // Connect output of the buffer to the buffered net.
            // And connect the input of the buffer to the source net.
            let buffer_output_pin = chip.pin_instance(&buffer_inst, &buffer_output);
            chip.connect_pin_instance(&buffer_output_pin, Some(buffer_net.clone()));

            // Set the location of the buffer.
            chip.set_transform(&buffer_inst, db::SimpleTransform::translate(center));

            // Connect the sinks to the buffer output.
            for sink in &sinks {
                let old_net = chip.connect_terminal(sink, Some(buffer_net.clone()));
                debug_assert_eq!(
                    old_net.as_ref(),
                    Some(&source_net),
                    "Sink is expected to be connected to the source net."
                );
            }

            // Connect the buffer input.
            let buffer_input_pin = chip.pin_instance(&buffer_inst, &buffer_input);
            chip.connect_pin_instance(&buffer_input_pin, Some(source_net.clone()));

            // Store the input pin to the new buffer.
            buffer_inputs.push(TerminalId::PinInstId(buffer_input_pin));

            // Remember created nets and instances.
            new_buffer_instances.push(buffer_inst);
            new_nets.push(buffer_net);
        }

        {
            // Recursively call this buffer insertion to create the next layer of buffers.
            // Make sure the signal is not inverted in the end.
            let (buffers, nets) = self.insert_buffers_recursive(
                chip,
                signal_source,
                &buffer_inputs,
                !need_inversion,
            )?;

            // Append created buffers and nets.
            new_buffer_instances.extend(buffers);
            new_nets.extend(nets);
        }
        Ok((new_buffer_instances, new_nets))
    }
}

impl<LN: NetlistEdit + LayoutEdit> SimpleBufferInsertion<LN> for SimpleBufferInsertionEngine
where
    LN::Coord: PrimInt,
{
    type Error = ();

    /// All terminals must be in the same parent cell.
    fn insert_buffers(
        &self,
        chip: &mut LN,
        signal_source: TerminalId<LN>,
        signal_sinks: &Vec<TerminalId<LN>>,
    ) -> Result<(Vec<LN::CellInstId>, Vec<LN::NetId>), Self::Error> {
        log::debug!("Create buffer tree with {} sinks.", signal_sinks.len());

        log::debug!("Find buffer cell: {}", &self.inverting_buffer_cell);

        // let parent_cell = chip.parent_cell(&chip.parent_of_pin_instance(&signal_source));

        // Find buffer cell to be used.
        let buffer_cell = chip
            .cell_by_name(&self.inverting_buffer_cell)
            .expect("Buffer cell not found.");

        // Detect the input and output of the buffer cell.
        let inputs: Vec<_> = chip
            .each_pin(&buffer_cell)
            .filter(|pin| chip.pin_direction(&pin).is_input())
            .collect();
        assert_eq!(inputs.len(), 1, "Buffer must have exactly one input pin.");
        let outputs: Vec<_> = chip
            .each_pin(&buffer_cell)
            .filter(|pin| chip.pin_direction(&pin).is_output())
            .collect();
        assert_eq!(outputs.len(), 1, "Buffer must have exactly one output pin.");

        // let buffer_input = inputs[0].clone();
        // let buffer_output = outputs[0].clone();

        // Sanity check: The source and all sinks must be connected to the same net.
        let source_net = chip.net_of_terminal(&signal_source);
        if source_net.is_none() {
            log::error!("Source terminal is not connected to any net.");
            // TODO: Pass a `Result::Err` instead o panicking?
            panic!("Source terminal is not connected to any net.");
        }
        {
            let correct_nets = signal_sinks
                .iter()
                .all(|sink| chip.net_of_terminal(sink) == source_net);
            if !correct_nets {
                log::error!("Sinks are not all connected to the same net as the source.");
                panic!("Sinks are not all connected to the same net as the source.");
            }
        }

        self.insert_buffers_recursive(
            chip,
            signal_source,
            signal_sinks,
            false, // No inversion needed. Signal has correct polarity.
        )
    }
}

/// Compute the center of mass of a list of points.
fn center_of_mass<C: PrimInt>(points: &Vec<db::Point<C>>) -> db::Point<C> {
    // Sum up all points and divide by number of points.
    points.iter().fold(db::Point::zero(), |a, b| a + b) / NumCast::from(points.len()).unwrap()
}

/// Group points together to clusters and return a vector with the cluster ID for each point and the number of clusters.
fn find_clusters<C: PrimInt>(points: &Vec<db::Point<C>>, max_cluster_size: u32) -> (Vec<u32>, u32) {
    // For each point store the original index and the cluster ID.
    // The index used to map back to the original points after sorting.
    let mut points_with_cluster_id: Vec<(db::Point<_>, usize, u32)> = points
        .iter()
        .enumerate()
        .map(|(idx, p)| (*p, idx, 0))
        .collect();
    // Sort by locations.
    points_with_cluster_id.sort_by_key(|(p, _, _)| (p.x, p.y));

    let mut cluster_id = 0u32;
    // While there is a point that does not belong to a cluster
    // take the first one (sorted by x and y) and use it as a start
    // of the cluster.

    let mut current_cluster_points = vec![];

    while let Some((first_point_idx, _)) = points_with_cluster_id
        .iter()
        .find_position(|(_, _, cluster)| *cluster == 0)
    {
        cluster_id += 1;

        points_with_cluster_id[first_point_idx].2 = cluster_id;
        let target_cluster_size = max_cluster_size;
        let start = points_with_cluster_id[first_point_idx].0;
        current_cluster_points.clear();
        current_cluster_points.push(start);
        // Find closest neighbour.
        for _ in 1..target_cluster_size {
            // Compute center of mass of the cluster.
            let center = center_of_mass(&current_cluster_points);

            let closest = points_with_cluster_id
                .iter()
                .enumerate()
                .filter(|(_, (_, _, cluster))| *cluster == 0) // Take only unassigned sinks.
                .min_by_key(|(_, (pos, _, _))| (center - *pos).norm1())
                .map(|(idx, _)| idx);
            if let Some(closest) = closest {
                points_with_cluster_id[closest].2 = cluster_id;
                current_cluster_points.push(points_with_cluster_id[closest].0);
            } else {
                // No unassigned node found.
                break;
            }
        }
    }

    let num_clusters = cluster_id;

    let mut result = vec![0u32; points_with_cluster_id.len()];
    for (_p, index, cluster_id) in points_with_cluster_id {
        result[index] = cluster_id - 1; // `0` meant 'no cluster'
    }

    (result, num_clusters)
}

#[test]
fn test_buffer_insertion() {
    use db::{Chip, Direction};
    use db::{HierarchyEdit, NetlistBase};

    let mut chip = Chip::new();
    let top = chip.create_cell("TOP".into());
    let sub = chip.create_cell("SUB".into());
    let inv = chip.create_cell("INV".into());

    let sub_data_in = chip.create_pin(&sub, "data_in".into(), Direction::Input);

    let _inv_in = chip.create_pin(&inv, "IN".into(), Direction::Input);
    let inv_out = chip.create_pin(&inv, "OUT".into(), Direction::Output);

    // Create a net where the driver and all the sinks will be attached to.
    let signal = chip.create_net(&top, Some("signal".into()));

    // Create a cell that drives the net 'signal'.
    let signal_source_instance =
        chip.create_cell_instance(&top, &inv, Some("signal_source".into()));
    chip.connect_pin_instance(
        &chip.pin_instance(&signal_source_instance, &inv_out),
        Some(signal.clone()),
    );

    let pseudo_rand = |i: i32| -> i32 {
        let i = i as i64;
        let a = i as i64 ^ 0x12345678;
        let b = a ^ (i * 1001 ^ 0xfedcba9i64);
        let c = b + a ^ (i * 661097 ^ 0xdeadbeefi64);
        (c % 1000) as i32
    };

    // Create sinks and attach them to the signal net.
    let num_sinks = 257;
    for i in 0..num_sinks {
        let inst = chip.create_cell_instance(&top, &sub, Some(format!("sink_{}", i).into()));
        chip.connect_pin_instance(
            &chip.pin_instance(&inst, &sub_data_in),
            Some(signal.clone()),
        );
        // Set instance location.
        let x = pseudo_rand(i);
        let y = pseudo_rand(i + num_sinks);

        chip.set_transform(&inst, db::SimpleTransform::translate((x, y)));
    }

    let buffering_engine = SimpleBufferInsertionEngine {
        inverting_buffer_cell: "INV".to_string(),
        max_fanout: 4,
    };

    let (_buffers, _nets) = buffering_engine
        .add_buffer_tree_on_net(&mut chip, &signal)
        .unwrap();
}