import pandas as pd
import argparse
import subprocess
import glob
import os

from .util import *

# optimised cutadapt args
CUTADAPT_ARGS = '-O 10 --match-read-wildcards'

def reverse_complement(seq):

    ambiguous_dna_complement = {
        "A": "T",
        "C": "G",
        "G": "C",
        "T": "A",
        "M": "K",
        "R": "Y",
        "W": "W",
        "S": "S",
        "Y": "R",
        "K": "M",
        "V": "B",
        "H": "D",
        "D": "H",
        "B": "V",
        "X": "X",
        "N": "N",
    }

    seq = seq.upper()

    rc_seq = ''
    for nt in seq[::-1]:
        rc_seq += ambiguous_dna_complement[nt]

    return rc_seq

def cutadapt_deplex(dada_table, adapters, cutadapt_args=CUTADAPT_ARGS, work_dir='work', rc=False):

    logging.info(f'cutadapt deplexing started at {work_dir}')

    os.makedirs(work_dir, exist_ok=True)

    # generate input fasta from DADA2 table

    dada_df = pd.read_csv(dada_table, sep='\t', index_col=0)

    fasta = f'{work_dir}/input_seqs.fasta'

    with open(fasta, 'w') as outfile:
        for seqid, seq in dada_df['sequence'].items():
            if rc:
                rc_seq = reverse_complement(seq)
                outfile.write(f'>{seqid}_rc\n{rc_seq}\n')
            else:
                outfile.write(f'>{seqid}\n{seq}\n')

    cmd = f"cutadapt {cutadapt_args} -g file:{adapters} "
    cmd += f"-o {work_dir}/ASV_{{name}}.fa {fasta}"

    process = subprocess.run(cmd.split(), capture_output=True, text=True)
    process.check_returncode()

    logging.info('cutadapt deplexing complete')

    return 0

def get_deplex_df(work_dir):
    '''
    Read demultiplexed fasta into dataframe with
    seqid: target, trimmed_sequence
    '''
    logging.info(f'parsing deplexed sequences from {work_dir}')

    deplex_dict = dict()
    # iterate over deplexed fasta files
    for fa in sorted(glob.glob(f'{work_dir}/ASV_*.fa')):
        target = fa.split('/')[-1].split('.')[0].split('_', maxsplit=1)[1]
        # basic parser
        with open(fa) as f:
            for line in f:
                line = line.strip()
                if line.startswith('>'):
                    seqid = line[1:]
                else:
                    deplex_dict[seqid] = {'target':target,'trimmed_sequence':line}
        # proper parser
        # for record in SeqIO.parse(fa, format='fasta'):
        #     deplex_dict[record.name] = {'target':target,'trimmed_sequence':str(record.seq)}
    deplex_df = pd.DataFrame(deplex_dict).T
    dup_seqid = deplex_df.index.duplicated()
    if dup_seqid.any():
        raise ValueError(f'duplicate seqids generated by dada2: {deplex_df.index[dup_seqid]}')
    return deplex_df

