core/aggregator.py

from __future__ import annotations

from collections import Counter, defaultdict
from statistics import mean

from core.schemas import (
    AggregatedReport,
    BrutalQuote,
    DeadliestStep,
    ExcludedPersona,
    FunnelStep,
    IssueCluster,
    Persona,
    SegmentFailure,
    SimulationResult,
    SummaryMetrics,
)


def _normalize_issue(issue: str) -> str:
    lowered = issue.lower()
    if "pricing" in lowered or "cost" in lowered or "free" in lowered:
        return "Pricing unclear"
    if "jargon" in lowered or "headline" in lowered or "word" in lowered:
        return "Headline jargon"
    if "visual" in lowered or "screenshot" in lowered or "demo" in lowered:
        return "No product visual"
    if "cta" in lowered or "button" in lowered or "sign" in lowered:
        return "CTA unclear"
    if "trust" in lowered or "proof" in lowered or "review" in lowered:
        return "Missing trust signals"
    if "contrast" in lowered or "accessibility" in lowered:
        return "Accessibility risk"
    return issue.strip().capitalize()


def aggregate_results(
    results: list[SimulationResult],
    personas: list[Persona],
    flow_steps: list[str],
) -> AggregatedReport:
    valid_count = len(results)
    total_personas = len(personas)
    valid_persona_ids = {result.persona_id for result in results}
    excluded_personas = [
        ExcludedPersona(
            persona_id=persona.persona_id,
            archetype=persona.archetype,
            reason="simulation result failed validation",
        )
        for persona in personas
        if persona.persona_id not in valid_persona_ids
    ]
    completed = sum(1 for result in results if result.completed_task)
    failed = valid_count - completed
    success_rate = completed / valid_count if valid_count else 0.0
    avg_confusion = mean([result.overall_confusion_score for result in results]) if results else 0.0
    avg_trust = mean([result.overall_trust_score for result in results]) if results else 0.0
    max_steps = max([result.total_steps for result in results], default=len(flow_steps) or 1)

    funnel: list[FunnelStep] = []
    for step in range(1, max_steps + 1):
        label = flow_steps[step - 1] if step <= len(flow_steps) else f"Step {step}"
        survival = (
            sum(1 for result in results if result.final_step_reached >= step) / valid_count
            if valid_count
            else 0.0
        )
        funnel.append(FunnelStep(step=step, label=label, survival_rate=survival))

    deadliest_step = None
    if len(funnel) > 1:
        drops = [
            (current.step, current.label, current.survival_rate - nxt.survival_rate)
            for current, nxt in zip(funnel, funnel[1:], strict=False)
        ]
        step, label, dropoff = max(drops, key=lambda item: item[2])
        deadliest_step = DeadliestStep(
            step=step,
            label=label,
            dropoff=max(dropoff, 0.0),
            primary_reason=_primary_reason(results, step),
        )

    by_segment: dict[str, list[SimulationResult]] = defaultdict(list)
    for result in results:
        by_segment[result.archetype].append(result)
    worst_segments = sorted(
        [
            SegmentFailure(
                archetype=archetype,
                failure_rate=1
                - (sum(r.completed_task for r in segment_results) / len(segment_results)),
            )
            for archetype, segment_results in by_segment.items()
        ],
        key=lambda item: item.failure_rate,
        reverse=True,
    )[:5]

    issue_counts = Counter(
        _normalize_issue(issue)
        for result in results
        for issue in result.top_issues
        if issue.strip()
    )
    top_issues = [
        IssueCluster(
            issue=issue,
            frequency=frequency,
            severity="high" if frequency >= 3 else "medium" if frequency == 2 else "low",
        )
        for issue, frequency in issue_counts.most_common(5)
    ]
    brutal_quotes = [
        BrutalQuote(
            quote=result.brutal_quote,
            persona=result.persona_id,
            archetype=result.archetype,
        )
        for result in sorted(results, key=lambda item: item.overall_confusion_score, reverse=True)
        if not result.completed_task
    ][:3]

    return AggregatedReport(
        summary=SummaryMetrics(
            total_personas=total_personas,
            valid_personas=valid_count,
            excluded=len(excluded_personas),
            completed=completed,
            failed=failed,
            success_rate=success_rate,
            avg_confusion=avg_confusion,
            avg_trust=avg_trust,
        ),
        failure_funnel=funnel,
        deadliest_step=deadliest_step,
        worst_segments=worst_segments,
        top_issues=top_issues,
        brutal_quotes=brutal_quotes,
        excluded_personas=excluded_personas,
        persona_results=results,
    )


def _primary_reason(results: list[SimulationResult], step: int) -> str:
    issues = Counter(
        _normalize_issue(issue)
        for result in results
        for log in result.step_log
        if log.step == step
        for issue in log.issues
    )
    return issues.most_common(1)[0][0] if issues else "No dominant issue"