import tensorflow as tf

    # [batch, 1, n_h_mixer]
    hidden = tf.nn.elu(tf.matmul(agent_qs_reshaped, w1_reshaped) + b1_reshaped)
    
    # Second layer
    w_final = tf.abs(tf.matmul(state, hyper_w_final))
    w_final_reshaped = tf.reshape(w_final, [-1, n_h_mixer, 1]) # reshape into batch of matrices
    b_final_reshaped = tf.reshape(hyper_b_final, [-1, 1, 1])

    # [batch, 1, 1]
    y = tf.matmul(hidden, w_final_reshaped) + b_final_reshaped

    q_tot = tf.reshape(y, [-1, 1])

    return q_tot


class QMix():
    def __init__(self, env, num_s, num_a, lr=0.0001, gamma=0.99, replace_target_iter=5000,
                 memory_size=200000, batch_size=256, epsilon=1, epsilon_decay=0.0001):
        self.n_agents = 2
        self.env = env
        self.name = "qmix"
