import sys
from pathlib import Path
from argparse import ArgumentParser

import h5py
import pandas as pd
import numpy as np
from tqdm import tqdm

from export import export_read_file


def get_args():
    parser = ArgumentParser(description="Parse sequencing_summary.txt files and .paf files to find split reads "
                                        "in an Oxford Nanopore Dataset",
                            add_help=False)
    general = parser.add_argument_group(title='General options')
    general.add_argument("-h", "--help",
                         action="help",
                         help="Show this help and exit"
                         )
    in_args = parser.add_argument_group(
        title='Input sources'
    )
    in_args.add_argument("-s", "--summary",
                         required=True,
                         nargs='+',
                         help='Sequencing summary file(s) generated by albacore or guppy. Can be compressed '
                              'using gzip, bzip2, xz, or zip')
    in_args.add_argument("--start-events",
                         help="start_events.csv file generated by event_finder.py",
                         default="",
                         required=True,
                         )
    in_args.add_argument("--end-events",
                         help="end_events.csv file generated by event_finder.py",
                         default="",
                         required=True,
                         )
    in_args.add_argument("--targets",
                         help="A text file of target read ids with one per line.",
                         default="",
                         required=True,
                         )
    in_args.add_argument("--bulk-files",
                         help="ONT bulk FAST5 files.",
                         nargs='+',
                         default="",
                         )
    in_args.add_argument("-o", "--output-name",
                         help="Name of the output folder, this will be generated if it does not exist",
                         required=True,
                         default=""
                         )
    in_args.add_argument("--extra-classifications",
                         help="Any extra MinKNOW classifications to include.",
                         nargs='*',
                         default="",
                         )
    return parser.parse_args()


def main():
    args = get_args()
    # debug(args)
    # # sys.exit()
    # Make folders
    for j in ['starts', 'ends']:
        Path('{i}/{j}/{k}'.format(i=args.output_name, j=j, k='fast5')).mkdir(parents=True, exist_ok=True)
    # Open files
    start_events = pd.read_csv(args.start_events, sep=',')
    end_events = pd.read_csv(args.end_events, sep=',')
    seq_sum_df = concat_files_to_df(file_list=args.summary, sep='\t')

    # Create end_time Series in seq_sum_df
    seq_sum_df['end_time'] = seq_sum_df['start_time'] + seq_sum_df['duration']
    # Sort and Groupby to segregate runs and channels
    seq_sum_df = seq_sum_df.sort_values(by=['run_id', 'channel', 'start_time'], ascending=True)
    seq_sum_df_1 = seq_sum_df.copy()
    gb = seq_sum_df.groupby(['run_id', 'channel'])
    gb1 = seq_sum_df_1.groupby(['run_id', 'channel'])
    # Get previous and next start times within groupby
    seq_sum_df['next_start'] = gb['start_time'].shift(-1)
    seq_sum_df_1['prev_start'] = gb1['start_time'].shift(1)

    target_read_ids = []
    with open(args.targets, 'r') as file:
        for line in file:
            target_read_ids.append(line.strip())

    classifications = ['pore', 'inrange', 'good_single', 'unblocking']
    if args.extra_classifications:
        classifications.extend(args.extra_classifications)

    # Get end_events for target_read_ids
    end_events = end_events[end_events['read_id'].isin(target_read_ids)]
    normal_ending_ids = end_events[end_events['time'].ge(0) &
                                   end_events['label'].isin(classifications)]['read_id'].unique()
    abnormally_ending_ids = end_events[~end_events['read_id'].isin(normal_ending_ids)]['read_id'].unique()
    end_target_ss = seq_sum_df[seq_sum_df['read_id'].isin(abnormally_ending_ids)]

    # Get start_events for target_read_ids
    start_events = start_events[start_events['read_id'].isin(target_read_ids)]
    normal_starting_ids = start_events[start_events['time'].le(0) &
                                       start_events['label'].isin(classifications)]['read_id'].unique()
    abnormally_starting_ids = start_events[~start_events['read_id'].isin(normal_starting_ids)]['read_id'].unique()
    start_target_ss = seq_sum_df_1[seq_sum_df_1['read_id'].isin(abnormally_starting_ids)]

    print('Collecting abnormally ending reads:')
    end_read_info = write_files(end_target_ss, args.bulk_files, 'start_time',
                                'next_start', '{i}/ends/fast5/'.format(i=args.output_name))
    end_read_info.to_csv('{}/ends_read_info.txt'.format(args.output_name), sep='\t', index=False, header=True)
    end_read_info.to_csv('{}/ends_filenames.txt'.format(args.output_name), sep='\t', index=False, header=False,
                         columns=['filename'])

    print('Collecting abnormally starting reads:')
    start_read_info = write_files(start_target_ss, args.bulk_files, 'prev_start',
                                  'end_time', '{i}/starts/fast5/'.format(i=args.output_name))
    start_read_info.to_csv('{}/starts_read_info.txt'.format(args.output_name), sep='\t', index=False, header=True)
    start_read_info.to_csv('{}/starts_filenames.txt'.format(args.output_name), sep='\t', index=False, header=False,
                           columns=['filename'])
    return


