import numpy as np
import tensorflow as tf
from tensorflow.python.keras.backend import dtype, print_tensor
from utils.thread import custom_Thread
from utils.randomTool import RandomHandler
from utils.npTcp import npTcpSender, npTcpReceiver
from utils.randomTool import get_random_mask, RandomHandler
from utils.coding import *
from utils.bnnLayers import sum_pool
from utils.np_ops import choose, extract_from_stream

def calculate_2_input_gate(a_masked: np.ndarray, b_masked: np.ndarray,
                           triple: np.ndarray, constant_computation: bool):

    ret = triple[2] ^ (a_masked & triple[1]) ^ (b_masked & triple[0])

    if constant_computation:
        ret ^= a_masked & b_masked

    return ret

def beaver_mult_bool(x_share: np.ndarray, y_share: np.ndarray, triple,
                     sender: npTcpSender, listener: npTcpReceiver,
                     constant_computation: bool):
    '''
    return z_share = x_share * y_share
    '''
    x_share_masked = x_share ^ triple[0]
    y_share_masked = y_share ^ triple[1]
    msg = np.array([x_share_masked, y_share_masked], dtype='bool')

    send_msg_thread = sender.send_ndarray_by_bytes_thread(msg, packbits=True)
    send_msg_thread.start()

    [x_other, y_other] = listener.receive_ndarray_by_bytes(size=msg.shape,
                                                           packbits=True)

    x_masked = x_share_masked ^ x_other
    y_masked = y_share_masked ^ y_other

    z_share = calculate_2_input_gate(x_masked, y_masked, triple,
                                     constant_computation)

    send_msg_thread.join()

    return z_share

def beaver_mult_bool_two_round_with_optimization(x_share: np.ndarray,
                                                 y_share: np.ndarray, triple,
                                                 sender: npTcpSender,
                                                 listener: npTcpReceiver,
                                                 constant_computation: bool):
    '''
    return z_share = x_share * y_share
    '''
    # modify d = x ^ u
    d_share = x_share ^ triple[0]
    e_share = y_share ^ triple[1]

    # print('d_share.shape:', d_share.shape)
    # print('e_share.shape:', e_share.shape)
    # print('x_share.shape:', x_share.shape)
    # print('y_share.shape:', y_share.shape)
    # print('triple[0].shape', triple[0].shape)
    # print('triple[1].shape:', triple[1].shape)
    # print(d.shape)
    # assert(np.sum(d_share[0]!=d_share[1])==0)
    # only send one element of d

    msg = np.array([d_share[0], e_share[0], e_share[1]], dtype='bool')

    send_msg_thread = sender.send_ndarray_by_bytes_thread(msg, packbits=True)
    send_msg_thread.start()
    rcv_thread = listener.receive_ndarray_by_bytes_thread(size=msg.shape,
                                                          packbits=True)
    rcv_thread.start()
    rcv_thread.join()

    [d_other_0, e_other_0, e_other_1] = rcv_thread.get_result()

    # recover the array
    d_other = np.array([d_other_0, d_other_0], dtype='bool')
    e_other = np.array([e_other_0, e_other_1], dtype='bool')

    d = d_share ^ d_other
    e = e_share ^ e_other
    z_share = triple[2] ^ (e & triple[0]) ^ (d & triple[1])

    if constant_computation:
        z_share ^= d & e

    send_msg_thread.join()

    return z_share

def calculate_3_input_gate(a_masked: np.ndarray, b_masked: np.ndarray,
                           c_masked: np.ndarray, help_tuple,
                           constant_computation: bool):
    # help_tuple: (u, v, w, uv, uw, vw, uvw)

    ret = (help_tuple[3] & c_masked) ^ (help_tuple[4] & b_masked) ^ (
        help_tuple[5] & a_masked) ^ (a_masked & b_masked & help_tuple[2]) ^ (
            a_masked & help_tuple[1] & c_masked) ^ (help_tuple[0] & b_masked
                                                    & c_masked) ^ help_tuple[6]

    if constant_computation:
        ret ^= a_masked & b_masked & c_masked

    return ret

def beaver_mult_bool_3_input(a_share: np.ndarray, b_share: np.ndarray,
                             c_share: np.ndarray, help_tuple,
                             sender: npTcpSender, listener: npTcpReceiver,
                             constant_computation: bool):
    '''
    return z_share = a_share * b_share * c_share
    help_tuple: (u, v, w, uv, uw, vw, uvw)
    '''
    a_share_masked = a_share ^ help_tuple[0]
    b_share_masked = b_share ^ help_tuple[1]
    c_share_masked = c_share ^ help_tuple[2]

    msg = np.array([a_share_masked, b_share_masked, c_share_masked],
                   dtype='bool')

    send_msg_thread = sender.send_ndarray_by_bytes_thread(msg, packbits=True)
    send_msg_thread.start()
    rcv_thread = listener.receive_ndarray_by_bytes_thread(size=msg.shape,
                                                          packbits=True)
    rcv_thread.start()

    rcv_thread.join()
    [a_other, b_other, c_other] = rcv_thread.get_result()

    a_masked = a_share_masked ^ a_other
    b_masked = b_share_masked ^ b_other
    c_masked = c_share_masked ^ c_other

    z_share = calculate_3_input_gate(a_masked, b_masked, c_masked, help_tuple,
                                     constant_computation)

    send_msg_thread.join()

    return z_share

def ot_3party(role,
              sender: list,
              listener: list,
              random_handler: RandomHandler,
              msg_0=None,
              msg_1=None,
              msg_dtype=None,
              msg_bit=None,
              msg_shape=None,
              choose_bit=None):
    '''
    role=0 means receiver
    role=1 means sender
    role=2 means role_ttp
    '''

    role_receiver = 0
    role_sender = 1
    role_ttp = 2

    ret = None

    if role == role_receiver:

        ori_shape = choose_bit.shape

        # get same random numbers with sender
        swap_bit = random_handler.get_same_random_with(role_sender,
                                                       bits=1,
                                                       size=ori_shape)
        random_handler.cnter += 1

        # Send choose_bit
        sender[role_ttp].send_ndarray_by_bytes(choose_bit ^ swap_bit,
                                               packbits=True)

        # receive the message
        ret = listener[role_ttp].receive_ndarray_by_bytes(size=ori_shape,
                                                          dtype=msg_dtype)

        # unmask the message
        mask = random_handler.get_same_random_with(role_sender,
                                                   bits=msg_bit,
                                                   size=(2, *ori_shape))
        random_handler.cnter += 1
        mask = ((True ^ choose_bit ^ swap_bit) * mask[0] +
                (choose_bit ^ swap_bit) * mask[1]).astype(ret.dtype)

        ret = ret ^ mask

    elif role == role_sender:

        # get same random numbers with receiver and swap the message
        swap_bit = random_handler.get_same_random_with(role_receiver,
                                                       bits=1,
                                                       size=msg_0.shape)
        random_handler.cnter += 1
        msg = np.zeros(shape=(2, *msg_0.shape), dtype=msg_0.dtype)
        msg[0] = choose(swap_bit, msg_0, msg_1)
        msg[1] = choose(True ^ swap_bit, msg_0, msg_1)

        # mask the message
        mask = random_handler.get_same_random_with(role_receiver,
                                                   bits=msg_bit,
                                                   size=(2, *msg_0.shape))
        random_handler.cnter += 1
        msg = msg ^ mask
        # send
        sender[role_ttp].send_ndarray_by_bytes(msg)

    elif role == role_ttp:

        # receive the choose bit from receiver
        choice_thread = listener[
            role_receiver].receive_ndarray_by_bytes_thread(msg_shape,
                                                           packbits=True)
        choice_thread.start()

        # receive the masked message from sender
        data_thread = listener[role_sender].receive_ndarray_by_bytes_thread(
            (2, *msg_shape), dtype=msg_dtype)
        data_thread.start()

        # keep the counter in sync
        random_handler.cnter += 1

        choice_thread.join()
        choose_bit = choice_thread.get_result()
        data_thread.join()
        data = data_thread.get_result()

        # select the message
        msg = ((True ^ choose_bit) * data[0] + choose_bit * data[1]).astype(
            data.dtype)

        # send the selected message to receiver
        sender[role_receiver].send_ndarray_by_bytes(msg)

    return ret

def extract_msb(role: int, listener: list, sender: list,
                random_handler: RandomHandler, a_share: np.ndarray,
                b_share: np.ndarray, total_bit: int, pos: int,
                triple: np.ndarray):
    data_owner = 0
    model_owner = 1
    ret = None
    if role == model_owner or role == data_owner:
        const_compute = True if role == model_owner else False
        p_triple = 0
        a_1 = extractbit(a_share, pos=0, n_bits=total_bit)
        b_1 = extractbit(b_share, pos=0, n_bits=total_bit)

        # 2)
        c_last = beaver_mult_bool(a_1, b_1, triple[:,
                                                   p_triple], sender[1 - role],
                                  listener[1 - role], const_compute)
        p_triple += 1

        for i in range(1, pos + 1):
            a_i = extractbit(a_share, i, total_bit)
            b_i = extractbit(b_share, i, total_bit)
            y_i = a_i ^ b_i
            # d)
            if i == pos:
                ret = y_i ^ c_last
            else:
                # a)
                d_i = beaver_mult_bool(a_i, b_i, triple[:, p_triple],
                                       sender[1 - role], listener[1 - role],
                                       const_compute)
                if const_compute:
                    d_i ^= True
                p_triple += 1

                # b)
                e_i = beaver_mult_bool(y_i, c_last, triple[:, p_triple],
                                       sender[1 - role], listener[1 - role],
                                       const_compute)
                p_triple += 1
                if const_compute:
                    e_i ^= True

                # c)
                c_i = beaver_mult_bool(e_i, d_i, triple[:, p_triple],
                                       sender[1 - role], listener[1 - role],
                                       const_compute)
                p_triple += 1
                if const_compute:
                    c_i ^= True
                c_last = c_i

        ret ^= random_handler.get_same_random_with(1 - role,
                                                   bits=1,
                                                   size=ret.shape)

    # all three party add the cnter
    random_handler.cnter += 1

    return ret

def count_ppa_n_triple(msb_pos: int):
    # msb_pos counts from zero
    ret = msb_pos
    l_bound = msb_pos - 1
    while l_bound > 0:
        is_even = False

        if l_bound % 2 == 0:
            is_even = True
            l_bound = l_bound - 1

        n_node = (l_bound + 1) // 2
        ret += n_node + n_node

        if is_even == True:
            # count the unpaired node
            l_bound = n_node + 1 - 1
        else:
            l_bound = n_node - 1
    return ret

