LCRConnector/kernel/lcrconn/src/resolver/bfs.rs

use std::cmp::Ordering;
use std::collections::BinaryHeap;
use std::iter::FusedIterator;

use super::Resolver;
use crate::common::{Circuit, CircuitValueTrait, DeviceKind, JointKind, LcrConnError};
use crate::dataset::{Dataset, DatasetCollection, DatasetItem};
use crate::query::{Request, Response};

// ============================================================================
// Lazy iterator structs for circuit generation
// ============================================================================

// YYC MARK:
// Some circuit are equivalent in topology.
// If we deduplicate these equaivalent circuit in building result,
// there are too complex works.
// So we should deduplicated these equivalent circuit at the beginning,
// i.e. when generating them.
// So following iterator structs are taking this job.

/// Iterator over all possible one-device circuits without repeating equivalent topology.
pub struct OneDeviceCircuitIter<'a> {
    items: &'a [DatasetItem],
    pos: usize,
}

impl<'a> OneDeviceCircuitIter<'a> {
    pub fn new(items: &'a [DatasetItem]) -> Self {
        Self { items, pos: 0 }
    }
}

impl Iterator for OneDeviceCircuitIter<'_> {
    type Item = Circuit;

    fn next(&mut self) -> Option<Self::Item> {
        if self.pos < self.items.len() {
            // Every single device is unique so we directly output them.
            // This feature is insured by dataset itself.
            let circuit = Circuit::from_one_device(self.items[self.pos].value);
            self.pos += 1;
            Some(circuit)
        } else {
            None
        }
    }
}

impl FusedIterator for OneDeviceCircuitIter<'_> {}

/// Iterator over all possible two-device circuits without repeating equivalent topology.
pub struct TwoDeviceCircuitIter<'a> {
    items: &'a [DatasetItem],
    i: usize,
    j: usize,
    joint_idx: usize,
}

impl<'a> TwoDeviceCircuitIter<'a> {
    pub fn new(items: &'a [DatasetItem]) -> Self {
        Self {
            items,
            i: 0,
            j: 0,
            joint_idx: 0,
        }
    }
}

impl Iterator for TwoDeviceCircuitIter<'_> {
    type Item = Circuit;

    fn next(&mut self) -> Option<Self::Item> {
        let n = self.items.len();
        if n == 0 {
            return None;
        }

        loop {
            if self.joint_idx < JointKind::ALL.len() {
                let jk = JointKind::ALL[self.joint_idx];
                self.joint_idx += 1;
                // The two devices in this circuit is always swapable,
                // so we iterate them without repeating.
                return Some(Circuit::from_two_devices(
                    self.items[self.i].value,
                    self.items[self.j].value,
                    jk,
                ));
            }

            // Advance to next combination
            self.joint_idx = 0;
            self.j += 1;
            if self.j >= n {
                self.i += 1;
                self.j = self.i;
                if self.i >= n {
                    return None;
                }
            }
        }
    }
}

impl FusedIterator for TwoDeviceCircuitIter<'_> {}

/// Iterator over three-device circuits where both joints share the same type.
///
/// In this case, all 3 devices are swapable and are iterated without repeating.
pub struct ThreeDeviceSameJointIter<'a> {
    items: &'a [DatasetItem],
    i: usize,
    j: usize,
    k: usize,
    joint_idx: usize,
}

impl<'a> ThreeDeviceSameJointIter<'a> {
    pub fn new(items: &'a [DatasetItem]) -> Self {
        Self {
            items,
            i: 0,
            j: 0,
            k: 0,
            joint_idx: 0,
        }
    }
}

impl Iterator for ThreeDeviceSameJointIter<'_> {
    type Item = Circuit;

    fn next(&mut self) -> Option<Self::Item> {
        let n = self.items.len();
        if n == 0 {
            return None;
        }

        loop {
            if self.joint_idx < JointKind::ALL.len() {
                let jk = JointKind::ALL[self.joint_idx];
                self.joint_idx += 1;
                return Some(Circuit::from_three_devices(
                    self.items[self.i].value,
                    self.items[self.j].value,
                    jk,
                    self.items[self.k].value,
                    jk,
                ));
            }

            self.joint_idx = 0;
            self.k += 1;
            if self.k >= n {
                self.j += 1;
                self.k = self.j;
                if self.j >= n {
                    self.i += 1;
                    self.j = self.i;
                    self.k = self.i;
                    if self.i >= n {
                        return None;
                    }
                }
            }
        }
    }
}

impl FusedIterator for ThreeDeviceSameJointIter<'_> {}

