from app.lookups.sqlite.base import ResultLookupInterface


class ProtPepTable(ResultLookupInterface):
    table_map = {'protein': {'fntable': 'protein_tables',
                             'feattable': 'proteins',
                             'isoqtable': 'protein_iso_quanted',
                             'isochtable': 'protquant_channels',
                             'prectable': 'protein_precur_quanted',
                             'fdrtable': 'protein_fdr',
                             'fullqpsmtable': 'protein_iso_fullpsms',
                             'psmcounts': 'prot_psmpep_counts',
                             },
                 'gene': {'fntable': 'gene_tables',
                          'feattable': 'genes',
                          'isoqtable': 'gene_iso_quanted',
                          'isochtable': 'genequant_channels',
                          'prectable': 'gene_precur_quanted',
                          'fdrtable': 'gene_fdr',
                          'fullqpsmtable': 'gene_iso_fullpsms',
                          'psmcounts': 'gene_psmpep_counts',
                          },
                 'assoc': {'fntable': 'gene_tables',
                           'feattable': 'associated_ids',
                           'isoqtable': 'assoc_iso_quanted',
                           'isochtable': 'genequant_channels',
                           'prectable': 'assoc_precur_quanted',
                           'fdrtable': 'assoc_fdr',
                           'fullqpsmtable': 'assoc_iso_fullpsms',
                           'psmcounts': 'gn_psmpep_counts',
                           },
                 'peptide': {'fntable': 'peptide_tables',
                             'feattable': 'peptide_sequences',
                             'isoqtable': 'peptide_iso_quanted',
                             'isochtable': 'pepquant_channels',
                             'prectable': 'peptide_precur_quanted',
                             'fdrtable': 'peptide_fdr',
                             'peptable': 'peptide_pep',
                             'fullqpsmtable': 'peptide_iso_fullpsms',
                             'flrtable': 'ptm_flr',
                             'psmcounts': 'pep_psmpep_counts',
                             }
                 }

    def __init__(self, fn=None):
        super().__init__(fn)

    def get_all_poolnames(self):
        cursor = self.get_cursor()
        cursor.execute(
            'SELECT DISTINCT set_name, set_id FROM biosets ORDER BY set_name')
        return cursor

    def store_table_files(self, tables):
        table = self.table_map[self.datatype]['fntable']
        self.store_many(
                'INSERT INTO {}(set_id, filename) VALUES(?, ?)'.format(table), tables)
        self.index_column('table_set_ix', table, 'set_id')

    def store_quant_channels(self, quantchannels, psmnrcols):
        table = self.table_map[self.datatype]['isochtable']
        if psmnrcols:
            sql = ('INSERT INTO {}({}, channel_name, amount_psms_name) VALUES'
                   '(?, ?, ?)'.format(table, self.colmap[table][1]))
        else:
            sql = ('INSERT INTO {}({}, channel_name) VALUES'
                   '(?, ?)'.format(table, self.colmap[table][1]))
        self.store_many(sql, quantchannels)

    def get_quantchannel_map(self):
        outdict = {}
        table = self.table_map[self.datatype]['isochtable']
        cursor = self.get_cursor()
        cursor.execute(
            'SELECT channel_id, {}, channel_name, amount_psms_name'
            ' FROM {}'.format(self.colmap[table][1], table))
        for channel_id, fnid, channel_name, amount_psms_name in cursor:
            try:
                outdict[fnid][channel_name] = (channel_id, amount_psms_name)
            except KeyError:
                outdict[fnid] = {channel_name: (channel_id, amount_psms_name)}
        return outdict

    def store_isobaric_quants(self, quants, psmnrcols):
        table = self.table_map[self.datatype]['isoqtable']
        if psmnrcols:
            sql = ('INSERT INTO {}({}, channel_id, quantvalue, amount_psms) ' 
                   'VALUES ' '(?, ?, ?, ?)'.format(table, self.colmap[table][1]))
        else:
            sql = ('INSERT INTO {}({}, channel_id, quantvalue) ' 
                   'VALUES ' '(?, ?, ?)'.format(table, self.colmap[table][1]))
        self.store_many(sql, quants)

    def index_iso_features(self):
        '''Index {feat}channel/{feat}id for feat isoquant table, and 
        {feat}table_id for feat channel table'''
        # FIXME not needed since sqlite creates auto covering index when merging, and
        # we only use this once
        table = self.table_map[self.datatype]['isoqtable']
        self.index_column('feat_iso_q_ix', table, f'{self.colmap[table][1]}, {self.colmap[table][2]}')
        table = self.table_map[self.datatype]['isochtable']
        self.index_column('feat_iso_ch_ix', table, self.colmap[table][1])

    def get_tablefnset_map(self):
        table = self.table_map[self.datatype]['fntable']
        cursor = self.get_cursor()
        cursor.execute(f'SELECT * FROM {table}')
        return {fn: [table_id, setid] for (table_id, setid, fn) in cursor}

    def get_feature_map(self, setids):
        '''Create map of all features in feature table, output:
        {accession_id: [accession_name, False]} where the False is a check to be set
        to True if an accession is in the input txt file (the features exist in the DB
        beforehand from the PSM table, but may for some reason not be in a protein table).
        '''
        table = self.table_map[self.datatype]['feattable']
        columns = self.colmap[table][0:2]
        cursor = self.get_cursor()
        cursor.execute(f'SELECT {columns[0]}, {columns[1]} FROM {table}')
        return {acc: [table_id, set(setids)] for (table_id, acc) in cursor}

    def delete_features(self, delete_ids):
        table = self.table_map[self.datatype]['psmcounts']
        acc_id = self.colmap[table][0]
        cursor = self.get_cursor()
        cursor.executemany(f'DELETE FROM {table} WHERE {acc_id} = ? AND set_id = ?', ((acc, setid) for acc, setids in delete_ids.items() for setid in setids))
        self.conn.commit()

    def store_singlecol(self, tablekey, vals):
        table = self.table_map[self.datatype][tablekey]
        cols = self.colmap[table]
        self.store_many('INSERT INTO {}({}, {}, {}) '
                        'VALUES (?, ?, ?)'.format(table, cols[0], cols[1],
                                                  cols[2]), vals)

    def store_precursor_quants(self, quants):
        self.store_singlecol('prectable', quants)
        table = self.table_map[self.datatype]['prectable']

    def store_fdr(self, fdr):
        self.store_singlecol('fdrtable', fdr)
        table = self.table_map[self.datatype]['fdrtable']

    def store_pep(self, pep):
        '''Overridden in peptide table lookup'''
        pass

    def store_fullq_psms(self, fullqpsms):
        self.store_singlecol('fullqpsmtable', fullqpsms)
        table = self.table_map[self.datatype]['fullqpsmtable']

    def store_ptm_flr(self, flrs):
        self.store_singlecol('flrtable', flrs)
        table = self.table_map[self.datatype]['flrtable']

    def check_isoquant_psmnrs(self):
        cursor = self.get_cursor()
        tables = self.table_map[self.datatype]
        psmschecksql = """SELECT EXISTS(SELECT amount_psms FROM {} 
                WHERE amount_psms IS NOT NULL)""".format(tables['isoqtable'])
        return cursor.execute(psmschecksql).fetchone()[0]

    def get_isoquant_headernames(self, stored_psmnrs):
        cursor = self.get_cursor()
        tables = self.table_map[self.datatype]
        if self.datatype == 'peptide':
            table_id_fieldname = 'peptable_id'
        else:
            table_id_fieldname = f"{tables['fntable'][:4]}table_id"
        if stored_psmnrs:
            sql = ('SELECT set_name || "_" || channel_name, set_name || "_" || amount_psms_name '
                    f'FROM {tables["isochtable"]} '
                    f'JOIN {tables["fntable"]} USING({table_id_fieldname}) '
                    'JOIN biosets USING(set_id)')
        else:
            sql = ('SELECT set_name || "_" || channel_name '
                    f'FROM {tables["isochtable"]} '
                    f'JOIN {tables["fntable"]} USING({table_id_fieldname}) '
                    'JOIN biosets USING(set_id)')
        return cursor.execute(sql)