def modify_triple_with_optimization(msb_pos, triple_do=None, triple_mo=None):
    # msb_pos is counted from zero
    # there are msb_pos + 1 bits

    p_triple = 0

    # calculate carry_bit
    p_triple += msb_pos

    # divide and conquer
    l_bound = msb_pos - 1

    while l_bound > 0:
        exist_unpaired_bits = False

        if l_bound % 2 == 0:
            l_bound = l_bound - 1
            exist_unpaired_bits = True

        # l_bound is odd
        l_idx = list(range(1, l_bound + 1, 2))
        len_l = len(l_idx)
        # len_l = l_bound // 2 + 1

        # data_1 = np.array([l_all_1, l_all_1], dtype='bool')
        # data_2 = np.array([r_all_1, carry_bit[r_idx]], dtype='bool')

        # [U, U'] -> [U, U]
        mid_p = p_triple + len_l
        nxt_p = mid_p + len_l

        # print(triple_do.shape)
        # print(triple_mo.shape)

        if triple_do is not None:
            triple_do[0, mid_p:nxt_p] = triple_do[0, p_triple:mid_p]
        if triple_mo is not None:
            triple_mo[0, mid_p:nxt_p] = triple_mo[0, p_triple:mid_p]
        p_triple = nxt_p

        if exist_unpaired_bits is True:
            len_l += 1

        l_bound = len_l - 1
        # print(l_bound, len_l)

    if (triple_do is not None) and (triple_mo is not None):
        delta = (triple_do[0] ^ triple_mo[0]) & (triple_do[1]
                                                 ^ triple_mo[1]) ^ triple_do[2]
        return delta
    elif triple_mo is not None:
        return triple_mo
    elif triple_do is not None:
        return triple_do

    return None

def extract_msb_ppa(role: int, listener: list, sender: list,
                    a_share: np.ndarray, b_share: np.ndarray, total_bit: int,
                    msb_pos: int, triple: np.ndarray):

    # msb_pos and l_bound counts from zero
    if role == 2:
        return None

    const_compute = True if role == 0 else False

    a_share_bit = list()
    b_share_bit = list()

    ret = extractbit(a_share, msb_pos, n_bits=total_bit) ^ extractbit(
        b_share, msb_pos, n_bits=total_bit)

    for i in range(0, msb_pos):
        a_share_bit.append(extractbit(a_share, i, n_bits=total_bit))
        b_share_bit.append(extractbit(b_share, i, n_bits=total_bit))

    a_share_bit = np.array(a_share_bit, dtype='bool')
    b_share_bit = np.array(b_share_bit, dtype='bool')

    p_triple = 0

    all_1 = a_share_bit ^ b_share_bit

    nxt_p = p_triple + len(a_share_bit)
    carry_bit = beaver_mult_bool(a_share_bit,
                                 b_share_bit,
                                 triple[:, p_triple:nxt_p],
                                 sender[1 - role],
                                 listener[1 - role],
                                 constant_computation=const_compute)
    p_triple = nxt_p

    l_bound = msb_pos - 1

    while l_bound > 0:
        # All the nodes are paired

        left_most_all_1 = None
        left_most_carry_bit = None

        if l_bound % 2 == 0:
            left_most_all_1 = all_1[l_bound]
            left_most_carry_bit = carry_bit[l_bound]
            l_bound = l_bound - 1

        l_idx = list(range(1, l_bound + 1, 2))
        r_idx = list(range(0, l_bound + 1, 2))
        l_all_1 = all_1[l_idx]
        r_all_1 = all_1[r_idx]

        data_1 = np.array([l_all_1, l_all_1], dtype='bool')
        data_2 = np.array([r_all_1, carry_bit[r_idx]], dtype='bool')

        nxt_p = p_triple + len(l_all_1) + len(l_all_1)
        data_ret = beaver_mult_bool(data_1,
                                    data_2,
                                    triple[:, p_triple:nxt_p].reshape(
                                        3, *data_1.shape),
                                    sender[1 - role],
                                    listener[1 - role],
                                    constant_computation=const_compute)
        p_triple = nxt_p

        all_1 = data_ret[0]
        carry_bit = data_ret[1] ^ carry_bit[l_idx]

        if left_most_all_1 is not None:
            all_1 = np.append(all_1, [left_most_all_1], axis=0)
            carry_bit = np.append(carry_bit, [left_most_carry_bit], axis=0)

        l_bound = len(all_1) - 1

    # print(p_triple, triple.shape[1],
    #         msb_pos, count_ppa_n_triple(msb_pos))
    return ret ^ carry_bit[0]

def extract_msb_ppa_with_optimize(role: int, listener: list, sender: list,
                                  a_share: np.ndarray, b_share: np.ndarray,
                                  total_bit: int, msb_pos: int,
                                  triple: np.ndarray):

    # msb_pos and l_bound counts from zero
    if role == 2:
        return None

    const_compute = True if role == 0 else False

    a_share_bit = list()
    b_share_bit = list()

    ret = extractbit(a_share, msb_pos, n_bits=total_bit) ^ extractbit(
        b_share, msb_pos, n_bits=total_bit)

    for i in range(0, msb_pos):
        a_share_bit.append(extractbit(a_share, i, n_bits=total_bit))
        b_share_bit.append(extractbit(b_share, i, n_bits=total_bit))

    a_share_bit = np.array(a_share_bit, dtype='bool')
    b_share_bit = np.array(b_share_bit, dtype='bool')

    p_triple = 0

    all_1 = a_share_bit ^ b_share_bit

    nxt_p = p_triple + len(a_share_bit)
    carry_bit = beaver_mult_bool(a_share_bit,
                                 b_share_bit,
                                 triple[:, p_triple:nxt_p],
                                 sender[1 - role],
                                 listener[1 - role],
                                 constant_computation=const_compute)
    p_triple = nxt_p

    l_bound = msb_pos - 1

    while l_bound > 0:

        # make nodes be paired
        left_most_all_1 = None
        left_most_carry_bit = None

        if l_bound % 2 == 0:
            left_most_all_1 = all_1[l_bound]
            left_most_carry_bit = carry_bit[l_bound]
            l_bound = l_bound - 1

        l_idx = list(range(1, l_bound + 1, 2))
        r_idx = list(range(0, l_bound + 1, 2))
        l_all_1 = all_1[l_idx]
        r_all_1 = all_1[r_idx]

        data_1 = np.array([l_all_1, l_all_1], dtype='bool')
        data_2 = np.array([r_all_1, carry_bit[r_idx]], dtype='bool')

        mid_p = p_triple + len(l_all_1)
        nxt_p = mid_p + len(l_all_1)
        cur_triple = triple[:, p_triple:nxt_p]
        cur_triple[0, mid_p:nxt_p] = cur_triple[0, p_triple:mid_p]
        data_ret = beaver_mult_bool_two_round_with_optimization(
            data_1,
            data_2,
            cur_triple.reshape(3, *data_1.shape),
            sender[1 - role],
            listener[1 - role],
            constant_computation=const_compute)
        p_triple = nxt_p

        all_1 = data_ret[0]
        carry_bit = data_ret[1] ^ carry_bit[l_idx]

        if left_most_all_1 is not None:
            all_1 = np.append(all_1, [left_most_all_1], axis=0)
            carry_bit = np.append(carry_bit, [left_most_carry_bit], axis=0)

        l_bound = len(all_1) - 1

    return ret ^ carry_bit[0]

def extract_msb_with_3_input_gate(role: int, listener: list, sender: list,
                                  a_share: np.ndarray, b_share: np.ndarray,
                                  total_bit: int, msb_pos: int,
                                  triple: np.ndarray, help_tuple: np.ndarray):
    '''
    help_tuple: (u, v, w, uv, uw, vw, uvw)
    '''
    # msb_pos and l_bound counts from zero
    if role == 2:
        return None

    const_compute = True if role == 0 else False

    # extract MSB
    ret = extractbit(a_share, msb_pos, n_bits=total_bit) ^ extractbit(
        b_share, msb_pos, n_bits=total_bit)

    a_share_bit = list()
    b_share_bit = list()

    # extract other bits
    for i in range(0, msb_pos):
        a_share_bit.append(extractbit(a_share, i, n_bits=total_bit))
        b_share_bit.append(extractbit(b_share, i, n_bits=total_bit))
    a_share_bit = np.array(a_share_bit, dtype='bool')
    b_share_bit = np.array(b_share_bit, dtype='bool')

    p_triple = 0
    p_help_tuple = 0

    # Start to solve the problem of the leaves

    all_1 = a_share_bit ^ b_share_bit
    p_triple_nxt = p_triple + len(a_share_bit)
    carry_bit = beaver_mult_bool(a_share_bit,
                                 b_share_bit,
                                 triple[:, p_triple:p_triple_nxt],
                                 sender[1 - role],
                                 listener[1 - role],
                                 constant_computation=const_compute)
    p_triple = p_triple_nxt

    # divide and conquer
    l_bound = msb_pos - 1

    while l_bound > 0:

        flag = (l_bound + 1) % 3
        n_paired = (l_bound + 1) // 3
        paired_l_bound = max(n_paired * 3 - 1, 0)

        l_idx = list(range(2, paired_l_bound + 1, 3))
        m_idx = list(range(1, paired_l_bound + 1, 3))
        r_idx = list(range(0, paired_l_bound + 1, 3))

        # the input of 3-input-and gate
        # print(l_bound, all_1.shape, carry_bit.shape, n_paired)
        l_all_1 = all_1[l_idx]
        m_all_1 = all_1[m_idx]
        r_all_1 = all_1[r_idx]

        l_carry_bit = carry_bit[l_idx]
        m_carry_bit = carry_bit[m_idx]
        r_carry_bit = carry_bit[r_idx]

        # the input of 2-input-and gate (the left most two bits)
        left_most_l_all_1 = None
        left_most_l_carry_bit = None
        left_most_r_all_1 = None
        left_most_r_carry_bit = None

        left_most_all_1 = None
        left_most_carry_bit = None

        if flag == 1:
            left_most_all_1 = all_1[l_bound:l_bound + 1]
            left_most_carry_bit = carry_bit[l_bound:l_bound + 1]

        elif flag == 2:
            left_most_l_all_1 = all_1[l_bound:l_bound + 1]
            left_most_l_carry_bit = carry_bit[l_bound:l_bound + 1]

            left_most_r_all_1 = all_1[l_bound - 1:l_bound]
            left_most_r_carry_bit = carry_bit[l_bound - 1:l_bound]

        cur_node_len = n_paired  # the number of 3-input gates
        p_help_tuple_nxt = p_help_tuple + cur_node_len

        # mask the input
        if n_paired > 0:
            masked_tuple_elements_sent = np.array([
                l_all_1 ^ help_tuple[0, p_help_tuple:p_help_tuple_nxt],
                m_all_1 ^ help_tuple[1, p_help_tuple:p_help_tuple_nxt],
                r_all_1 ^ help_tuple[2, p_help_tuple:p_help_tuple_nxt],
                r_carry_bit ^ help_tuple[7, p_help_tuple:p_help_tuple_nxt],
                m_carry_bit ^ help_tuple[11, p_help_tuple:p_help_tuple_nxt]
            ])
            data_sent = masked_tuple_elements_sent.reshape(-1)

        if flag == 2:
            masked_triple_elements_sent = np.array([
                left_most_l_all_1 ^ triple[0, p_triple:p_triple + 1],
                left_most_r_all_1 ^ triple[1, p_triple:p_triple + 1],
                left_most_l_all_1 ^ triple[0, p_triple + 1:p_triple + 2],
                left_most_r_carry_bit ^ triple[1, p_triple + 1:p_triple + 2]
            ],
                dtype='bool')
            if n_paired == 0:
                data_sent = masked_triple_elements_sent.reshape(-1)
            else:
                # print('l_bound:', l_bound)
                # print('data_sent.shape:', data_sent.shape)
                # print('masked_triple_elements_sent.shape:', masked_triple_elements_sent.shape)
                data_sent = np.append(masked_tuple_elements_sent.reshape(-1),
                                      masked_triple_elements_sent.reshape(-1),
                                      axis=0)

        # exchange the message
        send_msg_thread = sender[1 - role].send_ndarray_by_bytes_thread(
            data_sent, packbits=True)
        send_msg_thread.start()
        data_received = listener[1 - role].receive_ndarray_by_bytes(
            size=data_sent.shape, packbits=True)

        # recover the masked data

        data_recovered = data_sent ^ data_received

        p_stream = 0
        if n_paired > 0:
            p_stream_nxt = p_stream + np.prod(masked_tuple_elements_sent.shape)
            masked_tuple_elements = data_recovered[
                p_stream:p_stream_nxt].reshape(
                    masked_tuple_elements_sent.shape)
            p_stream = p_stream_nxt
            l_all_1_masked = masked_tuple_elements[0]
            m_all_1_masked = masked_tuple_elements[1]
            r_all_1_masked = masked_tuple_elements[2]
            r_carry_bit_masked = masked_tuple_elements[3]
            m_carry_bit_masked = masked_tuple_elements[4]

        if flag == 2:
            p_stream_nxt = p_stream + np.prod(
                masked_triple_elements_sent.shape)
            masked_triple_elements = data_recovered[
                p_stream:p_stream_nxt].reshape(
                    masked_triple_elements_sent.shape)
            p_stream = p_stream_nxt
            left_most_l_all_1_masked_1 = masked_triple_elements[0]
            left_most_r_all_1_masked = masked_triple_elements[1]
            left_most_l_all_1_masked_2 = masked_triple_elements[2]
            left_most_r_carry_bit_masked = masked_triple_elements[3]

        if n_paired > 0:
            # all_1 = l_all_1 & m_all_1 & r_all_1
            all_1 = calculate_3_input_gate(
                l_all_1_masked, m_all_1_masked, r_all_1_masked,
                help_tuple[[0, 1, 2, 3, 4, 5, 6],
                           p_help_tuple:p_help_tuple_nxt], const_compute)

            # carry_bit = (l_all_1 & m_all_1 & r_carry_bit) ^ (l_all_1 & m_carry_bit) ^ l_carry_bit
            carry_bit = calculate_3_input_gate(
                l_all_1_masked, m_all_1_masked, r_carry_bit_masked,
                help_tuple[[0, 1, 7, 3, 8, 9, 10],
                           p_help_tuple:p_help_tuple_nxt],
                const_compute) ^ calculate_2_input_gate(
                    l_all_1_masked, m_carry_bit_masked,
                    help_tuple[[0, 11, 12], p_help_tuple:p_help_tuple_nxt],
                    const_compute) ^ l_carry_bit

            p_help_tuple = p_help_tuple_nxt

        if flag == 2:
            # left_most_all_1 = left_most_l_all_1 & left_most_r_all_1
            left_most_all_1 = calculate_2_input_gate(
                left_most_l_all_1_masked_1, left_most_r_all_1_masked,
                triple[:, p_triple:p_triple + 1], const_compute)

            # left_most_carry_bit = left_most_l_all_1 & left_most_r_carry_bit ^ left_most_l_carry_bt
            left_most_carry_bit = calculate_2_input_gate(
                left_most_l_all_1_masked_2, left_most_r_carry_bit_masked,
                triple[:, p_triple + 1:p_triple + 2],
                const_compute) ^ left_most_l_carry_bit

            p_triple = p_triple + 2

        if left_most_all_1 is not None:
            # all_1 = np.append(all_1, left_most_all_1, axis=0)
            if n_paired == 0:
                carry_bit = left_most_carry_bit
                all_1 = left_most_all_1
            else:
                carry_bit = np.append(carry_bit, left_most_carry_bit, axis=0)
                all_1 = np.append(all_1, left_most_all_1, axis=0)

        l_bound = len(all_1) - 1
        send_msg_thread.join()

    return ret ^ carry_bit[0]