/// Iterator over three-device circuits where the two joint types differ.
///
/// In this case, the first 2 devices are swapable and are iterated without repeating,
/// while the third device iterates over all values independently.
pub struct ThreeDeviceDiffJointIter<'a> {
    items: &'a [DatasetItem],
    i: usize,
    j: usize,
    k: usize,
    joint_idx: usize,
}

impl<'a> ThreeDeviceDiffJointIter<'a> {
    pub fn new(items: &'a [DatasetItem]) -> Self {
        Self {
            items,
            i: 0,
            j: 0,
            k: 0,
            joint_idx: 0,
        }
    }
}

impl Iterator for ThreeDeviceDiffJointIter<'_> {
    type Item = Circuit;

    fn next(&mut self) -> Option<Self::Item> {
        let n = self.items.len();
        if n == 0 {
            return None;
        }

        loop {
            if self.joint_idx < JointKind::ALL.len() {
                let j = JointKind::ALL[self.joint_idx];
                self.joint_idx += 1;
                return Some(Circuit::from_three_devices(
                    self.items[self.i].value,
                    self.items[self.j].value,
                    j,
                    self.items[self.k].value,
                    j.flip(),
                ));
            }

            self.joint_idx = 0;
            self.k += 1;
            if self.k >= n {
                self.j += 1;
                self.k = 0;
                if self.j >= n {
                    self.i += 1;
                    self.j = self.i;
                    self.k = 0;
                    if self.i >= n {
                        return None;
                    }
                }
            }
        }
    }
}

impl FusedIterator for ThreeDeviceDiffJointIter<'_> {}

/// Type alias for the chained three-device circuit iterator.
pub type ThreeDeviceCircuitIter<'a> = std::iter::Chain<
    ThreeDeviceSameJointIter<'a>,
    ThreeDeviceDiffJointIter<'a>,
>;

// ============================================================================
// BfsItem
// ============================================================================

/// The entry used in BFS iteration storing circuit and value.
pub struct BfsItem {
    /// The circuit represented by this item.
    circuit: Circuit,
    /// The computed value of the circuit.
    value: f64,
    /// The unsigned difference between the target value and the value of this circuit.
    unsigned_difference: f64,
}

impl BfsItem {
    /// Create a new BFS item by computing values eagerly.
    ///
    /// # Errors
    ///
    /// See [`CircuitValueTrait::value`].
    pub fn new(circuit: Circuit, cv_trait: &CircuitValueTrait) -> Result<Self, LcrConnError> {
        let value = cv_trait.value(&circuit)?;
        let unsigned_difference = cv_trait.unsigned_difference(&circuit, Some(value))?;
        Ok(Self {
            circuit,
            value,
            unsigned_difference,
        })
    }

    /// The circuit represented by this item.
    pub fn circuit(&self) -> &Circuit {
        &self.circuit
    }

    /// The computed value of the circuit.
    pub fn value(&self) -> f64 {
        self.value
    }

    /// The unsigned difference between the target value and the value of this circuit.
    pub fn unsigned_difference(&self) -> f64 {
        self.unsigned_difference
    }

    /// Consume this item and return the inner circuit.
    pub fn into_circuit(self) -> Circuit {
        self.circuit
    }
}

// ============================================================================
// ResultBucket
// ============================================================================

/// An item stored in a [`ResultBucket`].
struct ResultBucketItem {
    /// The score associated with this item.
    score: f64,
    /// The underlying BfsItem.
    item: BfsItem,
    /// Monotonic counter used as a tiebreaker when scores are equal,
    /// ensuring that BinaryHeap never compares BfsItem directly.
    seq: usize,
}

impl ResultBucketItem {
    fn new(score: f64, item: BfsItem, seq: usize) -> Self {
        Self { score, item, seq }
    }
}

impl PartialEq for ResultBucketItem {
    fn eq(&self, other: &Self) -> bool {
        self.score == other.score && self.seq == other.seq
    }
}

impl Eq for ResultBucketItem {}

