//! Labelling-based semantics solver using subset enumeration.
//!
//! Computes extensions by enumerating conflict-free sets, filtering for
//! admissibility and completeness, and then applying the semantics-specific
//! selection criterion (maximality, stability, etc.).
//!
//! Suitable for small frameworks and as a correctness oracle; for larger
//! instances, use a SAT-based or SCC-decomposing solver.

use std::collections::BTreeSet;

use crate::framework::abstract_af::AbstractFramework;
use crate::framework::argument::ArgumentId;
use crate::framework::semantics::{Extension, Semantics};
use crate::pipeline::traits::SemanticsSolver;

/// Labelling-based semantics solver using subset enumeration.
///
/// Computes extensions by enumerating conflict-free sets, filtering for
/// admissibility and completeness, and then applying the semantics-specific
/// selection criterion (maximality, stability, etc.).
///
/// Suitable for small frameworks; for larger instances, use a SAT-based or
/// SCC-decomposing solver.
pub struct NaiveEnumerationSolver;

impl SemanticsSolver for NaiveEnumerationSolver {
    fn extensions(&self, framework: &AbstractFramework, semantics: Semantics) -> Vec<Extension> {
        match semantics {
            Semantics::Grounded => vec![grounded_extension(framework)],
            Semantics::Complete => complete_extensions(framework),
            Semantics::Preferred => preferred_extensions(framework),
            Semantics::Stable => stable_extensions(framework),
            Semantics::SemiStable => semi_stable_extensions(framework),
            Semantics::Ideal => vec![ideal_extension(framework)],
            Semantics::Eager => vec![eager_extension(framework)],
            Semantics::Admissible => {
                let args: Vec<ArgumentId> = framework.arguments.iter().copied().collect();
                enumerate_admissible_sets(framework, &args)
            }
        }
    }

    fn find_one(&self, framework: &AbstractFramework, semantics: Semantics) -> Option<Extension> {
        let args: Vec<ArgumentId> = framework.arguments.iter().copied().collect();
        match semantics {
            Semantics::Grounded => Some(grounded_extension(framework)),
            Semantics::Admissible => find_one_admissible_set(framework, &args),
            Semantics::Stable => find_one_stable_set(framework, &args),
            Semantics::Complete => find_one_complete_set(framework, &args),
            Semantics::Preferred => {
                let complete = complete_extensions(framework);
                complete
                    .iter()
                    .find(|ext| {
                        !complete
                            .iter()
                            .any(|other| other != *ext && ext.is_subset(other))
                    })
                    .cloned()
            }
            Semantics::SemiStable => semi_stable_extensions(framework).into_iter().next(),
            Semantics::Ideal => Some(ideal_extension(framework)),
            Semantics::Eager => Some(eager_extension(framework)),
        }
    }
}

// ---------------------------------------------------------------------------
// Extension computation algorithms
// ---------------------------------------------------------------------------

/// Compute the grounded extension as the least fixpoint of the Dung
/// characteristic function.
///
/// Each iteration accepts all currently-undefeated arguments (whose
/// attackers have all been marked defeated) and then propagates defeat to
/// their targets. The loop terminates when no new arguments are accepted --
/// guaranteed in at most |Args| iterations since the `grounded` set grows
/// monotonically and `framework.arguments` is finite.
pub(crate) fn grounded_extension(framework: &AbstractFramework) -> Extension {
    let mut grounded = BTreeSet::new();
    let mut defeated = BTreeSet::new();

    loop {
        let mut new_in = BTreeSet::new();

        for &arg in &framework.arguments {
            if grounded.contains(&arg) || defeated.contains(&arg) {
                continue;
            }

            let attackers = framework.attackers_of(arg);
            let all_defeated = attackers.iter().all(|a| defeated.contains(a));
            if all_defeated {
                new_in.insert(arg);
            }
        }

        if new_in.is_empty() {
            break;
        }

        for &arg in &new_in {
            for &target in framework.attacked_by(arg) {
                defeated.insert(target);
            }
        }

        grounded.extend(new_in);
    }

    grounded
}

