import os
from app.readers import tsv as tsvreader
from app.dataformats import prottable as prottabledata


def create_lookup(fns, pqdb, poolnames, featcolnr, ms1_qcolpattern,
        isobqcolpattern, fdrcolpattern, pepcolpattern, flrcolpattern):
    psmnrpattern = prottabledata.HEADER_NO_PSMS_SUFFIX
    fullpsmpattern = prottabledata.HEADER_NO_FULLQ_PSMS
    patterns = [ms1_qcolpattern, fdrcolpattern, pepcolpattern, fullpsmpattern,
            flrcolpattern]
    storefuns = [pqdb.store_precursor_quants, pqdb.store_fdr, pqdb.store_pep,
            pqdb.store_fullq_psms, pqdb.store_ptm_flr]
    poolmap = {name: set_id for (name, set_id) in pqdb.get_all_poolnames()}
    pqdb.store_table_files([(poolmap[pool], os.path.basename(fn))
                            for fn, pool in zip(fns, poolnames)])
    tablefnset_map = pqdb.get_tablefnset_map()
    feat_map = pqdb.get_feature_map(poolmap.values())
    for pattern, storefun in zip(patterns, storefuns):
        if pattern is None:
            continue
        colmap = get_colmap(fns, pattern, single_col=True)
        if colmap:
            store_single_col_data(fns, tablefnset_map, feat_map, storefun, colmap)
    if isobqcolpattern:
        isocolmap = get_colmap(fns, isobqcolpattern, antipattern=psmnrpattern)
        psmcolmap = get_colmap(fns, psmnrpattern)
        create_isobaric_quant_lookup(fns, tablefnset_map, feat_map, pqdb, featcolnr,
                isocolmap, psmcolmap)
    pqdb.store_psmpep_counts()
    delete_feats = {acc_id: setids for _, [acc_id, setids] in feat_map.items() if setids}
    pqdb.delete_features(delete_feats)


def get_colmap(fns, pattern, single_col=False, antipattern=False):
    """For table files, loops through headers and checks which column(s)
    match a passed pattern. Those column(s) names are returned in a map with
    filenames as keys"""
    colmap = {}
    for fn in fns:
        header = tsvreader.get_tsv_header(fn)
        basefn = os.path.basename(fn)
        try:
            cols = tsvreader.get_cols_in_file(pattern, header, single_col)
        except RuntimeError:
            # Columns are not in this file
            cols = []
        if antipattern:
            try:
                anticols = tsvreader.get_cols_in_file(antipattern, header,
                                                      single_col)
            except RuntimeError:
                # The filtering "anti"-columns are not in the file, 
                anticols = []
            cols = [col for col in cols if col not in anticols]
        if cols:
            colmap[basefn] = cols
        else:
            return False
    return colmap


def store_single_col_data(fns, prottable_id_set_map, pacc_map, pqdbmethod, colmap):
    """General method to store single column data from protein tables
    in lookup"""
    to_store = []
    for fn, header, pquant in tsvreader.generate_tsv_pep_protein_quants(fns):
        pacc_id = pacc_map[pquant[header[0]]][0]
        # Mark feature for used in this set
        pacc_map[pquant[header[0]]][1].discard(prottable_id_set_map[fn][1])
        pqdata = (pacc_id, prottable_id_set_map[fn][0], pquant[colmap[fn]])
        to_store.append(pqdata)
        if len(to_store) > 10000:
            pqdbmethod(to_store)
            to_store = []
    pqdbmethod(to_store)


def create_isobaric_quant_lookup(fns, tablefnset_map, featmap, pqdb,
                                 featcolnr, allquantcols, psmcolmap):
    """Creates a lookup dict from peptide/protein quant input files and some
    input parameters. This assumes the order of quant columns and
    number-of-PSM columns is the same."""
    pqdb.store_quant_channels(
            map_psmnrcol_to_quantcol(allquantcols, psmcolmap, tablefnset_map),
            psmcolmap)
    quantmap = pqdb.get_quantchannel_map()
    to_store = []
    for fn, header, pquant in tsvreader.generate_tsv_pep_protein_quants(fns):
        pqdata = get_isob_quant_data(pquant, header[featcolnr], tablefnset_map[fn][0],
                                     featmap, tablefnset_map[fn][1], quantmap)
        to_store.extend(pqdata)
        if len(to_store) > 10000:
            pqdb.store_isobaric_quants(to_store, psmcolmap)
            to_store = []
    pqdb.store_isobaric_quants(to_store, psmcolmap)
    #pqdb.index_iso_features()


