#!/usr/bin/env python

"""
A helper script for many matgendb functions.
"""

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "1.2"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Dec 1, 2012"

import os
import datetime
import logging
import multiprocessing
import json
import sys
import webbrowser
import time
import argparse

from pymatgen.apps.borg.queen import BorgQueen

from matgendb.query_engine import QueryEngine
from matgendb.creator import VaspToDbTaskDrone
from matgendb.util import get_settings, DEFAULT_SETTINGS, MongoJSONEncoder

logger = logging.getLogger(__name__)


def init_db(args):
    d = DEFAULT_SETTINGS
    doc = {}
    print("Please supply the following configuration values")
    print("(press Enter if you want to accept the defaults)\n")
    for k, v in d:
        val = raw_input("Enter {} (default: {}) : ".format(k, v))
        doc[k] = val if val else v
    doc["port"] = int(doc["port"])  # enforce the port as an int
    with open(args.config_file, "w") as f:
        json.dump(doc, f, indent=4, sort_keys=True)
    print("\nConfiguration written to {}!".format(args.config_file))


def update_db(args):
    FORMAT = "%(relativeCreated)d msecs : %(message)s"

    if args.logfile:
        logging.basicConfig(level=logging.INFO, format=FORMAT,
                            filename=args.logfile[0])
    else:
        logging.basicConfig(level=logging.INFO, format=FORMAT)

    d = get_settings(args.config_file)

    logger.info("Db insertion started at {}.".format(datetime.datetime.now()))
    additional_fields = {"author": args.author, "tags": args.tag}
    drone = VaspToDbTaskDrone(
        host=d["host"], port=d["port"],  database=d["database"],
        user=d["admin_user"], password=d["admin_password"], parse_dos=False,
        collection=d["collection"], update_duplicates=args.force_update_dupes,
        additional_fields=additional_fields, mapi_key=args.mapi_key)
    ncpus = multiprocessing.cpu_count() if not args.ncpus else args.ncpus
    logger.info("Using {} cpus...".format(ncpus))
    queen = BorgQueen(drone, number_of_drones=ncpus)
    queen.parallel_assimilate(args.directory)
    tids = map(int, filter(lambda x: x, queen.get_data()))
    logger.info("Db upate completed at {}.".format(datetime.datetime.now()))
    logger.info("{} new task ids inserted.".format(len(tids)))


def query_db(args):
    from prettytable import PrettyTable
    d = get_settings(args.config_file)
    qe = QueryEngine(host=d["host"], port=d["port"], database=d["database"],
                     user=d["readonly_user"], password=d["readonly_password"],
                     collection=d["collection"],
                     aliases_config=d.get("aliases_config", None))
    criteria = None
    if args.criteria:
        try:
            criteria = json.loads(args.criteria)
        except ValueError:
            print("Criteria {} is not a valid JSON string!".format(
                args.criteria))
            sys.exit(-1)

    # TODO: document this 'feature' --dang 4/4/2013
    is_a_file = lambda s: len(s) == 1 and s[0].startswith(':')
    if is_a_file(args.properties):
        with open(args.properties[0][1:], 'rb') as f:
            props = [s.strip() for s in f]
    else:
        props = args.properties

    if args.dump_json:
        for r in qe.query(properties=props, criteria=criteria):
            print(json.dumps(r, cls=MongoJSONEncoder))
    else:
        t = PrettyTable(props)
        t.float_format = "4.4"
        for r in qe.query(properties=props, criteria=criteria):
            t.add_row([r[p] for p in props])
        print(t)


