#!/usr/bin/env python

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

from __future__ import division

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

import os
import datetime
import logging
from matgendb.creator import VaspToDbTaskDrone
from pymatgen.apps.borg.queen import BorgQueen
import multiprocessing
import json
import sys
import webbrowser

logger = logging.getLogger(__name__)


def get_settings(args):
    if args.config_file:
        with open(args.config_file) as f:
            return json.load(f)
    elif os.path.exists("db.json"):
        with open("db.json") as f:
            return json.load(f)
    else:
        return {"host": "localhost",
                "port": 27017,
                "database": "vasp",
                "admin_user": None,
                "admin_password": None,
                "readonly_user": None,
                "readonly_password": None,
                "collection": "tasks",
                "aliases_config": None}


def init_db(args):
    if args.config_file:
        d = [("host", "localhost"),
             ("port", 27017),
             ("database", "vasp"),
             ("admin_user", None),
             ("admin_password", None),
             ("readonly_user", None),
             ("readonly_password", None),
             ("collection", "tasks"),
             ("aliases_config", None)]
        doc = {}
        for k, v in d:
            val = raw_input("Enter {} (default: {}) : ".format(k, v))
            doc[k] = val if val else v

        with open(args.config_file, "w") as f:
            json.dump(doc, f, indent=4, sort_keys=True)


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)

    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)
    ncpus = multiprocessing.cpu_count()
    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 matgendb.query_engine import QueryEngine
    from prettytable import PrettyTable
    d = get_settings(args)
    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)

    t = PrettyTable(args.properties)
    t.float_format = "4.4"
    for r in qe.query(properties=args.properties, criteria=criteria):
        t.add_row([r[p] for p in args.properties])
    print t


def run_server(args):
    import os
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "matgendb.webui.settings")
    os.environ["MGDB_CONFIG"] = json.dumps(get_settings(args))
    from django.core.management import call_command
    from multiprocessing import Process
    p1 = Process(target=call_command, args=("runserver",))
    p1.start()
    import time
    time.sleep(2)
    webbrowser.open_new_tab("127.0.0.1:8000")
    p1.join()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="""
    Command line db insertion script.

    Author: Shyue Ping Ong
    Version: 2.0
    Last updated: Jun 21 2012""")

    subparsers = parser.add_subparsers()

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

    pinit.add_argument("-c", "--config", dest="config_file", type=str,
                       nargs='?',
                       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.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,
                        nargs='+', required=True,
                        help="Desired properties. E.g., pretty_formula, "
                             "task_id, energy...")

    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.set_defaults(func=run_server)

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