import java.util.*;
import java.io.*;

/**
 * This class represents an MDP, and includes a constructor that can
 * read the MDP from a data file.  If you want to create your own
 * MDP's, you can do so by adding new constructors and filling in all
 * of the needed fields, or by creating data files that are read by
 * the provided constructor (see assignment instructions).
 **/
public class Mdp {

    /** total number of states */
    public int numStates;

    /** total number of actions */
    public int numActions;

    /** an array containing the names of all of the states */
    public String stateName[];

    /** an array containing the names of all of the actions */
    public String actionName[];

    /** an array containing the (immediate) reward associated with
     * every state **/
    public double reward[];

    /** a list of the possible states that can be reached from each
     * state under each action **/
    public int nextState[][][];

    /** a table of transition probabilities:
     * <tt>transProb[s][a][i]</tt> is the probability of transitioning
     * from state <tt>s</tt> under action <tt>a</tt> to state
     * <tt>nextState[s][a][i]</tt>. **/
    public double transProb[][][];

    /** the start state **/
    public int startState = -1;

    /** This constructor reads in data from <tt>filename</tt> and sets
     * up all of the public fields.  See assignment instructions for
     * information on the required format of this file.
     */
    public Mdp(String filename) throws FileNotFoundException, IOException {
	BufferedReader in;
	String line;
	Map state_map = new HashMap();
	Map action_map = new HashMap();
	Map reward_map = new HashMap();
	List trans_prob_list = new LinkedList();

	try {
	    in = new BufferedReader(new FileReader(filename));
	} catch (FileNotFoundException e) {
	    System.err.print("File "+filename+" not found.\n");
	    throw e;
	}

	while (true) {
	    try {
		line = in.readLine();
	    }
	    catch (IOException e) {
		System.err.println("Error reading file "+filename);
		throw e;
	    }

	    String[] words = null;

	    if (line == null)
		break;

	    line = line.trim( );

	    words = line.split("\\s+");

	    if (line.equals(""))
		continue;

	    if (words.length == 2) { // reward function entry
		int s = getId(state_map, words[0]);
		if (startState < 0)
		    startState = s;
		Double r;
		try {
		    r = new Double(words[1]);
		} catch (NumberFormatException e) {
		    System.err.println("Expected number in file "+filename+
				       " at line: " + line);
		    throw e;
		}
		reward_map.put(new Integer(s), r);
	    } else if (words.length >= 4 && words.length % 2 == 0) { // transition prob. entry
		int s = getId(state_map, words[0]);
		int a = getId(action_map, words[1]);

		for (int i = 2; i < words.length; i += 2) {
		    double p;
		    try {
			p = Double.parseDouble(words[i+1]);
		    } catch (NumberFormatException e) {
			System.err.println("Expected number in file "+filename+
					   " at line: " + line);
			throw e;
		    }
		    if (p < 0.) {
			String err = "Probabilities must be nonnegative in file " +
			    filename + " at line: " + line;
			System.err.println(err);
			throw new RuntimeException(err);
		    }
		    trans_prob_list.add(new TransProbItem(s, a, 
							  getId(state_map, words[i]),
							  p));
		}
	    } else {
		String err = "Badly formatted data in " +
		    filename + " at line: " + line;
		System.err.println(err);
		throw new RuntimeException(err);
	    }
	}

	stateName = mapToArray(state_map);
	numStates = stateName.length;

	actionName = mapToArray(action_map);
	numActions = actionName.length;

	reward = new double[numStates];

	Iterator it = reward_map.keySet().iterator();

	while (it.hasNext()) {
	    Object o = it.next();
	    reward[((Integer) o).intValue()] =
		((Double) reward_map.get(o)).doubleValue();
	}
	
	transProb = new double[numStates][numActions][];
	nextState = new int[numStates][numActions][];

	Map[][] next_state_map = new Map[numStates][numActions];

	for (int fs = 0; fs < numStates; fs++)
	    for (int a = 0; a < numActions; a++)
		next_state_map[fs][a] = new TreeMap();

	it = trans_prob_list.iterator();

	while (it.hasNext()) {
	    TransProbItem tpi = (TransProbItem) it.next();
	    Map m = next_state_map[tpi.from_state][tpi.action];
	    Integer ts = new Integer(tpi.to_state);
	    double old_val = (m.containsKey(ts)
			      ? ((Double) m.get(ts)).doubleValue()
			      : 0.0);
	    m.put(ts, new Double(old_val + tpi.prob));
	}

	for (int fs = 0; fs < numStates; fs++)
	    for (int a = 0; a < numActions; a++) {
		Map m = next_state_map[fs][a];
		int size = m.size();
		transProb[fs][a] = new double[size];
		nextState[fs][a] = new int[size];
		it = m.keySet().iterator();
		double sum = 0.0;
		for (int i = 0; i < size; i++) {
		    Integer ts = (Integer) it.next();
		    nextState[fs][a][i] = ts.intValue();
		    sum += transProb[fs][a][i] = ((Double) m.get(ts)).doubleValue();
		}
		if (sum <= 0.0) {
		    String err = "State " + stateName[fs] + " with action " +
			actionName[a] +
			" must have positive transition probability in file " +
			filename;
		    System.err.println(err);
		    throw new RuntimeException(err);
		}
		for (int i = 0; i < size; i++)
		    transProb[fs][a][i] /= sum;
	    }
    }

    private class TransProbItem {
	private int from_state;
	private int action;
	private int to_state;
	private double prob;

	private TransProbItem(int from_state,
			      int action,
			      int to_state,
			      double prob) {
	    this.from_state = from_state;
	    this.action = action;
	    this.to_state = to_state;
	    this.prob = prob;
	}
    }

    private int getId(Map m, String s) {
	if (!m.containsKey(s)) {
	    int n = m.size();
	    m.put(s, new Integer(n));
	}
	return ((Integer) m.get(s)).intValue();
    }

    private String[] mapToArray(Map m) {
	String a[] = new String[m.size()];
	Iterator it = m.keySet().iterator();
	while (it.hasNext()) {
	    String s = (String) it.next();
	    a[((Integer) m.get(s)).intValue()] = s;
	}
	return a;
    }


}
