code/self_healing_finetuning/eval/paired_eval_defect_ranker.py

from __future__ import annotations

import json
import re
from collections import Counter, defaultdict
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from n21.config import write_json
from n21.settings import SHFT_WORKSPACE_ROOT


ASSETS = ("equity", "fixed_income", "multi_asset")
ROLES = (
    "chief_investment_officer",
    "client_portfolio_manager",
    "performance_manager",
    "portfolio_manager",
    "researcher",
    "risk_manager",
)

DEFECT_TYPES = (
    "numeric_reasoning",
    "fact_inference_separation",
    "role_discipline",
    "risk_tradeoff_framing",
    "hallucination_unsupported_claim",
    "weak_source_grounding",
    "overfit_memorized_answer_style",
)


def utc_now() -> str:
    return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def read_json(path: Path) -> dict[str, Any] | None:
    if not path.exists():
        return None
    return json.loads(path.read_text(encoding="utf-8-sig"))


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    if not path.exists():
        return rows
    for line_no, line in enumerate(path.read_text(encoding="utf-8-sig").splitlines(), start=1):
        if not line.strip():
            continue
        item = json.loads(line)
        if not isinstance(item, dict):
            raise ValueError(f"{path}:{line_no} must contain a JSON object")
        rows.append(item)
    return rows


def latest_submitted_run(release_id: str) -> Path | None:
    runs_root = SHFT_WORKSPACE_ROOT / "runs"
    dirs = [
        item
        for item in runs_root.glob(f"run_{release_id}*")
        if item.is_dir() and not item.name.endswith("_paired_eval_code")
    ]
    dirs = sorted(dirs, key=lambda item: item.stat().st_mtime, reverse=True)
    for run_path in dirs:
        if (run_path / "trainer_state" / "train_handle.json").exists():
            return run_path
    return dirs[0] if dirs else None


def bool_check(score: dict[str, Any], name: str) -> bool | None:
    checks = score.get("checks")
    if not isinstance(checks, dict):
        return None
    value = checks.get(name)
    return value if isinstance(value, bool) else None


def answer_text(prediction: dict[str, Any]) -> str:
    return str(prediction.get("candidate_answer") or "")


def detect_role_discipline(prediction: dict[str, Any]) -> list[str]:
    text = answer_text(prediction).lower()
    reasons: list[str] = []
    if bool_check(prediction.get("candidate_score", {}), "no_investment_advice") is False:
        reasons.append("candidate gave buy/sell/hold/short-style investment advice")
    if re.search(r"\b(as an ai|i cannot|i do not have access|not financial advice)\b", text):
        reasons.append("candidate broke analyst-role voice")
    if len(text.split()) < 18:
        reasons.append("candidate answer too thin for role-grade analysis")
    return reasons


def detect_weak_source_grounding(prediction: dict[str, Any]) -> list[str]:
    text = answer_text(prediction).lower()
    prompt = str(prediction.get("prompt") or "").lower()
    reasons: list[str] = []
    grounding_terms = ("reported", "filing", "disclosed", "source", "data", "management", "statement", "fact")
    if not any(term in text for term in grounding_terms):
        reasons.append("candidate did not anchor answer to reported/source facts")
    prompt_numbers = set(re.findall(r"\$?\d+(?:\.\d+)?%?", prompt))
    if prompt_numbers:
        answer_numbers = set(re.findall(r"\$?\d+(?:\.\d+)?%?", text))
        if not prompt_numbers.intersection(answer_numbers):
            reasons.append("candidate omitted numeric evidence present in prompt")
    return reasons


