Created by: BladeSun

import paddle
import paddle.fluid as fluid
import numpy
from paddle.fluid.layers import create_parameter
from paddle.fluid import ParamAttr
import paddle.fluid.layers as layers

def forward_with_weight(input, weights):
    output = fluid.layers.matmul(input, weights[0]) + weights[1]
    return output

def gen_variables(area_name, input_dim, output_dim):
    var_list = list()
    names = ["w_1", "b_1"]
    is_biases = [False, True] #* 4
    shapes = [[input_dim, output_dim], [output_dim]]
    for name, is_bias, shape in zip(names, is_biases, shapes):
        var_list.append(create_parameter(name=area_name + name,
            shape=shape,
            dtype='float32', attr=ParamAttr(name=area_name + name), is_bias=is_bias))
    return var_list

def linear_regression_net(input, hidden):
    out = fluid.layers.fc(input, hidden)
    return out


batch_size = 5 
train_reader = fluid.io.batch(fluid.io.shuffle(paddle.dataset.uci_housing.train(), 500), batch_size)
test_reader = fluid.io.batch(paddle.dataset.uci_housing.test(), batch_size)

def reader():
    for i in range(100):
        yield (numpy.random.rand(4,5,13).astype('float32'), numpy.random.rand(4,5,1).astype('float32'))
)

max_epoch_num = 100
query_set_size = 4
num_update = 2
task_num = 4
support_x = fluid.data(name='s_x', shape=[None, None, 13], dtype='float32')
support_y = fluid.data(name='s_y', shape=[None, None, 1], dtype='float32')
query_x = fluid.data(name='q_x', shape=[None, None, 13], dtype='float32')
query_y = fluid.data(name='q_y', shape=[None, None, 1], dtype='float32')

def loss_func(x, y):
    return fluid.layers.square_error_cost(input=x, label=y)


weights = gen_variables('train', 13 , 1)




def gather_and_squeeze(inp_list, i):
    return [layers.squeeze(layers.gather(inp, index=i), axes=[0]) for inp in inp_list]

def cond(start_idx, end_idx, support_x, support_y, query_x, query_y, query_loss_array):   # 参数和loop_vars相对应
    return start_idx < end_idx 

def body(start_idx, end_idx, support_x, support_y, query_x, query_y, query_loss_array):  # 参数和loop_vars相对应
    query_outputs = []
    query_losses = []

    # support_x : task_num * support_set_num * dim
    task_support_x, task_support_y, task_query_x, task_query_y = \
        gather_and_squeeze([support_x, support_y, query_x, query_y], start_idx)
    #ret = fluid.layers.array_write(x_i, i, ret) 

    support_output = forward_with_weight(task_support_x, weights)
    support_loss = loss_func(support_output, task_support_y)
    grads =  fluid.backward.gradients(support_loss, weights)
    fast_weights = [param - 0.01 * grad for param, grad in zip(weights, grads)]
    query_output = forward_with_weight(task_query_x, fast_weights)
    query_loss = loss_func(query_output, task_query_y) 
    query_outputs.append(query_output)
    # query_loss: query_set_num * 1
    query_losses.append(query_loss)
    
    for j in range(num_update - 1):
        loss = loss_func(forward_with_weight(task_support_x, fast_weights), task_support_y) 
        grads = fluid.backward.gradients(loss, fast_weights)
        fast_weights = [param - 0.01 * grad for param, grad in zip(fast_weights, grads)]
        output = forward_with_weight(task_query_x, fast_weights)
        query_outputs.append(output)
        query_losses.append(loss_func(output, task_query_y))
    
    #each elem in array: num_update * query_set_size * 1
    query_loss_array = fluid.layers.array_write(layers.stack(query_losses, axis=0), start_idx, query_loss_array) 
    start_idx = start_idx + 1

    return start_idx, end_idx, support_x, support_y, query_x, query_y, query_loss_array  

arr = fluid.layers.create_array(dtype='float32')

start_idx = layers.fill_constant(shape=[1], dtype='int64', value=0)     # 循环计数器
end_idx = layers.fill_constant(shape=[1], dtype='int64', value=task_num)  # 循环次数
ret  = layers.while_loop(cond, body, [start_idx, end_idx, support_x, 
                                                    support_y, query_x, query_y, arr])
ret_loss_array_tensor = ret[6] 
#ret_loss_stacked_tensor: num_task * num_update * query_set_size * 1
ret_loss_stacked_tensor, array_idx = layers.tensor_array_to_tensor(input=ret_loss_array_tensor, axis=0, use_stack=True)
#stacked_array_tensor_query_loss_mean: num_task * num_update * 1
ret_loss_stacked_tensor_query_set_mean = layers.reduce_mean(ret_loss_stacked_tensor, dim=2) 
#ret_loss_stacked_task_avg: num_update * 1
ret_loss_stacked_task_avg = layers.reduce_sum(ret_loss_stacked_tensor_query_set_mean, dim= 0) / task_num 
meta_train_op = ret_loss_stacked_task_avg[-1]


sgd_optimizer = fluid.optimizer.SGD(learning_rate=0.001)
sgd_optimizer.minimize(meta_train_op)

place = fluid.CPUPlace()
exe = fluid.Executor(place)

num_epochs = 100
exe.run(fluid.default_startup_program())

for epoch in range(num_epochs):
    for batch_id, data_train in enumerate(reader()):
        avg_loss_value, = exe.run(feed={'s_x': data_train[0], 's_y': data_train[1],
                                         'q_x': data_train[0], 'q_y': data_train[1],}, fetch_list=[meta_train_op])

        if batch_id % 10 == 0 and batch_id is not 0:
            print("epoch: {}, batch_id: {}, loss is: {}".format(epoch, batch_id, avg_loss_value))