import tensorflow as tf
import numpy as np
import sys
from src.convolutional import set_convolutional
from src.crops_sr import extract_crops_z_sr, extract_crops_x_sr, pad_frame_sr, resize_images
import src.ConvNetValue_Class as cnvc

sys.path.append('../')


# the follow parameters *have to* reflect the design of the network to be imported
# TODO: <> instance version
def build_tracking_graph_from_instance(final_score_sz, design):
    pos_x_ph = tf.placeholder(tf.float64)
    pos_y_ph = tf.placeholder(tf.float64)

    z_sz_ph_w = tf.placeholder(tf.float64)
    z_sz_ph_h = tf.placeholder(tf.float64)

    # TODO: <> add placeholders for scale-ratio here
    # x_sz0_ph = tf.placeholder(tf.float64)
    # x_sz1_ph = tf.placeholder(tf.float64)
    # x_sz2_ph = tf.placeholder(tf.float64)

    x_s0r0_w = tf.placeholder(tf.float64)
    x_s0r0_h = tf.placeholder(tf.float64)
    x_s0r1_w = tf.placeholder(tf.float64)
    x_s0r1_h = tf.placeholder(tf.float64)
    x_s0r2_w = tf.placeholder(tf.float64)
    x_s0r2_h = tf.placeholder(tf.float64)

    x_s1r0_w = tf.placeholder(tf.float64)
    x_s1r0_h = tf.placeholder(tf.float64)
    x_s1r1_w = tf.placeholder(tf.float64)
    x_s1r1_h = tf.placeholder(tf.float64)
    x_s1r2_w = tf.placeholder(tf.float64)
    x_s1r2_h = tf.placeholder(tf.float64)

    x_s2r0_w = tf.placeholder(tf.float64)
    x_s2r0_h = tf.placeholder(tf.float64)
    x_s2r1_w = tf.placeholder(tf.float64)
    x_s2r1_h = tf.placeholder(tf.float64)
    x_s2r2_w = tf.placeholder(tf.float64)
    x_s2r2_h = tf.placeholder(tf.float64)

    # TODO <> set max scale
    x_max_ph_w = tf.placeholder(tf.float64)
    x_max_ph_h = tf.placeholder(tf.float64)

    # TODO: <> generate a instance of class ConvNetValue who has the latest value of params of model
    # matconvnet_value = cnvc.ConvNetValue()
    # TODO: <> OR I just import the value from file ConvNetValue_Class.py
    matconvnet_value = cnvc.ConvNetValue().update_value()

    filename = tf.placeholder(tf.string, [], name='filename')
    image_file = tf.read_file(filename)
    # Decode the image as a JPEG file, this will turn it into a Tensor
    image = tf.image.decode_jpeg(image_file)
    image = 255.0 * tf.image.convert_image_dtype(image, tf.float32)

    frame_sz = tf.shape(image)
    # used to pad the crops
    if design.pad_with_image_mean:
        avg_chan = tf.reduce_mean(image, axis=(0,1), name='avg_chan')
    else:
        avg_chan = None
    # pad with if necessary
    # TODO: <>
    #frame_padded_z, npad_z = pad_frame(image, frame_sz, pos_x_ph, pos_y_ph, z_sz_ph, avg_chan)
    frame_padded_z, npad_z = pad_frame_sr(image, frame_sz, pos_x_ph, pos_y_ph, z_sz_ph_w, z_sz_ph_h, avg_chan)
    frame_padded_z = tf.cast(frame_padded_z, tf.float32)
    # extract tensor of z_crops
    # TODO: <> introduce width and height here
    #z_crops = extract_crops_z(frame_padded_z, npad_z, pos_x_ph, pos_y_ph, z_sz_ph, design.exemplar_sz)
    z_crops = extract_crops_z_sr(frame_padded_z, npad_z, pos_x_ph, pos_y_ph, z_sz_ph_w, z_sz_ph_h, design.exemplar_sz)

    # TODO: <> set max scale
    #frame_padded_x, npad_x = pad_frame(image, frame_sz, pos_x_ph, pos_y_ph, x_sz2_ph, avg_chan)
    frame_padded_x, npad_x = pad_frame_sr(image, frame_sz, pos_x_ph, pos_y_ph, x_max_ph_w, x_max_ph_h, avg_chan)
    frame_padded_x = tf.cast(frame_padded_x, tf.float32)
    # extract tensor of x_crops (3 scales)
    # TODO: <> introduce 9 scale and ratio
    #x_crops = extract_crops_x(frame_padded_x, npad_x, pos_x_ph, pos_y_ph, x_sz0_ph, x_sz1_ph, x_sz2_ph, design.search_sz)
    x_crops = extract_crops_x_sr(frame_padded_x, npad_x, pos_x_ph, pos_y_ph,
                              x_s0r0_w, x_s0r0_h, x_s0r1_w, x_s0r1_h, x_s0r2_w, x_s0r2_h,
                              x_s1r0_w, x_s1r0_h, x_s1r1_w, x_s1r1_h, x_s1r2_w, x_s1r2_h,
                              x_s2r0_w, x_s2r0_h, x_s2r1_w, x_s2r1_h, x_s2r2_w, x_s2r2_h,
                              x_max_ph_w, x_max_ph_h,
                              design.search_sz)


    # TODO: <> create siamese from instance
    # use crops as input of (MatConvnet imported) pre-trained fully-convolutional Siamese net
    # template_z, templates_x, p_names_list, p_val_list = _create_siamese(os.path.join(env.root_pretrained,design.net), x_crops, z_crops)
    template_z, templates_x= _create_siamese_from_instance(z_crops, x_crops, matconvnet_value)
    template_z = tf.squeeze(template_z)
    # TODO: <>
    #templates_z = tf.stack([template_z, template_z, template_z])
    templates_z = tf.stack([template_z, template_z, template_z, template_z, template_z, template_z,
                            template_z, template_z, template_z])

    # TODO: <> match templates from instance
    # compare templates via cross-correlation
    #scores = _match_templates(templates_z, templates_x, p_names_list, p_val_list)
    scores = _match_templates_from_instance(templates_z, templates_x, matconvnet_value)

    # upsample the score maps
    scores_up = tf.image.resize_images(scores, [final_score_sz, final_score_sz],
        method=tf.image.ResizeMethod.BICUBIC, align_corners=True)
    return filename, image, templates_z, scores_up


