"""
corpkit: process CONLL formatted data
"""

def parse_conll(f,
                first_time=False,
                just_meta=False,
                usecols=None):
    """
    Make a pandas.DataFrame with metadata from a CONLL-U file
    
    Args:
        f (str): Filepath
        first_time (bool, optional): If True, add in sent index
        just_meta (bool, optional): Return only a metadata `dict`
        usecols (None, optional): Which columns must be parsed by pandas.read_csv
    
    Returns:
        pandas.DataFrame: DataFrame containing tokens and a ._metadata attribute
    """
    import pandas as pd
    try:
        from StringIO import StringIO
    except ImportError:
        from io import StringIO

    from collections import defaultdict

    # go to corpkit.constants to modify the order of columns if yours are different
    from corpkit.constants import CONLL_COLUMNS as head

    with open(f, 'r') as fo:
        data = fo.read().strip('\n')

    splitdata = []
    metadata = {}
    sents = data.split('\n\n')    
    for count, sent in enumerate(sents, start=1):
        metadata[count] = defaultdict(set)
        for line in sent.split('\n'):
            if line and not line.startswith('#') \
                and not just_meta:
                splitdata.append('\n%d\t%s' % (count, line))
            else:
                line = line.lstrip('# ')
                if '=' in line:
                    field, val = line.split('=', 1)
                    metadata[count][field].add(val)
        metadata[count] = {k: ','.join(v) for k, v in metadata[count].items()}
    if just_meta:
        return metadata

    # happens with empty files
    if not splitdata:
        return

    # head can only be as long as the list of cols in the df
    num_tabs = splitdata[0].strip('\t').count('\t')
    head = head[:num_tabs]
    
    # introduce sentence index for multiindex
    #for i, d in enumerate(splitdata, start=1):
    #    d = d.replace('\n', '\n%s\t' % str(i))
    #    splitdata[i-1] = d

    # turn into something pandas can read    
    data = '\n'.join(splitdata)
    data = data.replace('\n\n', '\n') + '\n'

    # remove slashes as early as possible
    data = data.replace('/', '-slash-')

    # open with sent and token as multiindex
    try:
        df = pd.read_csv(StringIO(data), sep='\t', header=None,
                         names=['s'] + head, index_col=['s', 'i'], usecols=usecols)
        #df.index = pd.MultiIndex.from_tuples([(1, i) for i in df.index])
    except ValueError:
        return
    df._metadata = metadata
    return df

def get_dependents_of_id(idx, df=False, repeat=False, attr=False, coref=False):
    """
    Get dependents of a token
    """
    sent_id, tok_id = getattr(idx, 'name', idx)
    deps = df.ix[sent_id, tok_id]['d'].split(',')
    out = []
    for govid in deps:
        if attr:
            # might not exist...
            try:
                tok = getattr(df.ix[sent_id,int(govid)], attr, False)
                if tok:
                    out.append(tok)
            except (KeyError, IndexError):
                pass
        else:
            out.append((sent_id, int(govid)))
    return out

def get_governors_of_id(idx, df=False, repeat=False, attr=False, coref=False):
    """
    Get governors of a token
    """
    
    # it can be a series or a tuple
    sent_id, tok_id = getattr(idx, 'name', idx)
    # get the governor id
    govid = df['g'].loc[sent_id, tok_id]
    if attr:
        return getattr(df.loc[sent_id,govid], attr, 'root')
    return [(sent_id, govid)]

def get_match(idx, df=False, repeat=False, attr=False, **kwargs):
    """
    Dummy function, for the most part
    """
    sent_id, tok_id = getattr(idx, 'name', idx)
    if attr:
        return df[attr].ix[sent_id, tok_id]
    return [(sent_id, tok_id)]

def get_head(idx, df=False, repeat=False, attr=False, **kwargs):
    """
    Get the head of a 'constituent'---'
    for 'corpus linguistics', if 'corpus' is searched, return 'linguistics'
    """

    sent_id, tok_id = getattr(idx, 'name', idx)
    #sent = df.ix[sent_id]
    token = df.ix[sent_id, tok_id]

    if not hasattr(token, 'c'):
        # this should error, because the data isn't there at all
        lst_of_ixs = [(sent_id, tok_id)]

    elif token['c'] == '_':
        lst_of_ixs = [(sent_id, tok_id)]
    # if it is the head, return it
    elif token['c'].endswith('*'):
        lst_of_ixs = [(sent_id, tok_id)]
    else:
        # should be able to speed this one up!
        just_same_coref = df.loc[sent_id][df.loc[sent_id]['c'] == token['c'] + '*']
        if not just_same_coref.empty:
            lst_of_ixs = [(sent_id, i) for i in just_same_coref.index]
        else:
            lst_of_ixs = [(sent_id, tok_id)]
    if attr:
        lst_of_ixs = [df.loc[i][attr] for i in lst_of_ixs]
    return lst_of_ixs

