import os
import sys
import numpy as np
from pie_intent import PIEIntent
from pie_data import PIE
import keras.backend as K
import tensorflow as tf
from prettytable import PrettyTable

dim_ordering = K.image_dim_ordering()


def count_ratio(data):
    pos = 0
    output = data['output']
    print(type(output))
    total = output.shape[0]
    pos_index=[]
    neg_index=[]
    for i in range(0,total):
        if output[i][0][0]==1:
            pos_index.append(i)
        else:
            neg_index.append(i)
    length = min(len(pos_index),len(neg_index))
    new_output=[]
    for i in range(0,length):
        new_output.append(output[pos_index[i]])
        new_output.append(output[neg_index[i]])
    new_output=np.array(new_output)
    lt=[]
    for i in range(0,new_output.shape[0]):
        lt.append(new_output[i][0][0])
    print(lt)
    for output_seq in new_output:
        pos += output_seq[0][0]
    return pos, new_output.shape[0] - pos

if __name__ == '__main__':
    ratio = float(sys.argv[1])
    data_opts = {
        'fstride': 1,
        'sample_type': 'all',
        'height_rng': [0, float('inf')],
        'squarify_ratio': 0,
        'data_split_type': 'default',  # kfold, random, default
        'seq_type': 'intention',  # crossing , intention
        'min_track_size': 0,  # discard tracks that are shorter
        'max_size_observe': 15,  # number of observation frames
        'max_size_predict': 5,  # number of prediction frames
        'seq_overlap_rate': 0.5,  # how much consecutive sequences overlap
        'balance': True,  # balance the training and testing samples
        'crop_type': 'context',  # crop 2x size of bbox around the pedestrian
        'crop_mode': 'pad_resize',  # pad with 0s and resize to VGG input
        'encoder_input_type': [],
        'decoder_input_type': ['bbox'],
        'output_type': ['intention_binary']
    }
    t = PIEIntent(num_hidden_units=128,
                  regularizer_val=0.001,
                  lstm_dropout=0.4,
                  lstm_recurrent_dropout=0.2,
                  convlstm_num_filters=64,
                  convlstm_kernel_size=2)
    saved_files_path = ''
    imdb = PIE(data_path=os.environ.copy()['PIE_PATH'])

    beh_seq_train = imdb.generate_data_trajectory_sequence('train', **data_opts)
    beh_seq_train = imdb.balance_samples_count(beh_seq_train, label_type='intention_binary',ratio=ratio)
    # beh_seq_val = imdb.generate_data_trajectory_sequence('val', **data_opts)
    # beh_seq_val = imdb.balance_samples_count(beh_seq_val, label_type='intention_binary',ratio=ratio)
    # beh_seq_test = imdb.generate_data_trajectory_sequence('test', **data_opts)
    # beh_seq_test = imdb.balance_samples_count(beh_seq_test, label_type='intention_binary', ratio=0)

    data_type = {'encoder_input_type': data_opts['encoder_input_type'],
                    'decoder_input_type': data_opts['decoder_input_type'],
                    'output_type': data_opts['output_type']}
    data_train=t.get_train_val_data(beh_seq_train, data_type=data_type, seq_length=15, overlap=0.5)
    # data_val=t.get_train_val_data(beh_seq_val, data_type=data_type, seq_length=15, overlap=0.5)
    # data_test = t.get_train_val_data(beh_seq_test, data_type=data_type, seq_length=15, overlap=1)

    pos_train,neg_train=count_ratio(data_train)
    # pos_val, neg_val = count_ratio(data_val)
    # pos_test, neg_test = count_ratio(data_test)

    pt = PrettyTable(['data_type','positive samples', 'negative samples','ratio'])
    pt.title = 'Proportion of positive and negative samples'
    pt.add_row(['train',pos_train, neg_train, round(pos_train/neg_train,2)])
    # pt.add_row(['val',pos_val, neg_val, round(pos_val/neg_val,2)])
    # pt.add_row(['test',pos_test, neg_test, round(pos_test/neg_test,2)])
    print(pt)

    print('seq_length=15, overlap=0.5')

    K.clear_session()
    tf.reset_default_graph()