#!/usr/bin/env python3
"""Adaptive selector for active birth-time rectification questions."""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Iterable


DEFAULT_DOMAIN_ORDER = [
    "education",
    "residence",
    "relationship",
    "career",
    "career_learning",
    "health_pressure",
    "public_work",
    "fine_timing",
]

DOMAIN_RELIABILITY = {
    "education": 1.0,
    "residence": 0.95,
    "relationship": 0.95,
    "career": 0.9,
    "career_learning": 0.88,
    "health_pressure": 0.72,
    "public_work": 0.82,
    "fine_timing": 0.45,
}

DOMAIN_EFFORT = {
    "education": 0.88,
    "residence": 0.9,
    "relationship": 0.85,
    "career": 0.82,
    "career_learning": 0.8,
    "health_pressure": 0.7,
    "public_work": 0.75,
    "fine_timing": 0.5,
}

QUESTION_LAYER_MAP = {
    "D1": "d1_ascendant",
    "D4": "divisional_ascendants",
    "D9": "divisional_ascendants",
    "D10": "divisional_ascendants",
    "D24": "divisional_ascendants",
    "D30": "divisional_ascendants",
    "D60": "divisional_ascendants",
    "UL": "arudha",
    "A7": "arudha",
    "A10": "arudha",
    "KP_cusp": "kp_cusps",
}


@dataclass(frozen=True)
class SelectorInput:
    questions: list[dict[str, Any]]
    answers: dict[str, str]
    candidate_count: int | None = None
    prior_question_ids: list[str] | None = None


def _is_question(value: Any) -> bool:
    return isinstance(value, dict) and bool(value.get("id"))