def get_hap_df(dada_table, work_dir, rc=False):

    logging.info('combining dada output with deplexing info')

    dada_df = pd.read_csv(dada_table, sep='\t', index_col=0)
    deplex_df = get_deplex_df(work_dir)
    if rc:
        logging.info('adding reverse complement sequences')
        rc_deplex_df = get_deplex_df(work_dir + '_rc')
        assert deplex_df.shape[0] == rc_deplex_df.shape[0], \
            'mismatch in deplex and reverse complement sequences number'
        n_rc = 0
        for seqid, r in deplex_df.iterrows():
            rc_seqid = f'{seqid}_rc'
            rc_r = rc_deplex_df.loc[rc_seqid]
            if (r.target == 'unknown') and (rc_r.target != 'unknown'):
                n_rc += 1
                for col in ('trimmed_sequence', 'target'):
                    deplex_df.loc[seqid, col] = rc_r[col]
        logging.info(f'added {n_rc} reverse complement sequences matching targets')
    dada_deplex_df = pd.merge(dada_df, deplex_df, left_index=True, right_index=True)
    dada_deplex_df.index.name = 'dada2_id'
    assert dada_deplex_df.shape[0] == dada_df.shape[0] == deplex_df.shape[0], \
        'lost some sequences in combining deplexing and original DADA2 data'
    
    hap_df = pd.melt(dada_deplex_df.reset_index(), 
            id_vars=['dada2_id', 'sequence', 'target', 'trimmed_sequence'],
            var_name='sample_id',
            value_name='reads')
    hap_df['target'] = hap_df.target.astype(str)
    if not hap_df['target'].isin(CUTADAPT_TARGETS).all():
        logging.warning('non-ANOSPP targets detected in deplexing')
    hap_df.rename(
        columns={
            'sequence':'untrimmed_sequence',
            'trimmed_sequence':'consensus'
            },
        inplace=True)
    # remove unsupported sequences
    hap_df = hap_df.query('reads > 0').copy()
    # collapse identical sequences
    pre_collapse_nseq = len(hap_df)
    hap_df = hap_df.groupby(['sample_id', 'target', 'consensus'])['reads'].sum().reset_index()
    post_collapse_nseq = len(hap_df)
    if post_collapse_nseq != pre_collapse_nseq:
        logging.info(
            f'collapsed {pre_collapse_nseq - post_collapse_nseq} sequences with identical inserts')
    
    hap_df['total_reads'] = hap_df.groupby(by=['sample_id', 'target']) \
        ['reads'].transform('sum')

    hap_df['reads_fraction'] = hap_df['reads'] / hap_df['total_reads']

    hap_df['nalleles'] = hap_df.groupby(by=['sample_id', 'target']) \
        ['consensus'].transform('nunique')

    hap_df = seqid_generator(hap_df)

    return hap_df

def prep_samples(samples_fn, run_id=None):
    '''
    Prepare sample manifest used for anospp pipeline
    '''

    # allow reading from tsv (new style) or csv (old style)
    if samples_fn.endswith('csv'):
        logging.info(f'preparing sample manifest from legacy file {samples_fn}')
        samples_df = pd.read_csv(samples_fn, sep=',', dtype='str')
        samples_df.rename(columns=({
            'Source_sample':'sample_id',
            'Run':'run_id',
            'Lane':'lane_index',
            'Tag':'tag_index',
            'Replicate':'replicate_id'
            }), 
            inplace=True)
    elif samples_fn.endswith('tsv'):
        logging.info(f'preparing sample manifest from new style file {samples_fn}')
        samples_df = pd.read_csv(samples_fn, sep='\t', dtype='str')
        samples_df.rename(columns=({'derived_sample_id':'sample_id'}), inplace=True)
        if run_id is not None:
            samples_df['run_id'] = run_id
            samples_df['lane_index'] = 1
            samples_df['tag_index'] = [i + 1 for i in range(len(samples_df))]
        else:
            assert samples_df.irods_path.str.match('/seq/\d{5}/\d{5}_\d#\d+.cram').all(), \
                ('tsv sample manifest input requires irods_path column to be present '
                'and match "/seq/12345/12345_1#123.cram"')
            samples_df['run_id'] = samples_df.irods_path.str.split('/').str.get(2)
            samples_df[['lane_index', 'tag_index']] = samples_df.irods_path \
                .str.split('/').str.get(3) \
                .str.split('_').str.get(1) \
                .str.split('.').str.get(0) \
                .str.split('#', expand=True)
    else:
        raise ValueError(f'Expected {samples_fn} to be in either tsv or csv format')    

    for col in ('sample_id',
                'run_id',
                'lane_index',
                'tag_index'):
        assert col in samples_df.columns, f'samples column {col} not found'
    samples_df['run_id'] = samples_df['run_id'].astype(int)
    samples_df['lane_index'] = samples_df['lane_index'].astype(int)
    samples_df['tag_index'] = samples_df['tag_index'].astype(int)
    
    # plate/well ids were recorded in legacy filetypes, keep as is
    if 'plate_id' in samples_df.columns and 'well_id' in samples_df.columns:
        pass
    else:
        # sample_id as `{plate_id}_{well_id}-{sanger_sample_id}` 
        try:
            plate_well_ids = samples_df['sample_id'].str.rsplit('-', n = 1).str.get(0)
            samples_df[['plate_id', 'well_id']] = plate_well_ids.str.rsplit('_', n = 1, expand=True)
            # do not allow for non-standard well IDs
            assert samples_df.well_id.isin(well_id_mapper().values()).all()
            # do not allow for duplicate plate IDs - issues with plotting
            assert (samples_df.plate_id.value_counts() <= 96).all()
            logging.info('inferring plate_id and well_id from sample_id')
        except:
            logging.info('inferring plate_id and well_id from tags')
            samples_df['plate_id'] = samples_df.apply(lambda r: f'p_{r.run_id}_{(r.tag_index - 1) // 96 + 1}',
                axis=1)
            samples_df['well_id'] = (samples_df.tag_index % 96).replace(well_id_mapper())
    
    assert ~samples_df.plate_id.isna().any(), 'Could not infer plate_id for all samples'
    assert ~samples_df.well_id.isna().any(), 'Could not infer well_id for all samples'
    assert samples_df.well_id.isin(well_id_mapper().values()).all(), 'Found well_id outside A1...H12'
    # id_library_lims as `{lims_plate_id}:{lims_well_id}`
    if ('id_library_lims' in samples_df.columns and
        samples_df.id_library_lims.str.contains(':').all()):
            logging.info('inferring lims_plate_id from id_library_lims')
            samples_df[['lims_plate_id', 'lims_well_id']] = samples_df.id_library_lims.str.split(
                ':', n = 1, expand=True
                )
    else:
        logging.info('inferring lims_plate_id from tags')
        samples_df['lims_plate_id'] = samples_df.apply(
            lambda r: f'lp_{r.run_id}_{(r.tag_index - 1) // 384 + 1}',
            axis=1
            )
        samples_df['lims_well_id'] = (samples_df.tag_index % 384).replace(lims_well_id_mapper())
    assert ~samples_df.lims_plate_id.isna().any(), 'Could not infer plate_id for all samples'
    assert ~samples_df.lims_well_id.isna().any(), 'Could not infer well_id for all samples'
    assert samples_df.lims_well_id.isin(lims_well_id_mapper().values()).all(), 'Found well_id outside A1...H12'
    
    # sample_name - short sample id for plotting
    # works if sanger_sample_id does not contain dashes
    # and if sample ID is a concatenation of sample_name and sanger_sample_id
    samples_df['sample_name'] = samples_df['sample_id'].str.rsplit('-', n=1).str.get(0)

    # first run record assumed to be the run ID for plot titles
    run_id = samples_df['run_id'].iloc[0]

    return run_id, samples_df

