"""
usage:
    load_all_sites.py [options] DU_DIR

arguments:
    DU_DIR      Directory with OEDesignUnit files

options:
    -v                    Verbose output
    -d                    Debug mode, no multiprocessing
    --ncpu N              Number of CPUs to use [default: 0]
"""
import docopt
import multiprocessing
import os
import random
import sys
from yaspin import yaspin
from itertools import repeat

from spruce.utils.file_utils import get_design_unit_files_by_code

from oeclient_utils.exceptions import ValidationError, BadResponse

import mmdsclient as mmdscli
from openeye import oechem


def get_current_sites():
    msg = "Gathering existing sites"
    with yaspin(text=msg, color="cyan") as sp:
        session = mmdscli.get_session(profile=os.environ["MMDS_PROFILE"])
        filters = {'limit': 1000, 'fields': 'id,title'}
        existing = set()
        try:
            for i, site in enumerate(session.list_resources(mmdscli.Site, filters=filters)):
                existing.add(site.title)
                if i % 1000:
                    sp.text = "{} : {}".format(msg, i+1)
        except BadResponse as e:
            print(str(e))
        sp.ok("✔")
        return existing


def get_images(design_unit_file):
    interaction = None
    density = None
    svg_file = design_unit_file.replace("_DU__", "_receptor__").replace(".oedu", ".svg")
    if os.path.isfile(svg_file):
        interaction = svg_file
    svg_file2 = design_unit_file.replace("_DU__", "_density__").replace(".oedu", ".svg")
    if os.path.isfile(svg_file2):
        density = svg_file2
    return interaction, density


def add_sites(input_list, existing):
    du_files = []
    codes = set()
    for du in input_list:
        basename = os.path.basename(du)
        code = basename.split('_')[0]
        codes.add(code)
        if 'biounit' not in du:
            du_files.append(du)

    if len(codes) != 1:
        return [], '--\nError: Found more than one code in same list: {}\n'.format(codes)

    code = codes.pop()
    session = mmdscli.get_session(profile=os.environ["MMDS_PROFILE"])
    experiments = list(session.list_resources(mmdscli.Experiment, filters={'code': code}))
    if len(experiments) != 1:
        return [], "--\nError: Found {} experiments for code {}".format(len(experiments), code)

    expt_id = experiments[0].id
    out = "--\nExperiment: {} code: {}\n".format(expt_id, code)
    sites = set()
    du = oechem.OEDesignUnit()
    for du_file in du_files:
        if not oechem.OEReadDesignUnit(du_file, du):
            out += "  Unable to read DU: {}\n".format(du_file)
            continue
        if du.GetTitle() in existing:
            out += "  {} is an existing site, skipping.\n".format(du.GetTitle())
            continue
        out += "  DU: {}\n".format(du_file)
        interaction_svg, density_svg = get_images(du_file)
        try:
            site = mmdscli.Site.add(session, expt_id, du_file,
                                    interaction_image=interaction_svg,
                                    density_image=density_svg)
            sites.add(site.id)
            out += "  Added site: {}\n".format(site.id)
        except (ValidationError, BadResponse) as e:
            out += "  {}\n".format(str(e))
    out += "added {} sites to {}\n".format(len(sites), code)
    return list(sites), out


def run(args):
    return add_sites(args[0], args[1])


def main():
    args = docopt.docopt(__doc__)
    verbose = args['-v']
    debug = args['-d']
    ncpu = int(args['--ncpu'])
    if ncpu == 0:
        ncpu = multiprocessing.cpu_count()
    if ncpu > 8:
        ncpu = 8

    existing = get_current_sites()

    data_dir = args["DU_DIR"]
    with yaspin(text="Gathering input DU files by code", color="cyan") as sp:
        du_files = get_design_unit_files_by_code(data_dir)
        sp.text = sp.text + " : {}".format(len(du_files))
        sp.ok("✔")

    if len(du_files) == 0:
        print("No new DUs to process. Exiting.")
        return 0

    random.shuffle(du_files)

    logfile = open('load_all_sites.log', 'w')

    msg = "Loading OEDesignUnit file sets."
    count = 0
    total_sites = 0
    with yaspin(text=msg, color="cyan") as sp:
        if not debug:
            pool = multiprocessing.Pool(ncpu)
            for sites, out in pool.imap_unordered(run, zip(du_files, repeat(existing))):
                count += 1
                sp.text = "{} : {}".format(msg, count)
                logfile.write(out)
                total_sites += len(sites)
            pool.terminate()
        else:
            for dus in du_files:
                sites, out = add_sites(dus, existing)
                sp.text = "{} : {}".format(msg, count)
                logfile.write(out)
                total_sites += len(sites)
        sp.ok("✔")
    print("Loaded {} sites".format(total_sites))


if __name__ == "__main__":
    sys.exit(main())