/// All complete extensions: admissible sets that are also fixpoints of the
/// characteristic function (contain every argument they defend).
fn complete_extensions(framework: &AbstractFramework) -> Vec<Extension> {
    let args: Vec<ArgumentId> = framework.arguments.iter().copied().collect();
    let admissible = enumerate_admissible_sets(framework, &args);

    admissible
        .into_iter()
        .filter(|set| framework.is_complete(set))
        .collect()
}

/// Preferred extensions: maximal (w.r.t. set inclusion) complete extensions.
fn preferred_extensions(framework: &AbstractFramework) -> Vec<Extension> {
    let complete = complete_extensions(framework);
    maximal_sets(&complete)
}

/// Stable extensions: conflict-free sets that attack every non-member.
fn stable_extensions(framework: &AbstractFramework) -> Vec<Extension> {
    let args: Vec<ArgumentId> = framework.arguments.iter().copied().collect();
    let cf_sets = enumerate_conflict_free_sets(framework, &args);

    cf_sets
        .into_iter()
        .filter(|set| framework.is_stable(set))
        .collect()
}

/// Semi-stable extensions: complete extensions whose range (S ∪ S⁺) is maximal.
fn semi_stable_extensions(framework: &AbstractFramework) -> Vec<Extension> {
    let complete = complete_extensions(framework);
    if complete.is_empty() {
        return vec![];
    }

    let with_undec: Vec<(Extension, BTreeSet<ArgumentId>)> = complete
        .into_iter()
        .map(|ext| {
            let attacked: BTreeSet<ArgumentId> = ext
                .iter()
                .flat_map(|&a| framework.attacked_by(a).iter().copied())
                .collect();
            let range: BTreeSet<ArgumentId> = ext.union(&attacked).copied().collect();
            let undec: BTreeSet<ArgumentId> =
                framework.arguments.difference(&range).copied().collect();
            (ext, undec)
        })
        .collect();

    let min_undec = with_undec.iter().map(|(_, u)| u.len()).min().unwrap_or(0);

    with_undec
        .into_iter()
        .filter(|(_, u)| u.len() == min_undec)
        .map(|(ext, _)| ext)
        .collect()
}

/// Ideal extension: the largest admissible subset of the intersection of
/// all preferred extensions (Dung, Mancarella & Toni 2007).
fn ideal_extension(framework: &AbstractFramework) -> Extension {
    let preferred = preferred_extensions(framework);
    if preferred.is_empty() {
        return BTreeSet::new();
    }

    let mut intersection = preferred[0].clone();
    for ext in &preferred[1..] {
        intersection = intersection.intersection(ext).copied().collect();
    }

    let args: Vec<ArgumentId> = intersection.iter().copied().collect();
    let admissible = enumerate_admissible_sets(framework, &args);

    admissible
        .into_iter()
        .max_by_key(std::collections::BTreeSet::len)
        .unwrap_or_default()
}

/// Eager extension: the largest admissible subset of the intersection of
/// all semi-stable extensions (Caminada 2007).
fn eager_extension(framework: &AbstractFramework) -> Extension {
    let semi_stable = semi_stable_extensions(framework);
    if semi_stable.is_empty() {
        return BTreeSet::new();
    }

    let mut intersection = semi_stable[0].clone();
    for ext in &semi_stable[1..] {
        intersection = intersection.intersection(ext).copied().collect();
    }

    let args: Vec<ArgumentId> = intersection.iter().copied().collect();
    let admissible = enumerate_admissible_sets(framework, &args);

    admissible
        .into_iter()
        .max_by_key(std::collections::BTreeSet::len)
        .unwrap_or_default()
}