# TODO: <> load the model from instance of class convnetvalue
def _create_siamese_from_instance(net_z, net_x, instance):
    # read mat file from net_path and start TF Siamese graph from placeholders X and Z
    # params_names_list, params_values_list = _import_from_matconvnet(net_path)

    # loop through the flag arrays and re-construct network, reading parameters of conv and bnorm layers
    for i in range(instance.num_layers):

        # set up conv "block" with bnorm and activation
        net_x = set_convolutional(net_x, instance.W_d[i], np.swapaxes(instance.b_d[i], 0, 1), instance.conv_stride_n[i],
                                  instance.bn_beta_d[i], instance.bn_gamma_d[i], instance.bn_moving_mean_d[i],
                                  instance.bn_moving_variance_d[i], filtergroup=instance.filtergroup_yn_n[i],
                                  batchnorm=instance.bnorm_yn_n[i], activation=instance.relu_yn_n[i],
                                  scope='conv' + str(i + 1), trainable=instance.trainable_n[i], reuse=tf.AUTO_REUSE)

        # notice reuse=True for Siamese parameters sharing
        net_z = set_convolutional(net_z, instance.W_d[i], np.swapaxes(instance.b_d[i], 0, 1), instance.conv_stride_n[i],
                                  instance.bn_beta_d[i], instance.bn_gamma_d[i], instance.bn_moving_mean_d[i],
                                  instance.bn_moving_variance_d[i], filtergroup=instance.filtergroup_yn_n[i],
                                  batchnorm=instance.bnorm_yn_n[i], activation=instance.relu_yn_n[i],
                                  scope='conv' + str(i + 1), trainable=instance.trainable_n[i], reuse=True)

        # add max pool if required
        if instance.pool_stride_n[i] > 0:
            print('\t\tMAX-POOL: size ' + str(instance.pool_sz) + ' and stride ' + str(instance.pool_stride_n[i]))
            net_x = tf.nn.max_pool(net_x, [1, instance.pool_sz, instance.pool_sz, 1],
                                   strides=[1, instance.pool_stride_n[i], instance.pool_stride_n[i], 1],
                                   padding='VALID', name='pool' + str(i + 1))
            net_z = tf.nn.max_pool(net_z, [1, instance.pool_sz, instance.pool_sz, 1],
                                   strides=[1, instance.pool_stride_n[i], instance.pool_stride_n[i], 1],
                                   padding='VALID', name='pool' + str(i + 1))

    return net_z, net_x



# TODO: <> instance version
def _match_templates_from_instance(net_z, net_x, instance):
    # finalize network
    # z, x are [B, H, W, C]
    net_z = tf.transpose(net_z, perm=[1,2,0,3])
    net_x = tf.transpose(net_x, perm=[1,2,0,3])
    # z, x are [H, W, B, C]
    Hz, Wz, B, C = tf.unstack(tf.shape(net_z))
    Hx, Wx, Bx, Cx = tf.unstack(tf.shape(net_x))
    # assert B==Bx, ('Z and X should have same Batch size')
    # assert C==Cx, ('Z and X should have same Channels number')
    net_z = tf.reshape(net_z, (Hz, Wz, B*C, 1))
    net_x = tf.reshape(net_x, (1, Hx, Wx, B*C))
    net_final = tf.nn.depthwise_conv2d(net_x, net_z, strides=[1,1,1,1], padding='VALID')
    # final is [1, Hf, Wf, BC]
    # TODO: <>
    # net_final = tf.concat(tf.split(net_final, 3, axis=3), axis=0)
    net_final = tf.concat(tf.split(net_final, 9, axis=3), axis=0)
    # final is [B, Hf, Wf, C]
    net_final = tf.expand_dims(tf.reduce_sum(net_final, axis=3), axis=3)
    # final is [B, Hf, Wf, 1]
    if instance.bnorm_adjust:
        bn_beta = instance.bn_beta_adj
        bn_gamma = instance.bn_gamma_adj
        bn_moments = instance.bn_moments_adj
        bn_moving_mean = instance.bn_moving_mean_adj
        bn_moving_variance = instance.bn_moving_variance_adj
        net_final = tf.layers.batch_normalization(net_final, beta_initializer=tf.constant_initializer(bn_beta),
                                                gamma_initializer=tf.constant_initializer(bn_gamma),
                                                moving_mean_initializer=tf.constant_initializer(bn_moving_mean),
                                                moving_variance_initializer=tf.constant_initializer(bn_moving_variance),
                                                # TODO: <> params' setting
                                                training=False, trainable=False)

    return net_final