def prep_stats(stats_fn):
    '''
    load DADA2 stats table from either DADA2_stats.tsv
    or overall_summary.txt file of ampliseq pipeline
    
    For legacy stats table, summarise across targets
    '''

    logging.info(f'preparing DADA2 statistics from {stats_fn}')

    stats_df = pd.read_csv(stats_fn, sep='\t')

    stats_df.rename(columns={
        # compatibility with legacy format
        's_Sample':'sample_id',
        'final':'dada2_postfilter_reads',
        # compatibility with new format 
        'sample':'sample_id',
        # generic renaming
        'DADA2_input':'dada2_input_reads',
        'filtered':'dada2_filtered_reads',
        'denoised':'dada2_denoised_reads',
        'denoisedF':'dada2_denoisedf_reads',
        'denoisedR':'dada2_denoisedr_reads',
        'merged':'dada2_merged_reads',
        'nonchim':'dada2_nonchim_reads'
        },
        inplace=True)
    
    for col in ('sample_id',
                'dada2_input_reads',
                'dada2_filtered_reads',
                'dada2_nonchim_reads'):
        assert col in stats_df.columns, f'stats column {col} not found'

    # denoising happens for F and R reads independently, we take minimum of those 
    # as an estimate for denoised read count
    if 'dada2_denoisedf_reads' in stats_df.columns and 'dada2_denoisedr_reads' in stats_df.columns:
        logging.info('found DADA2 stats for paired end reads')
        stats_df['dada2_denoised_reads'] = stats_df[[
            'dada2_denoisedf_reads',
            'dada2_denoisedr_reads'
            ]].min(axis=1)
        assert 'dada2_merged_reads' in stats_df.columns, f'stats column dada2_merged_reads not found'
    else:
        logging.info('found DADA2 stats for single end reads')
        assert 'dada2_denoised_reads' in stats_df.columns, f'stats column dada2_denoised_reads not found'
        # fake merged - same as denoised
        stats_df['dada2_merged_reads'] = stats_df['dada2_denoised_reads']
    # legacy stats calculated separately for each target, merging
    if 'target' in stats_df.columns:
        logging.info(f'summarising legacy DADA2 statistics across targets')
        stats_df = stats_df.groupby('sample_id').sum(numeric_only=True).reset_index()
    # overall_summary.txt
    if 'cutadapt_total_processed' in stats_df.columns:
        stats_df.rename(columns={
            'cutadapt_total_processed':'total_reads',
            'cutadapt_passing_filters':'readthrough_pass_reads'
        },
        inplace=True)
        # cutadapt stats recorded with thousands comma separator
        for col in ('total_reads', 'readthrough_pass_reads'):
            stats_df[col] = stats_df[col].astype(str).str.replace(',', '').astype(int)
        stats_df.drop(
            columns=['cutadapt_reverse_complemented', 'cutadapt_passing_filters_percent'],
            inplace=True
            )
    # DADA2_stats
    else:
        logging.warning(
            'DADA2_stats.tsv provided instead of overall_summary.txt, '
            'cutadapt readthrough stats will be missing'
            )
        stats_df['total_reads'] = stats_df['dada2_input_reads']
        stats_df['readthrough_pass_reads'] = stats_df['dada2_input_reads']
    
    return stats_df