def detect_overfit_style(prediction: dict[str, Any]) -> list[str]:
    text = answer_text(prediction)
    lower = text.lower()
    reasons: list[str] = []
    repeated_phrases = [
        "numeric anchor:",
        "reported facts:",
        "risk/tradeoff:",
        "required next evidence:",
    ]
    phrase_hits = sum(1 for phrase in repeated_phrases if phrase in lower)
    if phrase_hits >= 3:
        reasons.append("candidate uses repair-template wording heavily")
    sentences = [item.strip() for item in re.split(r"[.!?]\s+", lower) if item.strip()]
    if len(sentences) >= 4 and len(set(sentences)) <= max(2, len(sentences) // 2):
        reasons.append("candidate repeats sentence structure/content")
    return reasons


def classify_prediction(prediction: dict[str, Any]) -> dict[str, list[str]]:
    score = prediction.get("candidate_score", {})
    checks = score.get("checks") if isinstance(score.get("checks"), dict) else {}
    defects: dict[str, list[str]] = {name: [] for name in DEFECT_TYPES}

    if checks.get("numeric_reasoning_present") is False:
        defects["numeric_reasoning"].append("numeric_reasoning_present=false")
    if checks.get("fact_inference_separation") is False:
        defects["fact_inference_separation"].append("fact_inference_separation=false")
    if checks.get("risk_or_tradeoff_identified") is False or checks.get("neutral_language") is False:
        defects["risk_tradeoff_framing"].append("risk/tradeoff or neutral-language check failed")
    if checks.get("no_unsupported_certainty") is False:
        defects["hallucination_unsupported_claim"].append("no_unsupported_certainty=false")

    defects["role_discipline"].extend(detect_role_discipline(prediction))
    defects["weak_source_grounding"].extend(detect_weak_source_grounding(prediction))
    defects["overfit_memorized_answer_style"].extend(detect_overfit_style(prediction))

    return {key: value for key, value in defects.items() if value}


@dataclass
class RoleDefectResult:
    asset_class: str
    role: str
    release_id: str
    run_id: str | None
    status: str
    prediction_count: int
    failed_prediction_count: int
    defect_counts: dict[str, int]
    top_defects: list[dict[str, Any]]
    paired_predictions: str | None
    paired_eval_report: str | None


def rank_run_defects(*, run_path: Path, asset_class: str, role: str, release_id: str) -> RoleDefectResult:
    predictions_path = run_path / "eval" / "paired_predictions.jsonl"
    report_path = run_path / "eval" / "paired_eval_report.json"
    if not predictions_path.exists():
        return RoleDefectResult(
            asset_class=asset_class,
            role=role,
            release_id=release_id,
            run_id=run_path.name,
            status="proof_missing",
            prediction_count=0,
            failed_prediction_count=0,
            defect_counts={name: 0 for name in DEFECT_TYPES},
            top_defects=[],
            paired_predictions=str(predictions_path),
            paired_eval_report=str(report_path),
        )

    rows = read_jsonl(predictions_path)
    defect_counts: Counter[str] = Counter()
    examples: dict[str, list[dict[str, Any]]] = defaultdict(list)
    failed_count = 0
    for row in rows:
        defects = classify_prediction(row)
        if defects or row.get("candidate_score", {}).get("critical_pass") is False:
            failed_count += 1
        for defect, reasons in defects.items():
            defect_counts[defect] += 1
            if len(examples[defect]) < 3:
                examples[defect].append(
                    {
                        "id": row.get("id"),
                        "task": row.get("task"),
                        "delta": row.get("delta"),
                        "reasons": reasons,
                        "prompt_preview": str(row.get("prompt") or "")[:240],
                    }
                )

    counts = {name: int(defect_counts.get(name, 0)) for name in DEFECT_TYPES}
    top = [
        {"defect_type": name, "count": count, "examples": examples.get(name, [])}
        for name, count in sorted(counts.items(), key=lambda item: (-item[1], item[0]))
        if count > 0
    ]
    return RoleDefectResult(
        asset_class=asset_class,
        role=role,
        release_id=release_id,
        run_id=run_path.name,
        status="ranked" if rows else "no_predictions",
        prediction_count=len(rows),
        failed_prediction_count=failed_count,
        defect_counts=counts,
        top_defects=top,
        paired_predictions=str(predictions_path),
        paired_eval_report=str(report_path),
    )


def result_to_dict(result: RoleDefectResult) -> dict[str, Any]:
    return {
        "asset_class": result.asset_class,
        "role": result.role,
        "release_id": result.release_id,
        "run_id": result.run_id,
        "status": result.status,
        "prediction_count": result.prediction_count,
        "failed_prediction_count": result.failed_prediction_count,
        "defect_counts": result.defect_counts,
        "top_defects": result.top_defects,
        "paired_predictions": result.paired_predictions,
        "paired_eval_report": result.paired_eval_report,
    }


def rank_all_role_defects(*, output_path: Path | None = None) -> dict[str, Any]:
    role_results: list[RoleDefectResult] = []
    aggregate_counts: Counter[str] = Counter()

    for asset_class in ASSETS:
        for role in ROLES:
            release_id = f"linvest21_fingpt_{asset_class}_{role}_v1_001"
            run_path = latest_submitted_run(release_id)
            if run_path is None:
                result = RoleDefectResult(
                    asset_class=asset_class,
                    role=role,
                    release_id=release_id,
                    run_id=None,
                    status="run_missing",
                    prediction_count=0,
                    failed_prediction_count=0,
                    defect_counts={name: 0 for name in DEFECT_TYPES},
                    top_defects=[],
                    paired_predictions=None,
                    paired_eval_report=None,
                )
            else:
                result = rank_run_defects(run_path=run_path, asset_class=asset_class, role=role, release_id=release_id)
            role_results.append(result)
            aggregate_counts.update(result.defect_counts)

    payload = {
        "schema_version": "paired_eval_defect_ranker_v1",
        "created_at": utc_now(),
        "defect_types": list(DEFECT_TYPES),
        "summary": {
            "role_count": len(role_results),
            "ranked_role_count": sum(1 for item in role_results if item.status == "ranked"),
            "proof_missing_role_count": sum(1 for item in role_results if item.status == "proof_missing"),
            "run_missing_role_count": sum(1 for item in role_results if item.status == "run_missing"),
            "prediction_count": sum(item.prediction_count for item in role_results),
            "failed_prediction_count": sum(item.failed_prediction_count for item in role_results),
            "defect_counts": {name: int(aggregate_counts.get(name, 0)) for name in DEFECT_TYPES},
        },
        "roles": [result_to_dict(item) for item in role_results],
    }
    output = output_path or SHFT_WORKSPACE_ROOT / "verification" / "paired_eval_defect_ranking_latest.json"
    write_json(output, payload)
    markdown_output = output.with_suffix(".md")
    markdown_output.write_text(markdown_report(payload), encoding="utf-8")
    payload["output_path"] = str(output)
    payload["markdown_output_path"] = str(markdown_output)
    return payload


def markdown_report(payload: dict[str, Any]) -> str:
    summary = payload["summary"]
    widths = [12, 24, 14, 10, 10, 10, 10, 10, 10, 10]
    headers = [
        "Asset",
        "Role",
        "Status",
        "Numeric",
        "Fact/Inf",
        "Role",
        "Risk",
        "Unsup",
        "Ground",
        "Overfit",
    ]

    def sep(left: str, middle: str, right: str) -> str:
        return left + middle.join("─" * (width + 2) for width in widths) + right

    def line(values: list[Any]) -> str:
        return "│ " + " │ ".join(str(value).ljust(width) for value, width in zip(values, widths, strict=True)) + " │"

    lines = [
        "# Paired Eval Defect Ranking",
        "",
        f"Generated: {payload['created_at']}",
        "",
        "Defect taxonomy:",
        "",
    ]
    for defect_type in payload["defect_types"]:
        lines.append(f"- {defect_type}")
    lines.extend(
        [
            "",
            "Summary:",
            "",
            f"- role_count: {summary['role_count']}",
            f"- ranked_role_count: {summary['ranked_role_count']}",
            f"- proof_missing_role_count: {summary['proof_missing_role_count']}",
            f"- run_missing_role_count: {summary['run_missing_role_count']}",
            f"- prediction_count: {summary['prediction_count']}",
            f"- failed_prediction_count: {summary['failed_prediction_count']}",
            "",
            "Status By Role",
            "",
            sep("┌", "┬", "┐"),
            line(headers),
            sep("├", "┼", "┤"),
        ]
    )
    for row in payload["roles"]:
        counts = row["defect_counts"]
        lines.append(
            line(
                [
                    row["asset_class"],
                    row["role"],
                    row["status"],
                    counts["numeric_reasoning"],
                    counts["fact_inference_separation"],
                    counts["role_discipline"],
                    counts["risk_tradeoff_framing"],
                    counts["hallucination_unsupported_claim"],
                    counts["weak_source_grounding"],
                    counts["overfit_memorized_answer_style"],
                ]
            )
        )
    lines.append(sep("└", "┴", "┘"))
    return "\n".join(lines) + "\n"