def _normalize_questions(questions: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
    normalized = []
    for question in questions:
        if not _is_question(question):
            continue
        normalized.append(question)
    return normalized


def _question_answer_span(question: dict[str, Any]) -> int:
    scoring_map = question.get("scoring_map")
    if not isinstance(scoring_map, dict):
        return 0
    clusters = []
    for option in ("A", "B", "C", "D"):
        item = scoring_map.get(option)
        if isinstance(item, dict):
            clusters.append(str(item.get("cluster") or "neutral"))
    return len({cluster for cluster in clusters if cluster != "neutral"})


def _minute_feature_value(row: dict[str, Any], route: str, layer: str) -> Any:
    if route == "d1_ascendant":
        return {
            "sign": row.get("d1_ascendant"),
            "degree_in_sign": row.get("d1_degree_in_sign"),
            "longitude": row.get("d1_longitude"),
        }
    if route == "divisional_ascendants":
        return (row.get("divisional_ascendants") or {}).get(layer)
    if route == "arudha":
        return (row.get("arudha") or {}).get(layer)
    if route == "kp_cusps":
        return row.get("kp_cusps")
    return None


def _minute_relevance(question: dict[str, Any], minute_scan: dict[str, Any] | None) -> float:
    if not isinstance(minute_scan, dict):
        return 0.0
    rows = minute_scan.get("rows")
    if not isinstance(rows, list) or not rows:
        return 0.0
    layers = [str(layer) for layer in (question.get("sensitivity") or []) if str(layer)]
    if not any(QUESTION_LAYER_MAP.get(layer) for layer in layers):
        return 0.0
    signatures = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        signature = tuple(
            json.dumps(_minute_feature_value(row, route, layer), sort_keys=True, ensure_ascii=False, default=str)
            for layer in sorted(layers)
            for route in [QUESTION_LAYER_MAP.get(layer)]
            if route
        )
        signatures.append(signature)
    if not signatures:
        return 0.0
    distinct = len(set(signatures))
    transition_count = len(minute_scan.get("transitions") or [])
    spread = distinct / max(len(signatures), 1)
    return round(spread * 4.0 + min(transition_count, 6) * 0.35, 3)


def _question_effort(question: dict[str, Any]) -> float:
    prompt = str(question.get("prompt") or "")
    return max(0.35, min(1.0, 1.0 - (len(prompt) / 120.0)))


def _question_reliability(question: dict[str, Any]) -> float:
    domain = str(question.get("domain") or "")
    return float(question.get("factual_reliability") or DOMAIN_RELIABILITY.get(domain, 0.75))


def _question_domain_priority(question: dict[str, Any], prior_domain_counts: dict[str, int]) -> float:
    domain = str(question.get("domain") or "")
    base = float(question.get("domain_priority") or 0.0)
    if not base:
        try:
            base = float(len(DEFAULT_DOMAIN_ORDER) - DEFAULT_DOMAIN_ORDER.index(domain))
        except ValueError:
            base = 1.0
    penalty = prior_domain_counts.get(domain, 0) * 1.5
    return base - penalty


def _question_gating(question: dict[str, Any], candidate_count: int | None) -> tuple[bool, str | None]:
    if _question_answer_span(question) <= 0:
        return False, "no_discriminating_answer_clusters"
    if str(question.get("domain") or "") == "fine_timing" and (candidate_count is None or candidate_count > 12):
        return False, "fine_timing_requires_narrower_candidate_window"
    if str(question.get("domain") or "") == "health_pressure" and (candidate_count is None or candidate_count > 24):
        return False, "health_pressure_deferred_until_broader_context_is_resolved"
    return True, None


def _score_question(
    question: dict[str, Any],
    *,
    candidate_count: int | None,
    prior_domain_counts: dict[str, int],
    minute_scan: dict[str, Any] | None,
) -> dict[str, Any]:
    domain = str(question.get("domain") or "")
    reliability = _question_reliability(question)
    effort = _question_effort(question)
    discriminative_span = _question_answer_span(question)
    domain_priority = _question_domain_priority(question, prior_domain_counts)
    cluster_bias = float(question.get("question_gain_bias") or 0.0)
    minute_relevance = _minute_relevance(question, minute_scan)
    candidate_bonus = 0.0
    if candidate_count is not None:
        if candidate_count <= 9:
            candidate_bonus = 2.0
        elif candidate_count <= 21:
            candidate_bonus = 1.0
        else:
            candidate_bonus = 0.35
    score = (
        discriminative_span * 3.0
        + reliability * 2.5
        + effort * 1.5
        + domain_priority * 0.2
        + candidate_bonus
        + minute_relevance
        + cluster_bias
    )
    if str(question.get("domain") or "") == "fine_timing":
        score -= 2.5
    if prior_domain_counts.get(domain, 0):
        score -= prior_domain_counts[domain] * 1.0
    return {
        "question_id": question.get("id"),
        "domain": domain,
        "score": round(score, 3),
        "factual_reliability": round(reliability, 3),
        "user_effort": round(1.0 - effort, 3),
        "domain_diversity_penalty": prior_domain_counts.get(domain, 0),
        "candidate_separation": discriminative_span,
        "minute_relevance": round(minute_relevance, 3),
    }


def select_next_questions(questionnaire: dict[str, Any], answers: dict[str, str] | None = None, *, limit: int = 1) -> dict[str, Any]:
    questions = _normalize_questions(
        questionnaire.get("question_bank") if isinstance(questionnaire.get("question_bank"), list) else questionnaire.get("questions") or []
    )
    answered = {}
    prior_domain_counts: dict[str, int] = {}
    answers = answers or {}
    question_by_id = {str(question["id"]): question for question in questions}

    for question_id, choice in answers.items():
        question = question_by_id.get(str(question_id))
        if not question:
            continue
        answered[str(question_id)] = str(choice or "").strip().upper()
        domain = str(question.get("domain") or "")
        prior_domain_counts[domain] = prior_domain_counts.get(domain, 0) + 1

    remaining = [question for question in questions if str(question.get("id") or "") not in answered]
    candidate_count = None
    candidate_scan = questionnaire.get("candidate_scan")
    if isinstance(candidate_scan, dict):
        raw_count = candidate_scan.get("candidate_count")
        if isinstance(raw_count, int):
            candidate_count = raw_count
    minute_scan = candidate_scan.get("minute_scan") if isinstance(candidate_scan, dict) else None

    ranked = []
    for question in remaining:
        allowed, reason = _question_gating(question, candidate_count)
        if not allowed:
            ranked.append(
                {
                    "question_id": question.get("id"),
                    "domain": question.get("domain"),
                    "score": None,
                    "skipped": True,
                    "skip_reason": reason,
                    "factual_reliability": _question_reliability(question),
                    "candidate_separation": _question_answer_span(question),
                }
            )
            continue
        ranked.append(_score_question(question, candidate_count=candidate_count, prior_domain_counts=prior_domain_counts, minute_scan=minute_scan))

    usable = [item for item in ranked if not item.get("skipped")]
    usable.sort(key=lambda item: (-float(item["score"]), str(item["domain"]), str(item["question_id"])))

    if not usable:
        return {
            "selected_questions": [],
            "ranking": ranked,
            "stop": True,
            "stop_reason": "no_answer_can_improve_separation",
        }

    best = usable[0]
    if float(best["score"]) < 4.0:
        return {
            "selected_questions": [],
            "ranking": ranked,
            "stop": True,
            "stop_reason": "no_answer_can_improve_separation",
        }

    selected = []
    seen_domains: set[str] = set()
    for item in usable:
        question = question_by_id[str(item["question_id"])]
        domain = str(question.get("domain") or "")
        if domain in seen_domains:
            continue
        allowed, _ = _question_gating(question, candidate_count)
        if not allowed:
            continue
        selected.append({
            "id": question["id"],
            "domain": domain,
            "prompt": question.get("prompt"),
            "why_asked": list(question.get("why_this_question") or []),
            "candidate_ids_distinguished": [str(question.get("positive_cluster") or ""), str(question.get("negative_cluster") or "")],
            "technique_routes": list(question.get("sensitivity") or []),
            "answer_impact": dict(question.get("answer_impact") or {}),
            "selection_score": item["score"],
            "factual_reliability": item["factual_reliability"],
            "user_effort": item["user_effort"],
            "candidate_separation": item["candidate_separation"],
            "minute_relevance": item.get("minute_relevance", 0.0),
        })
        seen_domains.add(domain)
        if len(selected) >= max(1, limit):
            break

    return {
        "selected_questions": selected,
        "ranking": ranked,
        "stop": not bool(selected),
        "stop_reason": None if selected else "no_answer_can_improve_separation",
    }