from datetime import datetime
import re

from openelex.base.transform import Transform, registry
from openelex.models import Candidate, Contest, Office, Party, RawResult, Result

STATE = 'IL'
PLACE = 'Chicago'


meta_fields = ['source', 'election_id', 'state', 'place']

contest_fields = meta_fields + ['start_date',
                                'end_date',
                                'election_type',
                                'primary_type',
                                'result_type',
                                'special',
                                ]
candidate_fields = meta_fields + ['full_name', 'given_name',
                                  'family_name', 'additional_name']
result_fields = meta_fields + ['reporting_level', 'jurisdiction',
                               'votes', 'total_votes', 'vote_breakdowns']


class BaseTransform(Transform):

    district_offices = set([
        'U.S. Senate',
        'U.S. House of Representatives',
        'State Senate',
        'State House of Representatives',
    ])

    def __init__(self):
        super(BaseTransform, self).__init__()

    def get_raw_results(self):
        return RawResult.objects.filter(state=STATE, place=PLACE).no_cache()

class CreateContestsTransform(BaseTransform):
    name = 'chicago_create_unique_contests'

    def __call__(self):
        contests = []
        seen = set()

        for result in self.get_raw_results():
            key = self._contest_key(result)
            if key not in seen:
                fields = self.get_contest_fields(result)
                if fields:
                    fields['updated'] = datetime.now()
                    fields['created'] = datetime.now()
                    contest = Contest(**fields)
                    contests.append(contest)
                    seen.add(key)

        print "SEEN:", seen
        Contest.objects.insert(contests, load_bulk=False)
        logger.info("Created {} contests.".format(len(contests)))

    def _contest_key(self, raw_result):
        slug = raw_result.contest_slug
        return (raw_result.election_id, slug)

    def get_contest_fields(self, raw_result):
        fields = self._get_fields(raw_result, contest_fields)
        office = self._get_or_make_office(raw_result)
        if office:
            fields['office'] = office
            return fields
        else:
            return None

    def _get_fields(self, raw_result, field_names):
        return {k: getattr(raw_result, k) for k in field_names}

    def _get_or_make_office(self, raw_result):
        clean_name = self._clean_office(raw_result.office)

        if clean_name:

            office_query = {
                'state': STATE,
                'place': PLACE,
                'name': clean_name,
                'district': ''
            }

            if office_query['name'] is 'President':
                office_query['state'] = 'US'
                office_query['place'] = None

            if office_query['name'] in self.district_offices:
                if raw_result.district:
                    office_query['district'] = raw_result.district

            key = Office.make_key(**office_query)

            try:
                office = Office.objects.get(**office_query)
                return office
            except Office.DoesNotExist:
                office = Office(state=STATE, place=office_query['place'], name=office_query['name'], district=office_query['district'])
                Office.objects.insert([office])
                return office
        else:
            return None

    def _clean_office(self, office):
        """
        See: https://github.com/openelections/core/blob/dev/openelex/us/wa/load.py#L370

        """

        us_pres_regex = re.compile('president.+united\sstates|pres\sand\svice\spres', re.IGNORECASE)
        if re.search(us_pres_regex, office):
            return 'President'

        us_senator_regex = re.compile('senator.+u\.s\.|u\.s\..+senator|united\sstates\ssenator', re.IGNORECASE)
        if re.search(us_senator_regex, office):
            return 'U.S. Senate'

        return None



registry.register('il', CreateContestsTransform)


@@ -1918,74 +1918,8 @@
 lse)
-%0A        logger.info(%22Created %7B%7D contests.%22.format(len(contests)))
 %0A%0A  
@@ -2525,24 +2525,29 @@
 clean_office
+_name
 (raw_result.
@@ -2610,140 +2610,203 @@
 y = 
-%7B%0A                'state': STATE,%0A                'place': PLACE,%0A                'name': clean_name,%0A                'district': ''
+self._make_office_query(clean_name, raw_result)%0A%0A            try:%0A                office = Office.objects.get(**office_query)%0A                return office%0A            except Office.DoesNotExist:
 %0A   
@@ -2818,234 +2818,188 @@
     
-%7D%0A%0A
     
-        if office_query%5B'name'%5D is 'President':%0A                office_query%5B'state'%5D = 'US'%0A                office_query%5B'place'%5D = None%0A%0A            if office_query%5B'name'%5D in self.district_offices:%0A                if
+office = Office(**office_query)%0A                office.save()%0A                return office%0A        else:%0A            return None%0A%0A    def _make_office_query(self, office_name,
  raw
@@ -3005,25 +3005,17 @@
 w_result
-.district
+)
 :%0A      
@@ -3012,36 +3012,36 @@
 t):%0A        
- 
-   
+%22%22%22%0A
         
 office_query
@@ -3032,188 +3032,98 @@
     
-office_query%5B'district'%5D = raw_result.district%0A%0A            key = Office.make_key(**office_query)%0A%0A            try:%0A                office = Office.objects.get(**office_query)%0A    
+Gets the right state, place, district for an office%0A        %22%22%22%0A%0A        office_query = %7B%0A
     
@@ -3126,37 +3126,43 @@
             
-return
+'name':
  office
+_name
 %0A           
@@ -3162,272 +3162,324 @@
     
+%7D%0A%0A
     
-except Office.DoesNotExist:%0A                office = Office(state=STATE, place=office_query%5B'place'%5D, name=office_query%5B'name'%5D, district=office_query%5B'district'%5D)%0A                Office.objects.insert(%5Boffice%5D)%0A                return office%0A        else:%0A    
+    if office_name is 'President':%0A            office_query%5B'state'%5D = 'US'%0A            office_query%5B'place'%5D = None%0A%0A        if office_name in self.district_offices:%0A            office_query%5B'state'%5D = STATE%0A            if raw_result.district:%0A                office_query%5B'district'%5D = raw_result.district%0A%0A
     
@@ -3481,36 +3481,45 @@
 %0A        return 
-None
+office_query%0A
 %0A%0A    def _clean
@@ -3521,24 +3521,29 @@
 clean_office
+_name
 (self, offic

save offices, put all logic for office geographies in one place