/// Enumerate all conflict-free subsets of `args` via incremental powerset
/// construction, pruning branches that introduce a conflict.
fn enumerate_conflict_free_sets(
    framework: &AbstractFramework,
    args: &[ArgumentId],
) -> Vec<Extension> {
    let mut result = vec![BTreeSet::new()];

    for &arg in args {
        let mut new_sets = Vec::new();
        for set in &result {
            let has_conflict = framework
                .attacked_by(arg)
                .iter()
                .any(|&t| t == arg || set.contains(&t))
                || framework
                    .attackers_of(arg)
                    .iter()
                    .any(|&a| set.contains(&a));
            if !has_conflict {
                let mut candidate = set.clone();
                candidate.insert(arg);
                new_sets.push(candidate);
            }
        }
        result.extend(new_sets);
    }

    result
}

/// Enumerate all admissible sets of the framework via subset enumeration.
///
/// Pre-filters the argument pool by removing arguments defeated by the grounded
/// extension, since such arguments can never appear in any admissible set.
#[cfg(feature = "sat")]
pub(crate) fn admissible_extensions(framework: &AbstractFramework) -> Vec<Extension> {
    let grounded = grounded_extension(framework);
    let defeated: BTreeSet<ArgumentId> = grounded
        .iter()
        .flat_map(|&a| framework.attacked_by(a))
        .copied()
        .collect();
    let args: Vec<ArgumentId> = framework
        .arguments
        .iter()
        .filter(|a| !defeated.contains(a))
        .copied()
        .collect();
    enumerate_admissible_sets(framework, &args)
}

/// Enumerate all admissible sets: conflict-free sets where every member
/// is defended by the set.
fn enumerate_admissible_sets(framework: &AbstractFramework, args: &[ArgumentId]) -> Vec<Extension> {
    let cf_sets = enumerate_conflict_free_sets(framework, args);
    cf_sets
        .into_iter()
        .filter(|set| framework.defends_all_members(set))
        .collect()
}

/// Iterate through conflict-free subsets (same powerset logic as
/// `enumerate_conflict_free_sets`) and return the first set matching the predicate.
fn find_one_conflict_free_set_by<F>(
    framework: &AbstractFramework,
    args: &[ArgumentId],
    predicate: F,
) -> Option<Extension>
where
    F: Fn(&Extension) -> bool,
{
    let mut result = vec![BTreeSet::new()];

    for &arg in args {
        let mut new_sets = Vec::new();
        for set in &result {
            if predicate(set) {
                return Some(set.clone());
            }
            let has_conflict = framework
                .attacked_by(arg)
                .iter()
                .any(|&t| t == arg || set.contains(&t))
                || framework
                    .attackers_of(arg)
                    .iter()
                    .any(|&a| set.contains(&a));
            if !has_conflict {
                let mut candidate = set.clone();
                candidate.insert(arg);
                if predicate(&candidate) {
                    return Some(candidate);
                }
                new_sets.push(candidate);
            }
        }
        result.extend(new_sets);
    }

    result.into_iter().find(|set| predicate(set))
}

/// Return the first admissible set found during conflict-free enumeration.
fn find_one_admissible_set(
    framework: &AbstractFramework,
    args: &[ArgumentId],
) -> Option<Extension> {
    find_one_conflict_free_set_by(framework, args, |set| framework.defends_all_members(set))
}

/// Return the first stable set found during conflict-free enumeration.
fn find_one_stable_set(framework: &AbstractFramework, args: &[ArgumentId]) -> Option<Extension> {
    find_one_conflict_free_set_by(framework, args, |set| framework.is_stable(set))
}

/// Return the first complete set found during conflict-free enumeration.
fn find_one_complete_set(framework: &AbstractFramework, args: &[ArgumentId]) -> Option<Extension> {
    find_one_conflict_free_set_by(framework, args, |set| framework.is_complete(set))
}

/// Filter `sets` to retain only those not properly contained in another.
fn maximal_sets(sets: &[Extension]) -> Vec<Extension> {
    crate::algorithms::maximal::maximal_by_inclusion(sets)
}