def get_representative(idx,
                       df=False,
                       repeat=False,
                       attr=False,
                       **kwargs):
    """
    Get the representative coref head
    """
    sent_id, tok_id = getattr(idx, 'name', idx)

    token = df.ix[sent_id, tok_id]
    # if no corefs at all
    if not hasattr(token, 'c'):
        # this should error, because the data isn't there at all
        lst_of_ixs = [(sent_id, tok_id)]     
    # if no coref available
    elif token['c'] == '_':
        lst_of_ixs = [(sent_id, tok_id)]
    else:
        just_same_coref = df.loc[df['c'] == token['c'] + '*']
        if not just_same_coref.empty:
            lst_of_ixs = [just_same_coref.iloc[0].name]
        else:
            lst_of_ixs = [(sent_id, tok_id)]
    if attr:
        lst_of_ixs = [df.ix[i][attr] for i in lst_of_ixs]
    return lst_of_ixs

def get_all_corefs(s, i, df, coref=False):
    # if not in coref mode, skip
    if not coref:
        return [(s, i)]
    # if the word was not a head, forget it
    if not df.ix[s,i]['c'].endswith('*'):
        return [(s, i)]
    try:
        # get any other mention head for this coref chain
        just_same_coref = df.loc[df['c'] == df.ix[s,i]['c']]
        return list(just_same_coref.index)
    except:
        return [(s, i)]

def search_this(df, obj, attrib, pattern, adjacent=False, coref=False):
    """
    Search the dataframe for a single criterion
    """
    import re
    out = []

    # if searching by head, they need to be heads
    if obj == 'h':
        df = df.loc[df['c'].endswith('*')]

    # cut down to just tokens with matching attr
    # but, if the pattern is 'any', don't bother
    if hasattr(pattern, 'pattern') and pattern.pattern == r'.*':
        matches = df
    else:
        matches = df[df[attrib].fillna('').str.contains(pattern)]

    # functions for getting the needed object
    revmapping = {'g': get_dependents_of_id,
                  'd': get_governors_of_id,
                  'm': get_match,
                  'h': get_all_corefs,
                  'r': get_representative}

    getfunc = revmapping.get(obj)

    for idx in list(matches.index):

        if adjacent:
            if adjacent[0] == '+':
                tomove = -int(adj[1])
            elif adjacent[0] == '-':
                tomove = int(adj[1])
            idx = (idx[0], idx[1] + tomove)
        
        for mindex in getfunc(idx, df=df, coref=coref):

            if mindex:
                out.append(mindex)

    return list(set(out))

def show_fix(show):
    """show everything"""
    objmapping = {'d': get_dependents_of_id,
                  'g': get_governors_of_id,
                  'm': get_match,
                  'h': get_head}

    out = []
    for val in show:
        adj, val = determine_adjacent(val)
        obj, attr = val[0], val[-1]
        obj_getter = objmapping.get(obj)
        out.append(adj, val, obj, attr, obj_getter)
    return out

def dummy(x, *args, **kwargs):
    return x

def format_toks(to_process, show, df):
    """
    Format matches by show values
    """

    import pandas as pd

    objmapping = {'d': get_dependents_of_id,
                  'g': get_governors_of_id,
                  'm': get_match,
                  'h': get_head}

    sers = []

    dmode = any(x.startswith('d') for x in show)
    if dmode:
        from collections import defaultdict
        dicts = defaultdict(dict)

    for val in show:
        adj, val = determine_adjacent(val)
        if adj:
            if adj[0] == '+':
                tomove = int(adj[1])
            elif adj[0] == '-':
                tomove = -int(adj[1])

        obj, attr = val[0], val[-1]
        func = objmapping.get(obj, dummy)
        out = defaultdict(dict) if dmode else []
        for ix in list(to_process.index):
            piece = False
            if adj:
                ix = (ix[0], ix[1] + tomove)
                if ix not in df.index:
                    piece = 'none'
            if not piece:
                if obj == 'm':
                    piece = df.loc[ix][attr.replace('x', 'p')]
                    if attr == 'x':
                        from corpkit.dictionaries.word_transforms import taglemma
                        piece = taglemma.get(piece.lower(), piece.lower())
                    piece = [piece]
                else:
                    piece = func(ix, df=df, attr=attr)
                    if not isinstance(piece, list):
                        piece = [piece]
                if dmode:
                    dicts[ix][val] = piece
                else:
                    out.append(piece[0])

        if not dmode:
            ser = pd.Series(out, index=to_process.index)
            ser.name = val
            sers.append(ser)

    if not dmode:
        dx = pd.concat(sers, axis=1)
        if len(dx.columns) == 1:
            return dx.iloc[:,0]
        else:
            return dx.apply('/'.join, axis=1)
    else:
        index = []
        data = []
        for ix, dct in dicts.items():
            max_key, max_value = max(dct.items(), key=lambda x: len(x[1]))
            for val, pieces in dct.items():
                if len(pieces) == 1:
                    dicts[ix][val] = pieces * len(max_value)

            for tup in list(zip(*[i for i in dct.values()])):
                index.append(ix)
                data.append('/'.join(tup))
        return pd.Series(data, index=pd.MultiIndex.from_tuples(index))


