import pandas as pd
import logging

MOSQ_TARGETS = [str(i) for i in range(62)]
PLASM_TARGETS = ['P1', 'P2']
ANOSPP_TARGETS = MOSQ_TARGETS + PLASM_TARGETS
CUTADAPT_TARGETS = ANOSPP_TARGETS + ['unknown']

def setup_logging(verbose=False):
    try: 
        del logging.root.handlers[:]
    except:
        pass
    if verbose:
        logging.basicConfig(level=logging.INFO, format='[%(levelname)s] [%(asctime)s] %(message)s')
    else:
        logging.basicConfig(level=logging.WARNING, format='[%(levelname)s] [%(asctime)s] %(message)s')

def well_id_mapper():
    '''
    Yields mapping of tag_index 1,2...96 
    to well id A1,B1...H12.
    N.B. subsequent plates can be mapped using
    tag_index % 96, thus last well is inserted at 0
    '''

    well_ids = dict()
    tag = 1
    for col in range(1,13):
        for row in 'ABCDEFGH':
            well_ids[tag] = f'{row}{col}'
            tag += 1

    # edge case
    well_ids[0] = 'H12'
            
    return well_ids

def lims_well_id_mapper():
    '''
    Yields mapping of tag_index 1,2...384 
    to well id A1,C1...P24
    4 96-well plates order in quadrants is
    1 2
    3 4
    N.B. subsequent plates can be mapped using
    tag_index % 384, thus last well is inserted at 0
    '''
    lims_well_ids = dict()
    tag = 1
    # upper left quadrant
    for col in range(1,13):
        for row in 'ACEGIKMO':
            lims_well_ids[tag] = f'{row}{col * 2 - 1}'
            tag += 1
    # upper right quadrant
    for col in range(1,13):
        for row in 'ACEGIKMO':
            lims_well_ids[tag] = f'{row}{col * 2}'
            tag += 1
    # lower left quadrant
    for col in range(1,13):
        for row in 'BDFHJLNP':
            lims_well_ids[tag] = f'{row}{col * 2 - 1}'
            tag += 1
    # lower right quadrant
    for col in range(1,13):
        for row in 'BDFHJLNP':
            lims_well_ids[tag] = f'{row}{col * 2}'
            tag += 1

    # edge case
    lims_well_ids[0] = 'P24'

    return lims_well_ids

def seqid_generator(hap_df):
    '''
    assign identifiers to unique haplotypes
    '''

    seqids = dict()
    for tgt, group in hap_df.groupby('target'):
        for (i, cons) in enumerate(group['consensus'].unique()):
            seqids[tgt + cons] = '{}-{}'.format(tgt, i)
    hap_df['seqid'] = (hap_df.target + hap_df.consensus).replace(seqids)

    return hap_df

def human_format(num):
    if num is None:
        return ''
    magnitude = 0
    while abs(num) >= 1000:
        magnitude += 1
        num /= 1000.0
    # 1 though 999
    if magnitude == 0:
        return '%.0f' % (num)
    # 1.0k through 99.9k
    elif num < 100:
        return '%.1f%s' % (num, ['', 'K', 'M', 'G', 'T', 'P'][magnitude])
    # 100k through 999k
    else:
        return '%.0f%s' % (num, ['', 'K', 'M', 'G', 'T', 'P'][magnitude])

def well_ordering(well_series):

    # well order for plotting - A1, B1, ...
    well_sequence = []
    for col in range(1,13):
        for row in 'ABCDEFGH': 
            well_sequence.append(f'{row}{col}')

    return pd.Categorical(well_series, categories=well_sequence)

def prep_hap(hap_fn, anospp=True):
    '''
    load haplotypes table generated by prep
    '''

    logging.info(f'loading haplotypes table from {hap_fn}')

    hap_df = pd.read_csv(hap_fn, sep='\t', dtype={'target':'str'})

    # compatibility with old style haplotype column names
    hap_df.rename(columns=({
        's_Sample':'sample_id',
        'frac_reads':'reads_fraction'
        }), 
    inplace=True)

    for col in ('sample_id',
                'target',
                'consensus',
                'reads'):
        assert col in hap_df.columns, f'hap column {col} not found'

    # for legacy - same as in anospp_analysis.prep.get_hap_df
    if 'total_reads' not in hap_df.columns:
        hap_df['total_reads'] = hap_df.groupby(by=['sample_id', 'target']) \
            ['reads'].transform('sum')

    if 'reads_fraction' not in hap_df.columns:
        hap_df['reads_fraction'] = hap_df['reads'] / hap_df['total_reads']

    if 'seqid' not in hap_df.columns:
        hap_df = seqid_generator(hap_df)

    if 'nalleles' not in hap_df.columns:
        hap_df['nalleles'] = hap_df.groupby(by=['sample_id', 'target']) \
            ['seqid'].transform('nunique')

    hap_df['consensus'] = hap_df['consensus'].str.upper()
    
    if hap_df['target'].isin(CUTADAPT_TARGETS).all():
        hap_df['target'] = pd.Categorical(hap_df['target'], 
                                        categories=CUTADAPT_TARGETS, 
                                        ordered=True)
    else:
        logging.warning('non-ANOSPP targets detected in haps, targets order might be unstable')
    
    hap_df.sort_values(by=[
        'sample_id',
        'target'
        ], inplace=True)

    return hap_df

def prep_comb_stats(comb_stats_fn, anospp=True):

    logging.info(f'loading combined stats table from {comb_stats_fn}')

    comb_stats_df = pd.read_csv(comb_stats_fn, sep='\t')

    assert comb_stats_df['run_id'].nunique(dropna=False) == 1, f'multiple runs in {comb_stats_fn}'

    run_id = comb_stats_df['run_id'].iloc[0]

    return run_id, comb_stats_df