import numpy as np
import pandas as pd

HIT_FIELDS = ['query', 'target', 'pident', 'qcov', 'tcov', 'bits', 'evalue', 'qlen', 'tlen', 'alnlen']

FEATS = ['sum_bits', 'max_bits', 'n_sup', 'best_rank', 'pid_max', 'pid_best', 'qcov_max', 'tcov_max', 'nec_best',
         'top_bits', 'tot_bits', 'n_hits', 'top_pident', 'r_max', 'r_sum', 'r_n', 'log_top', 'ec_logfreq',
         'l3_support', 'l2_support', 'ec_share_l3', 'qlen', 'tlen_best', 'aln_frac_best', 'bits_per_qlen',
         'log_evalue_best']


def prefix(ec, k):
    return '.'.join(ec.split('.')[:k])


def l3_totals(freq):
    tot = {}
    for e, c in freq.items():
        p = prefix(e, 3)
        tot[p] = tot.get(p, 0) + c
    return tot


def build_candidates(h, ec_map, freq, kmax, l3_tot):
    """One row per (query, candidate EC) with features computed from the top-kmax alignment hits."""
    if len(h) == 0:
        return pd.DataFrame(columns=['query', 'ec'] + FEATS)
    H = h.sort_values(['query', 'bits'], ascending=[True, False], kind='mergesort')
    H = H.groupby('query', sort=False).head(kmax).copy()
    H['rank'] = H.groupby('query', sort=False).cumcount()
    qs = H.groupby('query', sort=False).agg(top_bits=('bits', 'max'), tot_bits=('bits', 'sum'),
                                            n_hits=('bits', 'size'), top_pident=('pident', 'first'),
                                            qlen=('qlen', 'first'))
    H['ecs'] = H['target'].map(ec_map)
    H['n_ec'] = H['ecs'].map(len)
    H = H.explode('ecs').rename(columns={'ecs': 'ec'})
    H['l3'] = H['ec'].map(lambda e: prefix(e, 3))
    H['l2'] = H['ec'].map(lambda e: prefix(e, 2))
    H['log_ev'] = np.log10(H['evalue'].astype(float) + 1e-200)
    a = H.groupby(['query', 'ec'], sort=False).agg(
        sum_bits=('bits', 'sum'), max_bits=('bits', 'max'), n_sup=('bits', 'size'), best_rank=('rank', 'min'),
        pid_max=('pident', 'max'), pid_best=('pident', 'first'), qcov_max=('qcov', 'max'), tcov_max=('tcov', 'max'),
        nec_best=('n_ec', 'first'), tlen_best=('tlen', 'first'), aln_best=('alnlen', 'first'),
        log_evalue_best=('log_ev', 'first'), l3=('l3', 'first'), l2=('l2', 'first')).reset_index()
    for lvl in ('l3', 'l2'):
        s = (H.drop_duplicates(['query', 'target', lvl]).groupby(['query', lvl], sort=False)['bits']
              .sum().rename(lvl + '_bits').reset_index())
        a = a.merge(s, on=['query', lvl], how='left')
    a = a.join(qs, on='query')
    a['r_max'] = a.max_bits / a.top_bits
    a['r_sum'] = a.sum_bits / a.tot_bits
    a['r_n'] = a.n_sup / a.n_hits
    a['log_top'] = np.log(a.top_bits)
    a['ec_logfreq'] = np.log1p(a['ec'].map(lambda e: freq.get(e, 0)))
    a['ec_share_l3'] = a['ec'].map(lambda e: freq.get(e, 0)) / a['l3'].map(lambda p: max(l3_tot.get(p, 0), 1))
    a['aln_frac_best'] = a.aln_best / a.qlen
    a['bits_per_qlen'] = a.max_bits / a.qlen
    a['l3_support'] = a.l3_bits / a.tot_bits
    a['l2_support'] = a.l2_bits / a.tot_bits
    a['tlen_best'] = a['tlen_best'].astype(float)
    return a[['query', 'ec', 'ec_logfreq'] + [f for f in FEATS if f != 'ec_logfreq']]


def select_from_scores(c, rel=0.5, max_ec=3):
    """c: DataFrame with query, ec, ec_logfreq, p. Returns {query: best-first EC list}."""
    out = {}
    if len(c) == 0:
        return out
    c = c.sort_values(['query', 'p', 'ec_logfreq', 'ec'], ascending=[True, False, False, True], kind='mergesort')
    for q, g in c.groupby('query', sort=False):
        ecs = g['ec'].tolist()
        ps = g['p'].tolist()
        top = ps[0]
        out[q] = [ecs[0]] + [e for e, pp in zip(ecs[1:max_ec], ps[1:max_ec]) if top > 0 and pp >= rel * top]
    return out


def rerank_select(cand, model, rel=0.5, max_ec=3):
    if len(cand) == 0:
        return {}
    c = cand[['query', 'ec', 'ec_logfreq']].copy()
    c['p'] = model.predict(cand[FEATS].to_numpy(dtype=float))
    return select_from_scores(c, rel, max_ec)