from ..segment import *
from ..util import *
from ..log import *

from extract import *
from filter import filter, juvenile_filter, sentence_match_arrest_filter

from transform.rulecontrollers.wvfbi_overlap import set_flag as set_wvfbi_overlap_flag
from transform.rulecontrollers.merge_cycles import transform as merge_cycles
from transform.rulecontrollers.ajuvenile import set_flag as set_juv_flag
from transform.rulecontrollers.id_filter import set_flag as set_id_flag
from transform.rulecontrollers.ajuvenile import transform as juv_transform
from transform.rulecontrollers.afollow_up import set_flag as set_followup_flag
from transform.rulecontrollers.pjp2 import set_pjp2
from transform.rulecontrollers.offense_deletion import set_flag as set_off_del_flag
from transform.rulecontrollers.chgseq import chgseq_transform
from transform.rulecontrollers.dates import convert as convert_dt
from transform.rulecontrollers.doa import set_doa
from transform.rulecontrollers.admrec import transform as set_admrec_flag
from transform.rulecontrollers.state_historic_flowchart import transform as flowchart_historic_transform
import transform.rulecontrollers.state_historic_intersegment as state_intersegment


__all__ = ['run_transform']

def validate_state_in_cycle(observations, **kw):
    state_var = '%sstate' % ('a' if kw['segment'] == 'arrest' else 'c')
    for cycnum, cycle_data in observations.iteritems():
        state = None
        for id,row in cycle_data.iteritems():
            if state is None:
                state = row[state_var]
                continue
            
            if state != row[state_var]:
                raise ValidationError("States do not match within cycle")

