# Based on: https://github.com/allenai/s2-folks/tree/main/examples/python
# 
# 

import os
import re

FORMAT_TEXT = 'Text'
FORMAT_DICT = 'Dict'
FORMAT_JSON = 'JSON'

FORMATS = (FORMAT_TEXT, FORMAT_DICT)

DS_WORMNEUROATLAS = 'WormNeuroAtlas'

import wormneuroatlas as wa
import pprint

pp = pprint.PrettyPrinter(depth=6)


def query_data_source(text, source, format):

    if source == DS_WORMNEUROATLAS:

        atlas = wa.NeuroAtlas()
        ds, ss = atlas.everything_about(text, return_values=True, return_text=True, print_text=False)

        if format==FORMAT_TEXT:
            info = '''Extracting info on neuron %s from %s...
'''%(text, source)
            info += """```%s```"""%ss
        elif format==FORMAT_DICT:
            info = """```
%s```"""%pp.pformat(ds)
        elif format==FORMAT_JSON:
            import json
            jj =  json.dumps(ds)
            info = """```
%s```"""%pp.pformat(jj)

    return info