/*******************************************************************************
 * DIANNE  - Framework for distributed artificial neural networks
 * Copyright (C) 2015  iMinds - IBCN - UGent
 *
 * This file is part of DIANNE.
 *
 * DIANNE is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Contributors:
 *     Tim Verbelen, Steven Bohez
 *******************************************************************************/
package be.iminds.iot.dianne.rl.agent.strategy;

import java.util.Map;

import be.iminds.iot.dianne.api.nn.NeuralNetwork;
import be.iminds.iot.dianne.api.rl.agent.ActionStrategy;
import be.iminds.iot.dianne.api.rl.agent.AgentProgress;
import be.iminds.iot.dianne.api.rl.environment.Environment;
import be.iminds.iot.dianne.nn.util.DianneConfigHandler;
import be.iminds.iot.dianne.rl.agent.strategy.config.BoltzmannConfig;
import be.iminds.iot.dianne.tensor.ModuleOps;
import be.iminds.iot.dianne.tensor.Tensor;
import be.iminds.iot.dianne.tensor.TensorOps;

/**
 * Discrete sampling ActionStrategy that takes the output of a (discrete) policy network
 * ending with a softmax/logsoftmax module and samples an actual discrete action from it.  
 * @author tverbele
 *
 */
public class DiscreteSamplingActionStrategy implements ActionStrategy {
	
	private BoltzmannConfig config;
	private NeuralNetwork nn;
	
	@Override
	public void setup(Map<String, String> config, Environment env, NeuralNetwork... nns) throws Exception {
		this.config = DianneConfigHandler.getConfig(config, BoltzmannConfig.class);
		this.nn = nns[0];
	}

	@Override
	public AgentProgress processIteration(long i, Tensor state) throws Exception {
		Tensor output = nn.forward(state);
		
		Tensor action = new Tensor(output.size());
		action.fill(0);
		
		if(TensorOps.max(output) < 0){
			// assume logsoftmax output, take exp
			output = TensorOps.exp(output, output);
		}
		
		double s = 0, r = Math.random();
		int a = 0;
		
		while((s += output.get(a)) < r && a < action.size())
			a++;
		
		action.set(1, a);
		
		return new AgentProgress(i, action);
	}

}

add discrete sampling strategy that samples from (log)softmax output