# ============================================================
# File:colorMap.py
# Description: program for making a color map of data
#              plotted on US map
# Author: Emily G. Allen
# Date: April 2008
# ============================================================  

# ============================================================
# Required Files
# ============================================================    
from state import *
from graphics import *
from math import *

# ============================================================
# ============================================================
# Creating/Drawing Colorized Map of US
# ============================================================
# ============================================================

# ============================================================
# Function: colorMap
# Description: creates a color map based upon range of data
#              within limits from minColor to maxColor
# Inputs: minColor (low values), maxColor (high values), data
#         minColor and maxColor are lists of [red, blue, green]
#         values
# Outputs: a color_rgb equivalency for each value in data
# ============================================================
def colorMap(minColor, maxColor, data):
        
    # calculate range of the data
    rnge = float(max(data) - min(data)) # to insure floating pt division later on
    if rnge == 0: # account for only one value
        rnge = 1

    # iterate over values in the data
    # calculating color equivalency for each value
    # equivalency is proportional to relative distance betweeen value and min(data)
    # minColor = [red, blue, green]
    # maxColor = [red, blue, green]
    colorList = []
    for value in data:
        difference = (value - min(data)) / rnge
        red = (maxColor[0] - minColor[0]) * difference + minColor[0]
        green = (maxColor[1] - minColor[1]) * difference + minColor[1]
        blue = (maxColor[2] - minColor[2]) * difference + minColor[2]

        colorList.append(color_rgb(red, green, blue)) # add rgb_color to list

    return colorList

# ============================================================
# Function: drawMap
# Description: draws (colorized) map of US
# Inputs: the drawing window, the statelist, the color list
# Outputs: none
# ============================================================  
def drawMap(window, stateList, capital, colorList):

    if len(colorList) == 0: # if no color, draw states in white
        for state in stateList:
            state.draw(window)
    else: # otherwise draw colored states
        for i in range(len(stateList)):
            stateList[i].draw(window, colorList[i])

    if capital: # if capital is True place capital dots
        for state in stateList:
            state.drawCapital(window)
            
# ============================================================
# Function: clearMap
# Description: clears the map to be redrawn
# Inputs: stateList
# Outputs: none
# ============================================================
def clearMap(stateList, window):
    for state in stateList:
        state.undraw()

    # erase legend
    rect = Rectangle(Point(100, 1100), Point(1400, 900))
    rect.setFill('white')
    rect.setOutline('white')
    rect.draw(window)

    # erase title
    rect = Rectangle(Point(0, 200), Point(1500,  0))
    rect.setFill('white')
    rect.setOutline('white')
    rect.draw(window)
        
# ============================================================
# Function: drawLegend
# Description: draws the colorbar legend
# Inputs: minColor, maxColor from which colormap was generated,
#         the data, categories, and the drawing window
# Outputs: none
# ============================================================  
def drawLegend(minColor, maxColor, data, categories, window):

    # determine the range being represented
    if categories == -1:
        increment =  range(int(min(data)), int(max(data)), int((max(data) - min(data)) / 10))
    else:
        increment = range(len(categories))

    # generated a color map for the legend out of the represented range
    legend = colorMap(minColor, maxColor, increment)

    # draw one rectangle for each color in legend
    x1, x2, y1, y2 = 100, 150, 1100, 1050 # initial rect. coords
    for c in range(len(legend)):
        r = Rectangle(Point(x1, y1), Point(x2, y2))
        r.setFill(legend[c])
        r.draw(window)

        # label the rectangle
        if categories == -1: # if quantitative values
            t = str(increment[c])
        else: 
            t = categories[c]  
        label = Text(Point(x1 + 25, y2 - 50), t)
        label.draw(window)
        
        x1 = x1 + 100 #update coordinates
        x2 = x2 + 100
          
# ============================================================
# ============================================================
# Processing Data to Be Plotted
# ============================================================
# ============================================================

