processing/scorer.py

from __future__ import annotations

import re
from datetime import datetime, timezone

from .query_intent import QueryIntent

def _score_recency(dataset: dict) -> float:
    raw = dataset.get("last_updated", "")
    if not raw:
        return 0.3

    try:
        for fmt in ("%Y-%m-%dT%H:%M:%S", "%Y-%m-%d", "%Y-%m", "%Y"):
            try:
                dt = datetime.strptime(str(raw)[:10], fmt[:len(fmt)])
                break
            except ValueError:
                continue
        else:
            return 0.3

        age_years = (datetime.now() - dt).days / 365.25
        return max(0.0, 1.0 - (age_years / 5.0))
    except Exception:
        return 0.3

_LICENSE_SCORES: dict[str, float] = {
    "public domain": 1.0,
    "cc0":           1.0,
    "cc-by":         0.9,
    "cc by":         0.9,
    "mit":           0.9,
    "apache":        0.85,
    "open":          0.8,
    "cc-by-nc":      0.6,
    "cc by-nc":      0.6,
    "research only": 0.5,
    "non-commercial":0.5,
    "proprietary":   0.2,
    "unknown":       0.4,
    "":              0.4,
}

def _score_license(dataset: dict) -> float:
    lic = str(dataset.get("license", "")).lower().strip()
    for key, score in _LICENSE_SCORES.items():
        if key and key in lic:
            return score
    return 0.4

def _score_size(dataset: dict) -> float:
    raw = str(dataset.get("size_estimate", "")).lower()
    if not raw or raw in ("unknown", "n/a", ""):
        return 0.4

    match = re.search(r"([\d,.]+)\s*(kb|mb|gb|tb|rows?|samples?|k|m)", raw)
    if not match:
        return 0.4

    num = float(match.group(1).replace(",", ""))
    unit = match.group(2)

    mb_map = {
        "kb": num / 1024,
        "mb": num,
        "gb": num * 1024,
        "tb": num * 1024 * 1024,
        "rows": num / 10_000,
        "row":  num / 10_000,
        "samples": num / 10_000,
        "sample":  num / 10_000,
        "k":  num / 10,
        "m":  num * 100,
    }
    size_mb = mb_map.get(unit, num)

    if size_mb < 0.1:
        return 0.1
    elif size_mb < 1:
        return 0.3
    elif size_mb < 10:
        return 0.5
    elif size_mb < 500:
        return 0.8
    else:
        return 1.0

_FORMAT_SCORES: dict[str, float] = {
    "csv":     1.0,
    "parquet": 0.95,
    "json":    0.9,
    "jsonl":   0.9,
    "tsv":     0.85,
    "xlsx":    0.7,
    "hdf5":    0.75,
    "h5":      0.75,
    "npy":     0.7,
    "zip":     0.6,
    "tar":     0.6,
    "api":     0.8,
}

def _score_format(dataset: dict) -> float:
    fmt = str(dataset.get("format", "")).lower()
    for key, score in _FORMAT_SCORES.items():
        if key in fmt:
            return score
    return 0.5

