"""
utils.py
------------------------------------------------
Contains general utilities for:
- Loading d3m dataset, writing d3m predictions
- Dealing with problem docs, protos
- Running functions with timelimit
"""

import sys
import glob
import pdb
import shutil
import json
import pandas
import numpy as np
import os
import traceback
import time
import git
import shutil
import pdb
from multiprocessing import Process

from d3m import metadata
from d3m.metadata import problem
from d3m.metadata.problem import *
from d3m.metadata.pipeline import Pipeline, PrimitiveStep
from d3m.metadata.base import Metadata
from d3m.container.dataset import D3MDatasetLoader, Dataset
from d3m import primitives
import d3m_outputs
from d3m.metadata import base as metadata_base
from d3m_outputs import Predictions

def get_git_root(path):
    git_repo = git.Repo(path, search_parent_directories=True)
    git_root = git_repo.git.rev_parse("--show-toplevel")
    return git_root

def load_problem_doc(problem_doc_uri: str):
    """
    Load problem_doc from problem_doc_uri

    Parameters
    ---------
    problem_doc_uri
        Uri where the problemDoc.json is located
    """
    with open(problem_doc_uri) as file:
        problem_doc = json.load(file)
    problem_doc_metadata = Metadata(problem_doc)
    return problem_doc_metadata

def add_target_columns_metadata(dataset: 'Dataset', problem_doc_metadata: 'Metadata'):
    """
    Add metadata to the dataset from problem_doc_metadata

    Parameters
    ---------
    dataset
        Dataset
    problem_doc_metadata:
        Metadata about the problemDoc
    """
    for data in problem_doc_metadata.query(())['inputs']['data']:
        targets = data['targets']
        for target in targets:
            semantic_types = list(dataset.metadata.query(
                (target['resID'], metadata_base.ALL_ELEMENTS, target['colIndex'])).get('semantic_types', []))
            if 'https://metadata.datadrivendiscovery.org/types/Target' not in semantic_types:
                semantic_types.append('https://metadata.datadrivendiscovery.org/types/Target')
                dataset.metadata = dataset.metadata.update(
                    (target['resID'], metadata_base.ALL_ELEMENTS, target['colIndex']), {'semantic_types': semantic_types})
            if 'https://metadata.datadrivendiscovery.org/types/TrueTarget' not in semantic_types:
                semantic_types.append('https://metadata.datadrivendiscovery.org/types/TrueTarget')
                dataset.metadata = dataset.metadata.update(
                    (target['resID'], metadata_base.ALL_ELEMENTS, target['colIndex']), {'semantic_types': semantic_types})
    return dataset

def load_data_from_dir(inputdir, mode="train"):
    """
    Returns problem_doc and dataset given input directory.
    """
    assert mode in ["train", "test", "score"]
    
    if mode == "train":
        problemdir = "%s/TRAIN/problem_TRAIN" % inputdir
        problem_doc_uri = "%s/problemDoc.json" % problemdir
        datasetdir = "%s/TRAIN/dataset_TRAIN" % inputdir
        dataset_uri = "%s/datasetDoc.json" % datasetdir
    if mode == "test":
        problemdir = "%s/TEST/problem_TEST" % inputdir
        problem_doc_uri = "%s/problemDoc.json" % problemdir
        datasetdir = "%s/TEST/dataset_TEST" % inputdir
        dataset_uri = "%s/datasetDoc.json" % datasetdir
    if mode == "score":
        problemdir = "%s/SCORE/problem_TEST" % inputdir
        problem_doc_uri = "%s/problemDoc.json" % problemdir
        datasetdir = "%s/SCORE/dataset_TEST" % inputdir
        dataset_uri = "%s/datasetDoc.json" % datasetdir
        
    # Problem doc and dataset loading
    problem_doc = load_problem_doc(problem_doc_uri)

    # Dataset    
    if 'file:' not in dataset_uri:
        dataset_uri = 'file://{dataset_uri}'.format(dataset_uri=os.path.abspath(dataset_uri))
        dataset = D3MDatasetLoader()
        dataset = dataset.load(dataset_uri=dataset_uri)
        dataset = add_target_columns_metadata(dataset, problem_doc)
        if mode == "train" or mode == "score":
            dataset = add_target_columns_metadata(dataset, problem_doc)

    return problem_doc, dataset

def is_higher_score_better(problem_doc):
    """
    Returns whether higher score is better for the given problem.

    Currently this is a hack: 
    If metric string has "err" in it, return False, otherwise True
    """

    metric_string = problem_doc.query(selector=()).get("inputs").get("performanceMetrics")[0].get("metric")
    return not "err" in metric_string.lower()

