#!/usr/bin/python3

"""Script to process poll results from Elections Canada."""

import psycopg2

# Ugly global db connection
conn = psycopg2.connect("dbname=electoralmap")

def create_cleanpoll():
    """
    Creates a new table called cleanpoll for cleaned polling data.
    """
    cur = conn.cursor()
    cur.execute("""CREATE TABLE cleanpoll ( fednum integer
                                          , emrp_name character varying(31)
                                          , pollname  character varying(255)
                                          , valid boolean DEFAULT False
                                          , electors integer DEFAULT 0
                                          , libvotes integer DEFAULT 0
                                          , convotes integer DEFAULT 0
                                          , ndpvotes integer DEFAULT 0
                                          , blqvotes integer DEFAULT 0
                                          , grnvotes integer DEFAULT 0
                                          , othvotes integer DEFAULT 0
                                          , nonvotes integer DEFAULT 0
                                          );""")
    conn.commit()

def add_votes(cur, fednum, emrp_name, party, votes):
    """Updates cleanpoll vote totals."""
    if not votes:
        return
    # SQL injection: use party as column name
    query = "UPDATE cleanpoll SET " + party + " = " + party + " + %(votes)s WHERE fednum = %(fednum)s AND emrp_name = %(emrp_name)s;"
    cur.execute(query,
                { 'fednum' : fednum
                , 'party' : party
                , 'votes' : votes
                , 'emrp_name' : emrp_name
                })

def get_psnums(cur, fednum, emrp_name):
    """
    Get all polling station numbers from pollresults which
    should be combined into the row for emrp_name.
    """
    cur.execute("""SELECT array_agg(DISTINCT psnum) FROM pollresults WHERE
                           ednum = %(fednum)s
                   AND NOT voidpoll
                   AND NOT nopoll
                   AND (   (   mergewith = ''
                           AND (psnum = %(psnum)s OR psnum LIKE %(psnum_letter)s)
                           )
                       OR  mergewith = %(psnum)s OR mergewith LIKE %(psnum_letter)s
                       );""",
                { 'psnum' : emrp_name
                , 'psnum_letter' : emrp_name + '_' # Matches, e.g., '26A'
                , 'fednum' : fednum
                })
    return cur.fetchone()[0]

def process_poll(cur, fednum, emrp_name):
    """
    Adds a new entry to the cleanpoll table for the specified poll.
    """
    # Create row
    cur.execute("INSERT INTO cleanpoll (fednum, emrp_name) VALUES (%s,%s)",
                (fednum, emrp_name))

    psnums = get_psnums(cur, fednum, emrp_name)
    # Map of party columns to search patterns
    parties = { 'libvotes' : '%Liberal%'
              , 'ndpvotes' : '%NDP%'
              , 'convotes' : '%Conservative%'
              , 'blqvotes' : '%Bloc%'
              , 'grnvotes' : '%Green%'
              }
    votetotals = []
    # For each party, calculate its vote total for all polls
    # merging into the given poll.
    for party, partypattern in parties.items():
        cur.execute("""SELECT sum(votes) FROM pollresults WHERE
                           ednum = %(fednum)s
                       AND psnum = ANY(%(psnums)s)
                       AND party_en LIKE %(partypattern)s;""",
                    { 'fednum' : fednum
                    , 'psnums' : psnums
                    , 'partypattern' : partypattern
                    })
        partyvotes, = cur.fetchone()
        add_votes(cur, fednum, emrp_name, party, partyvotes)
        votetotals.append(partyvotes)

    # Find the number of votes for independent candidates.
    params = parties
    params['fednum'] = fednum
    params['psnums'] = psnums
    cur.execute("""SELECT sum(votes) FROM pollresults WHERE
                       ednum = %(fednum)s
                   AND psnum = ANY(%(psnums)s)
                   AND NOT (   party_en LIKE %(libvotes)s
                           OR  party_en LIKE %(convotes)s
                           OR  party_en LIKE %(ndpvotes)s
                           OR  party_en LIKE %(blqvotes)s
                           OR  party_en LIKE %(grnvotes)s
                           )
                   ;""", params)
    othervotes, = cur.fetchone()
    add_votes(cur, fednum, emrp_name, 'othvotes', othervotes)
    votetotals.append(othervotes)

    # Count electors
    cur.execute("""SELECT sum(electors) FROM (
                        SELECT DISTINCT psnum, electors FROM pollresults
                        WHERE ednum = %(fednum)s AND psnum = ANY(%(psnums)s)
                   ) AS rows;""",
                params);
    electors, = cur.fetchone()
    add_votes(cur, fednum, emrp_name, 'electors', electors)

    # Calculate nonvoters
    nonvotes = electors - sum(votetotals)
    add_votes(cur, fednum, emrp_name, 'nonvotes', nonvotes)





@@ -1222,29 +1222,33 @@
 ()%0A%0Adef add_
-votes
+to_column
 (cur, fednum
@@ -1264,60 +1264,192 @@
 me, 
-party, votes):%0A    %22%22%22Updates cleanpoll vote totals.
+column, value):%0A    %22%22%22%0A    Adds 'value' to 'column' in cleanpoll table.%0A%0A    Note: This function uses the column parameter to build the query.%0A    Only call it with known inputs.%0A    
 %22%22%22%0A
@@ -1460,20 +1460,20 @@
 if not v
-otes
+alue
 :%0A      
@@ -1506,21 +1506,22 @@
 on: use 
-party
+column
  as colu
@@ -1562,29 +1562,30 @@
 oll SET %22 + 
-party
+column
  + %22 = %22 + p
@@ -1587,23 +1587,38 @@
 %22 + 
-party +
+column + %5C%0A           
  %22 + %25(v
 otes
@@ -1613,20 +1613,20 @@
  %22 + %25(v
-otes
+alue
 )s WHERE
@@ -1760,22 +1760,24 @@
  , '
-party' : party
+column' : column
 %0A   
@@ -1797,21 +1797,21 @@
 , 'v
-otes
+alue
 ' : v
-otes
+alue
 %0A   
@@ -3846,37 +3846,41 @@
 e()%0A        add_
-votes
+to_column
 (cur, fednum, em
@@ -4620,37 +4620,41 @@
 chone()%0A    add_
-votes
+to_column
 (cur, fednum, em
@@ -5029,37 +5029,41 @@
 chone()%0A    add_
-votes
+to_column
 (cur, fednum, em
@@ -5146,24 +5146,36 @@
 ctors - sum(
+(x for x in 
 votetotals)%0A
@@ -5172,16 +5172,22 @@
 tetotals
+ if x)
 )%0A    ad
@@ -5188,21 +5188,25 @@
     add_
-votes
+to_column
 (cur, fe

Rename add_votes to add_to_column