def _passes_constraint(dataset: dict, key: str, value: object) -> bool:
    text = " ".join([
        str(dataset.get("name", "")),
        str(dataset.get("description", "")),
        str(dataset.get("suitability_notes", "")),
    ]).lower()

    if key == "class_balance":
        if value == "imbalanced":
            return bool(re.search(
                r"\bimbalanced?\b|\bclass[\s_-]?imbalance\b|\bskewed\b"
                r"|\brare[\s_-]?class\b|\bminority[\s_-]?class\b"
                r"|\bunequal[\s_-]?class\b",
                text
            ))
        if value == "balanced":
            return bool(re.search(r"\bbalanced\b|\bequal[\s_-]?class\b", text))

    elif key == "labeled":
        if value is True:
            return bool(re.search(
                r"\blabell?ed\b|\bannotated\b|\bground[\s_-]?truth\b", text
            ))
        if value is False:
            return bool(re.search(r"\bunlabell?ed\b|\bunsupervised\b", text))

    elif key == "modality":
        modality_patterns = {
            "time_series": r"\btime[\s_-]?series\b|\btemporal\b|\bsequential\b",
            "image":       r"\bimage[s]?\b|\bvision\b|\bphoto[s]?\b|\bvisual\b",
            "text":        r"\btext\b|\bnlp\b|\bcorpus\b|\bsentence[s]?\b",
            "audio":       r"\baudio\b|\bspeech\b|\bsound\b",
            "tabular":     r"\btabular\b|\bcsv\b|\bspreadsheet\b|\bstructured\b",
            "graph":       r"\bgraph\b|\bnetwork\b|\bnode[s]?\b|\bedge[s]?\b",
        }
        pattern = modality_patterns.get(str(value), "")
        return bool(re.search(pattern, text)) if pattern else True

    elif key == "license_type":
        lic = str(dataset.get("license", "")).lower()
        if value == "commercial_friendly":
            return not bool(re.search(r"\bnon[\s_-]?commercial\b|\bnc\b|\bresearch[\s_-]?only\b", lic))
        if value in ("open", "public_domain"):
            return bool(re.search(
                r"\bopen\b|\bcc0\b|\bpublic[\s_-]?domain\b|\bmit\b|\bapache\b|\bcc[\s_-]?by\b", lic
            ))

    elif key == "min_size":
        return _score_size(dataset) >= 0.6

    elif key == "max_size":
        return _score_size(dataset) <= 0.5

    elif key == "min_recency":
        return _score_recency(dataset) >= 0.6

    elif key == "is_benchmark":
        return bool(re.search(r"\bbenchmark\b|\bsota\b|\bleaderboard\b", text))

    elif key == "multimodal":
        return bool(re.search(r"\bmultimodal\b|\bmulti[\s_-]?modal\b", text))

    return True

def hard_filter(
    datasets: list[dict],
    intent: QueryIntent,
) -> tuple[list[dict], list[dict]]:
    if not intent.hard_constraints:
        return datasets, []

    passed, rejected = [], []
    for ds in datasets:
        if all(
            _passes_constraint(ds, key, value)
            for key, value in intent.hard_constraints.items()
        ):
            passed.append(ds)
        else:
            rejected.append(ds)

    return passed, rejected

_BASE_WEIGHTS: dict[str, float] = {
    "llm_score":           0.80,
    "semantic_similarity": 0.05,
    "recency":             0.05,
    "license_openness":    0.05,
    "size_score":          0.025,
    "format_match":        0.025,
}

def resolve_weights(
    intent: QueryIntent,
    base_weights: dict[str, float] | None = None,
) -> dict[str, float]:
    weights = dict(base_weights or _BASE_WEIGHTS)

    for dim, boost in intent.weight_boosts.items():
        if dim in weights:
            weights[dim] = weights[dim] + boost

    weights = {k: max(0.0, min(1.0, v)) for k, v in weights.items()}

    total = sum(weights.values())
    return {k: v / total for k, v in weights.items()}

def score_results(
    datasets: list[dict],
    query: str,
    intent: QueryIntent,
    semantic_scores: dict[str, float] | None = None,
    llm_scores: dict[str, float] | None = None,
    base_weights: dict[str, float] | None = None,
) -> list[dict]:
    weights = resolve_weights(intent, base_weights)
    semantic_scores = semantic_scores or {}
    llm_scores = llm_scores or {}

    scored = []
    for ds in datasets:
        name = ds.get("name", "")

        dim_scores = {
            "llm_score":           llm_scores.get(name, 0.5),
            "semantic_similarity": semantic_scores.get(name, 0.5),
            "recency":             _score_recency(ds),
            "license_openness":    _score_license(ds),
            "size_score":          _score_size(ds),
            "format_match":        _score_format(ds),
        }

        relevance_score = sum(
            dim_scores[dim] * weights[dim]
            for dim in weights
        )

        scored.append({
            **ds,
            "relevance_score": round(relevance_score, 4),
            "dim_scores":      {k: round(v, 3) for k, v in dim_scores.items()},
            "active_weights":  {k: round(v, 3) for k, v in weights.items()},
            "active_constraints": intent.hard_constraints,
        })

    return sorted(scored, key=lambda x: x["relevance_score"], reverse=True)