def map_psmnrcol_to_quantcol(quantcols, psmcols, tablefnset_map):
    """This function yields tuples of table filename, isobaric quant column
    and if necessary number-of-PSM column"""
    if not psmcols:
        for fn in quantcols:
            for qcol in quantcols[fn]:
                yield (tablefnset_map[fn][0], qcol)
    else:
        for fn in quantcols:
            for qcol, psmcol in zip(quantcols[fn], psmcols[fn]):
                yield (tablefnset_map[fn][0], qcol, psmcol)


def get_isob_quant_data(pquant, featfield, fnid, featmap, set_id, qmap):
    """Turns a dict from a line of protein/peptide quant data into a set of
    tuples that can be stored in db"""
    feat_dbid = featmap[pquant[featfield]][0]
    featmap[pquant[featfield]][1].discard(set_id)
    for channel_name, (channel_id, psmfield) in qmap[fnid].items():
        if psmfield is None:
            yield (feat_dbid, channel_id, pquant[channel_name])
        else:
            yield (feat_dbid, channel_id, pquant[channel_name],
                   pquant[psmfield])


def build_proteintable(pqdb, mergecutoff, isobaric=False):
    """Fetches proteins and quants from joined lookup table, loops through
    them and when all of a protein's quants have been collected, yields the
    protein quant information."""
    # First get a dict with the data which is the same over sets (protein-gene mappings etc)
    pdmap = pqdb.create_pdata_map()
    previousfeat, outfeat = False, {}
    # Then loop over the SQLite output
    fdrfieldnr = pqdb.singlefields.index(prottabledata.HEADER_QVAL)
    # multi set output, means that a feature from DB can be delivered in
    # multiple lines (1/set). As soon as a new feat is received from db,
    # yield collected output and flush it for population with new protein
    for setname, feat_id, *setfeatvals in pqdb.merge_features():
        if feat_id != previousfeat and previousfeat and outfeat != {}:
            outfeat.update(pdmap[previousfeat])
            yield outfeat
            outfeat = {}
        if protein_pool_fdr_cutoff_ok(setfeatvals[fdrfieldnr], mergecutoff, feat_id, setname):
            for ix, field in enumerate(pqdb.singlefields):
                fieldname = '{}_{}'.format(setname, field)
                outfeat[fieldname] = setfeatvals[ix]
            if isobaric:
                channels = setfeatvals[len(pqdb.singlefields)].split(',')
                quants = setfeatvals[len(pqdb.singlefields)+1].split(',')
                try:
                    ampsms = setfeatvals[len(pqdb.singlefields)+2].split(',')
                except AttributeError:
                    ampsms = [0] * len(channels)
                for ch, quant, ampsm in zip(channels, quants, ampsms):
                    outfeat['{}_{}'.format(setname, ch)] = quant
                    if ampsms:
                        outfeat['{}_{}{}'.format(setname, ch, prottabledata.HEADER_NO_PSMS_SUFFIX)] = ampsm
        previousfeat = feat_id
    if outfeat != {}:
        outfeat.update(pdmap[feat_id])
        yield outfeat


def protein_pool_fdr_cutoff_ok(fdr, fdrcutoff, feature, setname):
    '''See if feature output meets the FDR cutoff, if any such cutoff exists.
    If there is no cutoff, we still need to be sure that the feature is in the input data,
    and not only in the input DB which is from a PSM table.'''
    if fdrcutoff is False:
        return True
    warning = False
    if fdr is False:
        warning = True
    try:
        fdr = float(fdr)
    except (TypeError, ValueError):
        # Catch 'NA', or other string or weird other type
        warning = True
    if warning:
        print(f'WARNING, filtering merge table on FDR but feat ID {feature} in '
                f'set {setname} has FDR value "{fdr}" in lookup.')
        return False
    return fdr < fdrcutoff
