import numpy as np
import os
import pickle
import model
import random
import ggtnn_graph_parse
import convert_story
import gzip
from enum import Enum
from ggtnn_graph_parse import MetadataList, PreppedStory
from graceful_interrupt import GracefulInterruptHandler
from pprint import pformat
import util
from train_exit_status import TrainExitStatus
from functools import reduce


def convert_answer(answer, num_words, format_spec, maxlen):
    Convert an answer into an appropriate answer matrix given
    a ModelOutputFormat.

    num_words should be after processing with get_effective_answer_words,
    so that the last word is the "stop" word
    assert format_spec in model.ModelOutputFormat
    if format_spec == model.ModelOutputFormat.subset:
        ans_mat = np.zeros((1,num_words), np.float32)
        for word in answer:
            ans_mat[0, word] = 1.0
    elif format_spec == model.ModelOutputFormat.category:
        ans_mat = np.zeros((1,num_words), np.float32)
        ans_mat[0,answer[0]] = 1.0
    elif format_spec == model.ModelOutputFormat.sequence:
        ans_mat = np.zeros((maxlen+1,num_words), np.float32)
        for i,word in enumerate(answer+[num_words-1]*(maxlen+1-len(answer))):
            ans_mat[i, word] = 1.0
    return ans_mat

def get_effective_answer_words(answer_words, format_spec):
    If needed, modify answer_words using format spec to add padding chars
    if format_spec == model.ModelOutputFormat.sequence:
        return answer_words + ["<stop>"]
        return answer_words

def sample_batch(matching_stories, batch_size, num_answer_words, format_spec):
    chosen_stories = [random.choice(matching_stories) for _ in range(batch_size)]
    return assemble_batch(chosen_stories, num_answer_words, format_spec)

def assemble_batch(story_fns, num_answer_words, format_spec):
    stories = []
    for sfn in story_fns:
        with,'rb') as f:
            cvtd_story, _, _, _ = pickle.load(f)
    sents, graphs, queries, answers = zip(*stories)
    cvtd_sents = np.array(sents, np.int32)
    cvtd_queries = np.array(queries, np.int32)
    max_ans_len = max(len(a) for a in answers)
    cvtd_answers = np.stack([convert_answer(answer, num_answer_words, format_spec, max_ans_len) for answer in answers])
    num_new_nodes, new_node_strengths, new_node_ids, next_edges = zip(*graphs)
    num_new_nodes = np.stack(num_new_nodes)
    new_node_strengths = np.stack(new_node_strengths)
    new_node_ids = np.stack(new_node_ids)
    next_edges = np.stack(next_edges)
    return cvtd_sents, cvtd_queries, cvtd_answers, num_new_nodes, new_node_strengths, new_node_ids, next_edges

def assemble_correct_graphs(story_fns):
    correct_strengths, correct_ids, correct_edges = [], [], []
    for sfn in story_fns:
        with,'rb') as f:
            cvtd_story, _, _, _ = pickle.load(f)
        strengths, ids, _, edges = convert_story.convert(cvtd_story)
    return tuple(np.concatenate(l,0) for l in (correct_strengths, correct_ids, correct_edges))

def visualize(m, story_buckets, wordlist, answerlist, output_format, outputdir, batch_size=1, seq_len=5, debugmode=False, snap=False):
    cur_bucket = random.choice(story_buckets)
    sampled_batch = sample_batch(cur_bucket, batch_size, len(answerlist), output_format)
    part_sampled_batch = sampled_batch[:3]
    with open(os.path.join(outputdir,'stories.txt'),'w') as f:
        ggtnn_graph_parse.print_batch(part_sampled_batch, wordlist, answerlist, file=f)
    with open(os.path.join(outputdir,'answer_list.txt'),'w') as f:
        f.write('\n'.join(answerlist) + '\n')
    if debugmode:
        args = sampled_batch
        fn = m.debug_test_fn
        args = part_sampled_batch[:2] + ((seq_len,) if output_format == model.ModelOutputFormat.sequence else ())
        fn = m.snap_test_fn if snap else m.fuzzy_test_fn
    results = fn(*args)
    for i,result in enumerate(results):,'result_{}.npy'.format(i)), result)

def test_accuracy(m, story_buckets, bucket_sizes, num_answer_words, format_spec, batch_size, batch_auto_adjust=None, test_graph=False):
    correct = 0
    out_of = 0
    for bucket, bucket_size in zip(story_buckets, bucket_sizes):
        cur_batch_size = adj_size(m, bucket_size, batch_size, batch_auto_adjust)
        for start_idx in range(0, len(bucket), cur_batch_size):
            stories = bucket[start_idx:start_idx+cur_batch_size]
            batch = assemble_batch(stories, num_answer_words, format_spec)
            answers = batch[2]
            args = batch[:2] + ((answers.shape[1],) if format_spec == model.ModelOutputFormat.sequence else ())

            if test_graph:
                _, batch_close, _ = m.eval(*batch, with_accuracy=True)
                out_answers, out_strengths, out_ids, out_states, out_edges = m.snap_test_fn(*args)
                close = np.isclose(out_answers, answers)
                batch_close = np.all(close, (1,2))


            batch_correct = np.sum(batch_close).tolist()
            batch_out_of = len(stories)
            correct +=  batch_correct
            out_of += batch_out_of

    return correct/out_of

