from BiddingKG.dl.table_head.models.model import get_model
from BiddingKG.dl.table_head.post_process import table_post_process
from BiddingKG.dl.table_head.pre_process import my_data_loader, table_pre_process

# init model
input_shape = (3, 10, 60)
output_shape = (3,)
model = get_model(input_shape, output_shape)

# load weights
model_path = "checkpoints/best.hdf5"
model.load_weights(model_path)


def predict(table_text_list):
    # 表格单元格数*2 即为单次预测batch_size
    data_list = table_pre_process(table_text_list, [], 0, is_train=False)
    batch_size = len(data_list)
    # print("batch_size", batch_size)

    # 数据预处理
    predict_x = my_data_loader(data_list, [], batch_size, is_train=False)

    # 预测
    predict_result = model.predict_generator(predict_x, steps=1)
    # print("predict_result", predict_result.shape)

    # 数据后处理
    table_label_list = table_post_process(table_text_list, predict_result)
    return table_label_list


if __name__ == '__main__':
    _str = "[['序号', '投标人名称', '价格得分', '技术得分', '商务得分', '综合得分', '排名'], " \
           "['序号', '投标人名称', '比例(20%)，', '比例(45%)，', '比例(35%)，', '100%', '排名'], " \
           "['1', '广州中科雅图信息技术有限公司', '19.71', '11.50', '11.00', '42.21', '3'], " \
           "['2', '核工业赣州工程勘察院', '19.64', '15.00', '11.00', '45.64', '2'], " \
           "['3', '广东晟腾地信科技有限公司', '20.00', '16.17', '14.00', '50.17', '1']]"

    data_list = eval(_str)
    print("len(data_list)", len(data_list))
    predict(data_list)