def run_transform(conn_db, arrest_file, sentence_file, arrest_del_queue, sentence_del_queue, **kw):
    conn_cursor = conn_db.cursor()
    
    arrest_records = get_records('arrest', conn_cursor, **kw)            
    sentence_records = get_records('sentence', conn_cursor, **kw)
    arrest_del_records = []
    sentence_del_records = []
    
    #handle date processing that was not completed during process stage  
    for segment, records in [('arrest', arrest_records), ('sentence', sentence_records)]:
        for row in records.itervalues():
            convert_dt(segment, row, kw['log'])
                
    #apply USSC id filter
    set_id_flag(sentence_records, **kw)
    sentence_del_records.extend(filter(sentence_records, ['id_flag']))

    #apply wvfbi filter
    set_wvfbi_overlap_flag(arrest_records, **kw)
    arrest_del_records.extend(filter(arrest_records, ['wvfbi_overlap_flag']))
                            
    #apply admrec filter
    set_admrec_flag('arrest', arrest_records)
    set_admrec_flag('sentence', sentence_records)
    arrest_del_records.extend(filter(arrest_records, ['admrec_flag']))
    sentence_del_records.extend(filter(sentence_records, ['admrec_flag']))
            
    #set date of arrest to flag bad date values
    set_doa(arrest_records, sentence_records)
    
    #filter out bad dates
    arrest_del_records.extend(filter(arrest_records, ['adateflag']))
            
    #flag juvenile records
    if kw.get('juv_flag'):
        set_juv_flag((kw['yrob'], kw['moob'], kw['daob']), arrest_records, kw['arrest_rulesets']['age'], **kw)
        #handle rule 3 exception
        juv_transform(arrest_records)
                
    #flag and filter arrest records outside of follow_up period
    set_followup_flag((kw['yrrel'], kw['morel'], kw['darel']), arrest_records, **kw)
    arrest_del_records.extend(filter(arrest_records, ['follow_up_flag']))
            
    #Set offense deletion flag
    set_off_del_flag('arrest', arrest_records, **kw)
    set_off_del_flag('sentence', sentence_records, **kw)
        
    #assign pjp2
    set_pjp2(sentence_records, kw['sentence_rulesets'], **kw)

    #apply age, offense flag filters
    if kw.get('juv_flag'):
        a_del_juv, c_del_juv = juvenile_filter(arrest_records, sentence_records, **kw)
        
        arrest_del_records.extend(a_del_juv)
        sentence_del_records.extend(c_del_juv)
    
    arrest_del_records.extend(filter(arrest_records, ['off_del_flag']))
    sentence_del_records.extend(filter(sentence_records, ['off_del_flag']))
    
    #apply sentence filter for matching casenum/cycnum
    sentence_del_records.extend(sentence_match_arrest_filter(arrest_records, sentence_records))
    
    #apply pre-cycle historical rules
    pre_cycle = {'arrest': {}, 'sentence': {}}
        
        #manipulate the records to conform to pre_cycle state historic rules
    for container in [('arrest', arrest_records), ('sentence', sentence_records)]:
        for key, row in container[1].iteritems():
            region = key_as_str(row['rapstate'])
            if region != 'wvfbi': 
                if (row['cycnum'], region) in pre_cycle[container[0]]:
                    pre_cycle[container[0]][(row['cycnum'], region)][key] = row
                else:
                    pre_cycle[container[0]][(row['cycnum'], region)] = {key:row}
            else:
                #add WVFBI records
                if container[0] == 'arrest': 
                    region = get_state_key(row['astate'])
                    #get wvfbi cases for arrest-ky
                    if region == 'ky':
                        if (row['cycnum'], region) in pre_cycle[container[0]]:
                            pre_cycle[container[0]][(row['cycnum'], region)][key] = row
                        else:
                            pre_cycle[container[0]][(row['cycnum'], region)] = {key:row}
                
    
    #apply the join flow chart rules
    for key, value in pre_cycle['sentence'].iteritems():
        region = key[1].lower()
        if region in ['ca','il','hi','mn','nc','tx','nb','ne', 'ok','nj', 'md', 'ut']:
            #apply the join flow chart rules
            flowchart_historic_transform(value, region)
        
    #apply the intersegement pre cycle rules       
    for key, arrest_cycle_obs in pre_cycle['arrest'].iteritems():
        region = key[1].lower()
        if region in ['ky', 'ma']:
            state_intersegment.pre_cycle(region, arrest_cycle_obs, arrest_records)
        elif region in ['ct', 'md']:
            state_intersegment.pre_cycle(region, arrest_cycle_obs, arrest_records, pre_cycle['sentence'].get(key, {}), sentence_records)
    
    #apply cycle rules
    merge_cycles(arrest_records, sentence_records)
    
    #apply post-cycle historical rules using old cycnums
    for key, arrest_cycle_obs in pre_cycle['arrest'].iteritems():
        region = key[1].lower()
        if region in ['ky']:
            state_intersegment.post_cycle(region, arrest_cycle_obs, arrest_records, pre_cycle['sentence'].get(key, {}), sentence_records)
    
    for key, sentence_cycle_obs in pre_cycle['sentence'].iteritems():
        region = key[1].lower()
        if region in ['mo', 'nb', 'nc']:
            state_intersegment.post_cycle(region, pre_cycle['arrest'].get(key, {}), arrest_records, sentence_cycle_obs, sentence_records)

        #manipulate the records to conform to pre_cycle state historic rules
    post_cycle = {'arrest': {}, 'sentence': {}}
    
    for container in [('arrest', arrest_records), ('sentence', sentence_records)]:
        for key, row in container[1].iteritems():
            region = key_as_str(row['rapstate'])
            if region != 'wvfbi':
                if (row['cycle'], region) in post_cycle[container[0]]:
                    post_cycle[container[0]][(row['cycle'], region)][key] = row
                else:
                    post_cycle[container[0]][(row['cycle'], region)] = {key:row}
            else:
                #add WVFBI records
                if container[0] == 'sentence': 
                    region = get_state_key(row['cstate'])
                    #get wvfbi cases for arrest-ky
                    if region == 'vt':
                        if (row['cycle'], region) in post_cycle[container[0]]:
                            post_cycle[container[0]][(row['cycle'], region)][key] = row
                        else:
                            post_cycle[container[0]][(row['cycle'], region)] = {key:row}

    for key, arrest_cycle_obs in post_cycle['arrest'].iteritems():
        region = key[1].lower()
        if region in ['de','ma']: #added MA rule 2 (I think this is the right place to put it)
            state_intersegment.post_cycle(region, arrest_cycle_obs, arrest_records, post_cycle['sentence'].get(key, {}), sentence_records)

    for key, sentence_cycle_obs in post_cycle['sentence'].iteritems():
        region = key[1].lower()
        if region in ['ia', 'sc', 'wi', 'vt', 'ar', 'ny']:
            state_intersegment.post_cycle(region, post_cycle['arrest'].get(key, {}), arrest_records, sentence_cycle_obs, sentence_records)

    #remove the delflag=1 records
    arrest_del_records.extend(filter(arrest_records, ['delflag']))
    sentence_del_records.extend(filter(sentence_records, ['delflag']))
    
    arrest_del_queue.put(arrest_del_records)
    sentence_del_queue.put(sentence_del_records)
    
    #apply achgseq
    chgseq_transform('a', arrest_records)
    #apply cchgseq
    chgseq_transform('c', sentence_records)

    #TODO dump the arrest_records to the write queue instead of file
    #dump the arrest_records to files   
    for record in arrest_records.itervalues():
        out = []
        
        for var in kw['variables']['arrest']:
            out.append(str(record.get(var.name)).replace('\\', '\\\\') if record.get(var.name) is not None else '\N')
        arrest_file.write('%s\n' % '|'.join(out))
    
    #dump the sentence_records to files
    for record in sentence_records.itervalues():
        out = []
        
        for var in kw['variables']['sentence']:
            out.append(str(record.get(var.name)).replace('\\', '\\\\') if record.get(var.name) is not None else '\N')
        sentence_file.write('%s\n' % '|'.join(out))

    conn_cursor.close()
    