"""Tiny helper for the TrialAutopsy open dataset. Python standard library only.

    import trialautopsy as ta
    rows = ta.load()                              # downloads https://trialautopsy.tech/data/trialautopsy.csv
    rows = ta.load("trialautopsy.csv")            # or read a local copy
    ta.counts(rows, "category")                   # Counter of labels
    ta.filter(rows, category="slow_accrual", sponsor="Pfizer")
    ta.similar(rows, "drug shipments spoiled")    # keyword similarity over the stated reasons

Labels are automatic and unverified. Records come from ClinicalTrials.gov.
"""
import csv
import io
import math
import re
import urllib.request
from collections import Counter

URL = "https://trialautopsy.tech/data/trialautopsy.csv"
_STOP = set("the and for was were with that this not due has have had are from but its been into than then they "
            "their there which will would could also after before during because study trial trials patient "
            "patients subject subjects".split())


def load(source=URL):
    """Rows as dicts. `source` is a path or an https URL."""
    if source.startswith("https://"):
        with urllib.request.urlopen(source, timeout=60) as r:
            text = r.read().decode("utf-8")
    else:
        with open(source, encoding="utf-8", newline="") as f:
            text = f.read()
    return list(csv.DictReader(io.StringIO(text)))


def counts(rows, field):
    return Counter(r[field] for r in rows)


def filter(rows, **where):
    """Rows whose fields equal the given values; a string value matches if it is contained, ignoring case."""
    def ok(r):
        return all(v.lower() in r.get(k, "").lower() if isinstance(v, str) else r.get(k) == v for k, v in where.items())
    return [r for r in rows if ok(r)]


def _tokens(text):
    return [re.sub(r"(ing|ed|es|s|ly)$", "", w) for w in re.findall(r"[a-z]{3,}", text.lower()) if w not in _STOP]


def similar(rows, text, k=10, min_score=0.2):
    """[(score, row)] for the stated reasons that share the most wording with `text` (TF-IDF cosine)."""
    docs = [_tokens(r["why_stopped"]) for r in rows]
    df = Counter(w for d in docs for w in set(d))
    n = len(rows)

    def vec(tokens):
        tf = Counter(tokens)
        v = {w: (1 + math.log(c)) * (math.log((n + 1) / (df.get(w, 0) + 1)) + 1) for w, c in tf.items()}
        norm = math.sqrt(sum(x * x for x in v.values())) or 1
        return {w: x / norm for w, x in v.items()}
    q = vec(_tokens(text))
    scored = []
    for r, d in zip(rows, docs):
        dv = vec(d)
        s = sum(x * dv.get(w, 0) for w, x in q.items())
        if s >= min_score:
            scored.append((s, r))
    scored.sort(key=lambda t: -t[0])
    return scored[:k]