from collections import defaultdict
import cardsharp as cs
from random import random
from datetime import datetime, timedelta
from dateutil.relativedelta import *
import time, os, json, sys, random
from ..configuration import config
from load import set_varcodes


__all__ = ['transform', 'transfrom_demographic']

def _set_rearrest(data, varcodes, var, time_rearrest):
    """Sets the most recent rearrest"""
    if (data[var.upper()] == float(varcodes[var.lower()]['missing']) or data[var.upper()] > time_rearrest):
        data[var.upper()] = time_rearrest

def transform(arrest_records, sentence_records, variables, maxcycles, varcodes, **kw):
    cycstr1 = "001"
    
    #sorts incoming cases by cycle, and identifies the case number in the first place it is found
    casenum = varcodes['casenum']['missing']
    acycles, ccycles = defaultdict(list), defaultdict(list)
    pjp2 = defaultdict(set)
    #defines dictionaries for the summary and full flat file variables, inserts the casenum
    data_full, data_full_all_cnts, num_chgcnts = {}, {}, {}
    data_sum = defaultdict(int)
    
    for row in arrest_records.itervalues(): 
        acycles[row['cycle']].append(row)
        if casenum == varcodes['casenum']['missing']: casenum = row['casenum']
        
    for row in sentence_records.itervalues(): 
        ccycles[row['cycle']].append(row)
        if casenum == varcodes['casenum']['missing']: casenum = row['casenum']
        pjp2[row['cycle']].add(row['pjp2'])
    
    data_full['CASENUM'], data_full_all_cnts['CASENUM'], data_sum['CASENUM'] = casenum, casenum, casenum    
    
    #creates the date-of-release as a datetime object
    if kw['yrrel'] != 9999 and kw['yrrel'] is not None: release_dt = datetime(kw['yrrel'],kw['morel'],kw['darel'])
    else: release_dt = datetime(9999,1,12)

    #FULL CALCULATIONS
    #defines dictionaries of tuples containing the record variable to be used in generating each output variable
    #along with the non-applicable code to be used in cases where data does not exist
    avars = {'OFF': 'achg16', 'OFD': 'achg', 'FED': 'afed', 'SVR': 'achgsvr', 
            'DSP': 'adsp', 'DOM': 'adom', 'WPN': 'aweap', 'CDV': 'acdv'}
    cvars = {'OFF': 'cchg16', 'OFD': 'cchg', 'FED': 'cfed', 'SVR': 'cchgsvr', 
             'DSP': 'cdsp', 'DOM': 'cdom', 'WPN': 'cweap', 'CDV': 'ccdv',
             'PJP': 'pjp2', 'CNF': 'cnf', 'PMX': 'jpmx', 'PRB': 'prb'}
    
    #iterates over the arrest and sentence segments with information unique to each segment
    for seg in [(acycles,'a','A','arr',avars),(ccycles,'c','J','cdisp',cvars)]:
        #iterates over the maximum number of cycles possible
        for i in range(1, maxcycles + 1):
            #generates a string for the cycle with the appropriate number of preceding zeroes
            cycstr = (3-len(str(i)))*"0" + str(i)
            #checks to see if the case contains the requested cycle
            if i in seg[0].keys(): 
                #defines variables to keep track of the charge codes
                prev_chg, chgs, all_chgs  = [], 1, 1
                #calculates the number of charges for the given cycle
                acnt = sum([row['%scnt' % seg[1]] if row['%scnt' % seg[1]] not in [9999, None, 0] else 1 for row in seg[0][i]])

                #calculates the total number of unique charge codes for the given cycle 
                #num_chgcnts[i] = chgcnts = len(set([row['%schg' % seg[1]] for row in seg[0][i]])) 
                chgcnts = len(set([row['%schg' % seg[1]] for row in seg[0][i]]))
                
                #calculate number of total charges
                num_chgcnts[i] = len([row['%schg' % seg[1]] for row in seg[0][i]])
                
                #sorts charges by their sequence number and iterates over them
                for row in sorted(seg[0][i], key=lambda k: k['%schgseq' % seg[1]]):
                    if chgs == 1: 
                        #if this is the first charge, record the date/state/count
                        data_full['%s%sYR' % (seg[2],cycstr)] = row['%syr' % seg[3]]
                        data_full['%s%sMO' % (seg[2],cycstr)] = row['%smo' % seg[3]]
                        data_full['%s%sDA' % (seg[2],cycstr)] = row['%sda' % seg[3]]
                        data_full['%s%sST' % (seg[2],cycstr)] = row['%sstate' % seg[1]]
                        data_full['%s%sTOT' % (seg[2],cycstr)] = acnt
                            
                    if chgs <= min([chgcnts,3]):
                        chg = row['%schg' % seg[1]]
                        #if the charge is not a duplicate of one previously recorded
                        if chg not in prev_chg:
                            prev_chg.append(chg)
                            #get the number of charges with this charge code in the cycle
                            chgcnt = sum([cnt if cnt not in [9999,None, 0] else 1 for cnt in [r['%scnt' % seg[1]] for r in seg[0][i] if r['%schg' % seg[1]] == chg]])
                            data_full['%s%sCNT%s' % (seg[2],cycstr,str(chgs))] = chgcnt
                            
                            #iterate over the variables for the segment and add them to data_full
                            for var in seg[4].keys():
                                data_full['%s%s%s%s' % (seg[2],cycstr,var,str(chgs))] = row[seg[4][var]]
                            if seg[2] == 'J':
                                #we need all DSP values for ADJ vars
                                for v in ['DSP', 'PMX']:
                                    data_full_all_cnts['J%s%s%s' % (cycstr, v, str(all_chgs))] = row[seg[4][v]]
                                                 
                            chgs += 1
                            all_chgs += 1
                        
                        else:
                            if seg[2] == 'J':
                                #we need all DSP values for ADJ vars
                                for v in ['DSP', 'PMX']:
                                    data_full_all_cnts['J%s%s%s' % (cycstr, v, str(all_chgs))] = row[seg[4][v]]
                                all_chgs += 1
                    else:
                        #we need all DSP values for ADJ vars
                        if seg[2] == 'A':
                            break
                        elif seg[2] == 'J':
                            for v in ['DSP', 'PMX']:
                                data_full_all_cnts['J%s%s%s' % (cycstr, v, str(all_chgs))] = row[seg[4][v]]
                            all_chgs += 1                        

                #if there are still less than 3 charges recorded, record non-applicable codes until the 3rd is reached
                while chgs <= 3:
                    _d_key = '%s%sCNT%s' % (seg[2],cycstr,str(chgs))
                    _v_key = '%scnt%i' % (seg[2], chgs)
                    
                    data_full[_d_key] = varcodes[_v_key.lower()]['na']
                    
                    for var in seg[4].keys():
                        _d_key = '%s%s%s%i' % (seg[2], cycstr, var, chgs)
                        _v_key = '%s%s%i' % (seg[2], var, chgs)
                        
                        data_full[_d_key] = varcodes[_v_key.lower()].get('na')

                    chgs += 1
            #if no data exists for this cycle, record non-applicable for all variables
            else:
                num_chgcnts[i] = 0
                s = seg[2]
                c = cycstr
                _s = s.lower()
                
                data_full['%s%sYR' % (s,c)] = varcodes['%syr' % _s]['na']
                data_full['%s%sMO' % (s,c)] = varcodes['%smo' % _s]['na']
                data_full['%s%sDA' % (s,c)] = varcodes['%sda' % _s]['na']
                data_full['%s%sST' % (s,c)] = varcodes['%sst' % _s]['na']
                data_full['%s%sTOT' % (s,c)] = varcodes['%stot' % _s]['na']
                
                for chgs in range(1,4):
                    _d_key = '%s%sCNT%s' % (seg[2],cycstr,str(chgs))
                    _v_key = '%scnt%i' % (seg[2], chgs)
                    
                    data_full[_d_key] = varcodes[_v_key.lower()]['na']
                    for var in seg[4].keys():
                        _d_key = '%s%s%s%i' % (s, cycstr, var, chgs)
                        _v_key = '%s%s%i' % (s, var, chgs)
                        
                        data_full[_d_key] = varcodes[_v_key.lower()]['na']
                           

    #SUMMARY CALCULATIONS
    for var in ['TMREARR', 'TMREARRC', 'TMREARRA', 'TMREARRI', 'TMREARRP']:
        if var.lower() in varcodes: #only add to data_sum if variable is present
            data_sum[var] = float(varcodes[var.lower()]['missing'])
    
    #pjp2 codes that indicate a conviction
    conviction_set = set([1, 2, 3, 4, 5, 6])  
    #set the disposition set for ADJ variables
    dsp_set = range(31,48) + [47,50,91,93]
    
    ncrp_state_map= {58:6}
    
    ncrp_state = ncrp_state_map[kw['state']] if kw['state'] in ncrp_state_map else kw['state']
    
    dsp = defaultdict(set)
    jpmx = defaultdict(set)
    
    #iterate over all cycles to get list of dispositions for a given cycle
    for i in range(1, maxcycles + 1):
        cycstr = (3-len(str(i)))*"0" + str(i)
        find_next_dsp = True
        
        _c = 1
        while _c <= num_chgcnts[i]:
            dsp[i].add(data_full_all_cnts['J%sDSP%s' % (cycstr,str(_c))] )
            jpmx[i].add(data_full_all_cnts['J%sPMX%s' % (cycstr,str(_c))])
            _c += 1
    
    for i in range(1, maxcycles + 1):
        #set flag variables to False
        arrdt_flag, arrdt_follow_up_flag, anytype_flag, dsp_flag, firsttype_flag, state_flag, prs_flag = False, False, False, False, False, False, False
        cycstr = (3-len(str(i)))*"0" + str(i)  
        
        if data_full['A%sYR' % cycstr] != int(varcodes['ayr']['na']) and release_dt.year != 9999:
            arr_dt = datetime(data_full['A%sYR' % cycstr], data_full['A%sMO' % cycstr], data_full['A%sDA' % cycstr])
            
            #set arr_dt flag to true if arrest date < release date
            if arr_dt < release_dt: arrdt_flag = True
            
            #set arrdt_follow_up_flag to True if relase date after day of release but less than 5 years later 
            if release_dt <= arr_dt <= release_dt + timedelta(days=(365.25*kw['follow_period'])): arrdt_follow_up_flag = True
            
            #get the number of dayfrom release to FIRST rearrest
            if arr_dt >= release_dt:
                tmrearr = (arr_dt - release_dt).days
                _set_rearrest(data_sum, varcodes, 'TMREARR', tmrearr) 
                #calculate TMREARRC if 1 <= pjp2 <= 6
                #get the length of time between release and first rearrest resulting in a conviction
                if len(pjp2[i].intersection(conviction_set)) > 0:
                    _set_rearrest(data_sum, varcodes, 'TMREARRC', tmrearr)
                #get the length of time between release and first rearrest let lead to an adjudication
                if len(dsp[i].intersection(dsp_set)) > 0:
                    _set_rearrest(data_sum, varcodes, 'TMREARRA', tmrearr)                    
                
                if 'TMREARRI' in data_sum: #only set rearrest if variable is present in varcodes
                    if len(pjp2[i].intersection((1,2,3))) > 0:
                        _set_rearrest(data_sum, varcodes, 'TMREARRI', tmrearr)
                if 'TMREARRP' in data_sum: #only set rearrest if variable is present in varcodes
                    if len(pjp2[i].intersection((1,))) > 0:
                        _set_rearrest(data_sum, varcodes, 'TMREARRP', tmrearr)
                    else:
                        for x in jpmx[i]:
                            if 365 <= x < 99799799.77:
                                _set_rearrest(data_sum, varcodes, 'TMREARRP', tmrearr)
                                break
                    
        for j in range(1, 4):
            #set first most serios charge flag
            if j == 1 and data_full['J%sPJP%s' % (cycstr,str(j))] in [1, 2, 3]: firsttype_flag = True
            #set any charge flag
            if data_full['J%sPJP%s' % (cycstr,str(j))] in [1, 2, 3, 4, 5, 6]: anytype_flag = True
        
        _c = 1
        
        while _c <= num_chgcnts[i]:
            #set any dsp flag
            if data_full_all_cnts['J%sDSP%s' % (cycstr,str(_c))] in dsp_set: dsp_flag = True 
            #set prison flag, only set flag if variable is present in varcodes
            if 'prprs' in varcodes:
                if data_full.get('J%sPJP%s' % (cycstr,str(_c)), 0) == 1 or 365 <= data_full_all_cnts.get('J%sPMX%s' % (cycstr,str(_c)), 0) < 99799799.77: prs_flag = True
            _c+=1
        
            
        #if arrest state != ncrp_state and is not unknown set out of state flag to true
        if data_full['A%sST' % cycstr] != ncrp_state and data_full['A%sST' % cycstr] != 99: state_flag = True
            
        if arrdt_flag:
            data_sum['PRARR'] += 1
            if anytype_flag:
                data_sum['PRCNV'] += 1
                if state_flag:
                    data_sum['PROSCNV'] += 1
            if firsttype_flag:
                data_sum['PRCNF'] += 1
                if state_flag:
                    data_sum['PROSCNF'] += 1
            if dsp_flag:
                data_sum['PRADJ'] += 1
                if state_flag:
                    data_sum['PROSADJ'] += 1
            if state_flag:
                data_sum['PROSARR'] += 1
            if prs_flag: 
                data_sum['PRPRS'] += 1
                if state_flag:
                    data_sum['PROSPRS'] += 1
                
        if arrdt_follow_up_flag:
            data_sum['REARR'] += 1
            if anytype_flag:
                data_sum['RECNV'] += 1
                if state_flag: 
                    data_sum['REOSCNV'] += 1
            if firsttype_flag:
                data_sum['RECNF'] += 1
                if state_flag:
                    data_sum['REOSCNF'] += 1
            if dsp_flag:
                data_sum['READJ'] += 1
                if state_flag:
                    data_sum['REOSADJ'] += 1
            if state_flag:
                data_sum['REOSARR'] += 1
            if prs_flag:
                data_sum['REPRS'] += 1
                if state_flag:
                    data_sum['REOSPRS'] += 1
                    

    #COMPILE OUTPUT LINES
    full_output = [data_full['CASENUM']]
    for i in range(1, maxcycles + 1):
        cycstr = (3-len(str(i)))*"0" + str(i)  
        full_output.extend([data_full['A%sYR' % cycstr], data_full['A%sMO' % cycstr], data_full['A%sDA' % cycstr], data_full['A%sST' % cycstr], data_full['A%sTOT' % cycstr]])
        for j in range(1, 4):
            full_output.extend([data_full['A%sOFF%s' % (cycstr,str(j))], data_full['A%sOFD%s' % (cycstr,str(j))], data_full['A%sCNT%s' % (cycstr,str(j))]])
            full_output.extend([data_full['A%sFED%s' % (cycstr,str(j))], data_full['A%sSVR%s' % (cycstr,str(j))], data_full['A%sDSP%s' % (cycstr,str(j))]])
            full_output.extend([data_full['A%sDOM%s' % (cycstr,str(j))], data_full['A%sWPN%s' % (cycstr,str(j))], data_full['A%sCDV%s' % (cycstr,str(j))]])
    for i in range(1, maxcycles + 1):
        cycstr = (3-len(str(i)))*"0" + str(i)  
        full_output.extend([data_full['J%sYR' % cycstr], data_full['J%sMO' % cycstr], data_full['J%sDA' % cycstr], data_full['J%sST' % cycstr], data_full['J%sTOT' % cycstr]])
        for j in range(1, 4):
            full_output.extend([data_full['J%sOFF%s' % (cycstr,str(j))], data_full['J%sOFD%s' % (cycstr,str(j))], data_full['J%sCNT%s' % (cycstr,str(j))]])
            full_output.extend([data_full['J%sFED%s' % (cycstr,str(j))], data_full['J%sSVR%s' % (cycstr,str(j))], data_full['J%sDSP%s' % (cycstr,str(j))]])
            full_output.extend([data_full['J%sDOM%s' % (cycstr,str(j))], data_full['J%sWPN%s' % (cycstr,str(j))], data_full['J%sCDV%s' % (cycstr,str(j))]])
            full_output.extend([data_full['J%sPJP%s' % (cycstr,str(j))], data_full['J%sCNF%s' % (cycstr,str(j))], data_full['J%sPMX%s' % (cycstr,str(j))], data_full['J%sPRB%s' % (cycstr,str(j))]])

    sum_output = [data_sum['CASENUM']]
    var_list = ['PRARR', 'PRCNV', 'PRCNF', 'PROSARR', 'PROSCNV', 'PROSCNF', 'PRADJ', 'PROSADJ', 'PRPRS', 'PROSPRS',
                'REARR', 'RECNV', 'RECNF', 'REOSARR', 'REOSCNV', 'REOSCNF', 'READJ', 'REOSADJ', 'RECNFTV', 'REPRS', 'REOSPRS',
                'TMREARR', 'TMREARRC', 'TMREARRA', 'TMREARRI', 'TMREARRP']

    sum_output.extend([data_sum[var] for var in var_list])
    
    return (full_output, sum_output)

def transfrom_demographic(casenum, demo_records, max_profiles, varcodes):
    row = [casenum]
    #TODO make gender match sex in metadata 
    c = 0
    #TODO vars should load from metadata
    _vars = ['dstate', 'gender', 'race', 'hispanic', 'plcbth', 'ctzshp', 'dead']
    
    for _r in demo_records.itervalues():
        c += 1
        row.extend([_r[var] if _r[var] else varcodes[var]['missing'] for var in _vars])
    
    while c < max_profiles:
        row.extend([varcodes[v]['na'] for v in _vars])
        c += 1
    
    return row
    
    
    
