"""General text quality on the full MTEB(eng, v2) benchmark (41 tasks).

Give each model its own --output-dir, or local models will reuse each other's cached scores.

    uv run mteb_tasks.py jaisawhney/SonicEmbed --output-dir mteb-results/sonicembed
"""
import argparse
import csv
import json
from pathlib import Path
from typing import cast

import mteb
import torch
from mteb import BenchmarkResults
from sentence_transformers import SentenceTransformer
from sentence_transformers.base.modules.transformer import Transformer

BENCHMARK = "MTEB(eng, v2)"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--output-dir", default="mteb-results")
    parser.add_argument("--device", default="cuda", help="cuda, mps or cpu")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Only the text side is needed.
    model =SentenceTransformer(
        args.model,
        device=args.device,
        model_kwargs={
            "torch_dtype": torch.bfloat16 if args.device == "cuda" else torch.float32,
        },
        config_kwargs={
            "vision_config": None,
            "audio_config": None,
        },
    )

    # Otherwise a text that's just a link gets fetched as a video, image or audio file.
    transformer = cast(Transformer, model[0])
    for modality in transformer.modalities:
        if modality != "text":
            del transformer.modality_config[modality]
    transformer.input_formatter.supported_modalities = transformer.modalities

    # MTEB calls this task type STS.
    if "SentenceSimilarity" in model.prompts:
        model.prompts.setdefault(
            "STS",
            model.prompts["SentenceSimilarity"],
        )

    benchmark = mteb.get_benchmark(BENCHMARK)

    cache =mteb.ResultCache(str(output_dir / "cache"))

    print(f"Evaluating {args.model} on {BENCHMARK}")
    print(f"Tasks: {len(benchmark.tasks)}")

    results = mteb.evaluate(
        model,
        tasks=benchmark,
        encode_kwargs={
            "batch_size": args.batch_size,
        },
        cache=cache,
        overwrite_strategy="only-missing",
    )

    scores ={
        result.task_name: 100 * result.get_score()
        for result in results.task_results
    }

    summary =benchmark.get_score(
        BenchmarkResults(model_results=[results]),
        raise_error=True,
    )[results.model_name]

    rows = []

    for task in benchmark.tasks:
        name = task.metadata.name
        score = scores.get(name)

        rows.append({
            "Task": name,
            "Score": round(score, 2) if score is not None else None,
        })

    with open(output_dir / "task_scores.csv", "w", newline="") as file:
        writer = csv.DictWriter(
            file,
            fieldnames=["Task", "Score"],
        )
        writer.writeheader()
        writer.writerows(rows)

    with open(output_dir / "summary.json", "w") as file:
        json.dump(
            {
                "model": results.model_name,
                "benchmark": BENCHMARK,
                "summary": summary,
            },
            file,
            indent=2,
        )

    print(f"\n{args.model} ({BENCHMARK})")
    print("| Task | Score |")
    print("|---|---:|")

    for row in rows:
        score = row["Score"]
        value = f"{score:.2f}" if score is not None else "N/A"
        print(f"| {row['Task']} | {value} |")

    print("\nAggregate scores:")

    for name, score in summary.items():
        if name == "Rank":
            continue

        value = f"{score * 100:.2f}" if score is not None else "N/A"
        print(f"| {name} | {value} |")

    print(f"\nResults saved to: {output_dir.resolve()}")


if __name__ == "__main__":
    main()