def verify_triple(triple: np.ndarray, sender, listener):
    send_msg_thread = sender.send_ndarray_by_bytes_thread(triple,
                                                          packbits=True)
    send_msg_thread.start()
    rcv_thread = listener.receive_ndarray_by_bytes_thread(size=triple.shape,
                                                          packbits=True)
    rcv_thread.start()

    send_msg_thread.join()
    rcv_thread.join()

    triple_recieved = rcv_thread.get_result()

    tmp = triple ^ triple_recieved
    assert (np.sum((tmp[0] & tmp[1]) != tmp[2]) == 0)

def count_tuple_and_triple(msb_pos: int):
    '''
    return (n_tuple, n_triple)
    '''
    n_tuple, n_triple = 0, msb_pos

    l_bound = msb_pos - 1

    while l_bound > 0:
        flag = (l_bound + 1) % 3

        if flag == 2:
            n_triple += 2

        n_pair = (l_bound + 1) // 3
        n_tuple += n_pair
        l_bound = n_pair - 1

        if flag != 0:
            l_bound += 1

    return n_tuple, n_triple

def get_triple_delta(triple_mo: np.ndarray, triple_do: np.ndarray):
    '''
    get the delta sent to model owner
    '''
    return (triple_mo[0] ^ triple_do[0]) & (triple_mo[1]
                                            ^ triple_do[1]) ^ triple_do[2]

def get_tuple_delta(tuple_mo: np.ndarray, tuple_do: np.ndarray):
    '''
    get the delta sent to model owner
    tuple:
    u,  v,  w,  uv, uw, vw, uvw
    0,  1,  2,   3,  4,  5, 6

    u,  v,  w1, uv, uw1,    vw1,    uvw1
    0,  1,  7,  3,  8,      9,      10

    u,  v2, uv2
    0,  11, 12
    '''
    u = tuple_mo[0] ^ tuple_do[0]
    v = tuple_mo[1] ^ tuple_do[1]
    w = tuple_mo[2] ^ tuple_do[2]
    w1 = tuple_mo[7] ^ tuple_do[7]
    v2 = tuple_mo[11] ^ tuple_do[11]

    ret = np.array([
        u & v ^ tuple_do[3], u & w ^ tuple_do[4], v & w ^ tuple_do[5],
        u & v & w ^ tuple_do[6], u & w1 ^ tuple_do[8], v & w1 ^ tuple_do[9],
        u & v & w1 ^ tuple_do[10], u & v2 ^ tuple_do[12]
    ],
        dtype='bool')
    return ret

