data_analysis.py

"""
Phase 1 — Data Analysis
========================
Generate comprehensive dataset statistics before training.
Outputs a JSON report and prints a human-readable summary.
"""
import json
import statistics
from collections import Counter
from tqdm import tqdm

import config
import utils


def run():
    coll = utils.get_training_collection()
    report = {}

    # ── 1. Total Articles ────────────────────────────────────────────────
    total = coll.count_documents({})
    report["total_articles"] = total
    print(f"Total articles: {total:,}")

    # ── 2. Articles Missing Labels / Title / Abstract ────────────────────
    no_mesh = coll.count_documents({"mesh_ids": {"$size": 0}})
    no_abstract = coll.count_documents({"abstract": {"$in": ["", None]}})
    no_title = coll.count_documents({"title": {"$in": ["", None]}})
    report["articles_missing_labels"] = no_mesh
    report["articles_without_abstract"] = no_abstract
    report["articles_without_title"] = no_title
    print(f"Articles missing labels: {no_mesh}")
    print(f"Articles without abstract: {no_abstract}")
    print(f"Articles without title: {no_title}")

    # ── 3. Unique MeSH Terms ─────────────────────────────────────────────
    pipeline_unique = [
        {"$unwind": "$mesh_ids"},
        {"$group": {"_id": "$mesh_ids"}},
        {"$count": "count"},
    ]
    res = list(coll.aggregate(pipeline_unique))
    unique_terms = res[0]["count"] if res else 0
    report["unique_mesh_terms"] = unique_terms
    print(f"Unique MeSH terms: {unique_terms:,}")

    # ── 4. Average MeSH Terms per Article ────────────────────────────────
    pipeline_avg = [
        {"$project": {"n": {"$size": "$mesh_terms"}}},
        {"$group": {"_id": None, "avg": {"$avg": "$n"}, "min": {"$min": "$n"}, "max": {"$max": "$n"}}},
    ]
    res = list(coll.aggregate(pipeline_avg))
    if res:
        report["avg_mesh_terms_per_article"] = round(res[0]["avg"], 2)
        report["min_mesh_terms_per_article"] = res[0]["min"]
        report["max_mesh_terms_per_article"] = res[0]["max"]
        print(f"Avg MeSH terms/article: {res[0]['avg']:.2f}  (min={res[0]['min']}, max={res[0]['max']})")

    # ── 5. Major vs Minor Distribution ───────────────────────────────────
    pipeline_mm = [
        {"$unwind": "$mesh_terms"},
        {"$group": {"_id": "$mesh_terms.is_major_topic", "count": {"$sum": 1}}},
    ]
    major, minor = 0, 0
    for r in coll.aggregate(pipeline_mm):
        if r["_id"] is True:
            major = r["count"]
        else:
            minor = r["count"]
    report["major_topics_count"] = major
    report["minor_topics_count"] = minor
    report["major_pct"] = round(100 * major / (major + minor), 2) if (major + minor) else 0
    print(f"Major topics: {major:,} ({report['major_pct']}%)")
    print(f"Minor topics: {minor:,} ({100 - report['major_pct']}%)")

    # ── 6. Top 100 Most Frequent MeSH Terms ──────────────────────────────
    pipeline_top = [
        {"$unwind": "$mesh_terms"},
        {"$group": {"_id": "$mesh_terms.mesh_id", "name": {"$first": "$mesh_terms.mesh_name"}, "count": {"$sum": 1}}},
        {"$sort": {"count": -1}},
        {"$limit": 100},
    ]
    top_terms = []
    for r in coll.aggregate(pipeline_top):
        top_terms.append({"mesh_id": r["_id"], "mesh_name": r["name"], "count": r["count"]})
    report["top_100_mesh_terms"] = top_terms
    print(f"\nTop 10 MeSH terms:")
    for i, t in enumerate(top_terms[:10], 1):
        print(f"  {i:2d}. {t['mesh_id']:<12} {t['mesh_name']:<45} {t['count']:>8,}")

    # ── 7. Abstract Length Distribution (character-level) ────────────────
    print("\nComputing abstract length distribution (sampling 50k docs)...")
    sample_cursor = coll.find(
        {"abstract": {"$ne": ""}},
        {"abstract": 1},
    ).limit(50_000)
    lengths = [len(doc.get("abstract", "")) for doc in sample_cursor]

    if lengths:
        report["abstract_length_distribution"] = {
            "sample_size": len(lengths),
            "mean": round(statistics.mean(lengths), 1),
            "median": round(statistics.median(lengths), 1),
            "stdev": round(statistics.stdev(lengths), 1) if len(lengths) > 1 else 0,
            "min": min(lengths),
            "max": max(lengths),
            "p25": round(sorted(lengths)[len(lengths) // 4], 1),
            "p75": round(sorted(lengths)[3 * len(lengths) // 4], 1),
        }
        dist = report["abstract_length_distribution"]
        print(f"  Mean: {dist['mean']:,.0f}  Median: {dist['median']:,.0f}  "
              f"Stdev: {dist['stdev']:,.0f}  Min: {dist['min']}  Max: {dist['max']}")

    # ── Save Report ──────────────────────────────────────────────────────
    utils.save_json(report, config.ANALYSIS_REPORT_PATH)
    print(f"\n[OK] Full report saved to {config.ANALYSIS_REPORT_PATH}")
    return report


if __name__ == "__main__":
    run()