use pyo3::prelude::*;
use std::collections::HashMap;

type Transaction = (Vec<u32>, Vec<f32>);
type ProjectedDB = Vec<(usize, usize, f32)>;

fn hupm_mine_recursive(
    transactions: &[Transaction],
    pdb: &ProjectedDB,
    min_utility: f32,
    max_len: Option<usize>,
    current_pattern: &mut Vec<u32>,
    results: &mut Vec<(f32, Vec<u32>)>,
) {
    if let Some(ml) = max_len {
        if current_pattern.len() >= ml {
            return;
        }
    }

    let mut item_twu: HashMap<u32, f32> = HashMap::new();
    let mut item_exact_utility: HashMap<u32, f32> = HashMap::new();

    for &(tx_idx, start_pos, prefix_util) in pdb {
        let (items, utils) = &transactions[tx_idx];

        let mut remaining_util = 0.0;
        for j in start_pos..items.len() {
            remaining_util += utils[j];
        }

        let local_tu = prefix_util + remaining_util;

        for j in start_pos..items.len() {
            let item = items[j];
            let item_u = utils[j];

            *item_twu.entry(item).or_insert(0.0) += local_tu;
            *item_exact_utility.entry(item).or_insert(0.0) += prefix_util + item_u;
        }
    }

    let mut promising_items: Vec<u32> = item_twu
        .into_iter()
        .filter(|&(_, twu)| twu >= min_utility)
        .map(|(item, _)| item)
        .collect();

    promising_items.sort_unstable();

    for item in promising_items {
        current_pattern.push(item);

        let exact_u = *item_exact_utility.get(&item).unwrap_or(&0.0);
        if exact_u >= min_utility {
            results.push((exact_u, current_pattern.clone()));
        }

        let mut new_pdb: ProjectedDB = Vec::with_capacity(pdb.len());
        for &(tx_idx, start_pos, prefix_util) in pdb.iter() {
            let (items, utils) = &transactions[tx_idx];

            if let Some(offset) = items[start_pos..].iter().position(|&x| x == item) {
                let actual_pos = start_pos + offset;
                if actual_pos + 1 < items.len() {
                    let new_prefix_util = prefix_util + utils[actual_pos];
                    new_pdb.push((tx_idx, actual_pos + 1, new_prefix_util));
                }
            }
        }

        if !new_pdb.is_empty() {
            hupm_mine_recursive(
                transactions,
                &new_pdb,
                min_utility,
                max_len,
                current_pattern,
                results,
            );
        }

        current_pattern.pop();
    }
}

pub fn hupm_simple(
    transactions: &[Transaction],
    min_utility: f32,
    max_len: Option<usize>,
) -> Vec<(f32, Vec<u32>)> {
    let mut pdb = Vec::with_capacity(transactions.len());
    for i in 0..transactions.len() {
        if !transactions[i].0.is_empty() {
            pdb.push((i, 0, 0.0));
        }
    }

    let mut results = Vec::new();
    let mut current_pattern = Vec::new();

    hupm_mine_recursive(
        transactions,
        &pdb,
        min_utility,
        max_len,
        &mut current_pattern,
        &mut results,
    );

    results
}

#[pyfunction]
#[pyo3(signature = (items_list, utils_list, min_utility, max_len=None))]
pub fn hupm_mine_py(
    items_list: Vec<Vec<u32>>,
    utils_list: Vec<Vec<f32>>,
    min_utility: f32,
    max_len: Option<usize>,
) -> PyResult<(Vec<f32>, Vec<Vec<u32>>)> {
    if items_list.len() != utils_list.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "items_list and utils_list must have the same length",
        ));
    }

    let mut transactions = Vec::with_capacity(items_list.len());
    for (items, utils) in items_list.into_iter().zip(utils_list.into_iter()) {
        if items.len() != utils.len() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "Each inner list of items and utils must have the same length",
            ));
        }

        let mut pairs: Vec<(u32, f32)> = items.into_iter().zip(utils.into_iter()).collect();
        pairs.sort_unstable_by(|a, b| a.0.cmp(&b.0));

        let sorted_items: Vec<u32> = pairs.iter().map(|&(i, _)| i).collect();
        let sorted_utils: Vec<f32> = pairs.iter().map(|&(_, u)| u).collect();

        transactions.push((sorted_items, sorted_utils));
    }

    let raw_res = hupm_simple(&transactions, min_utility, max_len);

    let mut utilities = Vec::with_capacity(raw_res.len());
    let mut patterns = Vec::with_capacity(raw_res.len());

    for (util, pat) in raw_res {
        utilities.push(util);
        patterns.push(pat);
    }

    Ok((utilities, patterns))
}