def SecureBiNN(role, random_handler, listener, sender, data_share,
               model_share):
    # role 0: data owner (share_0, share_1)
    # role 1: model owner (share_1, share_2)
    # role 2: trusted third party (share_2, share_0)
    data_owner = 0
    model_owner = 1
    ttp = 2

    role_last = (role - 1) % 3
    role_next = (role + 1) % 3

    ret = data_share

    for i in range(len(model_share)):

        random_handler.step_synchronous()

        cur_layer = model_share[i]
        name = cur_layer['name']
        # print(name)
        # print('Len ret:', len(ret))

        if name.startswith('dense'):
            ret = tf.matmul(ret[0], cur_layer['W'][0]) + \
                tf.matmul(ret[0], cur_layer['W'][1]) + \
                tf.matmul(ret[1], cur_layer['W'][0]) + \
                cur_layer['b']
            ret += random_handler.get_3_out_of_3_shares(bits=int(
                cur_layer['dtype'][3:]),
                size=ret.shape)
            random_handler.cnter += 1

        elif name.startswith('conv2d'):
            ret = tf.nn.conv2d(
                ret[0], cur_layer['W'][0], strides=cur_layer['strides'], padding=cur_layer['padding']) + \
                tf.nn.conv2d(
                    ret[0], cur_layer['W'][1], strides=cur_layer['strides'], padding=cur_layer['padding']) + \
                tf.nn.conv2d(
                    ret[1], cur_layer['W'][0], strides=cur_layer['strides'], padding=cur_layer['padding'])
            ret = tf.nn.bias_add(ret, cur_layer['b'])
            ret += random_handler.get_3_out_of_3_shares(bits=int(
                cur_layer['dtype'][3:]),
                size=ret.shape)
            random_handler.cnter += 1

        elif name.startswith('batch_normalization'):
            ret += cur_layer['val']

        elif name.startswith('activation') or name.startswith('max_pooling'):
            if name.startswith('max_pooling'):
                ret = ret[0]
                ret = sum_pool(ret, kernel_size=cur_layer['pool_size'])
                if role == data_owner:
                    ret = ret - 1

            input_dtype = cur_layer['dtype']
            output_dtype = cur_layer['nxt_dtype']

            total_bit = 32
            input_bit = int(input_dtype[3:])
            pos = cur_layer['msb']

            output_bit = int(output_dtype[3:])
            ori_shape = ret.shape
            ret = ret.numpy()

            # n_triple = 3 * (pos - 1) + 1
            # n_triple = count_ppa_n_triple(msb_pos=pos)
            n_tuple, n_triple = count_tuple_and_triple(msb_pos=pos)

            # Prepare beaver's triple
            '''
            tuple:
            u,  v,  w,  uv, uw, vw, uvw
            0,  1,  2,   3,  4,  5, 6

            u,  v,  w1, uv, uw1,    vw1,    uvw1
            0,  1,  7,  3,  8,      9,      10

            u,  v2, uv2
            0,  11, 12

            generate tuple, then generate triple
            '''

            if role == ttp:
                if name.startswith('activation'):
                    send_datashare_thread = sender[
                        data_owner].send_ndarray_by_bytes_thread(ret)
                    send_datashare_thread.start()

                tuple_do = random_handler.get_same_random_with(
                    data_owner, bits=1, size=(13, n_tuple, *ret.shape))
                tuple_mo = random_handler.get_same_random_with(
                    model_owner, bits=1, size=(13, n_tuple, *ret.shape))
                random_handler.cnter += 1

                triple_do = random_handler.get_same_random_with(
                    data_owner, bits=1, size=(3, n_triple, *ret.shape))
                triple_mo = random_handler.get_same_random_with(
                    model_owner, bits=1, size=(3, n_triple, *ret.shape))
                random_handler.cnter += 1

                # delta = modify_triple_with_optimization(
                #     pos, triple_do, triple_mo)
                # delta = (triple_do[0] ^ triple_mo[0]) & (
                #     triple_do[1] ^ triple_mo[1]) ^ triple_do[2]

                delta_tuple = get_tuple_delta(tuple_mo, tuple_do)
                delta_triple = get_triple_delta(triple_mo, triple_do)
                delta = np.append(delta_tuple.reshape(-1),
                                  delta_triple.reshape(-1))

                send_delta_thread = sender[
                    model_owner].send_ndarray_by_bytes_thread(delta,
                                                              packbits=True)
                send_delta_thread.start()
                a_share = None
                b_share = None
                triple = None
                tuple = None
                send_delta_thread.join()

                if name.startswith('activation'):
                    send_datashare_thread.join()

            elif role == data_owner:
                if name.startswith('activation'):
                    ttp_share = listener[ttp].receive_ndarray_by_bytes(
                        size=ret.shape, dtype=ret.dtype)
                    ret += ttp_share

                # triple = random_handler.get_same_random_with(
                #     ttp, bits=1, size=(3, n_triple, *ret.shape))
                # triple = modify_triple_with_optimization(
                #     pos, triple_do=triple)
                # random_handler.cnter += 1

                tuple = random_handler.get_same_random_with(ttp,
                                                            bits=1,
                                                            size=(13, n_tuple,
                                                                  *ret.shape))
                random_handler.cnter += 1

                triple = random_handler.get_same_random_with(ttp,
                                                             bits=1,
                                                             size=(3, n_triple,
                                                                   *ret.shape))
                random_handler.cnter += 1

                a_share = ret
                b_share = np.zeros_like(ret)

            elif role == model_owner:

                # delta = listener[ttp].receive_ndarray_by_bytes(
                #     size=(n_triple, *ret.shape), packbits=True)
                # triple = random_handler.get_same_random_with(
                #     ttp, bits=1, size=(3, n_triple, *ret.shape))
                # triple = modify_triple_with_optimization(
                #     pos, triple_mo=triple)
                # random_handler.cnter += 1
                # triple[2] = delta

                tuple = random_handler.get_same_random_with(ttp,
                                                            bits=1,
                                                            size=(13, n_tuple,
                                                                  *ret.shape))
                random_handler.cnter += 1

                triple = random_handler.get_same_random_with(ttp,
                                                             bits=1,
                                                             size=(3, n_triple,
                                                                   *ret.shape))
                random_handler.cnter += 1

                delta = listener[ttp].receive_ndarray_by_bytes(
                    size=((8 * n_tuple + n_triple) * np.prod(ret.shape)),
                    packbits=True)

                len_delta_triple = np.prod(triple[2].shape)
                len_delta_tuple = len(delta) - len_delta_triple

                delta_tuple = np.reshape(delta[:len_delta_tuple],
                                         (-1, n_tuple, *ret.shape))
                delta_triple = np.reshape(delta[len_delta_tuple:],
                                          triple[2].shape)

                tuple[3] = delta_tuple[0]
                tuple[4] = delta_tuple[1]
                tuple[5] = delta_tuple[2]
                tuple[6] = delta_tuple[3]
                tuple[8] = delta_tuple[4]
                tuple[9] = delta_tuple[5]
                tuple[10] = delta_tuple[6]
                tuple[12] = delta_tuple[7]

                triple[2] = delta_triple

                a_share = np.zeros_like(ret)
                b_share = ret

            # if role != ttp:
            #     verify_triple(triple, sender[1-role], listener[1-role])
            # Extract MSB
            # ret = extract_msb_ppa(
            #     role=role,
            #     listener=listener,
            #     sender=sender,
            #     a_share=a_share,
            #     b_share=b_share,
            #     total_bit=total_bit,
            #     msb_pos=pos,
            #     triple=triple
            # )
            # ret = extract_msb_ppa_with_optimize(
            #     role=role,
            #     listener=listener,
            #     sender=sender,
            #     a_share=a_share,
            #     b_share=b_share,
            #     total_bit=total_bit,
            #     msb_pos=pos,
            #     triple=triple
            # )
            ret = extract_msb_with_3_input_gate(role=role,
                                                listener=listener,
                                                sender=sender,
                                                a_share=a_share,
                                                b_share=b_share,
                                                total_bit=total_bit,
                                                msb_pos=pos,
                                                triple=triple,
                                                help_tuple=tuple)

            # Here ret_mo ^ ret_do = MSB

            # 3P - Oblivious Transfer

            if role == data_owner:
                ret = ot_3party(role=role,
                                sender=sender,
                                listener=listener,
                                random_handler=random_handler,
                                msg_dtype='u' + output_dtype,
                                msg_bit=output_bit,
                                choose_bit=ret)

            elif role == model_owner:
                msg_0 = ret
                msg_1 = True ^ ret
                msg_0 = ((True ^ msg_0) +
                         msg_0 * cur_layer['deact_val']).astype(output_dtype)
                msg_1 = ((True ^ msg_1) +
                         msg_1 * cur_layer['deact_val']).astype(output_dtype)
                msg_0 = int2uint(msg_0)
                msg_1 = int2uint(msg_1)

                # add r to two messages
                r = get_random_mask(msg_0.shape, bits=output_bit)
                msg_0 -= r
                msg_1 -= r

                # send the message
                ot_3party(role=role,
                          sender=sender,
                          listener=listener,
                          random_handler=random_handler,
                          msg_0=msg_0,
                          msg_1=msg_1,
                          msg_bit=output_bit)
                ret = r

            elif role == ttp:
                ot_3party(role=role,
                          sender=sender,
                          listener=listener,
                          random_handler=random_handler,
                          msg_dtype='u' + output_dtype,
                          msg_shape=ori_shape)
                ret = np.zeros(shape=ori_shape, dtype='u' + output_dtype)

            # Reshare
            if cur_layer['deact_val'] == -1:
                random_handler.step_synchronous()
                mask = random_handler.get_3_out_of_3_shares(output_bit,
                                                            size=ret.shape)
                ret += mask
                send_thread = sender[role_last].send_ndarray_by_bytes_thread(
                    ret)
                send_thread.start()
                share_rcv = listener[role_next].receive_ndarray_by_bytes(
                    size=ret.shape, dtype=ret.dtype)
                ret = tf.constant([ret, share_rcv], dtype=tf.int32)
                send_thread.join()
            else:
                ret = tf.constant([
                    ret,
                ], dtype=tf.int32)

        elif name.startswith('bnn__dense'):
            ret = tf.matmul(ret[0], cur_layer['W'][0]) + \
                tf.matmul(ret[0], cur_layer['W'][1]) + \
                tf.matmul(ret[1], cur_layer['W'][0])
            ret += random_handler.get_3_out_of_3_shares(bits=int(
                cur_layer['dtype'][3:]),
                size=ret.shape)
            random_handler.cnter += 1

        elif name.startswith('bnn__conv2d'):
            ret = tf.nn.conv2d(
                ret[0], cur_layer['W'][0], strides=cur_layer['strides'], padding=cur_layer['padding']) + \
                tf.nn.conv2d(
                    ret[0], cur_layer['W'][1], strides=cur_layer['strides'], padding=cur_layer['padding']) + \
                tf.nn.conv2d(
                    ret[1], cur_layer['W'][0], strides=cur_layer['strides'], padding=cur_layer['padding'])
            ret += random_handler.get_3_out_of_3_shares(bits=int(
                cur_layer['dtype'][3:]),
                size=ret.shape)
            random_handler.cnter += 1

        elif name.startswith('flatten'):
            share_0 = tf.reshape(ret[0], (ret[0].shape[0], -1))
            share_1 = tf.reshape(ret[1], (ret[1].shape[0], -1))
            ret = [share_0, share_1]

    return ret

def test_secure_activation(n_bit: int, n_batch: int, role: int, listener: list, sender: list, config):
    data_owner = 0
    model_owner = 1
    ttp = 2

    role_last = (role - 1) % 3
    role_next = (role + 1) % 3

    ret = np.zeros(n_batch, dtype='int32')
    cur_layer = {
        'dtype': 'int32',
        'output_dtype': 'int32',
        'msb': n_bit-1,
        'nxt_dtype': 'int32',
        'deact_val': -1
    }

    input_dtype = cur_layer['dtype']
    output_dtype = cur_layer['nxt_dtype']

    total_bit = 32
    input_bit = int(input_dtype[3:])
    pos = cur_layer['msb']

    output_bit = int(output_dtype[3:])
    ori_shape = ret.shape

    # n_triple = 3 * (pos - 1) + 1
    # n_triple = count_ppa_n_triple(msb_pos=pos)
    n_tuple, n_triple = count_tuple_and_triple(msb_pos=pos)

    random_handler = RandomHandler(role, config['r'][role],
                                   config['r'][role_next])
    # Prepare beaver's triple
    '''
    tuple:
    u,  v,  w,  uv, uw, vw, uvw
    0,  1,  2,   3,  4,  5, 6

    u,  v,  w1, uv, uw1,    vw1,    uvw1
    0,  1,  7,  3,  8,      9,      10

    u,  v2, uv2
    0,  11, 12

    generate tuple, then generate triple
    '''

    if role == ttp:
        send_datashare_thread = sender[
            data_owner].send_ndarray_by_bytes_thread(ret)
        send_datashare_thread.start()

        tuple_do = random_handler.get_same_random_with(
            data_owner, bits=1, size=(13, n_tuple, *ret.shape))
        tuple_mo = random_handler.get_same_random_with(
            model_owner, bits=1, size=(13, n_tuple, *ret.shape))
        random_handler.cnter += 1

        triple_do = random_handler.get_same_random_with(
            data_owner, bits=1, size=(3, n_triple, *ret.shape))
        triple_mo = random_handler.get_same_random_with(
            model_owner, bits=1, size=(3, n_triple, *ret.shape))
        random_handler.cnter += 1

        # delta = modify_triple_with_optimization(
        #     pos, triple_do, triple_mo)
        # delta = (triple_do[0] ^ triple_mo[0]) & (
        #     triple_do[1] ^ triple_mo[1]) ^ triple_do[2]

        delta_tuple = get_tuple_delta(tuple_mo, tuple_do)
        delta_triple = get_triple_delta(triple_mo, triple_do)
        delta = np.append(delta_tuple.reshape(-1),
                          delta_triple.reshape(-1))

        send_delta_thread = sender[
            model_owner].send_ndarray_by_bytes_thread(delta,
                                                      packbits=True)
        send_delta_thread.start()
        a_share = None
        b_share = None
        triple = None
        tuple = None
        send_delta_thread.join()

        send_datashare_thread.join()

    elif role == data_owner:
        ttp_share = listener[ttp].receive_ndarray_by_bytes(
            size=ret.shape, dtype=ret.dtype)
        ret += ttp_share

        tuple = random_handler.get_same_random_with(ttp,
                                                    bits=1,
                                                    size=(13, n_tuple,
                                                          *ret.shape))
        random_handler.cnter += 1

        triple = random_handler.get_same_random_with(ttp,
                                                     bits=1,
                                                     size=(3, n_triple,
                                                           *ret.shape))
        random_handler.cnter += 1

        a_share = ret
        b_share = np.zeros_like(ret)

    elif role == model_owner:

        tuple = random_handler.get_same_random_with(ttp,
                                                    bits=1,
                                                    size=(13, n_tuple,
                                                          *ret.shape))
        random_handler.cnter += 1

        triple = random_handler.get_same_random_with(ttp,
                                                     bits=1,
                                                     size=(3, n_triple,
                                                           *ret.shape))
        random_handler.cnter += 1

        delta = listener[ttp].receive_ndarray_by_bytes(
            size=((8 * n_tuple + n_triple) * np.prod(ret.shape)),
            packbits=True)

        len_delta_triple = np.prod(triple[2].shape)
        len_delta_tuple = len(delta) - len_delta_triple

        delta_tuple = None
        delta_triple = None

        if len_delta_tuple > 0:
            delta_tuple = np.reshape(delta[:len_delta_tuple],
                                     (-1, n_tuple, *ret.shape))
            tuple[3] = delta_tuple[0]
            tuple[4] = delta_tuple[1]
            tuple[5] = delta_tuple[2]
            tuple[6] = delta_tuple[3]
            tuple[8] = delta_tuple[4]
            tuple[9] = delta_tuple[5]
            tuple[10] = delta_tuple[6]
            tuple[12] = delta_tuple[7]

        if len_delta_triple > 0:
            delta_triple = np.reshape(delta[len_delta_tuple:],
                                      triple[2].shape)
            triple[2] = delta_triple

        a_share = np.zeros_like(ret)
        b_share = ret

    ret = extract_msb_with_3_input_gate(role=role,
                                        listener=listener,
                                        sender=sender,
                                        a_share=a_share,
                                        b_share=b_share,
                                        total_bit=total_bit,
                                        msb_pos=pos,
                                        triple=triple,
                                        help_tuple=tuple)

    # Here ret_mo ^ ret_do = MSB

    # 3P - Oblivious Transfer

    if role == data_owner:
        ret = ot_3party(role=role,
                        sender=sender,
                        listener=listener,
                        random_handler=random_handler,
                        msg_dtype='u' + output_dtype,
                        msg_bit=output_bit,
                        choose_bit=ret)

    elif role == model_owner:
        msg_0 = ret
        msg_1 = True ^ ret
        msg_0 = ((True ^ msg_0) +
                 msg_0 * cur_layer['deact_val']).astype(output_dtype)
        msg_1 = ((True ^ msg_1) +
                 msg_1 * cur_layer['deact_val']).astype(output_dtype)
        msg_0 = int2uint(msg_0)
        msg_1 = int2uint(msg_1)

        # add r to two messages
        r = get_random_mask(msg_0.shape, bits=output_bit)
        msg_0 -= r
        msg_1 -= r

        # send the message
        ot_3party(role=role,
                  sender=sender,
                  listener=listener,
                  random_handler=random_handler,
                  msg_0=msg_0,
                  msg_1=msg_1,
                  msg_bit=output_bit)
        ret = r

    elif role == ttp:
        ot_3party(role=role,
                  sender=sender,
                  listener=listener,
                  random_handler=random_handler,
                  msg_dtype='u' + output_dtype,
                  msg_shape=ori_shape)
        ret = np.zeros(shape=ori_shape, dtype='u' + output_dtype)

    # Reshare
    random_handler.step_synchronous()
    mask = random_handler.get_3_out_of_3_shares(output_bit,
                                                size=ret.shape)
    ret += mask
    send_thread = sender[role_last].send_ndarray_by_bytes_thread(
        ret)
    send_thread.start()
    share_rcv = listener[role_next].receive_ndarray_by_bytes(
        size=ret.shape, dtype=ret.dtype)
    ret = tf.constant([ret, share_rcv], dtype=tf.int32)
    send_thread.join()

    return None