def make_series(ser, df=False, obj=False,
                att=False, adj=False):
    """
    To apply to a DataFrame to add complex criteria, like 'gf'
    """
    # distance mode
    if att == 'a':
        count = 0
        if obj == 'g':
            if ser[obj] == 0:
                return '-1'
            ser = df.loc[ser.name[0], ser['g']]
        while count < 20:
            if ser['mf'].lower() == 'root':
                return str(count)
            ser = df.loc[ser.name[0], ser['g']]
            count += 1
        return '20+'

    # h is head of this particular group
    if obj == 'h':
        cohead = ser['c']
        if cohead.endswith('*'):
            return ser['m' + att]
        elif cohead == '_':
            return 'none'
        else:
            sent = df.loc[ser.name[0]]
            just_cof = sent[sent['c'] == cohead + '*']
            if just_cof.empty:
                return ser['m' + att]
            else:
                return just_cof.iloc[0]['m' + att]
        
    # r is the representative mention head
    if obj == 'r':
        cohead = ser['c']
        if cohead == '_':
            return 'none'
        if not cohead.endswith('*'):
            cohead = cohead + '*'
        # iterrows is slow, but we only need the first instance
        just_cof = df[df['c'] == cohead]
        if just_cof.empty:
            return ser['m' + att]
        else:
            return just_cof.iloc[0]['m' + att]

    if obj == 'g':
        if ser[obj] == 0:
            return 'root'
        else:
            try:
                return df[att][ser.name[0], ser[obj]]
            # this keyerror can happen if governor is punctuation, for example
            except KeyError:
                return

    # if dependent, we need to return a df-like thing instead
    elif obj == 'd':
        #import pandas as pd
        idxs = [(ser.name[0], int(i)) for i in ser[obj].split(',')]
        dat = df[att].ix[idxs]
        return dat

    # todo: fix everything below here
    elif obj == 'r': # get the representative
        cohead = ser['c'].rstrip('*')
        refs = df[df['c'] == cohead + '*']
        return refs[att].ix[0]

    elif obj == 'h': # get head
        cohead = ser['c']
        if cohead.endswith('*'):
            return ser[att]
        else:
            sent = df[att].loc[ser.name[0]]
            return sent[sent['c'] == cohead + '*']

    # potential naming conflict with sent index ...
    elif obj == 's': # get whole phrase"
        cohead = ser['c']
        sent = df[att].loc[ser.name[0]]
        return sent[sent['c'] == cohead.rstrip('*')].values
    
def joiner(ser):
    return ser.str.cat(sep='/') 

def make_new_for_dep(dfmain, dfdep, name):
    """
    If showind dependent, we have to make a whole new dataframe

    :param dfmain: dataframe with everything in it
    :param dfdep: dataframe with just dependent
    """
    import pandas as pd
    import numpy as np
    new = []
    newd = []
    index = []
    for (i, ml), (_, dl) in zip(dfmain.iterrows(), dfdep.iterrows()):
        if all(pd.isnull(i) for i in dl.values):
            index.append(i)
            new.append(ml)
            newd.append('none')
            continue
        else:
            for bit in dl:
                if pd.isnull(bit):
                    continue
                index.append(i)
                new.append(ml)
                newd.append(bit)
    
    #todo: account for no matches
    index = pd.MultiIndex.from_tuples(index, names=['s', 'i'])
    newdf = pd.DataFrame(new, index=index)
    newdf[name] = newd
    return newdf

def turn_pos_to_wc(ser, showval):
    if not showval:
        return ser
    import pandas as pd
    from corpkit.dictionaries.word_transforms import taglemma   
    vals = [taglemma.get(piece.lower(), piece.lower())
                  for piece in ser.values]
    news = pd.Series(vals, index=ser.index)
    news.name = ser.name[:-1] + 'x'
    return news

