#!python

from PW_explorer.export import *
from PW_explorer.pwe_helper import (
    load_from_temp_pickle,
    set_current_project_name,
    get_current_project_name,
    get_save_folder,
    save_to_txt_file,
)

import argparse
import os


def __main__():
    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--project_name", type=str, help="provide session/project name used while parsing")
    parser.add_argument("-s", "--schema", action="store_true", help="generate sql schemas", default=False)
    parser.add_argument("-sql", action="store_true", help="include if you want to export a sql db", default=False)
    parser.add_argument("-csv", action="store_true", help="include if you want to export in csv", default=False)
    parser.add_argument("-h5", action="store_true", help="include if you want to export in hdf5 format", default=False)
    parser.add_argument("-msg", action="store_true", help="include if you want to export in msgpack format",
                        default=False)
    parser.add_argument("-pkl", action="store_true", help="include if you want to export in pickle format",
                        default=False)
    parser.add_argument("-asp_simple", action='store_true', help="Output the facts as asp simple tuples.",
                        default=False)
    parser.add_argument("-asp_triples", action='store_true', help="Output the facts as asp facts in form of triples",
                        default=False)
    parser.add_argument("-rels_to_use", type=str, nargs='*', default=[],
                        help="To be used with -asp_simple or -asp_triples. Enlist the relation names to output")
    parser.add_argument("-pws_to_use", type=int, nargs='*', default=[],
                        help="To be used with -asp_simple or -asp_triples. Enlist the pw_ids to output")
    parser.add_argument("--rel_facts_with_arity", dest='rel_facts_with_arity', action='store_true',
                        help="Use with -asp_simple or -asp_triples to include arity information in exported asp "
                             "relation names.")
    parser.add_argument("--no-rel_facts_with_arity", dest='rel_facts_with_arity', action='store_false',
                        help="Use with -asp_simple or -asp_triples to include arity information in exported asp "
                             "relation names.")
    parser.set_defaults(rel_facts_with_arity=True)
    parser.add_argument("--include_pw_ids", dest='include_pw_ids', action='store_true',
                        help="Use with -asp_simple or -asp_triples to include pw_ids in exported asp "
                             "relation facts.")
    parser.add_argument("--no-include_pw_ids", dest='include_pw_ids', action='store_false',
                        help="Use with -asp_simple or -asp_triples to include pw_ids in exported asp "
                             "relation facts.")
    parser.set_defaults(include_pw_ids=True)
    args = parser.parse_args()

    project_name = ''
    if args.project_name is None:
        project_name = get_current_project_name()
        if project_name is None:
            print("Couldn't find current project. Please provide a project name.")
            exit(1)
    else:
        project_name = args.project_name

    export_to_sql = args.sql
    export_to_csv = args.csv
    export_to_hdf = args.h5
    export_to_msg = args.msg
    export_to_pkl = args.pkl

    dfs = load_from_temp_pickle(project_name, 'dfs')

    if export_to_sql:
        if export_to_sqlite_db(get_save_folder(project_name, 'sql_export'), dfs, project_name):
            print("Successfully exported to sql")
        else:
            print("SQL Export Failed")
    if export_to_csv:
        if export('csv', get_save_folder(project_name, 'csv_export'), dfs):
            print("Successfully exported to csv")
        else:
            print("CSV Export Failed")
    if export_to_hdf:
        if export('h5', get_save_folder(project_name, 'h5_export'), dfs):
            print("Successfully exported to hdf")
        else:
            print("HDF Export Failed")
    if export_to_msg:
        if export('msg', get_save_folder(project_name, 'msg_export'), dfs):
            print("Successfully exported to msg")
        else:
            print("MSGPACK Export Failed")
    if export_to_pkl:
        if export('pkl', get_save_folder(project_name, 'pkl_export'), dfs):
            print("Successfully exported to pkl")
        else:
            print("PICKLE Export Failed")

    if args.schema:
        schemas = get_sqlite_schema(dfs)
        print('\n'.join(schemas))

    if args.asp_simple:
        pws = load_from_temp_pickle(project_name, 'pws')
        attr_defs = load_from_temp_pickle(project_name, 'attr_defs')
        asp_facts = export_as_asp(pws=pws, simple_or_triples='simple', rel_facts_with_arity=args.rel_facts_with_arity,
                                  include_pw_ids=args.include_pw_ids, attr_defs=attr_defs,
                                  pw_ids_to_output=args.pws_to_use, rels_to_output=args.rels_to_use)
        save_to_txt_file(asp_facts, os.path.join(get_save_folder(project_name, 'asp_simple'), 'asp_simple.lp4'))
        print("Succesfully exported as asp simple facts.")

    if args.asp_triples:
        pws = load_from_temp_pickle(project_name, 'pws')
        asp_facts = export_as_asp(pws=pws, simple_or_triples='triples', rel_facts_with_arity=args.rel_facts_with_arity,
                                  include_pw_ids=args.include_pw_ids, pw_ids_to_output=args.pws_to_use,
                                  rels_to_output=args.rels_to_use)
        save_to_txt_file(asp_facts, os.path.join(get_save_folder(project_name, 'asp_triples'), 'asp_triples.lp4'))
        print("Successfully exported in asp triples format.")

    set_current_project_name(project_name)


if __name__ == '__main__':
    __main__()