use alloc::collections::BTreeMap;

use crate::Felt;

#[cfg(test)]
mod tests;

// RANGE CHECKER
// ================================================================================================

/// Range checker for the VM.
///
/// This component is responsible for building an execution trace for all 16-bit range checks
/// performed by the VM. Thus, the [RangeChecker] doesn't actually check if a given value fits
/// into 16-bits, but rather keeps track of all 16-bit range checks performed by the VM.
///
/// ## Execution trace
/// The execution trace generated by the range checker consists of 2 columns. Conceptually, the
/// table starts with value 0 and ends with value 65535.
///
/// The layout is illustrated below.
///
///    m     v
/// ├─────┴─────┤
///
/// In the above, the meaning of the columns is as follows:
/// - Column `v` contains the value being range-checked where `v` must be a 16-bit value. The values
///   must be in increasing order and the jump allowed between two values should be a power of 3
///   less than or equal to 3^7, and duplicates are allowed.
/// - Column `m` specifies the lookup multiplicity, which is how many lookups are to be included for
///   a given value.
///
/// Thus, for example, if a value was range-checked just once, we'll need to add a single row to
/// the table with (m, v) set to (1, v), where v is the value. If the value was range-checked 5
/// times, we'll need to specify the row (5, v).
pub struct RangeChecker {
    /// Tracks lookup count for each checked value.
    lookups: BTreeMap<u16, usize>,
}

impl RangeChecker {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------
    /// Returns a new [RangeChecker] instantiated with an empty lookup table.
    pub fn new() -> Self {
        let mut lookups = BTreeMap::new();
        // we need to make sure that the first row after the padded section and the last row of the
        // range checker table are initialized. this simplifies trace table building later on.
        lookups.insert(0, 0);
        lookups.insert(u16::MAX, 0);
        Self { lookups }
    }

    // TRACE MUTATORS
    // --------------------------------------------------------------------------------------------

    /// Adds the specified value to the trace of this range checker's lookups.
    pub fn add_value(&mut self, value: u16) {
        self.lookups.entry(value).and_modify(|v| *v += 1).or_insert(1);
    }

    /// Adds a batch of range-check lookups to this [RangeChecker] instance.
    pub fn add_range_checks(&mut self, values: &[u16]) {
        // Callers provide two Merkle-depth, U32DIV, or memory-delta values; four u32 helper
        // values; or five Merkle-index witness values.
        debug_assert!(matches!(values.len(), 2 | 4 | 5));

        for value in values.iter() {
            // add the specified value to the trace of this range checker's lookups.
            self.add_value(*value);
        }
    }

    // EXECUTION TRACE GENERATION (INTERNAL)
    // --------------------------------------------------------------------------------------------

    /// Writes the range-checker `(m, v)` columns in-place into the row-major core buffer
    /// `core_data` (physical row width `stride`), at column offsets `m_off` and `v_off`.
    ///
    /// The two range slots in `core_data` are assumed to be zero on entry, so the front
    /// padding rows (`[0, num_padding_rows)`, all zero) are left untouched.
    pub fn write_range_into_core(
        self,
        core_data: &mut [Felt],
        stride: usize,
        m_off: usize,
        v_off: usize,
        range_table_len: usize,
        core_height: usize,
    ) {
        let num_padding_rows = core_height - range_table_len;

        let end = {
            let mut sink = |step: usize, m: u64, v: u64| {
                core_data[step * stride + m_off] = Felt::new_unchecked(m);
                core_data[step * stride + v_off] = Felt::new_unchecked(v);
            };
            self.emit_table_rows(&mut sink, num_padding_rows)
        };
        assert_eq!(end, core_height, "range checker table length mismatch vs core height");
    }

    /// Emits the range-checker table rows (bridge rows + value rows + the trailing extra
    /// `u16::MAX` row) through `sink`, starting at row `start_step`. `sink(step, m, v)` is
    /// invoked once per row with monotonically increasing `step`.
    fn emit_table_rows<F: FnMut(usize, u64, u64)>(&self, sink: &mut F, start_step: usize) -> usize {
        let mut step = start_step;
        let mut prev_value = 0u16;
        for (&value, &num_lookups) in self.lookups.iter() {
            write_rows(sink, &mut step, num_lookups, value, prev_value);
            prev_value = value;
        }

        // Preserve the trailing zero-multiplicity u16::MAX row in the range-table layout. It
        // contributes zero to the RangeCheck sum, and its zero delta is permitted by the range AIR.
        write_trace_row(sink, &mut step, 0, (u16::MAX).into());
        step
    }

    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns the number of rows needed to support all 16-bit lookups requested by the VM.
    pub fn get_number_range_checker_rows(&self) -> usize {
        // Account for the trailing zero-multiplicity u16::MAX row emitted by `emit_table_rows`.
        let mut num_rows = 1;

        let mut prev_value = 0u16;
        for value in self.lookups.keys() {
            // add one row for each value in the range checker table
            num_rows += 1;
            // determine the delta between this and the previous value. we need to know this delta
            // to determine if we need to insert any "bridge" rows to the  table, this is needed
            // since the gap between two values in the range checker can only be a power of 3 less
            // than or equal to 3^7.
            let delta = value - prev_value;
            num_rows += get_num_bridge_rows(delta);
            prev_value = *value;
        }
        num_rows
    }

    // TEST HELPERS
    // --------------------------------------------------------------------------------------------

    /// Returns length of execution trace required to describe all 16-bit range checks performed
    /// by the VM.
    #[cfg(test)]
    pub(super) fn trace_len(&self) -> usize {
        self.get_number_range_checker_rows()
    }
}

impl Default for RangeChecker {
    fn default() -> Self {
        Self::new()
    }
}

// HELPER FUNCTIONS
// ================================================================================================

/// Calculates the number of bridge rows that are need to be added to the trace between two values
/// to be range checked.
pub fn get_num_bridge_rows(delta: u16) -> usize {
    let mut gap = delta;
    let mut bridge_rows = 0_usize;
    let mut stride = 3_u16.pow(7);
    while gap != stride {
        if gap > stride {
            bridge_rows += 1;
            gap -= stride;
        } else {
            stride /= 3;
        }
    }
    bridge_rows
}

/// Adds a row for the values to be range checked. In case the difference between the current and
/// next value is not a power of 3, this function will add additional bridge rows to the trace.
fn write_rows<F: FnMut(usize, u64, u64)>(
    sink: &mut F,
    step: &mut usize,
    num_lookups: usize,
    value: u16,
    prev_value: u16,
) {
    let mut gap = value - prev_value;
    let mut prev_val = prev_value;
    let mut stride = 3_u16.pow(7);
    while gap != stride {
        if gap > stride {
            gap -= stride;
            prev_val += stride;
            write_trace_row(sink, step, 0, prev_val as u64);
        } else {
            stride /= 3;
        }
    }
    write_trace_row(sink, step, num_lookups, value as u64);
}

/// Populates a single row at the specified step via `sink`.
fn write_trace_row<F: FnMut(usize, u64, u64)>(
    sink: &mut F,
    step: &mut usize,
    num_lookups: usize,
    value: u64,
) {
    sink(*step, num_lookups as u64, value);
    *step += 1;
}