from __future__ import absolute_import

import numpy as np
import random

from brainfuck.core import Instruction

from . import brainfuck_common as common

print_interval = 5


def preset_flags(FLAGS):
    FLAGS.max_length = 1
    FLAGS.test_max_length = 1
    FLAGS.input_dim = 8 + 1 + 3 + 1 + 2  # 13 + 2 code 8 nonzero 1 skip 3 direction 1
    FLAGS.output_dim = 3 * (3 + 3 + 3 + 1 + 3) + 2  # 13 + 2 head 3 value 3 skip 3 direction 1 interaction 3


def run(ntm, seq_length, sess, idx=None, print_=True):
    return common.meta_run(ntm, seq_length, sess, generate_training_sequence=generate_training_sequence, idx=idx, seq_to_inst=seq_to_inst, print_=print_)


def train(ntm, config, sess):
    return common.meta_train(ntm, config, sess, generate_training_sequence=generate_training_sequence)


tf_set = [False, True]
order = '><+-.,[]'


def generate_training_sequence(length, config, idx=None):
    if not idx:
        idx = random.randint(0, 2 ** 30)
    context, instruction = common.context_products[idx % len(common.context_products)]

    token = context[0]
    if token is None:
        token = random.choice(order)
    token_id = order.index(token)
    nonzero = context[1]
    if nonzero is None:
        nonzero = random.choice(tf_set)
    skip = context[2]
    if skip is None:
        skip = random.choice((0, 1, 2))
    direction = context[3]
    direction_id = direction > 0

    seq_in = np.zeros([1, config.input_dim], dtype=np.float32)
    # marker 2 bits
    seq_in[0][2 + token_id] = 1
    # token 8 bits -> 10
    seq_in[0][10] = nonzero
    # nonzero 1 bit -> 11
    seq_in[0][11 + skip] = 1
    # skip 3 bits -> 14
    seq_in[0][14] = direction_id

    direction_id = instruction.direction > 0
    interaction_id = [None, 'i', 'o'].index(instruction.interaction)

    seq_out = np.zeros([1, config.output_dim], dtype=np.float32)
    idx = 2
    subidx = (instruction.head_diff + 1) * 3
    seq_out[0][idx + subidx:idx + subidx + 3] = 1
    idx += 9
    subidx = (instruction.value_diff + 1) * 3
    seq_out[0][idx + subidx:idx + subidx + 3] = 1
    idx += 9
    subidx = (instruction.skip_diff + 1) * 3
    seq_out[0][idx + subidx:idx + subidx + 3] = 1
    idx += 9
    seq_out[0][idx:idx + 3] = direction_id
    idx += 3
    subidx = (interaction_id * 3)
    seq_out[0][idx + subidx:idx + subidx + 3] = 1

    return (token, nonzero, skip, direction), instruction, list(seq_in), list(seq_out)


def seq_to_inst(seq):
    bits = seq[0]

    def get_diff(base):
        diff = sorted([
            (bits[base + 0:base + 3].sum(), -1),
            (bits[base + 3:base + 6].sum(), 2),
            (bits[base + 6:base + 9].sum(), 1),
        ], reverse=True)[0][1]
        if diff == 2:
            diff = 0
        return diff

    head_diff = get_diff(2)
    value_diff = get_diff(11)
    skip_diff = get_diff(20)
    direction = 1 if bits[29:32].sum() > 1.5 else -1
    interaction_id = get_diff(32)
    interaction = (interaction_id == 0 and 'i') or (interaction_id == 1 and 'o') or None
    instruction = Instruction(head_diff, value_diff, skip_diff, direction, interaction)
    return instruction
