import sys

sys.path.append('.')
sys.path.insert(0, "..")

import yaml
import tensorflow as tf
from checkpoint_tracker import Tracker
from data import data_loader, vocabulary
from meta_model import VarMisuseModel
import DD
import sm_helper as hp

###############################################################


g_model = None
g_data = None
g_cnt_pass = [0, 0, 0]
g_all_data = []
g_cnt_dict = {}


###############################################################


def check_dd_results(d_result):
    if hp.g_buggy_type == hp.g_buggy_types[0]:  # buggy
        return d_result[-1][0] == 0 and d_result[-1][-1] == 1.0
    elif hp.g_buggy_type == hp.g_buggy_types[1]:  # nonbuggy
        return d_result[-1][0] == 1.0 and d_result[-1][1] == 0 and d_result[-1][2] == 0
    else:
        return False


class MyDD(DD.DD):
    def __init__(self):
        DD.DD.__init__(self)

    def _test(self, deltas):
        if not deltas:
            return self.PASS

        try:
            g_cnt_pass[0] = g_cnt_pass[0] + 1
            d_tokens = [d[1] for d in deltas]
            d_sample = hp.update_marker_positions(d_tokens[:])
            if len(d_sample) > 0:  # TODO: Is parsable?
                g_cnt_pass[1] = g_cnt_pass[1] + 1
                hp.set_eval_tmp_json(d_sample)
                d_result = evaluate_single_data(g_data, g_model)
                time = hp.get_current_time()
                print('time = {}, result = {}'.format(time, d_result))
                if check_dd_results(d_result):
                    g_cnt_pass[2] = g_cnt_pass[2] + 1
                    hp.g_sample = d_sample.copy()
                    j_data = hp.get_dd_json_data(time, g_cnt_pass, d_result, d_sample.copy())
                    for j in j_data:
                        g_all_data.append(j)
                    g_all_data.append('\n')
                    return self.FAIL
        except:
            pass

        return self.PASS


###############################################################


def evaluate_single_data(data, model):
    losses, accs, preds = [], [], []  # get_metrics()
    for batch in data.batcher(mode='eval'):
        tokens, edges, error_loc, repair_targets, repair_candidates = batch
        token_mask = tf.clip_by_value(tf.reduce_sum(tokens, -1), 0, 1)
        pointer_preds, _ = model(tokens, token_mask, edges, training=False)
        batch_loss, batch_acc, batch_pred = model.get_loss(pointer_preds, token_mask, error_loc, repair_targets,
                                                           repair_candidates)
        losses, accs, preds = batch_loss, batch_acc, batch_pred
        break  # single sample only
    accs = [a.numpy().tolist() for a in accs]
    losses = [l.numpy().tolist() for l in losses]
    preds = [p.numpy().tolist() for p in preds]
    return [preds, losses, accs]


###############################################################


if __name__ == '__main__':
    config = yaml.safe_load(open(hp.config_path))
    print("Configuration:", config)

    hp.set_model_json_from_configuration(config["model"]["configuration"])
    if hp.vm_model_path is None:
        raise ValueError("Must provide a path to pre-trained models when running final evaluation")

    data = data_loader.DataLoader(hp.data_path, config["data"], vocabulary.Vocabulary(hp.vocabulary_path))
    model = VarMisuseModel(config['model'], data.vocabulary.vocab_dim)
    model.run_dummy_input()
    tracker = Tracker(model, hp.vm_model_path)
    tracker.restore_single_ckpt()

    # g_data/g_model for DD
    g_data, g_model = data, model

    # apply_dd_eval(g_data, g_model)
    js_count, dd_count = 0, 0
    with open(hp.vm_json_path) as file:
        for line in file:
            if not line.strip(): continue
            js_count += 1
            try:
                print("\nStart: {}\n".format(js_count))
                g_all_data.clear()
                g_cnt_pass = [0, 0, 0]

                sample = hp.check_eval_tmp(line.strip())
                assert sample["results"]["model"] == config["model"]["configuration"]
                results = evaluate_single_data(data, model)

                if hp.g_buggy_type == hp.g_buggy_types[0]:  # buggy
                    assert results[-1][0] == 0 and results[-1][-1] == 1.0
                else:  # nonbuggy
                    assert results[-1][0] == 1.0 and results[-1][1] == 0 and results[-1][2] == 0

                # original sample
                g_all_data.append("\nOriginal sample:\n")
                g_all_data.append(hp.get_pretty_sample(sample.copy()))

                # save filename
                save_name = "{}___L{}.txt".format(sample["txt_file"], sample["js_count"])

                # create deltas by tokens
                tokens = list(sample["marker_tokens"])
                deltas = list(zip(range(len(tokens)), tokens))

                try:
                    # run ddmin
                    mydd = MyDD()
                    print("Simplifying failure-inducing input...")
                    g_all_data.append("\n\nTrace of simplified code(s):\n")
                    c = mydd.ddmin(deltas)  # Invoke DDMIN
                    print("The 1-minimal failure-inducing input is", c)
                    print("Removing any element will make the failure go away.")
                    # TODO: c to code?
                    c = [d[1] for d in c]
                    s = hp.update_marker_positions(c[:])
                    g_all_data.append("\n\nMinimal simplified tokens:\n")
                    g_all_data.append(str(s["source_tokens"]))
                    dd_count += 1
                except Exception as e:
                    g_all_data.append("\n\nException:\n{}".format(str(e)))

                # print/save all simplified code
                output_file = hp.dd_file.format(save_name)
                hp.save_simplified_code(g_all_data, output_file)
                print("\nDone: {}-{}\n".format(js_count, dd_count))

                if dd_count >= hp.g_dd_count:
                    quit()
            except Exception as e:
                print("\nError: {}\n{}".format(js_count, str(e)))