def concline_generator(matches, idxs, df, metadata,
                       add_meta, category, fname, preserve_case=False):
    """
    Get all conclines

    :param matches: a list of formatted matches
    :param idxs: their (sent, word) idx
    """
    conc_res = []
    # potential speedup: turn idxs into dict
    from collections import defaultdict
    mdict = defaultdict(list)
    # if remaking idxs here, don't need to do it earlier
    idxs = list(matches.index)
    for mid, (s, i) in zip(matches, idxs):
    #for s, i in matches:
        mdict[s].append((i, mid))
    # shorten df to just relevant sents to save lookup time
    df = df.loc[list(mdict.keys())]
    # don't look up the same sentence multiple times
    for s, tup in sorted(mdict.items()):
        sent = df.loc[s]
        if not preserve_case:
            sent = sent.str.lower()
        meta = metadata[s]
        sname = meta.get('speaker', 'none')
        for i, mid in tup:
            if not preserve_case:
                mid = mid.lower()
            ix = '%d,%d' % (s, i)
            start = ' '.join(sent.loc[:i-1].values)
            end = ' '.join(sent.loc[i+1:].values)
            lin = [ix, category, fname, sname, start, mid, end]
            if add_meta:
                for k, v in sorted(meta.items()):
                    if k in ['speaker', 'parse', 'sent_id']:
                        continue
                    if isinstance(add_meta, list):
                        if k in add_meta:
                            lin.append(v)
                    elif add_meta is True:
                        lin.append(v)
            conc_res.append(lin)
    return conc_res

def p_series_to_x_series(val):
    return taglemma.get(val.lower(), val.lower())

def fast_simple_conc(dfss, idxs, show,
                     metadata=False,
                     add_meta=False, 
                     fname=False,
                     category=False,
                     only_format_match=True,
                     conc=False,
                     preserve_case=False,
                     gramsize=1,
                     window=None):
    """
    Fast, simple concordancer, heavily conditional
    to save time.
    """
    if dfss.empty:
        return [], []
        
    import pandas as pd

    # best case, the user doesn't want any gov-dep stuff
    simple = all(i.startswith('m') and not i.endswith('a') for i in show)
    # worst case, the user wants something from dep
    dmode = any(x.startswith('d') for x in show)
    # make a quick copy if need be because we modify the df
    df = dfss.copy() if not simple else dfss
    # add text to df columns so that it resembles 'show' values
    lst = ['s', 'i', 'w', 'l', 'e', 'p', 'f']

    # for ner, change O to 'none'
    if 'e' in df.columns:
        df['e'] = df['e'].str.replace('^O$', 'none')

    df.columns = ['m' + i if len(i) == 1 and i in lst \
                  else i for i in list(df.columns)]

    # this is the data needed for concordancing
    df_for_lr = df['mw'] if only_format_match else df

    just_matches = df.loc[idxs]
    
    # if the showing can't come straight out of the df, 
    # we can add columns with the necessary information
    if not simple:
        formatted = []
        import numpy as np

        for ind, i in enumerate(show):
            # nothing to do if it's an m feature
            if i.startswith('m') and not i.endswith('a'):
                continue
            # defaults for adjacent work

            adj, tomove, adjname = False, False, ''
            adj, i = determine_adjacent(i)
            adjname = ''.join(adj) if hasattr(adj, '__iter__') else ''
            
            # get number of places to shift left or right
            if adj:
                if adj[0] == '+':
                    tomove = -int(adj[1])
                elif adj[0] == '-':
                    tomove = int(adj[1])

            # cut df down to just needed bits for the sake of speed
            # i.e. if we want gov func, get only gov and func cols
            ob, att = i[0], i[-1]
            xmode = att == 'x'
            if xmode:
                att = 'p'
                show[ind] = show[ind][:-1] + 'p'
            # for corefs, we also need the coref data
            if ob in ['h', 'r']:
                dfx = df[['c', 'm' + att]]
            else:
                lst = ['s', 'i', 'w', 'l', 'f', 'p']
                if att in lst and ob != 'm':
                    att = 'm' + att
                if ob == 'm' and att != 'a':
                    dfx = df[['m' + att]]
                elif att == 'a':
                    dfx = df[['mf', 'g']]
                else:
                    dfx = df[[ob, att]]
            # decide if we need to format everything
            if (not conc or only_format_match) and not adj:
                to_proc = just_matches
            else:
                to_proc = df
            # now we get or generate the new column
            if ob == 'm' and att != 'a':
                ser = to_proc['m' + att]
            else:
                ser = to_proc.apply(make_series, df=dfx, obj=ob, att=att, axis=1)
            if xmode:
                ser = ser.apply(p_series_to_x_series)

            # adjmode simply shifts series and index
            if adj:
                #todo: this shifts next sent into previous sent!
                ser = ser.shift(tomove)
                ser = ser.fillna('none')

            # dependent mode produces multiple matches
            # so, we have to make a new dataframe with duplicate indexes
            # todo: what about when there are two dep options?
            ser.name = adjname + i
            if ob != 'd':
                df[ser.name] = ser
            else:
                df = make_new_for_dep(df, ser, i)

        df = df.fillna('none')

    # x is wordclass. so, we just get pos and translate it
    nshow = [(i.replace('x', 'p'), i.endswith('x')) for i in show]

    # generate a series of matches with slash sep if multiple show vals
    if len(nshow) > 1:

        if conc and not only_format_match:
            first = turn_pos_to_wc(df[nshow[0][0]], nshow[0][1])
            llist = [turn_pos_to_wc(df[sho], xmode) for sho, xmode in nshow[1:]]
            df = first.str.cat(others=llist, sep='/')
            matches = df[idxs]
        else:
            justm = df.loc[idxs]
            first = turn_pos_to_wc(justm[nshow[0][0]], nshow[0][1])
            llist = [turn_pos_to_wc(justm[sho], xmode) for sho, xmode in nshow[1:]]
            matches = first.str.cat(others=llist, sep='/')
            if conc:
                df = df_for_lr
    else:
        if conc and not only_format_match:
            df = turn_pos_to_wc(df[nshow[0][0]], nshow[0][1])
            matches = df[idxs]
        else:
            matches = turn_pos_to_wc(df[nshow[0][0]][idxs], nshow[0][1])
            if conc:
                df = df_for_lr
    
    # get rid of (e.g.) nan caused by no_punct=True
    matches = matches.dropna(axis=0, how='all')

    if not preserve_case:
        matches = matches.str.lower()

    if not conc:
        # todo: is matches.values faster?
        return list(matches), []
    else:
        conc_res = concline_generator(matches, idxs, df,
                                      metadata, add_meta,
                                      category, fname,
                                      preserve_case=preserve_case)

    return list(matches), conc_res