def get_global_score(local_score, problem_doc):
    """
    Convert local_score to global_score. A global_score has the property
    the a lower value always means a better solution.

    The algorithm is as follows:
    If problem_doc indicates higher scores are better:
      return LARGE_EPS - local_score
    else
      return LARGE_EPS + local_score

    LARGE_EPS should be chosen to be bigger than any possible |local_score|.
    However if this is not the case and |local_score| > LARGE_EPS, we return
    a global score LARGE_EPS (which is basically 0) to indicate an error.
    """
    LARGE_EPS = 1e12

    if np.abs(local_score) > LARGE_EPS:
        return LARGE_EPS

    higher_is_better = is_higher_score_better(problem_doc)
    sign = -1 if higher_is_better else 1

    return LARGE_EPS + sign*local_score

def save_pipeline(pipeline_file_path, pipeline, pipeline_rank=None):
    """
    Save pipeline to pipeline_file_path. pipeline is a d3m object.
    """
    to_json_structure = pipeline.to_json_structure()
    if pipeline_rank is not None:
        to_json_structure["pipeline_rank"] = pipeline_rank
    with open(pipeline_file_path, "w") as f:
        f.write(json.dumps(to_json_structure,
                           sort_keys=True, indent=4, separators=(',', ': ')))

def save_supporting_files(supporting_files_path):
    """
    Save all supporting files to the given directory path.
    Simply copy over the whole repository to this path.
    """
    git_root_path = get_git_root(os.path.dirname(os.path.abspath(__file__)))

    def ign_func(d, files):
        keep_files_with_py = [f for f in files if os.path.isfile(os.path.join(d, f)) and not f.endswith(".py")]
        ignore_files_with_leading_dot = [f for f in files if os.path.isfile(os.path.join(d,f)) and f[0] == "."]
        return keep_files_with_py + ignore_files_with_leading_dot
    
    shutil.copytree(git_root_path, supporting_files_path, ignore=ign_func)

def save_executable(exe_path, pipeline_id):
    """
    Save a bash script with a command that will run the pipeline
    against a target directory and generate predictions.
    """
    # TODO(maxlam): FINISH
    command = "testing"
    with open(exe_path, "w") as f:
        f.write(command)

def clear_all_files_but_top(output_dir, n_to_keep=5):
    """
    Clears all files but the top n_to_keep.
    Assumes that lower scores are better.
    """    
    pipelines_files_path = "%s/pipelines_ranked" % output_dir
    files = os.listdir(pipelines_files_path)
    files_and_scores = [(x, float(x.replace(".json",""))) for x in files]
    files_and_scores.sort(key=lambda x:x[1])
    files_and_scores_to_delete = files_and_scores[n_to_keep:]
    for f, score in files_and_scores_to_delete:
        pipeline_id = str(score)
        pipelines_files_path_to_delete = "%s/pipelines_ranked/%s.json" % (output_dir, pipeline_id)
        for to_delete in [pipelines_files_path_to_delete]:
            if os.path.isdir(to_delete):
                shutil.rmtree(to_delete)
            else:
                if os.path.exists(to_delete):
                    os.remove(to_delete)

def save_scored_pipeline(outputdir, local_score, pipeline, problem_doc, keep_top=5, clear_directories=False):
    """
    Save scored pipeline, executables and supporting files to the given output directory as specified by NIST. 
    Keep only the top keep_top configurations.    
    
    Creates the following directory structure if it doesn't exist:
    outputdir/
      pipelines_ranked
      pipelines_scored
      pipelines_searched
      subpipelines
      pipeline_runs
      additional_inputs

    The process is as follows:

    1. Given local_score (e.g: 68.5), converts this to a global score (e.g: 1000-68.5 = 931.5)
       in which a lower value implies a better prediction. This global score is the id of the pipeline.
       Set the pipeline's id as such.
    2. Given the id of the pipeline, create `outputdir`/pipelines_ranked/`pipeline_id`.json containing
       the pipeline json.
    3. Collect all pipelines saved; sort them; delete all but the top keep_top pipelines (in pipelines_ranked. 
       pipelines_scored should contain every single pipeline ever scored).
    
    ########################################################
    # Note this function is not necessarily thread safe!!! #
    ########################################################

    Furthermore: It is possible that this function may be in the middle of writing to output
    when the program is terminated, leading to nasty bugs. Please give enough time for this 
    function before the deadline of the evaluation.
    
    Arguments
    - outputdir - Full path to output to save pipelines, supporting_files, executables
    - local_score - The local score of the pipeline
    - pipeline - The pipeline object (not json) representing pipeline to save
    - problem_doc - Problem doc of the problem of the pipeline being saved
    - keep_top - Number of total pipelines / supporting files to keep
    - clear_directories - Should clear the directories or not
    """

    # Create the directory structure
    pipelines_ranked = "%s/pipelines_ranked" % outputdir
    pipelines_scored = "%s/pipelines_scored" % outputdir
    pipelines_searched = "%s/pipelines_searched" % outputdir
    subpipelines = "%s/subpipelines" % outputdir
    pipeline_runs = "%s/pipeline_runs" % outputdir
    additional_inputs = "%s/additional_inputs" % outputdir

    all_paths = [pipelines_ranked, pipelines_scored,
                 pipelines_searched, subpipelines,
                 pipeline_runs, additional_inputs]
    
    # Clear the path outputs if specified
    if clear_directories:
        for path_to_clear in all_paths:
            if os.path.exists(path_to_clear):
                shutil.rmtree(path_to_clear)
                
    # Create the dirs
    for necessary_path in all_paths:
        if not os.path.exists(necessary_path):
            os.makedirs(necessary_path)

    # Get pipeline id
    pipeline_id = str(get_global_score(local_score, problem_doc))
    pipeline.id = pipeline_id

    target_pipeline_file_path = "%s/%s.json" % (pipelines_ranked, pipeline_id)
    
    # If pipeline json with pipeline id already exists, don't do anything
    if os.path.exists(target_pipeline_file_path):
        return

    paths_to_save = ["%s/%s.json" % (d, pipeline_id) for d in [pipelines_ranked, pipelines_searched, pipelines_scored]]

    # Save the pipeline
    for p in paths_to_save:
        save_pipeline(p, pipeline, pipeline_rank=pipeline_id)

    # The new D3M evaluation doesn't require saving supporting files or executables!
    #save_supporting_files(target_supporting_files_path)
    #save_executable(target_executable_file_path, pipeline_id)

    # Clear all but the top
    clear_all_files_but_top(outputdir, n_to_keep=keep_top)

