#!/usr/bin/python
# -*- coding: utf-8 -*-

import psycopg2
import sys
import csv
import glob
import re


def load_names():
    """Import names from SSN data dump into database"""

    con = psycopg2.connect(database='ssn_names')
    cur = con.cursor()

    # for filename in glob.iglob('*1880.txt'):
    for filename in glob.iglob('*.txt'):
        year = int(re.findall('\d+', filename)[0])
        print('starting {0}'.format(year))

        try:
            with open(filename, 'r') as f:
                reader = csv.reader(f, delimiter=',', quotechar='"')
                for row in reader:
                    data = tuple(row[0:3]) + (year,)
                    load_statement = "INSERT INTO raw_names " + \
                                     "(name, gender, count, year) " + \
                                     "VALUES ('%s', '%s', '%s', '%s')" % data
                    cur.execute(load_statement)
        except psycopg2.DatabaseError as e:
            print('Error {0}'.format(e))
            sys.exit(1)

        finally:
            if con:
                con.commit()


def calculate_rank():
    rank_statement = """
UPDATE raw_names
       set rank = r.rnk
FROM (
     SELECT name, year, gender, rank() OVER (partition by year, gender ORDER BY count DESC) AS rnk 
     FROM raw_names
     ) r
WHERE raw_names.name = r.name and raw_names.gender = r.gender and raw_names.year = r.year
"""

    con = psycopg2.connect(database='ssn_names')
    cur = con.cursor()
    cur.execute(rank_statement)
    con.commit()

def generate_list():
    unique_names_statement_1 = """
INSERT INTO names(name, gender)
SELECT name, 
       gender
  FROM raw_names
 GROUP BY name, gender
"""

    unique_names_statement_2 = """
SELECT name,
       gender,
       count(name) as num
  INTO temp_max
  FROM raw_names
 WHERE rank <= 50
 GROUP by name, gender
"""

    unique_names_statement_3 = """
UPDATE names n
   SET num_above_max = (SELECT num
                               FROM temp_max t
                              WHERE n.name = t.name
                                AND n.gender = t.gender)
"""

    unique_names_statement_4 = """
drop table temp_max
"""

    unique_names_statement_5 = """
SELECT name,
       gender,
       count(name) as num
  INTO temp_min
  FROM raw_names
 WHERE rank > 500
 GROUP by name, gender
"""
    
    unique_names_statement_6 = """
UPDATE names n
   SET num_below_min = (SELECT num
                               FROM temp_min t
                              WHERE n.name = t.name
                                AND n.gender = t.gender)
"""
    
    unique_names_statement_7 = """
drop table temp_min
"""

    unique_names_statement_8 = """
SELECT name,
       gender,
       count(name)
  INTO temp_count
  FROM raw_names
 GROUP by name, gender
"""

    unique_names_statement_9 = """
UPDATE names n
   SET total_count = t.count
  FROM temp_count t
 WHERE t.name = n.name
   AND t.gender = n.gender
"""

    unique_names_statement_10 = """
UPDATE names
   SET total_count_below = (136 - (total_count)) + num_below_min
"""

    unique_names_statement_11 = """
drop table temp_count
"""

    con = psycopg2.connect(database='ssn_names')
    cur = con.cursor()
    cur.execute(unique_names_statement_1)
    cur.execute(unique_names_statement_2)
    cur.execute(unique_names_statement_3)
    cur.execute(unique_names_statement_4)
    cur.execute(unique_names_statement_5)
    cur.execute(unique_names_statement_6)
    cur.execute(unique_names_statement_7)
    cur.execute(unique_names_statement_8)
    cur.execute(unique_names_statement_9)
    cur.execute(unique_names_statement_10)
    cur.execute(unique_names_statement_11)
    con.commit()


def main():
    # load_names()
    print('calculating rank')
    calculate_rank()
    print('generating list')
    generate_list()


if __name__ == "__main__":
    main()