def adj_size(m, cur_bucket_size, batch_size, batch_auto_adjust):
    if batch_auto_adjust is not None:
        # Adjust batch size for this bucket
        edge_size = (cur_bucket_size**3) * (m.new_nodes_per_iter**2) * m.num_edge_types
        if m.sequence_representation:
            # In sequence representation mode, we are doing stuff with all objects at the same time
            # so add a multiple of the edge size to get a nice bound
            edge_size = edge_size * 4
        max_batch_size = batch_auto_adjust//edge_size
        return min(batch_size, max_batch_size)
        return batch_size

def train(m, story_buckets, bucket_sizes, len_answers, output_format, num_updates, outputdir, start=0, batch_size=BATCH_SIZE, validation_buckets=None, validation_bucket_sizes=None, stop_at_accuracy=None, stop_at_loss=None, stop_at_overfitting=None, save_params=1000, validation_interval=1000, batch_auto_adjust=None, interrupt_file=None):
    with GracefulInterruptHandler() as interrupt_h:
        for i in range(start+1,num_updates+1):
            exit_with = None
            cur_bucket, cur_bucket_size = random.choice(list(zip(story_buckets, bucket_sizes)))
            cur_batch_size = adj_size(m, cur_bucket_size, batch_size, batch_auto_adjust)
            sampled_batch = sample_batch(cur_bucket, cur_batch_size, len_answers, output_format)
            loss, info = m.train(*sampled_batch)
            if np.any(np.isnan(loss)):
                print("Loss at timestep {} was nan! Aborting".format(i))
                return TrainExitStatus.nan_loss # Don't bother saving
            with open(os.path.join(outputdir,'data.csv'),'a') as f:
                if i == 1:
                    keylist = "iter, loss, " + ", ".join(k for k,v in sorted(info.items())) + "\n"
                    if validation_buckets is not None:
                        with open(os.path.join(outputdir,'valid.csv'),'w') as f2:
                f.write("{}, {},".format(i,loss) + ", ".join(str(v) for k,v in sorted(info.items())) + "\n")
            if i % 1 == 0:
                print("update {}: {}\n{}".format(i,loss,pformat(info)))
            if i % validation_interval == 0:
                if validation_buckets is not None:
                    cur_bucket, cur_bucket_size = random.choice(list(zip(validation_buckets, validation_bucket_sizes)))
                    cur_batch_size = adj_size(m, cur_bucket_size, batch_size, batch_auto_adjust)
                    sampled_batch = sample_batch(cur_bucket, cur_batch_size, len_answers, output_format)
                    valid_loss, valid_info = m.eval(*sampled_batch)
                    print("validation at {}: {}\n{}".format(i,valid_loss,pformat(valid_info)))
                    with open(os.path.join(outputdir,'valid.csv'),'a') as f:
                        f.write("{}, {}, ".format(i,valid_loss) + ", ".join(str(v) for k,v in sorted(valid_info.items())) + "\n")
                    valid_accuracy = test_accuracy(m, validation_buckets, validation_bucket_sizes, len_answers, output_format, batch_size, batch_auto_adjust, (not m.train_with_query))
                    print("Best-choice accuracy at {}: {}".format(i,valid_accuracy))
                    with open(os.path.join(outputdir,'valid_acc.csv'),'a') as f:
                        f.write("{}, {}\n".format(i,valid_accuracy))
                    if stop_at_accuracy is not None and valid_accuracy >= stop_at_accuracy:
                        print("Accuracy reached threshold! Stopping training")
                        exit_with = TrainExitStatus.success
                    if stop_at_loss is not None and valid_loss <= stop_at_loss:
                        print("Loss reached threshold! Stopping training")
                        exit_with = TrainExitStatus.success
                    if stop_at_overfitting is not None and valid_loss/loss > stop_at_overfitting:
                        print("Model appears to be overfitting! Stopping training")
                        exit_with = TrainExitStatus.overfitting
            if exit_with is None and (interrupt_h.interrupted or (interrupt_file is not None and os.path.isfile(interrupt_file))):
                exit_with = TrainExitStatus.interrupted
            if (save_params is not None and i % save_params == 0) or (exit_with is not None) or (i==num_updates):
                util.save_params(m.params, open(os.path.join(outputdir, 'params{}.p'.format(i)), 'wb'))
            if exit_with is not None:
                return exit_with
    return TrainExitStatus.reached_update_limit