from typeguard import typechecked
from ..models.error_models import GuideCountErrorType, GuideCountError
from ..models.quality_control_models import SingleInferenceQualityControlResult, MatchSetSingleInferenceQualityControlResult, QualityControlResult
from ..models.mapping_models import GeneralMappingInferenceDict
from typing import DefaultDict, Tuple, Optional, Union
from collections import defaultdict
from ..processing.crispr_editing_processing import get_non_error_dict

# -----------------------------
# QC helpers
# -----------------------------

def iter_all_obs_values(d, contains_sample_barcode: bool):
    """
    Iterate over all observed values regardless of nesting.
    If sample barcodes are present, d is nested: {sample: {guide: InferenceResult}}
    Otherwise, d is flat: {guide: InferenceResult}
    """
    if contains_sample_barcode:
        for sample_dict in d.values():
            for obs_value in sample_dict.values():
                yield obs_value
    else:
        for obs_value in d.values():
            yield obs_value


def get_umi_noncollapsed_counts(d, contains_guide_umi: bool, contains_sample_barcode: bool):
    return sum(
        sum(v.observed_value.values()) if contains_guide_umi else v.observed_value
        for v in iter_all_obs_values(d, contains_sample_barcode)
    )


def get_umi_collapsed_counts(d, contains_guide_umi: bool, contains_sample_barcode: bool):
    return sum(
        len(v.observed_value) if contains_guide_umi else v.observed_value
        for v in iter_all_obs_values(d, contains_sample_barcode)
    )


# -----------------------------
# Update num counts in QC results
# -----------------------------

def set_num_non_error_counts(qc_result, counts_inferred_dict, contains_guide_umi: bool, contains_sample_barcode: bool):
    if contains_guide_umi:
        qc_result.num_non_error_umi_noncollapsed_counts = get_umi_noncollapsed_counts(
            counts_inferred_dict, contains_guide_umi=True, contains_sample_barcode=contains_sample_barcode
        )
        qc_result.num_non_error_umi_collapsed_counts = get_umi_collapsed_counts(
            counts_inferred_dict, contains_guide_umi=True, contains_sample_barcode=contains_sample_barcode
        )
    else:
        qc_result.num_non_error_counts = get_umi_noncollapsed_counts(
            counts_inferred_dict, contains_guide_umi=False, contains_sample_barcode=contains_sample_barcode
        )
    return qc_result


def set_num_guide_count_error_types(qc_result, attribute_name: str, observed_guide_reporter_umi_counts_inferred, contains_guide_umi: bool, contains_sample_barcode: bool):
    guide_count_error_type_umi_noncollapsed_count: DefaultDict = defaultdict(int)
    guide_count_error_type_umi_collapsed_count: DefaultDict = defaultdict(int)
    guide_count_error_type_count: DefaultDict = defaultdict(int)

    for obs_value in iter_all_obs_values(observed_guide_reporter_umi_counts_inferred, contains_sample_barcode):
        attr = getattr(obs_value.inferred_value, attribute_name, None)
        error = getattr(attr, "error", None)

        # Only count errors, skip NO_ERROR
        if error is not None and error.guide_count_error_type.name != "NO_ERROR":
            if contains_guide_umi:
                guide_count_error_type_umi_noncollapsed_count[error.guide_count_error_type] += sum(obs_value.observed_value.values())
                guide_count_error_type_umi_collapsed_count[error.guide_count_error_type] += len(obs_value.observed_value)
            else:
                guide_count_error_type_count[error.guide_count_error_type] += obs_value.observed_value

    qc_result.guide_count_error_type_umi_noncollapsed_count = guide_count_error_type_umi_noncollapsed_count
    qc_result.guide_count_error_type_umi_collapsed_count = guide_count_error_type_umi_collapsed_count
    qc_result.guide_count_error_type_count = guide_count_error_type_count
    return qc_result


# -----------------------------
# Main QC
# -----------------------------

def perform_counts_quality_control(
    observed_guide_reporter_umi_counts_inferred,
    contains_guide_umi: bool,
    contains_guide_surrogate: bool,
    contains_guide_barcode: bool,
    contains_sample_barcode: bool,
):
    qc_result = QualityControlResult()

    # Helper to process a single attribute
    def set_match_set_single_inference_quality_control_results(attribute_name: str):
        qc = MatchSetSingleInferenceQualityControlResult()
        qc.non_error_dict = get_non_error_dict(observed_guide_reporter_umi_counts_inferred, attribute_name)
        qc = set_num_non_error_counts(qc, qc.non_error_dict, contains_guide_umi, contains_sample_barcode)
        qc = set_num_guide_count_error_types(qc, attribute_name, observed_guide_reporter_umi_counts_inferred, contains_guide_umi, contains_sample_barcode)
        return qc

    # Assign QC results
    qc_result.protospacer_match = set_match_set_single_inference_quality_control_results("protospacer_match")
    if contains_guide_barcode:
        qc_result.protospacer_match_barcode_match = set_match_set_single_inference_quality_control_results("protospacer_match_barcode_match")
    if contains_guide_surrogate:
        qc_result.surrogate_match = set_match_set_single_inference_quality_control_results("surrogate_match")
    if contains_sample_barcode:
        qc_result.sample_barcode_match = set_match_set_single_inference_quality_control_results("sample_barcode_match")

    # Shared totals
    if contains_guide_umi:
        qc_result.total_umi_noncollapsed_counts = get_umi_noncollapsed_counts(
            observed_guide_reporter_umi_counts_inferred, contains_guide_umi, contains_sample_barcode
        )
        qc_result.total_umi_collapsed_counts = get_umi_collapsed_counts(
            observed_guide_reporter_umi_counts_inferred, contains_guide_umi, contains_sample_barcode
        )

    return qc_result