"""rio-tiler models."""

import itertools
import warnings
from enum import Enum
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

import attr
import numpy
from affine import Affine
from pydantic import BaseModel
from rasterio.coords import BoundingBox
from rasterio.crs import CRS
from rasterio.dtypes import dtype_ranges
from rasterio.plot import reshape_as_image
from rasterio.transform import from_bounds
from rio_color.operations import parse_operations
from rio_color.utils import scale_dtype, to_math_type

from .errors import InvalidDatatypeWarning
from .types import ColorMapType, GDALColorMapType, IntervalTuple, NumType
from .utils import linear_rescale, render, resize_array

class NodataTypes(str, Enum):
    """rio-tiler Nodata types."""

    Alpha = "Alpha"
    Mask = "Mask"
    Internal = "Internal"
    Nodata = "Nodata"
    Empty = "None"

class RioTilerBaseModel(BaseModel):
    """Base Model for rio-tiler models."""

    def __getitem__(self, item):
        """Keep `getter` access for compatibility."""
        return self.__dict__[item]

class Bounds(RioTilerBaseModel):
    """Dataset Bounding box"""

    bounds: BoundingBox

class SpatialInfo(Bounds):
    """Dataset SpatialInfo"""

    minzoom: int
    maxzoom: int

class Info(SpatialInfo):
    """Dataset Info."""

    band_metadata: List[Tuple[str, Dict]]
    band_descriptions: List[Tuple[str, str]]
    dtype: str
    nodata_type: NodataTypes
    colorinterp: Optional[List[str]]
    scale: Optional[float]
    offset: Optional[float]
    colormap: Optional[GDALColorMapType]

    class Config:
        """Config for model."""

        extra = "allow"
        use_enum_values = True

class BandStatistics(RioTilerBaseModel):
    """Image statistics"""

    min: float
    max: float
    mean: float
    count: float
    sum: float
    std: float
    median: float
    majority: float
    minority: float
    unique: float
    histogram: List[List[NumType]]
    valid_percent: float
    masked_pixels: float
    valid_pixels: float

    class Config:
        """Config for model."""

        extra = "allow"  # We allow extra values for `percentiles_{}`

def to_coordsbbox(bbox) -> Optional[BoundingBox]:
    """Convert bbox to CoordsBbox nameTuple."""
    return BoundingBox(*bbox) if bbox else None

def rescale_image(
    data: numpy.ndarray,
    mask: numpy.ndarray,
    in_range: Sequence[IntervalTuple],
    out_range: Sequence[IntervalTuple] = ((0, 255),),
    out_dtype: Union[str, numpy.number] = "uint8",
):
    """Rescale image data."""
    if len(data.shape) < 3:
        data = numpy.expand_dims(data, axis=0)

    nbands = data.shape[0]

    if len(in_range) != nbands:
        in_range = ((in_range[0]),) * nbands

    if len(out_range) != nbands:
        out_range = ((out_range[0]),) * nbands

    for bdx in range(nbands):
        data[bdx] = numpy.where(
            mask,
            linear_rescale(data[bdx], in_range=in_range[bdx], out_range=out_range[bdx]),
            0,
        )

    return data.astype(out_dtype)

@attr.s
class ImageData:
    """Image Data class.

    Attributes:
        data (numpy.ndarray): pixel values.
        mask (numpy.ndarray): rasterio mask values.
        assets (list, optional): list of assets used to construct the data values.
        bounds (BoundingBox, optional): bounding box of the data.
        crs (rasterio.crs.CRS, optional): Coordinates Reference System of the bounds.
        metadata (dict, optional): Additional metadata. Defaults to `{}`.
        band_names (list, optional): name of each band. Defaults to `["1", "2", "3"]` for 3 bands image.

    """

    data: numpy.ndarray = attr.ib()
    mask: numpy.ndarray = attr.ib()
    assets: Optional[List] = attr.ib(default=None)
    bounds: Optional[BoundingBox] = attr.ib(default=None, converter=to_coordsbbox)
    crs: Optional[CRS] = attr.ib(default=None)
    metadata: Optional[Dict] = attr.ib(factory=dict)
    band_names: List[str] = attr.ib()

    @data.validator
    def _validate_data(self, attribute, value):
        """ImageData data has to be a 3d array in form of (count, height, width)"""
        if not len(value.shape) == 3:
            raise ValueError(
                "ImageData data has to be an array in form of (count, height, width)"
            )

    @band_names.default
    def _default_names(self):
        return [f"{ix + 1}" for ix in range(self.count)]

    @mask.default
    def _default_mask(self):
        return numpy.zeros((self.height, self.width), dtype="uint8") + 255

    def __iter__(self):
        """Allow for variable expansion (``arr, mask = ImageData``)"""
        for i in (self.data, self.mask):
            yield i

    @classmethod
    def create_from_list(cls, data: Sequence["ImageData"]):
        """Create ImageData from a sequence of ImageData objects.

        Args:
            data (sequence): sequence of ImageData.

        """
        h, w = zip(*[(img.height, img.width) for img in data])
        if len(set(h)) > 1 or len(set(w)) > 1:
            warnings.warn(
                "Cannot concatenate images with different size. Will resize using max width/heigh",
                UserWarning,
            )
            max_h, max_w = max(h), max(w)
            for img in data:
                if img.height == max_h and img.width == max_w:
                    continue
                img.data = resize_array(img.data, max_h, max_w)
                img.mask = resize_array(img.mask, max_h, max_w)

        arr = numpy.concatenate([img.data for img in data])
        mask = numpy.all([img.mask for img in data], axis=0).astype(numpy.uint8) * 255
        assets = list(
            dict.fromkeys(
                itertools.chain.from_iterable(
                    [img.assets for img in data if img.assets]
                )
            )
        )

        bounds_values = [img.bounds for img in data if img.bounds]
        bounds = bounds_values[0] if bounds_values else None

        crs_values = [img.crs for img in data if img.crs]
        crs = crs_values[0] if crs_values else None

        band_names = list(
            itertools.chain.from_iterable(
                [img.band_names for img in data if img.band_names]
            )
        )

        return cls(
            arr, mask, assets=assets, crs=crs, bounds=bounds, band_names=band_names
        )

    def as_masked(self) -> numpy.ma.MaskedArray:
        """return a numpy masked array."""
        data = numpy.ma.array(self.data)
        data.mask = self.mask == 0
        return data

    def data_as_image(self) -> numpy.ndarray:
        """Return the data array reshaped into an image processing/visualization software friendly order.

        (bands, rows, columns) -> (rows, columns, bands).

        """
        return reshape_as_image(self.data)

    @property
    def width(self) -> int:
        """Width of the data array."""
        return self.data.shape[2]

    @property
    def height(self) -> int:
        """Height of the data array."""
        return self.data.shape[1]

    @property
    def count(self) -> int:
        """Number of band."""
        return self.data.shape[0]

    @property
    def transform(self):
        """Returns the affine transform."""
        return (
            from_bounds(*self.bounds, self.width, self.height)
            if self.bounds
            else Affine.scale(self.width, -self.height)
        )

    def rescale(
        self,
        in_range: Sequence[IntervalTuple],
        out_range: Sequence[IntervalTuple] = ((0, 255),),
        out_dtype: Union[str, numpy.number] = "uint8",
    ):
        """Rescale data in place."""
        self.data = rescale_image(self.data, self.mask, in_range, out_range, out_dtype)

    def apply_color_formula(self, color_formula: Optional[str]):
        """Apply rio-color formula in place."""
        self.data[self.data < 0] = 0
        for ops in parse_operations(color_formula):
            self.data = scale_dtype(ops(to_math_type(self.data)), numpy.uint8)

    def post_process(
        self,
        in_range: Optional[Sequence[IntervalTuple]] = None,
        out_dtype: Union[str, numpy.number] = "uint8",
        color_formula: Optional[str] = None,
        **kwargs: Any,
    ) -> "ImageData":
        """Post-process image data.

        Args:
            in_range (tuple): input min/max bounds value to rescale from.
            out_dtype (str, optional): output datatype after rescaling. Defaults to `uint8`.
            color_formula (str, optional): rio-color formula (see: https://github.com/mapbox/rio-color).
            kwargs (optional): keyword arguments to forward to `rio_tiler.utils.linear_rescale`.

        Returns:
            ImageData: new ImageData object with the updated data.

        Examples:
            >>> img.post_process(in_range=((0, 16000), ))

            >>> img.post_process(color_formula="Gamma RGB 4.1")

        """
        data = self.data.copy()
        mask = self.mask.copy()

        if in_range:
            data = rescale_image(data, mask, in_range, out_dtype=out_dtype, **kwargs)

        if color_formula:
            data[data < 0] = 0
            for ops in parse_operations(color_formula):
                data = scale_dtype(ops(to_math_type(data)), numpy.uint8)

        return ImageData(
            data,
            mask,
            crs=self.crs,
            bounds=self.bounds,
            assets=self.assets,
            metadata=self.metadata,
        )

    def render(
        self,
        add_mask: bool = True,
        img_format: str = "PNG",
        colormap: Optional[ColorMapType] = None,
        **kwargs,
    ) -> bytes:
        """Render data to image blob.

        Args:
            add_mask (bool, optional): add mask to output image. Defaults to `True`.
            img_format (str, optional): output image format. Defaults to `PNG`.
            colormap (dict or sequence, optional): RGBA Color Table dictionary or sequence.
            kwargs (optional): keyword arguments to forward to `rio_tiler.utils.render`.

        Returns:
            bytes: image.

        """
        img_format = img_format.upper()

        if img_format == "GTIFF":
            if "transform" not in kwargs:
                kwargs.update({"transform": self.transform})
            if "crs" not in kwargs and self.crs:
                kwargs.update({"crs": self.crs})

        data = self.data.copy()
        datatype_range = (dtype_ranges[str(data.dtype)],)

        if not colormap:
            if img_format in ["PNG"] and data.dtype not in ["uint8", "uint16"]:
                warnings.warn(
                    f"Invalid type: `{data.dtype}` for the `{img_format}` driver. Data will be rescaled using min/max type bounds.",
                    InvalidDatatypeWarning,
                )
                data = rescale_image(data, self.mask, in_range=datatype_range)

            elif img_format in ["JPEG", "WEBP"] and data.dtype not in ["uint8"]:
                warnings.warn(
                    f"Invalid type: `{data.dtype}` for the `{img_format}` driver. Data will be rescaled using min/max type bounds.",
                    InvalidDatatypeWarning,
                )
                data = rescale_image(data, self.mask, in_range=datatype_range)

            elif img_format in ["JP2OPENJPEG"] and data.dtype not in [
                "uint8",
                "int16",
                "uint16",
            ]:
                warnings.warn(
                    f"Invalid type: `{data.dtype}` for the `{img_format}` driver. Data will be rescaled using min/max type bounds.",
                    InvalidDatatypeWarning,
                )
                data = rescale_image(data, self.mask, in_range=datatype_range)

        if add_mask:
            return render(
                data, self.mask, img_format=img_format, colormap=colormap, **kwargs
            )

        return render(data, img_format=img_format, colormap=colormap, **kwargs)