def make_collocate_show(show, current):
    """
    Turn show into a collocate showing thing
    """
    out = []
    for i in show:
        out.append(i)
    for i in show:
        newn = '%s%s' % (str(current), i)
        if not newn.startswith('-'):
            newn = '+' + newn
        out.append(newn)
    return out

def show_this(df, matches, show, metadata, conc=False,
              coref=False, category=False, show_conc_metadata=False, **kwargs):

    only_format_match = kwargs.pop('only_format_match', True)
    ngram_mode = kwargs.get('ngram_mode', True)
    preserve_case = kwargs.get('preserve_case', False)
    gramsize = kwargs.get('gramsize', 1)
    window = kwargs.get('window', None)

    matches = sorted(list(matches))

    # add index as column if need be
    if any(i.endswith('s') for i in show):
        df['ms'] = [str(i) for i in df.index.labels[0]]
    if any(i.endswith('i') for i in show):
        df['mi'] = [str(i) for i in df.index.labels[1]]
    
    # attempt to leave really fast
    if kwargs.get('countmode'):
        return len(matches), {}
    if len(show) == 1 and not conc and gramsize == 1 and not window:
        if show[0] in ['ms', 'mi', 'mw', 'ml', 'mp', 'mf']:
            get_fast = df.loc[matches][show[0][-1]]
            if not preserve_case:
                get_fast = get_fast.str.lower()
            return list(get_fast), {}

    # todo: make work for ngram, collocate and coref
    if all(i[0] in ['m', 'g', '+', '-', 'd', 'h', 'r'] for i in show):
        if gramsize == 1 and not window:
            return fast_simple_conc(df,
                                matches,
                                show,
                                metadata,
                                show_conc_metadata,
                                kwargs.get('filename', ''),
                                category,
                                only_format_match,
                                conc=conc,
                                preserve_case=preserve_case,
                                gramsize=gramsize,
                                window=window)
        else:
            resbit = []
            concbit = []
            iterab = range(1, gramsize + 1) if gramsize > 1 else range(-window, window+1)
            for i in iterab:
                if i == 0:
                    continue
                if window:
                    nnshow = make_collocate_show(show, i)
                else:
                    nnshow = show
                r, c = fast_simple_conc(df,
                                matches,
                                nnshow,