from ....errors import ConvertError
from datetime import datetime
from ....util import get_state_key
from collections import defaultdict
from dateutil.relativedelta import *

def set_flag(dob, arrest_records, ruleset, **kw):    
    if dob[0] is not None:
        yrob, moob, daob = dob
    else:
        yrob = 9999
                
    for row in arrest_records.itervalues():
        try:
            doa = datetime(row['arryr'], row['arrmo'], row['arrda'])
        except ValueError:
            #log error
            row['juv_flag'] = False
            continue
        
        #check to make sure we got the date of birth, otherwise log an error
        if yrob == 9999 or yrob < 1900:
            err_msg = "No date of birth for subject casenum=%s" % row['casenum'] if yrob == 9999 else "Date of birth < 1900 for subject casenum=%s" % row['casenum'] 
            kw['log'].write('ajuvenile', [row['id'], row['astate'], yrob, err_msg], store=True)
            row['juv_flag'] = False
            row['bad_age_flag'] = 1
            continue
        
        if doa.year == 9999:
            kw['log'].write('ajuvenile', [row['id'], row['astate'], str(doa), 
                                          "No date of arrest for casenum=%s" % row['casenum']], store=True)
            row['juv_flag'] = False
            continue
        
        dob = datetime(yrob, moob, daob)
        
        #calculate the subjects age at the date of arrest
        age = relativedelta(doa, dob).years
        
        try:
            #get the juvenile age of arrest for a given state
            juv_age = int(ruleset.convert(get_state_key(row['astate'])))
            
        except ConvertError:
            #if state was not found default to 17
            juv_age = 17
        
        '''Exception Rule 2: If age at arrest is < than 7, assume that either DOB or Arrest date is wrong:
                - keep all entries as adult, but include in bad arrest age flag'''
        if age < 7:
            kw['log'].write('ajuvenile', [row['id'], row['astate'], str(doa), str(dob), 
                                          "Either date of arrest or date of birth wrong for casenum=%s" % row['casenum']], store=True)
            row['juv_flag'] = False
            row['bad_age_flag'] = 1
            row['under_seven_flag'] = 1
            continue
        
        '''if subjects age at the time of the arrest is less than or equal to the states juvenile age 
        than flag case for possible deletion in recidivism file'''
        row['juv_flag'] = True if age <= juv_age else False
            
            
def transform(arrest_records):
    '''This transform implements RULE 3:
        
        If atype != 2 (for all records) Then:
            If age at arrest is < than 7, and 2 or more other cycle in the file are marked as underage; assume DOB is wrong:
        
                1) Set juv_flag = False and set bad_age_flag
    '''
    
    arrest_obs = arrest_records.copy()
    juv_cycles = defaultdict(set)
    atype = dict()
    under_seven = dict()
    
    for r in arrest_records.itervalues():
        #if atype = 2 then we delete all juv records records for a case
        if r['atype'] == 2:
            atype[r['casenum']] = True
        
        #get cases that have an under seven flag set
        if r.get('under_seven_flag'):
            under_seven[r['casenum']] = True
        
        #otherwise we want to get the set of cycles that has juv_flags
        if r['juv_flag']:
            juv_cycles[r['casenum']].add(r['cycnum']) 
            
        
    for id, row in arrest_records.iteritems():
        #check for atype = 2 for any record. If found do nothing for the given casenum
        if row['casenum'] in atype:
            continue
        
        #if the current case is under seven years at date of arrest and th 
        if under_seven.get(row['casenum']) and len(juv_cycles[row['casenum']]) > 1 and len(under_seven) > 1:
            row['juv_flag'] = False
            row['bad_age_flag'] = 1
        