def run_server(args):
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "matgendb.webui.settings")
    os.environ["MGDB_CONFIG"] = json.dumps(get_settings(args.config_file))
    from django.core.management import call_command
    from multiprocessing import Process
    p1 = Process(target=call_command,
                 args=("runserver",  "{}:{}".format(args.host, args.port)))
    p1.start()
    if not args.nobrowser:
        time.sleep(2)
        webbrowser.open("http://{}:{}".format(args.host, args.port))
    p1.join()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="""
    mgdb is a complete command line db management script for pymatgen-db. It
    provides the facility to insert vasp runs, perform queries, and run a web
    server for exploring databases that you create. Type mgdb -h to see the
    various options.

    Author: Shyue Ping Ong
    Version: 3.0
    Last updated: Mar 23 2013""")

    subparsers = parser.add_subparsers()

    pinit = subparsers.add_parser("init", help="Initialization tools.")

    pinit.add_argument("-c", "--config", dest="config_file", type=str,
                       nargs='?', default="db.json",
                       help="Creates an db config file for the database. "
                            "Default filename is db.json.")
    pinit.set_defaults(func=init_db)

    pinsert = subparsers.add_parser("insert", help="Insert vasp runs.")

    pinsert.add_argument("directory", metavar="directory", type=str,
                         default=".", help="Root directory for runs.")
    pinsert.add_argument("-c", "--config", dest="config_file", type=str,
                         help="Config file to use. Generate one using mgdb "
                              "init --config filename.json if necessary. "
                              "Otherwise, the code searches for a db.json. If"
                              "none is found, an no-authentication "
                              "localhost:27017/vasp database and tasks "
                              "collection is assumed.")

    pinsert.add_argument("-l", "--logfile", dest="logfile", type=str,
                         help="File to log db insertion. Defaults to stdout.")
    pinsert.add_argument("-t", "--tag", dest="tag", type=str, nargs=1,
                         default=[],
                         help="Tag your runs for easier search."
                              " Accepts multiple tags")
    pinsert.add_argument("-f", "--force", dest="force_update_dupes",
                         action="store_true",
                         help="Force update duplicates. This forces the "
                              "analyzer to reanalyze already inserted data.")
    pinsert.add_argument("-a", "--author", dest="author", type=str, nargs=1,
                         default=None,
                         help="Enter a *unique* author field so that you can "
                              "trace back what you ran.")
    pinsert.add_argument("-n", "--ncpus", dest="ncpus", type=int,
                         default=None,
                         help="Number of CPUs to use in inserting. If "
                              "not specified, multiprocessing will use "
                              "the number of cpus detected.")

    pinsert.add_argument("-k", "--mapi_key", dest="mapi_key", type=str,
                         default=None,
                         help="Materials API key to use. If supplied, "
                              "the insertion process will attempt to "
                              "determine stability information using the "
                              "Materials REST API.")
    pinsert.set_defaults(func=update_db)

    pquery = subparsers.add_parser("query",
                                   help="Query tools. Requires the "
                                        "use of pretty_table.")

    pquery.add_argument("-c", "--config", dest="config_file", type=str,
                        help="Config file to use. Generate one using mgdb "
                             "init --config filename.json if necessary. "
                             "Otherwise, the code searches for a db.json. If"
                             "none is found, an no-authentication "
                             "localhost:27017/vasp database and tasks "
                             "collection is assumed.")

    pquery.add_argument("--crit", dest="criteria", type=str, default=None,
                        help="Query criteria in typical json format. E.g., "
                             "{\"task_id\": 1}.")

    pquery.add_argument("--props", dest="properties", type=str, default=[],
                        nargs='+', required=True,
                        help="Desired properties. Repeatable. E.g., pretty_formula, "
                             "task_id, energy...")

    pquery.add_argument("--dump", dest="dump_json", action='store_true', default=False,
                        help="Simply dump results to JSON instead of a tabular view")

    pquery.set_defaults(func=query_db)

    pserve = subparsers.add_parser("runserver",
                                   help="Run a server to the database.")

    pserve.add_argument("-c", "--config", dest="config_file", type=str,
                        help="Config file to use. Generate one using mgdb "
                             "init --config filename.json if necessary. "
                             "Otherwise, the code searches for a db.json. If"
                             "none is found, an no-authentication "
                             "localhost:27017/vasp database and tasks "
                             "collection is assumed.")

    pserve.add_argument("--nobrowser", dest="nobrowser", action="store_true",
                        help="Don't automatically open a browser window "
                             "(e.g. for shell-based servers)")

    pserve.add_argument("--port", dest="port", type=int, default=8000,
                        help="Port to run the server on (default: 8000)")

    pserve.add_argument("--host", dest="host", type=str, default="127.0.0.1",
                        help="Host to run the server on (default: 127.0.0.1)")

    pserve.set_defaults(func=run_server)

    args = parser.parse_args()
    args.func(args)