def write_files(target_ss, bulkfiles, read_start_col, read_end_col, export_path, remove_pore=True):
    """Abstraction for export_read_file for collecting read info

    Parameters
    ----------
    target_ss : pd.DataFrame
        DataFrame of reads to generate reads for
    bulkfiles: list
        list of bulk FAST5 files
    read_start_col : str
        Column in the target_ss that start index is derived from
    read_end_col : str
        Column in the target_ss that end index is derived from
    export_path : str
        The folder where read files will be written
    remove_pore : bool
        Remove pore-like signal from trace (>1500)

    Returns
    -------
    pd.DataFrame
        DataFrame of read info about reads that have been written
    """
    d = {
        'read_id': [],
        'channel': [],
        'start_index': [],
        'end_index': [],
        'bv_read_id': [],
        'filename': [],
        'bv_filename': []
    }
    files_written = 0
    for bf in tqdm(bulkfiles):
        f = h5py.File(bf, 'r')
        run_id = f['UniqueGlobalKey']["tracking_id"].attrs["run_id"].decode('utf8')
        sf = int(f["UniqueGlobalKey"]["context_tags"].attrs["sample_frequency"].decode('utf8'))
        t = target_ss[target_ss['run_id'] == run_id]
        t = t.dropna()
        f.close()
        file = h5py.File(bf, 'r')
        for idx, row in tqdm(t.iterrows(), total=t.shape[0], desc=run_id):
            si = int(np.floor(row[read_start_col] * sf))
            ei = int(np.floor(row[read_end_col] * sf))
            d['read_id'].append(row['read_id'])
            d['channel'].append(row['channel'])
            d['start_index'].append(si)
            d['end_index'].append(ei)
            d['bv_read_id'].append("{ch}-{start}-{end}".format(ch=row['channel'], start=si, end=ei))
            d['filename'].append(row['filename'])
            d['bv_filename'].append(export_read_file(row['channel'],
                                                     si,
                                                     ei,
                                                     file,
                                                     export_path,
                                                     remove_pore=remove_pore))
            files_written += 1
    print('{} reads written'.format(files_written))
    return pd.DataFrame(d)


def concat_files_to_df(file_list, **kwargs):
    """Return a pandas.DataFrame from a list of files
    """
    df_list = []
    for f in file_list:
        try:
            df_list.append(pd.read_csv(filepath_or_buffer=f, **kwargs))
        except pd.errors.ParserError as e:
            print('{}\nThis is usually caused by an input file not being the expected format'.format(repr(e)))
            sys.exit(1)
        except Exception as e:
            sys.exit(1)
    return pd.concat(df_list, ignore_index=True)


def debug(args):
    dirs = dir(args)
    for attr in dirs:
        if attr[0] != '_':
            print('{a:<25} {b}'.format(a=attr, b=getattr(args, attr)))


if __name__ == '__main__':
    main()