import argparse
import sys
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
from torch.nn.utils.rnn import pad_packed_sequence

from torch.nn.utils.rnn import pack_padded_sequence
import preprocess
from model import EncoderRNN, DecoderRNN, VanillaDecoderRNN, BAttnDecoderRNN, LAttnDecoderRNN
import dataloader
from gensim.models import word2vec
from gensim import models

def getDataLoader(data, batch_size=100, num_workers=4):
    dataset = dataloader.TestDataset(data)
    loader = Data.DataLoader(
        dataset=dataset,  # torch TensorDataset format
        batch_size=batch_size,  # mini batch size
        shuffle=False,
        num_workers=num_workers,
        collate_fn=dataloader.test_collate_fn
    )
    return loader

def to_var(x, volatile=False):
    if torch.cuda.is_available():
        x = x.cuda()
    return Variable(x, volatile=volatile)

def getwv():
    i2v = preprocess.load_obj('index2vocab')
    model = models.Word2Vec.load('w2v200.model.bin')
    embedding = np.zeros((len(i2v), model.wv.vector_size))
    j = 0
    for i in range(len(i2v)):
        try:
            embedding[i] = model.wv[i2v[i]]
        except:
            print(i2v[i], 'not in w2v')
            j += 1
            continue
    print(j)
    return embedding

def main(args):
    fout = open(args.output, 'w')
    maxlen = 20
    num_layers = 1
    dim = 512
    quotes_set = set(['[',']','{','}','!','?','。'])
    sentences, v_size = preprocess.label2onehot(args.input)
    data = getDataLoader(sentences)
    v = preprocess.load_obj('vocabindex')
    i2v = preprocess.load_obj('index2vocab')
    print(sentences)
    v_size = len(v)
    print('sentences shape =', sentences.shape)
    we = getwv()
    encoder = EncoderRNN(v_size, dim, we, num_layers=num_layers, bidirectional=False)
    #decoder = VanillaDecoderRNN(dim, v_size, num_layers=num_layers)
    #decoder = BAttnDecoderRNN(dim, v_size, num_layers=num_layers)
    decoder = LAttnDecoderRNN(dim, v_size, we, num_layers=num_layers)
    encoder.load_state_dict(torch.load(args.encoder))
    decoder.load_state_dict(torch.load(args.decoder))
    if torch.cuda.is_available():
        encoder.cuda()
        decoder.cuda()
    print(encoder)
    print(decoder)
    i = 0
    for sentence in sentences:
        print(sentence)
        s = to_var(torch.LongTensor(sentence).view(1, -1))
        encoder_hidden = encoder.initHidden(num_layers, 1)
        encoder_o, encoder_hidden = encoder(s, [s.size(1)], encoder_hidden)
        encoder_o = pad_packed_sequence(encoder_o, batch_first=True)[0]
        decoder_hidden = encoder_hidden[:num_layers]
        inputword = v['<start>']
        flag = True
        sentence = ''
        words = []
        length = 0
        while flag:
            inputword = to_var(torch.LongTensor([inputword]).view(1,-1))
            output, decoder_hidden = decoder(inputword,decoder_hidden, encoder_o)
            maxkey = np.argmax(output[0].data)
            inputword = maxkey
            word = i2v[maxkey.item()]
            length += 1
            if length > maxlen:
                flag = False
            if word == '<end>' or word == '<pad>':
                flag = False
            elif word == '<unk>':
                continue
            else:
                if word == '.' or word == '。':
                    #sentence = sentence[:-1]
                    #sentence += word
                    flag = False
                else:
                    if word in words and word in quotes_set:
                        continue
                    if (len(words) == 0 or word != words[-1]):
                        words.append(word)
        sentence = ' '.join(words)
        print(sentence)
        fout.write(sentence+'\n')
        i += 1

    fout.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, required=True, help='input file')
    parser.add_argument("--encoder", type=str, required=True, help='encoder model')
    parser.add_argument("--decoder", type=str, required=True, help='decoder model')
    parser.add_argument("--output", type=str, required=True, help='output file')
    args = parser.parse_args()
    main(args)


from predictionserver.futureconventions.activityconventions import (
    Activity, ActivityContext
)
from predictionserver.futureconventions.memoconventions import Memo
from predictionserver.serverhabits.sortedsethabits import (
    SortedSetType, SortedSetNameGranularity, SortedSetKeyGranularity
)
from predictionserver.servermixins.ownershipserver import OwnershipServer
from predictionserver.servermixins.memoserver import MemoServer
from collections import OrderedDict
from typing import Union
from copy import deepcopy

class SortedSetServer(MemoServer, OwnershipServer):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def _authorize_sortedset_change(
            self,
            sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            **kwargs
    ):
        if name_granularity in [
                SortedSetNameGranularity.name,
                SortedSetNameGranularity.name_and_delay]:
            return self._authorize(name=kwargs['name'], write_key=kwargs['write_key'])
        elif name_granularity == SortedSetNameGranularity.write_key:
            return self.is_valid_key(write_key=kwargs['write_key'])
        elif name_granularity in [
            SortedSetNameGranularity.code, SortedSetNameGranularity.code_and_genus
        ]:
            return self.shash(write_key=kwargs['write_key']) == kwargs['code']
        elif name_granularity in [SortedSetNameGranularity.memory]:
            return False  # System only
        else:
            return True

    def get_sortedset_values(
            self,
            sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            **name_kwargs
    ):
        """ Retrieve entire sortedset_type """
        location = self.sortedset_location(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            **name_kwargs
        )
        raw = list(reversed(
            self.client.zrange(
                name=location,
                start=0,
                end=-1,
                withscores=True)
        ))
        return OrderedDict([(k, float(v) if v is not None else 0.0) for k, v in raw])

    def add_sortedset_value(
            self, sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            key_granularity: SortedSetKeyGranularity,
            name_kwargs: dict,
            key_kwargs: dict,
            value: float,
            write_key=None,
            verbose=False
    ):
        return self._add_sortedset_value_implementation(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            key_granularity=key_granularity,
            name_kwargs=name_kwargs,
            write_key=write_key,
            key_kwargs=key_kwargs,
            value=value,
            verbose=verbose
        )

    def multiply_sortedset(
            self,
            sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            name_kwargs: dict,
            weight: float,
            write_key=None,
            verbose=False
    ):
        return self._multiply_sortedset_implementation(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            name_kwargs=name_kwargs,
            weight=weight,
            write_key=write_key,
            verbose=verbose
        )

    def delete_sortedset(
            self,
            sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            name_kwargs: dict,
            write_key: str,
            verbose=False
    ):
        """ If permission, confirmed delete of sortedset_type by user """
        return self._delete_sortedset_implementation(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            name_kwargs=name_kwargs,
            write_key=write_key,
            verbose=verbose
        )

    # ------------------ #
    #   Implementation   #
    # ------------------ #

    def _delete_sortedset_implementation(
            self,
            sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            name_kwargs: dict,
            verbose: Union[bool, int],
            write_key=None
    ):
        """ If permission, confirmed delete of sortedset_type by user """
        write_key = write_key or name_kwargs.get('write_key')
        allowed = self._authorize_sortedset_change(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            write_key=write_key,
            code=name_kwargs.get('code'),
            name=name_kwargs.get('name')
        )
        location = self.sortedset_location(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            **name_kwargs
        )
        memo_kwargs = deepcopy(name_kwargs)
        memo_kwargs.update({'write_key': write_key})
        memo = Memo(
            activity=Activity.delete,
            sortedset_type=sortedset_type,
            allowed=int(allowed),
            success=int(allowed),
            **memo_kwargs
        )
        if allowed:
            execution = self.client.delete(location)
            assert self.add_memo_as_owner_confirm(
                memo=memo, success=1, execution=execution
            )
        else:
            assert self.add_memo_as_owner_error(memo=memo, success=0, execution=0)
        return allowed if not verbose else memo.as_dict()

    def _add_sortedset_value_implementation(
            self,
            sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            key_granularity: SortedSetKeyGranularity,
            name_kwargs: dict,
            key_kwargs: dict,
            verbose: Union[bool, int],
            value: float,
            write_key=None,
            context: ActivityContext = None
    ):
        """ If permission, confirmed setting of sortedset_type value by user """
        write_key = write_key or name_kwargs.get('write_key') or key_kwargs.get(
            'write_key')
        allowed = self._authorize_sortedset_change(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            write_key=write_key,
            code=name_kwargs.get('code'),
            name=name_kwargs.get('name')
        )
        location = self.sortedset_location(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            **name_kwargs
        )

        key = self.sortedset_key(key_granularity=key_granularity, **key_kwargs)
        memo_kwargs = deepcopy(name_kwargs)
        memo_kwargs.update({'write_key': write_key})
        memo = Memo(
            activity=Activity.add,
            sortedset_type=sortedset_type,
            context=context,
            allowed=int(allowed),
            success=int(allowed),
            **memo_kwargs
        )
        if allowed:
            mapping = {key: value}
            execution = self.client.zadd(name=location, mapping=mapping)
            assert self.add_memo_as_owner_confirm(
                memo=memo, success=1, execution=execution
            )
        else:
            assert self.add_memo_as_owner_error(memo=memo, success=0, execution=0)
        return allowed if not verbose else memo.as_dict()

    def _multiply_sortedset_implementation(
            self,
            sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            name_kwargs: dict,
            weight: float,
            write_key=None,
            verbose=False
    ):
        """ If permission, multiply all values in sortedset_type by float multiplier """
        write_key = write_key or name_kwargs.get('write_key')
        allowed = self._authorize_sortedset_change(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            write_key=write_key,
            code=name_kwargs.get('code'),
            name=name_kwargs.get('name')
        )
        location = self.sortedset_location(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            **name_kwargs
        )
        memo_kwargs = deepcopy(name_kwargs)
        memo_kwargs.update({'write_key': write_key})
        memo = Memo(
            activity=Activity.multiply,
            sortedset_type=sortedset_type,
            allowed=int(allowed),
            success=int(allowed),
            value=weight,
            **memo_kwargs
        )
        if allowed:
            import random
            temporary_key = 'temporary_' + \
                ''.join([random.choice(['a', 'b', 'c']) for _ in range(20)])
            shrink_pipe = self.client.pipeline(transaction=True)
            shrink_pipe.zunionstore(dest=temporary_key, keys={location: weight})
            shrink_pipe.zunionstore(dest=location, keys={temporary_key: 1})
            shrink_pipe.expire(name=temporary_key, time=1)
            _ = shrink_pipe.execute()
            assert self.add_memo_as_owner_confirm(memo=memo, success=1, execution=1)
        else:
            assert self.add_memo_as_owner_error(memo=memo, success=0, execution=0)
        return allowed if not verbose else memo.as_dict()

    def _pipe_incr_sortedset_value_implementation(
            self,
            pipe,
            sortedset_type: SortedSetType,
            name_granularity: SortedSetNameGranularity,
            key_granularity: SortedSetKeyGranularity,
            name_kwargs: dict,
            key_kwargs: dict,
            amount=float
    ):
        """ System use only """
        location = self.sortedset_location(
            sortedset_type=sortedset_type,
            name_granularity=name_granularity,
            **name_kwargs
        )
        key = self.sortedset_key(key_granularity=key_granularity, **key_kwargs)
        pipe = pipe.zincrby(name=location, value=key, amount=amount)
        return pipe


