# Name: Yucheng Qian
# NetID: yq791
# Homework 3: Election prediction

import csv
import os
import time

def read_csv(path):
    """
    Reads the CSV file at path, and returns a list of rows. Each row is a
    dictionary that maps a column name to a value in that column, as a string.
    """
    output = []
    for row in csv.DictReader(open(path)):
        output.append(row)
    return output


################################################################################
# Problem 1: State edges
################################################################################

def row_to_edge(row):
    """
    Given an election result row or poll data row, returns the Democratic edge
    in that state.
    """
    return float(row["Dem"]) - float(row["Rep"])  

def state_edges(election_result_rows):
    """
    Given a list of election result rows, returns state edges.
    The input list does has no duplicate states;
    that is, each state is represented at most once in the input list.
    """
    # declare a dict to put states and their corresponding state edges
    state_edges_dict = {}
    # traverse all rows in the input
    for row in election_result_rows:
    	# since there is no duplicate states in the input, we can directly update
    	# the dict by using row_to_edge() above
    	state_edges_dict[row['State']] = row_to_edge(row)
    return state_edges_dict


################################################################################
# Problem 2: Find the most recent poll row
################################################################################

def earlier_date(date1, date2):
    """
    Given two dates as strings (formatted like "Oct 06 2012"), returns True if 
    date1 is after date2.
    """
    return (time.strptime(date1, "%b %d %Y") < time.strptime(date2, "%b %d %Y"))

def most_recent_poll_row(poll_rows, pollster, state):
    """
    Given a list of poll data rows, returns the most recent row with the
    specified pollster and state. If no such row exists, returns None.
    """
    # in order to return None if there is no answer, we initialize res with None.
    most_recent_row = None
    # traverse all rows to find rows with indicated pollster and state
    for row in poll_rows:
    	# if pollster and state match
    	if pollster == row['Pollster'] and state == row['State']:
    		# if we find one target for the first time, we directly assign the
    		# answer with it since there is nothing to compare with
    		if not most_recent_row:
    			most_recent_row = row.copy()
    		# if we find a new row, we need to compare the dates
    		else :
    			if earlier_date(most_recent_row['Date'], row['Date']):
    				most_recent_row = row.copy()
    return most_recent_row


################################################################################
# Problem 3: Pollster predictions
################################################################################

def unique_column_values(rows, column_name):
    """
    Given a list of rows and the name of a column (a string), returns a set
    containing all values in that column.
    """
    # declare a set that guarantees no duplicates in the answer
    value_set = set()
    # for all rows, add the value of indicated column to the set
    for row in rows:
    	value_set.add(row[column_name])
    return value_set

def pollster_predictions(poll_rows):
    """
    Given a list of poll data rows, returns pollster predictions.
    """
    # declare a dict to maintain answers
    predictions = dict()
    # get all pollsters and states
    pollsters = unique_column_values(poll_rows, 'Pollster')
    states = unique_column_values(poll_rows, 'State')
    # for each pollster
    for pollster in pollsters:
    	# declare a list to keep state edges
    	state_egdes_list = []
    	# for each state
    	for state in states:
    		# get most recent row
    		most_recent_row = most_recent_poll_row(poll_rows, pollster, state)
    		if most_recent_row:
    			state_egdes_list.append(most_recent_row)
    	# assign the state edges with each pollster
    	predictions[pollster] = state_edges(state_egdes_list)
    return predictions

################################################################################
# Problem 4: Pollster errors
################################################################################

def average_error(state_edges_predicted, state_edges_actual):
    """
    Given predicted state edges and actual state edges, returns
    the average error of the prediction.
    """
    # initialize the error with a decimal number 0.0
    total_error = 0.0
    # traverse all keys
    for key in state_edges_predicted.keys():
    	# increment the total error
    	total_error += abs(state_edges_predicted[key] - state_edges_actual[key])
    # edge case - if state_edges_predicted is empty, return -1.0
    if len(state_edges_predicted) == 0:
    	return -1.0
    return total_error / len(state_edges_predicted)

def pollster_errors(pollster_predictions, state_edges_actual):
    """
    Given pollster predictions and actual state edges, retuns pollster errors.
    """
    # declare errors to be the return object which is a dict
    errors = dict()
    # traverse the predictions to calculate the error
    for key in pollster_predictions:
    	errors[key] = average_error(pollster_predictions[key], state_edges_actual)
    return errors


################################################################################
# Problem 5: Pivot a nested dictionary
################################################################################

def pivot_nested_dict(nested_dict):
    """
    Pivots a nested dictionary, producing a different nested dictionary
    containing the same values.
    The input is a dictionary d1 that maps from keys k1 to dictionaries d2,
    where d2 maps from keys k2 to values v.
    The output is a dictionary d3 that maps from keys k2 to dictionaries d4,
    where d4 maps from keys k1 to values v.
    For example:
      input = { "a" : { "x": 1, "y": 2 },
                "b" : { "x": 3, "z": 4 } }
      output = {'y': {'a': 2},
                'x': {'a': 1, 'b': 3},
                'z': {'b': 4} }
    """
    # declare res as the return object which should be a dict
    res = dict()
    # traverse the pollsters
    for pollster in nested_dict:
    	# travserse the states
    	for state in nested_dict[pollster]:
    		# if first meet a state, we need to create a new dict
    		if state not in res:
    			res[state] = dict()
    		# put the pollster value in the state dict
    		res[state][pollster] = nested_dict[pollster][state]
    return res

