import logging
import pathlib

from IPython.display import display
import pandas as pd
import rdflib

ENGINE = 'postgresql://postgres@/glottolog3'

QUERY = '''
SELECT
  l.id AS glottocode,
  l.name,
  ll.level,
  ll.category,
  substring(wikidata_link->>'url' FROM '/([^/]+)$') AS qid,
  substring(wikipedia_link->>'url' FROM '/([^/]+)$') AS title,
  i.name AS iso639_3
FROM language AS l
JOIN languoid AS ll USING (pk)
CROSS JOIN jsonb_path_query(l.jsondata::jsonb,
  '$.links[*] ? (@.url starts with "https://www.wikidata.org/entity/")') AS wikidata_link
LEFT JOIN jsonb_path_query(l.jsondata::jsonb,
  '$.links[*] ? (@.url starts with "https://en.wikipedia.org/wiki/")') AS wikipedia_link ON TRUE
LEFT JOIN (
  languageidentifier AS li
  JOIN identifier AS i
  ON li.identifier_pk = i.pk AND i.type = 'iso639-3'
) ON li.language_pk = l.pk
ORDER BY l.id
'''.strip()


logging.basicConfig(format='[%(levelname)s@%(name)s] %(message)s', level=logging.INFO)
logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)


gf = pd.read_sql_query(QUERY, ENGINE, index_col='glottocode', dtype='string')

gf.info(memory_usage='deep')
assert gf.index.is_unique
assert gf.index.is_monotonic_increasing
gf.head()

%%time

ENDPOINT = 'https://query.wikidata.org/sparql'

class SCHEMA(rdflib.SDO):
    """https://github.com/RDFLib/rdflib/issues/1120"""
    _NS = rdflib.Namespace(rdflib.SDO._NS.replace('https://', 'http://'))

PREFIXES = {'schema': SCHEMA}

SPARQL_QUERY = '''
SELECT
  ?glottocode
  (strafter(str(?languoid), str(wd:)) AS ?qid)
  (?languoidLabel AS ?name)
  (strafter(str(?siteLink), "https://en.wikipedia.org/wiki/") AS ?title)
WHERE {
  ?languoid wdt:P1394 ?glottocode.
  FILTER (REGEX(?glottocode, "^[a-z0-9]{4}[0-9]{4}$")).
  OPTIONAL {
    ?siteLink schema:about ?languoid;
              schema:inLanguage "en";
              schema:isPartOf <https://en.wikipedia.org/>.
  }
  SERVICE wikibase:label {
    bd:serviceParam wikibase:language "en".
    ?languoid rdfs:label ?languoidLabel.
  }
}
ORDER BY
  ?glottocode
  xsd:integer(strafter(str(?languoid), str(wd:Q)))
'''.strip()

CSV_PATH = pathlib.Path('wikidata.csv')


def open_sparql_graph(endpoint=ENDPOINT, *, prefixes=PREFIXES):
    logging.info('endpoint: %r', endpoint)
    graph = rdflib.ConjunctiveGraph('SPARQLStore')
    graph.open(endpoint)
    logging.info('graph: %s', graph)
    logging.info('prefixes: %r', prefixes)
    for prefix, namespace in prefixes.items():
        graph.namespace_manager.bind(prefix, namespace, replace=True)
    logging.debug('namespaces: %r', list(graph.namespaces()))
    return graph


def iterrows(query, *, prefixes=PREFIXES,
             limit=None, verbose: bool = False,
             per_request: int = 100_000):
    if limit is None:
        limit = float('inf')

    graph = open_sparql_graph()

    offset = 0
    while offset < limit:
        request_limit = min(limit - offset, per_request)
        request_query = (f'{query}\n'
                         f'OFFSET {offset:d}\n'
                         f'LIMIT {request_limit:d}')
        logging.info("graph.query('''\n%s\n''')", request_query)
        result = graph.query(request_query)
        if not offset:
            yield (v.toPython().removeprefix('?') for v in result.vars)
        for n, values in enumerate(result, 1):
            yield (v.toPython() if v is not None else None for v in values)
        if n < request_limit:
            return
        offset += n


def read_sparql_query(query, *, limit=None, **kwargs):
    rows = iterrows(query, limit=limit)
    columns = list(next(rows))
    return pd.DataFrame.from_records(rows, columns=columns, **kwargs)


def read_glottocodes(*, path=CSV_PATH, encoding='utf-8'):
    if not path.exists():
        df = read_sparql_query(SPARQL_QUERY, index='glottocode').astype('string')
        df.to_csv(path, encoding=encoding)
    return pd.read_csv(path, index_col='glottocode', encoding=encoding).astype('string')


wf = read_glottocodes()

wf.info(memory_usage='deep')
assert wf.index.is_monotonic_increasing
wf.head(10)

KEEP_ALL = False

wf[wf.index.duplicated(keep=KEEP_ALL)]

df = gf.join(wf, on='glottocode', rsuffix='_wd').drop(['category', 'name_wd', 'iso639_3'], axis='columns')

df.info(memory_usage='deep')
df.head()

df[df.index.duplicated(keep=KEEP_ALL)]

mismatch = df.loc[df['qid'] != df['qid_wd'], ['name', 'level', 'title', 'title_wd', 'qid', 'qid_wd']]

mismatch.info(memory_usage='deep')
mismatch[mismatch.index.duplicated(keep=KEEP_ALL)]

with pd.option_context('display.max_rows', 150):
    display(mismatch)