impl PartialOrd for ResultBucketItem {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for ResultBucketItem {
    fn cmp(&self, other: &Self) -> Ordering {
        // BinaryHeap is a max-heap: the greatest element is at the top.
        // We want the entry with the largest score at the top.
        match self.score.partial_cmp(&other.score) {
            Some(Ordering::Equal) | None => self.seq.cmp(&other.seq),
            Some(ord) => ord,
        }
    }
}

/// A bounded bucket that keeps up to N entries with the smallest scores.
///
/// When the bucket is full, inserting a new item only succeeds if its score
/// is less than the current maximum; the maximum is then evicted.
pub struct ResultBucket {
    /// Maximum number of items the bucket can hold.
    n: usize,
    /// Max-heap of [`ResultBucketItem`].
    /// The entry with the largest score sits at index 0.
    heap: BinaryHeap<ResultBucketItem>,
    /// Monotonic counter fed to each [`ResultBucketItem`] as a tiebreaker,
    /// preventing BinaryHeap from comparing BfsItem on score collisions.
    counter: usize,
}

impl ResultBucket {
    /// Create a new bucket that holds at most `n` items.
    pub fn new(n: usize) -> Self {
        Self {
            n,
            heap: BinaryHeap::new(),
            counter: 0,
        }
    }

    /// The number of items currently in the bucket.
    pub fn len(&self) -> usize {
        self.heap.len()
    }

    /// Whether the bucket is empty.
    pub fn is_empty(&self) -> bool {
        self.heap.is_empty()
    }

    /// Insert a [`BfsItem`] with the given score.
    ///
    /// If the bucket is not yet full the item is always inserted.
    /// Otherwise the item is only inserted when `score` is smaller
    /// than the largest score currently in the bucket; the entry
    /// with the largest score is then evicted.
    ///
    /// # Returns
    ///
    /// `true` if the item was inserted, `false` otherwise.
    pub fn insert(&mut self, item: BfsItem, score: f64) -> bool {
        let entry = ResultBucketItem::new(score, item, self.counter);
        if self.heap.len() < self.n {
            self.heap.push(entry);
            self.counter += 1;
            true
        } else if score >= self.heap.peek().unwrap().score {
            false
        } else {
            *self.heap.peek_mut().unwrap() = entry;
            self.counter += 1;
            true
        }
    }

    /// Consume the bucket and return all stored items.
    pub fn into_items(self) -> Vec<BfsItem> {
        self.heap.into_iter().map(|entry| entry.item).collect()
    }
}

// ============================================================================
// BfsResolver
// ============================================================================

/// A resolver that uses brute-force search to find the best matching circuits.
pub struct BfsResolver {
    /// The datasets for all device kinds.
    datasets: DatasetCollection,
}

impl BfsResolver {
    /// Create a new BFS resolver with the given datasets.
    pub fn new(datasets: DatasetCollection) -> Self {
        Self { datasets }
    }

    /// Iterate all possible circuits with one device without repeating equivalent topology.
    pub fn iter_one_device_circuit(dataset: &Dataset) -> OneDeviceCircuitIter<'_> {
        OneDeviceCircuitIter::new(dataset.items())
    }

    /// Iterate all possible circuits with two devices without repeating equivalent topology.
    pub fn iter_two_devices_circuit(dataset: &Dataset) -> TwoDeviceCircuitIter<'_> {
        TwoDeviceCircuitIter::new(dataset.items())
    }

    /// Iterate all possible circuits with three devices without repeating equivalent topology.
    pub fn iter_three_devices_circuit(dataset: &Dataset) -> ThreeDeviceCircuitIter<'_> {
        ThreeDeviceSameJointIter::new(dataset.items())
            .chain(ThreeDeviceDiffJointIter::new(dataset.items()))
    }

    fn pick_dataset(&self, device_kind: DeviceKind) -> &Dataset {
        match device_kind {
            DeviceKind::Resistor => self.datasets.resistor(),
            DeviceKind::Capacitor => self.datasets.capacitor(),
            DeviceKind::Inductor => self.datasets.inductor(),
        }
    }
}

impl Resolver for BfsResolver {
    fn resolve(&self, request: &Request) -> Result<Response, LcrConnError> {
        // Pick dataset from collection
        let dataset = self.pick_dataset(request.device_kind);

        // Iterate circuit item one by one
        let mut bucket = ResultBucket::new(request.count_limit);
        let cv_trait = CircuitValueTrait::new(request.device_kind, request.target_value);

        let circuits = Self::iter_one_device_circuit(dataset)
            .chain(Self::iter_two_devices_circuit(dataset))
            .chain(Self::iter_three_devices_circuit(dataset));

        for circuit in circuits {
            let item = BfsItem::new(circuit, &cv_trait)?;
            // If circuit absolute difference is out of tolerance, skip it directly.
            if item.unsigned_difference() > request.tolerance {
                continue;
            }
            // Put it into bucket
            bucket.insert(item, item.unsigned_difference());
        }

        // Return result
        let circuits: Vec<Circuit> = bucket
            .into_items()
            .into_iter()
            .map(BfsItem::into_circuit)
            .collect();
        Response::new(request, circuits)
    }
}