tools/discovery.py

"""Discovery tools for finding similar datasets and suggesting ML tasks."""

from typing import Optional, List, Dict, Any
from utils.hf_client import get_client
from utils.formatting import format_similar_datasets, format_task_suggestions, format_comparison


# Common ML task patterns based on column names and types
TASK_PATTERNS = {
    "text-classification": {
        "columns": ["text", "label", "sentence", "review", "comment", "content"],
        "name": "Text Classification",
        "target_hints": ["label", "class", "category", "sentiment", "target"]
    },
    "question-answering": {
        "columns": ["question", "answer", "context", "response"],
        "name": "Question Answering",
        "target_hints": ["answer", "response"]
    },
    "summarization": {
        "columns": ["article", "summary", "document", "highlights", "abstract"],
        "name": "Text Summarization",
        "target_hints": ["summary", "highlights", "abstract"]
    },
    "translation": {
        "columns": ["source", "target", "en", "de", "fr", "es", "translation"],
        "name": "Machine Translation",
        "target_hints": ["target", "translation"]
    },
    "image-classification": {
        "columns": ["image", "label", "img", "photo"],
        "name": "Image Classification",
        "target_hints": ["label", "class", "category"]
    },
    "named-entity-recognition": {
        "columns": ["tokens", "ner_tags", "tags", "entities"],
        "name": "Named Entity Recognition",
        "target_hints": ["ner_tags", "tags", "entities", "labels"]
    },
    "token-classification": {
        "columns": ["tokens", "labels", "tags", "pos_tags"],
        "name": "Token Classification",
        "target_hints": ["labels", "tags"]
    },
    "text-generation": {
        "columns": ["prompt", "completion", "input", "output", "instruction"],
        "name": "Text Generation / Instruction Following",
        "target_hints": ["completion", "output", "response"]
    },
    "tabular-classification": {
        "columns": ["target", "label", "class"],
        "name": "Tabular Classification",
        "target_hints": ["target", "label", "class", "y"]
    },
    "tabular-regression": {
        "columns": ["target", "price", "value", "score", "rating"],
        "name": "Tabular Regression",
        "target_hints": ["target", "price", "value", "score", "rating"]
    }
}


def find_similar(
    dataset_id: str,
    top_k: int = 5
) -> str:
    """
    Find datasets similar to a given dataset based on tags, domain, and structure.

    Use this tool to discover alternative or complementary datasets for your task.
    Similarity is based on shared tags, similar column structures, and domain overlap.

    Args:
        dataset_id: The dataset to find similar datasets for (e.g., "imdb", "squad")
        top_k: Number of similar datasets to return (1-10, default: 5)

    Returns:
        List of similar datasets with:
        - Dataset ID and download count
        - Similarity score (0-1)
        - Reason for similarity (shared tags, similar structure, etc.)

    How similarity is computed:
        - Tag overlap (same task categories, languages, domains)
        - Similar column names and structures
        - Same author/organization
        - Related task types
    """
    top_k = max(1, min(10, top_k))

    client = get_client()

    # Get info about the source dataset
    source_info = client.get_dataset_info(dataset_id)
    if "error" in source_info:
        return f"Error: Could not fetch info for dataset '{dataset_id}': {source_info.get('error')}"

    source_tags = set(source_info.get('tags', []))

    # Get schema for structure comparison
    source_schema = client.get_schema(dataset_id)
    source_columns = set(source_schema.get('columns', [])) if "error" not in source_schema else set()

    # Extract key tags for search
    search_terms = []
    for tag in source_tags:
        if ':' in tag:
            # Task category tags like "task_categories:text-classification"
            if tag.startswith('task_categories:'):
                search_terms.append(tag.split(':')[1])
            elif tag.startswith('language:'):
                search_terms.append(tag.split(':')[1])
        elif len(tag) > 2:
            search_terms.append(tag)

    # Search for candidates
    candidates = []
    for term in search_terms[:3]:  # Use top 3 terms
        results = client.search_datasets(term, limit=20)
        candidates.extend(results)

    # Remove duplicates and source dataset
    seen = {dataset_id}
    unique_candidates = []
    for ds in candidates:
        if ds['id'] not in seen:
            seen.add(ds['id'])
            unique_candidates.append(ds)

    # Score candidates
    scored = []
    for ds in unique_candidates[:30]:  # Limit processing
        try:
            ds_info = client.get_dataset_info(ds['id'])
            ds_tags = set(ds_info.get('tags', []))

            # Compute similarity score
            tag_overlap = len(source_tags & ds_tags)
            tag_score = tag_overlap / max(len(source_tags), 1)

            # Check column similarity
            ds_schema = client.get_schema(ds['id'])
            ds_columns = set(ds_schema.get('columns', [])) if "error" not in ds_schema else set()
            col_overlap = len(source_columns & ds_columns)
            col_score = col_overlap / max(len(source_columns), 1) if source_columns else 0

            # Combined score
            similarity = (tag_score * 0.6) + (col_score * 0.4)

            # Determine reason
            reasons = []
            if tag_overlap > 0:
                common_tags = list(source_tags & ds_tags)[:3]
                reasons.append(f"Shared tags: {', '.join(common_tags)}")
            if col_overlap > 0:
                common_cols = list(source_columns & ds_columns)[:3]
                reasons.append(f"Similar columns: {', '.join(common_cols)}")
            if ds_info.get('author') == source_info.get('author'):
                reasons.append("Same author")
                similarity += 0.1

            if similarity > 0.1:
                scored.append({
                    "id": ds['id'],
                    "downloads": ds.get('downloads', 0),
                    "similarity_score": min(1.0, similarity),
                    "reason": "; ".join(reasons) if reasons else "Related domain"
                })
        except Exception:
            continue

    # Sort by similarity and return top_k
    scored.sort(key=lambda x: x['similarity_score'], reverse=True)
    return format_similar_datasets(scored[:top_k])


