Evaluations Framework Hill-Climbing

import Evaluations
import FoundationModels
import Testing

// MARK: - Score Dimensions

struct TagScoreDimensions: ScoreDimensions {
    @ScoreDimension(scale: 1...4, description: "How well tags represent plot, theme, and book info")
    var relevance: Score

    @ScoreDimension(scale: 1...4, description: "How useful the tags are as search terms")
    var usefulness: Score
}

// MARK: - Dataset Entry

struct TagAlignmentEntry: EvaluationInput {
    let bookSummary: String
    let generatedTags: [String]
    let expertRelevance: Int   // Human expert rating 1-4
    let expertUsefulness: Int  // Human expert rating 1-4
}

// MARK: - Cohen's Kappa Helper

func cohensKappa(expertRatings: [Int], modelRatings: [Int]) -> Double {
    guard expertRatings.count == modelRatings.count, !expertRatings.isEmpty else { return 0 }
    let n = Double(expertRatings.count)
    // Observed accuracy
    let matches = zip(expertRatings, modelRatings).filter { $0.0 == $0.1 }.count
    let po = Double(matches) / n
    // Expected accuracy (chance agreement)
    let scale = 1...4
    let pe = scale.reduce(0.0) { sum, k in
        let expertFrac = Double(expertRatings.filter { $0 == k }.count) / n
        let modelFrac  = Double(modelRatings.filter  { $0 == k }.count) / n
        return sum + expertFrac * modelFrac
    }
    guard pe < 1.0 else { return 1.0 }
    return (po - pe) / (1.0 - pe)
}

// MARK: - Alignment Evaluation

struct TagJudgeAlignmentEvaluation: Evaluation {
    typealias Input  = TagAlignmentEntry
    typealias Output = [String]
    typealias Dimensions = TagScoreDimensions

    let dataset: [TagAlignmentEntry] = [
        TagAlignmentEntry(
            bookSummary: "A young boy joins pirates on a treasure hunt, facing loyalty vs betrayal.",
            generatedTags: ["adventure", "pirates", "treasure", "coming-of-age"],
            expertRelevance: 4,
            expertUsefulness: 2
        ),
        TagAlignmentEntry(
            bookSummary: "Four sisters navigate love and loss in Civil War-era New England.",
            generatedTags: ["poignant", "quiet-steadiness", "family", "women"],
            expertRelevance: 4,
            expertUsefulness: 2
        )
    ]

    func subject(for input: TagAlignmentEntry) async throws -> [String] {
        // Return pre-generated tags from the dataset (no model call needed here)
        return input.generatedTags
    }

    var evaluators: [any Evaluator<[String], TagScoreDimensions>] {
        [
            ModelJudgeEvaluator(
                model: SystemLanguageModel.default,
                prompt: """Rate these book tags on relevance (how well they capture plot/theme) \
and usefulness (how good they are as search terms). Scale: 1=poor, 4=excellent."""
            )
        ]
    }

    func aggregate(
        results: [EvaluationResult<TagAlignmentEntry, [String], TagScoreDimensions>]
    ) -> [String: Double] {
        let expertRel  = results.map { $0.input.expertRelevance }
        let expertUse  = results.map { $0.input.expertUsefulness }
        let modelRel   = results.compactMap { $0.scores.relevance.map { Int($0.value.rounded()) } }
        let modelUse   = results.compactMap { $0.scores.usefulness.map { Int($0.value.rounded()) } }

        let kappaRel = cohensKappa(expertRatings: expertRel, modelRatings: modelRel)
        let kappaUse = cohensKappa(expertRatings: expertUse, modelRatings: modelUse)
        let overallKappa = (kappaRel + kappaUse) / 2.0

        return [
            "kappa.relevance":  kappaRel,
            "kappa.usefulness": kappaUse,
            "kappa.overall":    overallKappa
        ]
    }
}

// MARK: - Swift Testing Integration

@Test func judgeAlignmentMeetsThreshold() async throws {
    let evaluation = TagJudgeAlignmentEvaluation()
    let metrics = try await evaluation.run()
    let overallKappa = metrics["kappa.overall"] ?? 0
    // 0.6+ = substantial agreement per Cohen (1960)
    #expect(overallKappa >= 0.6, "Judge drift detected: kappa=\(overallKappa). Refine judge prompt.")
}