# ============================================================
# Function: loadData
# Description: loads a list of data to be plotted
# Inputs: stateList (data from file; in format, state, data)
#         first line indicates whether or not data are quantitative
#         if qualitative, list number of categories
# Outputs: list of data
# ============================================================
def loadData(stateList):
    print '======================================================'
    print 'YOU HAVE ELECTED TO LOAD DATA FROM A FILE'
    print '======================================================\n'
    print 'Make sure your data file is in the following format: '
    print 'DATATYPE'
    print 'Map Title'
    print 'State Name1, Value1'
    print 'State Name2, Value2'
    print '...\n'
    print 'DATATYPE should be either QUANT (quantitative) or'
    print 'CAT (qualitative or categories).  If DATATYPE is CAT'
    print 'follow specifier with list of categories, e.g., :'
    print '     CAT=Bush, Kerry\n'
    print 'See \'election2004.txt\' for example.\n\n'

    # get file name
    filename = raw_input('Please enter data file name: ')
    infile = open(filename, 'r')

    # figure out the data type
    datatype = infile.readline()
    if datatype.find('CAT') != -1: # if qualitative
        # determine the categories and assign #s
        datatype = datatype.split('=') # segregate category list 
        categories = datatype[1].split(',') # segregate categories
        for c in range(len(categories)): # clean up string (remove spaces and newlines
            categories[c] = categories[c].strip()

    else: # otherwise quantitative
        categories = -1
        
    # read in title
    title = infile.readline()
    
    # read in remainder of data
    data = []
    for line in infile:
        line = line.split(',')
        stateName = line[0] # first thing is state name
        value = line[1] # second thing on line is the value

        # find the right state and associate value
        for state in stateList:
            if stateName == state.getName():
                if categories == -1:
                    # quantitative, evaluate value and
                    # associate with the state
                    state.setValue(eval(value))
                else:
                    # qualitative, find the index of the value
                    # in the category list and associate that
                    # with the state
                    state.setValue(categories.index(value.strip()))   

    infile.close()
    return title, categories, stateList

# ============================================================
# Function: loadEnergyData
# Description: loads energy data and makes available for plotting
# Inputs: none
# Outputs: set of lists of energy data, titles
# ============================================================        
def loadEnergyData():
    infile = open('stateEnergy.txt', 'r')

    # column titles
    titles = infile.readline()
    titles = titles.split(',')

    data = []
    # resut of data
    for line in infile:
        data.append(line.split(','))

    return titles, data
        
# ============================================================
# Function: setEnergyData
# Description: stores appropriate energy data subset in states
# Inputs: selected data to store, energy data table, state list
# Outputs: state list
# ============================================================  
def setEnergyData(choice, data, stateList):
    for currentState in data:
        # iterate thru states
        for state in stateList:
            # find saved state with same name
            if currentState[0] == state.getName():
                state.setValue(eval(currentState[choice]))
    return stateList

# ============================================================
# ============================================================
# Main Program
# ============================================================
# ============================================================         
def main():
    # create drawing window
    win = GraphWin('States', 1000, 700)
    win.setCoords(0, 1200, 1500 ,0)
    win.setBackground('white')# draw the map and legend 
  
    # load state polygons and draw initial map
    states = loadStateData()
    drawMap(win, states, False, [])

    # load energy file into menu
    energyTitles, energyData = loadEnergyData()
  
    # as long as user wants to update plot
    plotAgain = 'Y'
    while plotAgain == 'Y': #
        # create menu to choose personal data file or energy file
        print '\n\n======================================================'
        print 'SELECT DATA'
        print '======================================================'
        print 'Energy Data Sets: '
        i = 1
        while i < len(energyTitles):      
            if i == 2:
                print '\nPercent Total Energy Consumption per Source\n'
            if i == 9:
                print '\nPercent Total Energy Consumption per Consumer\n'
                
            print '[%d]: %s' % (i, energyTitles[i])
            i = i + 1
            
        print '\n[%d]: Load Data File' % i
        print '======================================================'

        dataChoice = input('Enter choice (1 - %d): ' % len(energyTitles))
        while dataChoice > len(energyTitles):
              dataChoice = input('Please enter a value in the correct range: ')

        # select data to be plotted
        if dataChoice == len(energyTitles): # load other data
            title, categories, states = loadData(states)
        else: # energy data
            categories = -1
            if dataChoice <= 8:
                title = 'Source: '
            else:
                title = 'Consumer: '
            title = title + energyTitles[dataChoice] + ' Consumption (BTU)'
            states = setEnergyData(dataChoice, energyData, states)

        print '\n\n======================================================'
        print 'Please enter lower and upper values for color map.'
        print 'For example, to make a scale that maps from blue to red: '
        print 'minColor = [0, 50, 255]'
        print 'maxColor = [255, 50, 0]'
        print '======================================================s'
        minColor = input('Please enter minColor [r, g, b]: ')
        maxColor = input('Please enter maxColor [r, g, b]: ')

        # generate color map
        colors = colorMap(minColor, maxColor, getStateData(states))
           
        clearMap(states, win) # clear items
        drawMap(win, states, False, colors)
        drawLegend(minColor, maxColor, getStateData(states), categories, win)

        # draw Title
        titleText = Text(Point(750, 100), title)
        titleText.draw(win)

        print '\n\n======================================================'
        plotAgain = raw_input('Generate another plot (Y/N)? ')
        plotAgain = plotAgain[0].upper() # deal with y/Y yes/Yes
                              
main()