def suggest_tasks(dataset_id: str) -> str:
    """
    Analyze a dataset and suggest suitable machine learning tasks.

    Use this tool to discover what ML tasks a dataset could be used for,
    based on its column structure, data types, and metadata.

    Args:
        dataset_id: The dataset to analyze (e.g., "imdb", "squad", "cnn_dailymail")

    Returns:
        List of suggested ML tasks with:
        - Task name and confidence level (high/medium/low)
        - Reasoning for the suggestion
        - Recommended target column
        - Recommended feature columns

    Task types detected:
        - Text Classification (sentiment, topic, intent)
        - Question Answering
        - Summarization
        - Translation
        - Image Classification
        - Named Entity Recognition
        - Token Classification
        - Text Generation
        - Tabular Classification/Regression
    """
    client = get_client()

    # Get schema
    schema = client.get_schema(dataset_id)
    if "error" in schema:
        return format_task_suggestions({"error": f"Could not load schema: {schema['error']}"})

    columns = [c.lower() for c in schema.get('columns', [])]
    features = schema.get('features', {})

    # Get dataset info for tags
    info = client.get_dataset_info(dataset_id)
    tags = [t.lower() for t in info.get('tags', [])] if "error" not in info else []

    suggestions: List[Dict[str, Any]] = []

    for task_id, pattern in TASK_PATTERNS.items():
        # Check column name matches
        pattern_cols = [c.lower() for c in pattern['columns']]
        matching_cols = [c for c in columns if any(p in c for p in pattern_cols)]

        # Check tag matches
        tag_match = any(task_id in t for t in tags)

        if matching_cols or tag_match:
            # Determine confidence
            if tag_match and len(matching_cols) >= 2:
                confidence = "high"
            elif tag_match or len(matching_cols) >= 2:
                confidence = "medium"
            else:
                confidence = "low"

            # Find target column
            target_hints = [h.lower() for h in pattern['target_hints']]
            target_col = None
            for col in columns:
                if any(hint in col for hint in target_hints):
                    target_col = col
                    break

            # Feature columns (all except target)
            feature_cols = [c for c in columns if c != target_col][:5]

            # Build reason
            reasons = []
            if matching_cols:
                reasons.append(f"Found columns: {', '.join(matching_cols[:3])}")
            if tag_match:
                reasons.append("Dataset tags indicate this task")

            suggestions.append({
                "name": pattern['name'],
                "confidence": confidence,
                "reason": "; ".join(reasons),
                "target_column": target_col,
                "feature_columns": feature_cols
            })

    # Sort by confidence
    confidence_order = {"high": 0, "medium": 1, "low": 2}
    suggestions.sort(key=lambda x: confidence_order.get(x['confidence'], 3))

    if not suggestions:
        # Generic suggestion based on column types
        has_text = any('string' in str(features.get(c, '')).lower() for c in schema.get('columns', []))
        has_numeric = any('int' in str(features.get(c, '')).lower() or 'float' in str(features.get(c, '')).lower()
                         for c in schema.get('columns', []))

        if has_text:
            suggestions.append({
                "name": "Text Analysis (Generic)",
                "confidence": "low",
                "reason": "Dataset contains text columns",
                "target_column": None,
                "feature_columns": columns[:5]
            })
        if has_numeric:
            suggestions.append({
                "name": "Regression/Classification (Generic)",
                "confidence": "low",
                "reason": "Dataset contains numeric columns",
                "target_column": columns[-1] if columns else None,
                "feature_columns": columns[:-1] if len(columns) > 1 else columns
            })

    return format_task_suggestions({
        "dataset_id": dataset_id,
        "tasks": suggestions[:5]  # Return top 5 suggestions
    })