def combine_stats(stats_df, hap_df, samples_df):
    '''
    Combined per-sample statistics
    '''

    logging.info('preparing combined stats')

    # some samples can be missing from stats due to DADA2 filtering
    if set(stats_df.sample_id) - set(samples_df.sample_id) != set():
        logging.error('sample_id mismatch between samples and stats, QC results will be compromised')
    # some samples can be missing from haps due to DADA2 & post-processing filtering
    elif set(hap_df.sample_id) - set(samples_df.sample_id) != set():
        logging.error('sample_id mismatch between haps and samples, QC results will be compromised')

    comb_stats_df = pd.merge(stats_df, samples_df, on='sample_id', how='outer')
    for col in comb_stats_df.columns:
        if col.endswith('_reads'):
            comb_stats_df[col] = comb_stats_df[col].fillna(0).astype(int)
    comb_stats_df.set_index('sample_id', inplace=True)
    
    comb_stats_df['target_reads'] = hap_df[hap_df.target != 'unknown'] \
        .groupby('sample_id')['reads'].sum()
    comb_stats_df['target_reads'] = comb_stats_df['target_reads'].fillna(0).astype(int)

    comb_stats_df['overall_filter_rate'] = comb_stats_df['target_reads'] / comb_stats_df['total_reads']
    comb_stats_df['overall_filter_rate'] = comb_stats_df['overall_filter_rate'].fillna(0)

    comb_stats_df['unassigned_asvs'] = hap_df[hap_df.target == 'unknown'] \
        .groupby('sample_id')['consensus'].nunique()
    comb_stats_df['unassigned_asvs'] = comb_stats_df['unassigned_asvs'].fillna(0).astype(int)
    
    comb_stats_df['targets_recovered'] = hap_df[hap_df.target != 'unknown'] \
        .groupby('sample_id')['target'].nunique()
    comb_stats_df['targets_recovered'] = comb_stats_df['targets_recovered'].fillna(0).astype(int)
    
    comb_stats_df['raw_mosq_targets_recovered'] = hap_df[hap_df.target.isin(MOSQ_TARGETS)] \
        .groupby('sample_id')['target'].nunique()
    comb_stats_df['raw_mosq_targets_recovered'] = comb_stats_df['raw_mosq_targets_recovered'].fillna(0).astype(int)

    comb_stats_df['raw_multiallelic_mosq_targets'] = (
        hap_df[hap_df.target.isin(MOSQ_TARGETS)].groupby('sample_id')['target'].value_counts() > 2
        ).groupby(level='sample_id').sum()
    comb_stats_df['raw_multiallelic_mosq_targets'] = comb_stats_df['raw_multiallelic_mosq_targets'].fillna(0).astype(int)
    
    comb_stats_df['raw_mosq_reads'] = hap_df[hap_df.target.isin(MOSQ_TARGETS)] \
        .groupby('sample_id')['reads'].sum()
    comb_stats_df['raw_mosq_reads'] = comb_stats_df['raw_mosq_reads'].fillna(0).astype(int)
    
    for pt in PLASM_TARGETS:
        ptl = pt.lower()
        comb_stats_df[f'{ptl}_reads'] = hap_df[hap_df.target == pt] \
            .groupby('sample_id')['reads'].sum()
        comb_stats_df[f'{ptl}_reads'] = comb_stats_df[f'{ptl}_reads'].fillna(0).astype(int)
    
    comb_stats_df.reset_index(inplace=True)
    comb_stats_df.sort_values(by='tag_index', inplace=True)
        
    return comb_stats_df
    # [[
    #     'sample_id',
    #     'total_reads',
    #     'readthrough_pass_reads',
    #     'dada2_input_reads',
    #     'dada2_filtered_reads',
    #     'dada2_denoised_reads',
    #     'dada2_merged_reads',
    #     'dada2_nonchim_reads',
    #     'target_reads',
    #     'overall_filter_rate',
    #     'unassigned_asvs',
    #     'targets_recovered',
    #     'raw_mosq_targets_recovered',
    #     'raw_multiallelic_mosq_targets',
    #     'raw_mosq_reads',
    #     'p1_reads',
    #     'p2_reads'
    # ]]