################################################################################
# Problem 6: Average the edges in a single state
################################################################################

def average_error_to_weight(error):
    """
    Given the average error of a pollster, returns that pollster's weight.
    The error must be a positive number.
    """
    return error ** (-2)

# The default average error of a pollster who did no polling in the
# previous election.
DEFAULT_AVERAGE_ERROR = 5.0

def pollster_to_weight(pollster, pollster_errors):
    """"
    Given a pollster and a pollster errors, return the given pollster's weight.
    """
    if pollster not in pollster_errors:
        weight = average_error_to_weight(DEFAULT_AVERAGE_ERROR)
    else:
        weight = average_error_to_weight(pollster_errors[pollster])
    return weight


def weighted_average(items, weights):
    """
    Returns the weighted average of a list of items.
    
    Arguments:
    items is a list of numbers.
    weights is a list of numbers, whose sum is nonzero.
    
    Each weight in weights corresponds to the item in items at the same index.
    items and weights must be the same length.
    """
    assert len(items) > 0
    assert len(items) == len(weights)
    # declare total as the return value which is a decimal
    total = 0.0
    # for all pairs from two lists
    for i in range(len(items)):
    	# we increment the total for the product of both value
    	total += items[i] * weights[i]
    # we return the total divided by sum of weights
    return total / sum(weights)


def average_edge(pollster_edges, pollster_errors):
    """
    Given pollster edges and pollster errors, returns the average of these edges
    weighted by their respective pollster errors.
    """
    # in order to make use of weighted_average() function, we need to create 2 lists
    pollsters = []
    weights = []
    # traverse all pollster edges
    for key in pollster_edges:
    	# we update two lists with same order
    	pollsters.append(pollster_edges[key])
    	weights.append(pollster_to_weight(key, pollster_errors))
    # simply return the result of weighted_average() function
    return weighted_average(pollsters, weights)

    
################################################################################
# Problem 7: Predict the 2012 election
################################################################################

def predict_state_edges(pollster_predictions, pollster_errors):
    """
    Given pollster predictions from a current election and pollster errors from
    a past election, returns the predicted state edges of the current election.
    """
    # firstly we get the pivot dict of predictions
    pivot_dict = pivot_nested_dict(pollster_predictions)
    # declare a dict as the return object
    res = dict()
    # traverse all states, calculate the average edge and put that in the dict
    for key in pivot_dict.keys():
    	res[key] = average_edge(pivot_dict[key], pollster_errors)
    return res

################################################################################
# Electoral College, Main Function, etc.
################################################################################

def electoral_college_outcome(ec_rows, state_edges):
    """
    Given electoral college rows and state edges, returns the outcome of
    the Electoral College, as a map from "Dem" or "Rep" to a number of
    electoral votes won.  If a state has an edge of exactly 0.0, its votes
    are evenly divided between both parties.
    """
    ec_votes = {}               # maps from state to number of electoral votes
    for row in ec_rows:
        ec_votes[row["State"]] = float(row["Electors"])

    outcome = {"Dem": 0, "Rep": 0}
    for state in state_edges:
        votes = ec_votes[state]
        if state_edges[state] > 0:
            outcome["Dem"] += votes
        elif state_edges[state] < 0:
            outcome["Rep"] += votes
        else:
            outcome["Dem"] += votes/2.0
            outcome["Rep"] += votes/2.0
    return outcome


def print_dict(dictionary):
    """
    Given a dictionary, prints its contents in sorted order by key.
    Rounds float values to 8 decimal places.
    """
    for key in sorted(dictionary.keys()):
        value = dictionary[key]
        if type(value) == float:
            value = round(value, 8)
        print key, value


def main():
    """
    Main function, which is executed when election.py is run as a Python script.
    """
    # Read state edges from the 2008 election
    edges_2008 = state_edges(read_csv("data/2008-results.csv"))
    
    # Read pollster predictions from the 2008 and 2012 election
    polls_2008 = pollster_predictions(read_csv("data/2008-polls.csv"))
    polls_2012 = pollster_predictions(read_csv("data/2012-polls.csv"))
    
    # Compute pollster errors for the 2008 election
    error_2008 = pollster_errors(polls_2008, edges_2008)
    
    # Predict the 2012 state edges
    prediction_2012 = predict_state_edges(polls_2012, error_2008)
    
    # Obtain the 2012 Electoral College outcome
    ec_2012 = electoral_college_outcome(read_csv("data/2012-electoral-college.csv"),
                                        prediction_2012)
    
    print "Predicted 2012 election results:"
    print_dict(prediction_2012)
    print
    
    print "Predicted 2012 Electoral College outcome:"
    print_dict(ec_2012)
    print    


# If this file, election.py, is run as a Python script (such as by typing
# "python election.py" at the command shell), then run the main() function.
if __name__ == "__main__":
    main()