def compare_datasets(
    dataset_a: str,
    dataset_b: str
) -> str:
    """
    Compare two datasets side-by-side to understand their differences.

    Use this tool when deciding between similar datasets or understanding
    how datasets differ in structure, size, and content.

    Args:
        dataset_a: First dataset ID to compare (e.g., "imdb")
        dataset_b: Second dataset ID to compare (e.g., "rotten_tomatoes")

    Returns:
        Comparison table showing:
        - Download and like counts
        - Number of columns
        - Column names (common and unique)
        - License information
        - Tags comparison
        - Data types comparison

    Use cases:
        - Choosing between similar datasets for a task
        - Understanding dataset versions or variants
        - Finding complementary datasets
    """
    client = get_client()

    # Get info for both datasets
    info_a = client.get_dataset_info(dataset_a)
    info_b = client.get_dataset_info(dataset_b)

    if "error" in info_a:
        return f"Error loading dataset A ({dataset_a}): {info_a.get('error')}"
    if "error" in info_b:
        return f"Error loading dataset B ({dataset_b}): {info_b.get('error')}"

    # Get schemas
    schema_a = client.get_schema(dataset_a)
    schema_b = client.get_schema(dataset_b)

    cols_a = set(schema_a.get('columns', [])) if "error" not in schema_a else set()
    cols_b = set(schema_b.get('columns', [])) if "error" not in schema_b else set()

    comparison = {
        "dataset_a": dataset_a,
        "dataset_b": dataset_b,
        "comparison": {
            "Downloads": {
                "a": f"{info_a.get('downloads', 0):,}",
                "b": f"{info_b.get('downloads', 0):,}"
            },
            "Likes": {
                "a": str(info_a.get('likes', 0)),
                "b": str(info_b.get('likes', 0))
            },
            "License": {
                "a": info_a.get('license') or "N/A",
                "b": info_b.get('license') or "N/A"
            },
            "Columns": {
                "a": str(len(cols_a)),
                "b": str(len(cols_b))
            },
            "Author": {
                "a": info_a.get('author') or "N/A",
                "b": info_b.get('author') or "N/A"
            }
        },
        "common_columns": list(cols_a & cols_b),
        "unique_to_a": list(cols_a - cols_b),
        "unique_to_b": list(cols_b - cols_a)
    }

    # Compare tags
    tags_a = set(info_a.get('tags', []))
    tags_b = set(info_b.get('tags', []))
    common_tags = tags_a & tags_b

    if common_tags:
        comparison["comparison"]["Common Tags"] = {
            "a": str(len(common_tags)),
            "b": ", ".join(list(common_tags)[:3])
        }

    return format_comparison(comparison)