def prep_dada2(args):

    setup_logging(verbose=args.verbose)
    
    logging.info('ANOSPP data prep started')

    os.makedirs(args.outdir, exist_ok=True)

    cutadapt_deplex(args.dada_table, args.adapters, CUTADAPT_ARGS, args.work_dir)

    if args.rc:
        cutadapt_deplex(args.dada_table, args.adapters, CUTADAPT_ARGS, args.work_dir + '_rc', rc=True)

    hap_df = get_hap_df(args.dada_table, args.work_dir, args.rc)

    hap_fn = f'{args.outdir}/haps.tsv'

    hap_df[[
        'sample_id',
        'target',
        'consensus',
        'reads',
        'total_reads',
        'reads_fraction',
        'nalleles',
        'seqid'
    ]].to_csv(hap_fn, sep='\t', index=False)

    logging.info('Combining DADA2 stats with prep stats')

    hap_df = prep_hap(hap_fn)
    run_id, samples_df = prep_samples(args.manifest, args.run_id)
    stats_df = prep_stats(args.stats)

    comb_stats_df = combine_stats(stats_df, hap_df, samples_df)

    comb_stats_fn = f'{args.outdir}/comb_stats.tsv'

    comb_stats_df.to_csv(comb_stats_fn, sep='\t', index=False)

    logging.info('ANOSPP data prep complete')

def main():
    
    parser = argparse.ArgumentParser("Convert DADA2 output to ANOSPP haplotypes tsv")
    parser.add_argument('-t', '--dada_table', 
                        help='dada2 output table, in ampliseq pipeline it is called DADA2_table.tsv', 
                        required=True)
    parser.add_argument('-a', '--adapters', 
                        help='adapters fasta file for deplexing with cutadapt', 
                        required=True)
    parser.add_argument('-m', '--manifest', help='Samples manifest tsv file', required=True)
    parser.add_argument('-s', '--stats', 
                        help='ampliseq overall_summary.tsv file or DADA2_stats.tsv file', required=True)
    parser.add_argument('-o', '--outdir', 
                        help='output directory for haplotypes and stats tsv files. Default: prep', 
                        default='prep')
    parser.add_argument('-w', '--work_dir', 
                        help='working directory for intermediate files. Default: work',
                        default='work')
    parser.add_argument('-i', '--run_id',
                        help=('use this run_id and sample order to infer run, lane and tag instead of `irods_path` column.' 
                              'Default: do not override'),
                        default=None)
    parser.add_argument('-r', '--rc', 
                        help=('Also run deplexing on reverse complement ASVs, '
                              'useful when amplicon orientation is not known'), 
                        action='store_true', default=False)
    parser.add_argument('-v', '--verbose', 
                        help='Include INFO level log messages', action='store_true', default=False)

    args = parser.parse_args()

    args.work_dir=args.work_dir.rstrip('/')
    for fn in args.adapters, args.dada_table:
        assert os.path.isfile(fn), f'{fn} not found'

    prep_dada2(args)

if __name__ == '__main__':
    main()