def get_entrypoint_resource_id(dataset):
    """
    Gets the entrypoint resource id of a dataset.
    
    TODO(maxlam): Maybe put this in train utils?    
    """
    for resource_id in dataset.keys():
        if ('https://metadata.datadrivendiscovery.org/types/DatasetEntryPoint' in
            dataset.metadata.query(selector=(resource_id,)).get("semantic_types")):
            return resource_id
    return resource_id

def write_predictions_to_file(preds, fpath, problemdoc, dataset=None):
    """
    Writes predictions to file path fpath.
    
    Args:
    - preds - Dataframe of predictions
    
    TODO(maxlam): CLEAN THIS UP!
    """
    preds.to_csv(fpath, index=False, header=True)

def run_with_limited_time(func, args, kwargs, time):
    """Runs a function with time limit

    :param func: The function to run
    :param args: The functions args, given as tuple
    :param kwargs: The functions keywords, given as dict
    :param time: The time limit in seconds
    :return: True if the function ended successfully. False if it was terminated.
    """
    p = Process(target=func, args=args, kwargs=kwargs)
    p.start()
    p.join(time)
    if p.is_alive():
        p.terminate()
        return (False, None)
    return True

def convert_problem_doc_proto_to_metadata(problem_doc):
    # Convert problem_doc proto to json
    json_problem_doc = {}

    # General info
    json_problem_doc["about"] = {
        "problemID": problem_doc.problem.id,
        "problemVersion": problem_doc.problem.version,
        "problemName": problem_doc.problem.name,
        "taskType": problem_doc.problem.task_type,
        "taskSubType": problem_doc.problem.task_subtype,
        "problemSchemaVersion": "3.1.1",
        "problemDescription": problem_doc.problem.description
    }

    # Set inputs
    inputs_data = []
    for input_proto in problem_doc.inputs:
        targets_data = []
        for target_data in input_proto.targets:
            targets_data.append({
                "targetIndex": target_data.target_index,
                "resID": target_data.resource_id,
                "colIndex": target_data.column_index,
                "colName": target_data.column_name,
            })
        inputs_data.append({
            "datasetID": input_proto.dataset_id,
            "targets": targets_data
        })

    # Set metrics
    perf_metrics = []
    for performance_metric in problem_doc.problem.performance_metrics:
        print(performance_metric.metric)
        reverse_map = {v:k for k,v in PerformanceMetric.get_map().items()}
        metric_string = reverse_map[PerformanceMetric(performance_metric.metric)]
        perf_metrics.append({            
            "metric": metric_string
        })
    json_problem_doc["inputs"] = {
        "data": inputs_data,
        "performanceMetrics": perf_metrics
    }        

    # From json use standard methods to parse to metadata
    return Metadata(json_problem_doc)

def convert_dataset_uri_to_dataset(problem_doc, dataset_uri, mode="train"):
    if "file://" in dataset_uri:
        dataset_uri = dataset_uri[len("file://"):]

    # Dataset    
    dataset_uri = 'file://{dataset_uri}'.format(dataset_uri=os.path.abspath(dataset_uri))
    dataset = D3MDatasetLoader()
    dataset = dataset.load(dataset_uri=dataset_uri)
    dataset = add_target_columns_metadata(dataset, problem_doc)
    if mode == "train" or mode == "score":
        dataset = add_target_columns_metadata(dataset, problem_doc)
    return dataset