from django.shortcuts import render, redirect, get_object_or_404
from django.contrib import messages
from django.contrib.auth.decorators import login_required
from .forms import UserRegisterForm, UserUpdateForm, ProfileUpdateForm
from django.http import HttpResponse
from django.shortcuts import render, redirect
from django.contrib.auth import login, authenticate
from django.contrib.sites.shortcuts import get_current_site
from django.utils.encoding import force_bytes, force_text
from django.utils.http import urlsafe_base64_encode, urlsafe_base64_decode
from django.template.loader import render_to_string
from .tokens import account_activation_token
from django.contrib.auth.models import User
from django.core.mail import EmailMessage
from blog.models import Post

def register(request):
    if request.method == 'POST':
        form = UserRegisterForm(request.POST)
        if form.is_valid():
            user = form.save(commit=False)
            user.is_active = False
            user.save()
            current_site = get_current_site(request)
            mail_subject = 'Активируйте свой аккаунт'
            message = render_to_string('users/acc_active_email.html', {
                'user': user,
                'domain': current_site.domain,
                'uid': urlsafe_base64_encode(force_bytes(user.pk)),
                'token':account_activation_token.make_token(user),
            })
            to_email = form.cleaned_data.get('email')
            email = EmailMessage(
                        mail_subject, message, to=[to_email]
            )
            email.send()
            return HttpResponse('
Подтвердите свой аккаунт через почту
</body')
    else:
        form = UserRegisterForm()
    context = {
        'title': 'Register',
        'form': form
    }
    return render(request, 'users/register.html', context)

def activate(request, uidb64, token):
    try:
        uid = force_text(urlsafe_base64_decode(uidb64))
        user = User.objects.get(pk=uid)
    except(TypeError, ValueError, OverflowError, User.DoesNotExist):
        user = None
    if user is not None and account_activation_token.check_token(user, token):
        user.is_active = True
        user.save()
        login(request, user)
        # return redirect('home')
        return HttpResponse('
Спасибо! Вы подтвердили свой аккаунт !
</body')
    else:
        return HttpResponse('
Что-то пошло не так !!!
</body')

@login_required
def profile(request):
    context = {
        'title': 'Profile',
        'posts': Post.objects.all()
    }
    return render(request, 'users/profile.html', context)

@login_required
def profile_edit(request):
    if request.method == "POST":
        u_form = UserUpdateForm(request.POST, instance=request.user)
        p_form = ProfileUpdateForm(request.POST, request.FILES, instance=request.user.profile)
        if u_form.is_valid() and p_form.is_valid():
            u_form.save()
            p_form.save()
            messages.success(request, f'Profile of changed!')
            return redirect('profile')
    else:
        u_form = UserUpdateForm(instance=request.user)
        p_form = ProfileUpdateForm(instance=request.user.profile)
    context = {
        'title': 'Profile Edit',
        'u_form': u_form,
        'p_form': p_form
    }
    return render(request, 'users/profile_edit.html', context)


import librosa
import pickle
import numpy as np
import random
import torch
from torch import nn
import os
import soundfile as sf
from shutil import copyfile
import threading
import matplotlib.pyplot as plt
from ConvAE import ConvAE, create_network, accuracy_1_min_mab, normalized_loss
from pipeline_whole import SampleLibrary
sampler_dataset = pickle.load(open("simple_dataset.pkl", "rb"))
print("done loading data")

def worker_init_fn(worker_id):
    seed = np.random.get_state()[1][0] + worker_id
    np.random.seed(seed)
    random.seed(seed)

sampler_dataset_iterable = iter(torch.utils.data.DataLoader(sampler_dataset,
                                                      batch_size=8,
                                                      num_workers=0,
                                                      # pin_memory=True,
                                                      worker_init_fn=worker_init_fn)) # what is pin memory?
fig, (ax1, ax2, ax3) = plt.subplots(ncols=3)
plt.ion() #needed to prevent show() from blocking
# plt.figure(figsize=(15,25))

def visualize(content, style, output, save = False, name = None, step = 'no-step', visible = True):
    ax1.clear()
    ax2.clear()
    ax3.clear()
    ax1.set_axis_off()
    ax2.set_axis_off()
    ax3.set_axis_off()
    ax1.title.set_text(f"Ground Truth")
    ax1.imshow(content)
    ax2.title.set_text(f"Style ({name})")
    ax2.imshow(style)
    ax3.title.set_text(f"Output")
    ax3.imshow(output)
    if visible:
        plt.show()
    plt.pause(1)
    if save and name is None:
        plt.savefig("test.png")
    elif save:
        plt.savefig(f"{name}_{step}.png", bbox_inches='tight',pad_inches = 0)

def phase_reconstructor(output, filename):
    N_FFT = 1024
    output = np.squeeze(output)
    a = np.zeros_like(output)
    a = np.exp(output) - 1
    print("phase reconstruction")
    # This code is supposed to do phase reconstruction
    p = 2 * np.pi * np.random.random_sample(a.shape) - np.pi
    for i in range(500):
        if i % 100 == 0:
            print("\t", i)
        S = a * np.exp(1j * p)
        x = librosa.istft(S)
        p = np.angle(librosa.stft(x, N_FFT))
    sf.write(filename, x, 22050)

def make_generator():
    feature_maps = 64
    depth = 8
    pooling_freq = 1e100  # large number to disable pooling layers
    strided_conv_freq = 2
    strided_conv_feature_maps = 64
    input_dim = (1, 513, 513)
    style_hidden = 32
    num_style_interferences = 3

    CONV_ENC_BLOCK = [("conv1", feature_maps), ("relu1", None)]
    CONV_ENC_LAYERS = create_network(CONV_ENC_BLOCK, depth,
                                        pooling_freq=pooling_freq,
                                        strided_conv_freq=strided_conv_freq,
                                        strided_conv_channels=strided_conv_feature_maps,
#                                         batch_norm_freq = 1e100
                                        )
    CONV_ENC_NW = CONV_ENC_LAYERS
    model = ConvAE(input_dim, enc_config=CONV_ENC_NW, num_instruments = 5, style_hidden_size = style_hidden, num_style_interferences =
                   num_style_interferences)
    return model

def make_classifier():
    feature_maps = 8
    depth = 12
    pooling_freq = 1e100  # large number to disable pooling layers
    strided_conv_freq = 2
    strided_conv_feature_maps = 8
    code_size = 5
    input_dim = (1, 513, 513)
    CONV_ENC_BLOCK = [("conv1", feature_maps), ("relu1", None)]
    CONV_ENC_LAYERS = create_network(CONV_ENC_BLOCK, depth,
                                     pooling_freq=pooling_freq,
                                     strided_conv_freq=strided_conv_freq,
                                     strided_conv_channels=strided_conv_feature_maps,
                                     )

    CONV_ENC_NW = CONV_ENC_LAYERS + [("flatten1", None), ("linear1", code_size), ("softmax1", None)]

    model = ConvAE(input_dim, enc_config=CONV_ENC_NW, disable_decoder = True)
    return model

reconstructor_thread_list = list()

def test_evaluate(generator_model, classifier_frozen, num_tests_per_instrument, step, save = True, diff_target = False):
    instruments = ["distortion", "harp", "harpsichord", "piano", "timpani"]
    averages_dict = {}
    style = None
#     input(generator_model.state_dict().keys())
    generator_model.eval()
    if save:
        print("saving model!")
        generator_model.save_feature_maps(step)
        torch.save(generator_model.state_dict(), f"ae_params_{step}.pt")

    soft_make_dir(f"evaluation_{step}")
    os.chdir(f"evaluation_{step}")
    if torch.cuda.is_available():
        device = "cuda"
    else:
        device = "cpu"

    for instrument in instruments:
        print(instrument)
        for i in range(num_tests_per_instrument):
            print(f"\t{i}")
            if diff_target:
                data, target, style = sampler_dataset.sampleTest(test_instrument=instrument)
                style = np.expand_dims(style, axis = 0)
                style = torch.as_tensor(style, device=device)

            else:
                data, target, one_hot = sampler_dataset.samplePair(test = True, test_instrument = instrument)

            data = np.expand_dims(data, axis = 0)
            target = np.expand_dims(target, axis = 0)
            data = torch.as_tensor(data, device=device)
            target = torch.as_tensor(target, device=device)

            style = style if diff_target else target

            style_vec = classifier_frozen.forward(style, None)

            out, code = generator_model.forward(data, style_vec)

            if torch.cuda.is_available():
                out = out.cpu().detach()
                target = target.cpu().detach()
                data = data.cpu().detach()
                style = style.cpu().detach()

            encoding_loss = ((out - target)**2).mean()
            if instrument not in averages_dict.keys():
                averages_dict[instrument] = 0
            averages_dict[instrument] += encoding_loss.numpy()
            print("saving plot!")
            visualize(data[0][0], style[0][0], out[0][0], save=True, name=instrument, step = i, visible = False)

            print("saving generating audios")
            t_gen = threading.Thread(target=phase_reconstructor, args=(out.numpy()[0], f"generated_{instrument}_{i}.wav"))
            t_truth = threading.Thread(target=phase_reconstructor, args=(target.numpy()[0], f"truth_{instrument}_{i}.wav"))
            t_gen.start()
            t_truth.start()
            reconstructor_thread_list.append(t_gen)
            reconstructor_thread_list.append(t_truth)
        averages_dict[instrument] /= num_tests_per_instrument

        for t in reconstructor_thread_list: #multithreading per trial
            t.join()
    print(averages_dict)
    import json
    json = json.dumps(averages_dict)
    f = open("losses.json", "w")
    f.write(json)
    f.close()
    os.chdir("../")
    generator_model.train()

def soft_make_dir(path):
    try:
        os.mkdir(path)
    except:
        print("directory already exists!")

def load_and_evaluate(generator_model, classifier_frozen, file_name, num_tests_per_instrument, step):
    generator_model.load_state_dict(torch.load(file_name))
    generator_model.load_feature_maps(step)
#     for key, value in generator_model.state_dict().items():
#         print(f"key {key}, value {value.shape}")
#     input("asdf")
    print("done loading!")
    soft_make_dir("evaluation")
    os.chdir("evaluation")
    test_evaluate(generator_model, classifier_frozen, num_tests_per_instrument = num_tests_per_instrument, step = step, save = False,
                  diff_target = True) #turn to true
    print("done!!")

def load_and_freeze_classifier(classifier, load_dir):
    classifier.load_state_dict(torch.load(load_dir))
    for p in classifier.parameters():
        p.requires_grad = False
    print("done loading and freezing classifier")
    return classifier

if __name__ == "__main__":
    #TODO: these are the parameters you can modify
    classifier_load_dir = "experiments_classifier/audio_classifier_smaller/evaluation_15000/classifier_params_15000.pt"
    experiment = "piano_fullmodel_unsupervised_truthstyle"
    num_training_steps = 15000
    path = f"experiments/{experiment}"
    load_test = True

    generator_model = make_generator()
    classifier = make_classifier()
    # input(classifier)
    classifier_frozen = load_and_freeze_classifier(classifier, classifier_load_dir)
    print("done generating and loading models")

    if torch.cuda.is_available():
        print("cuda available!")
        device = "cuda"
        classifier_frozen = classifier_frozen.cuda()
        generator_model = generator_model.cuda()
    else:
        device = "cpu"

    AE_optimizer = torch.optim.Adam(generator_model.parameters(), lr=1e-3)
    loss = nn.MSELoss()

    if load_test:
        exp_dir =  "piano_fullmodel_unsupervised_truthstyle"
        step = 15000
        file_name = f"ae_params_{step}.pt"
        tests_per_instrument = 25 #chanve back
        try:
            os.chdir(f"experiments/{exp_dir}")
        except:
            print("this experiment doesn't exist!")
            quit()
        load_and_evaluate(generator_model, classifier_frozen, file_name, num_tests_per_instrument = tests_per_instrument, step = step)
        quit()

    soft_make_dir(path)
    copyfile("train_unsupervised.py", f"{path}/train_unsupervised.py")
    copyfile("ConvAE.py", f"{path}/ConvAE.py")
    copyfile("pipeline_whole.py", f"{path}/pipeline_whole.py")
    os.chdir(path)

    norm_mult = 1e-7
    for i in range(num_training_steps + 1):
        x, target, one_hot = sampler_dataset_iterable.next()
        one_hot = torch.as_tensor(one_hot, device=device, dtype = torch.float32)
        x = torch.as_tensor(x, device=device)
        target = torch.as_tensor(target, device=device)
        style_vec = classifier_frozen.forward(target, None)
#         out, code = generator_model.forward(x, one_hot) #CHANGE BACK
        out, code = generator_model.forward(x, style_vec)
        if i % 1500 == 0:
            print("eval time!")
            test_evaluate(generator_model, classifier_frozen, num_tests_per_instrument = 5, step = i, save = True, diff_target = True)
        # print("l1 loss: ", norm_mult * torch.sum(torch.abs(out)).cpu().detach().numpy())
        # print("norm loss: " , norm_mult * torch.norm(out).cpu().detach().numpy())
        # encoding_loss = normalized_loss(out, target)
        encoding_loss = loss(out, target) #+ norm_mult * torch.sum(torch.abs(out))
        print(i, " ", encoding_loss.cpu().detach().numpy())
        AE_optimizer.zero_grad()
        encoding_loss.backward()
        AE_optimizer.step()


import datetime
import logging
import os
import unittest

import numpy as np

from config import Config
from src.data.dataset import FaceDataset
from src.model.knn_model import KnnModel
from src.model.sequencial_model import SequentialModel
from src.problem2 import score

logging.config.fileConfig(Config.LOGGING_FILE)
logger = logging.getLogger()

class TestModel(unittest.TestCase):
    """
    Test the sequential model and the knn model.
    """

    def setUp(self):
        data1 = FaceDataset()
        data1.read_dataset('tests/fixtures/faces/same_faces')

        data2 = FaceDataset()
        data2.read_dataset('tests/fixtures/faces/newdata')

        data1.save(os.path.join('tests/fixtures/faces/same_faces/testset.pkl'))
        data2.save(os.path.join('tests/fixtures/faces/newdata/testset.pkl'))

    def test_sequencial_model(self):
        dataset = FaceDataset.load('tests/fixtures/faces/same_faces/testset.pkl')
        X, y = dataset.get_distance_vectors()

        modeldir = os.path.join(Config.SAVED_MODELS, datetime.datetime.now().strftime('%Y-%m-%d'))

        kmodel = SequentialModel()
        kmodel.load(modeldir)

        score = kmodel.evaluate(X, y)
        self.assertGreater(score, 0.8)

    def test_sequencial_model_newdata(self):
        dataset = FaceDataset.load('tests/fixtures/faces/newdata/testset.pkl')
        dataset = dataset.create_unknown_class()

        modeldir = os.path.join(Config.SAVED_MODELS, datetime.datetime.now().strftime('%Y-%m-%d'))
        trainset = FaceDataset.load(os.path.join(modeldir, Config.TRAINSET_FILE))

        seqmodel = SequentialModel()
        seqmodel.load(modeldir)

        results = score(seqmodel, dataset, trainset)
        self.assertGreater(results, 0.8)

    def test_knn_model(self):
        """
        Tests the data loading from a json file.
        :return:
        """
        dataset = FaceDataset.load('tests/fixtures/faces/same_faces/testset.pkl')
        modeldir = os.path.join(Config.SAVED_MODELS, datetime.datetime.now().strftime('%Y-%m-%d'))

        knnmodel = KnnModel.load(modeldir)

        X, y = dataset.get_embeddings()
        score = knnmodel.evaluate(X, y)
        self.assertGreater(score, 0.8)

    def test_knn_model_newdata(self):
        """

        :return:
        """
        dataset = FaceDataset.load('tests/fixtures/faces/newdata/testset.pkl')
        dataset = dataset.create_unknown_class()

        modeldir = os.path.join(Config.SAVED_MODELS, datetime.datetime.now().strftime('%Y-%m-%d'))
        trainset = FaceDataset.load(os.path.join(modeldir, Config.TRAINSET_FILE))

        knnmodel = KnnModel.load(modeldir)

        result = score(knnmodel, dataset, trainset)
        self.assertGreater(result, 0.8)

    def test_0(self):
        modeldir = os.path.join(Config.SAVED_MODELS, datetime.datetime.now().strftime('%Y-%m-%d'))

        kmodel = SequentialModel()
        kmodel.load(modeldir)

        X, y = np.zeros((1, 512)), np.array([1])
        score = kmodel.evaluate(X, y)
        self.assertGreater(score, 0.8)

if __name__ == '__main__':
    unittest.main()


# Recursive function with memoization
class Solution:
    def isInterleave(self, s1: str, s2: str, s3: str) -> bool:

        L1, L2, L3 = len(s1), len(s2), len(s3)
        visited = set()

        def interleavingString(idx1, idx2, idx3):
            if idx1 == L1 and idx2 == L2 and idx3 == L3:
                return True

            a, b = False, False

            if idx1 < L1 and idx3 < L3:
                if s1[idx1] == s3[idx3]:
                    if (idx1+1,idx2,idx3+1) not in visited:
                        visited.add((idx1+1, idx2, idx3+1))
                        a = interleavingString(idx1+1, idx2, idx3+1)
            if idx2 < L2 and idx3 < L3:
                if s2[idx2] == s3[idx3]:
                    if (idx1, idx2+1, idx3+1) not in visited:
                        visited.add((idx1, idx2+1, idx3+1))
                        b = interleavingString(idx1, idx2+1, idx3+1)
            return a or b

        if interleavingString(0, 0, 0):
            return True
        return False

# Using dynamic programming with a queue and memoization
import collections

class Solution:
    def isInterleave(self, s1: str, s2: str, s3: str) -> bool:

        L1, L2, L3 = len(s1), len(s2), len(s3)

        queue = collections.deque([(0, 0, 0)])

        visited = set()

        while queue:
            idx1, idx2, idx3 = queue.popleft()

            if (idx1, idx2, idx3) in visited:
                continue
            visited.add((idx1, idx2, idx3))

            if idx1==L1 and idx2==L2 and idx3==L3:
                return True

            if idx1<L1 and idx3<L3:
                if s1[idx1]==s3[idx3]:
                    queue.append((idx1+1, idx2, idx3+1))
            if idx2<L2 and idx3<L3:
                if s2[idx2]==s3[idx3]:
                    queue.append((idx1, idx2+1, idx3+1))
        return False

# Recursion(dfs) with cache
class Solution:
    def isInterleave(self, s1: str, s2: str, s3: str) -> bool:
        if len(s1) + len(s2) != len(s3):
            return False

        dp = {}
        def dfs(i1, i2):
            if i1 == len(s1) and i2 == len(s2):
                return True
            if (i1, i2) in dp:
                return dp[(i1,i2)]

            if i1 < len(s1) and s1[i1] == s3[i1+i2] and dfs(i1+1, i2):
                return True
            if i2 < len(s2) and s2[i2] == s3[i1+i2] and dfs(i1, i2+1):
                return True

            dp[(i1, i2)] = False
            return False

        return dfs(0, 0)