# import scimap as sm
from .. import tools as tl
from .. import preprocessing as pp
from .. import plotting as pl
from .. import helpers as hl

import sys
import argparse
import pathlib
import textwrap
from joblib import Parallel, delayed

def mcmicro_wrap(argv=sys.argv):
    
    parser = argparse.ArgumentParser()

    parser.add_argument(
        'csv',
        help='single-cell quantification table (CSV file) generated by mcmicro',
    )
    parser.add_argument(
        '-o', '--output', default='.',
        help='output directory'
    )

    args = parser.parse_args(argv[1:])

    run_result = Parallel(n_jobs=-1, verbose=1)(
        delayed(clustering)([None, args.csv, '-o', args.output, '--clustering-method', m]) 
        for m in ['spatial', 'kmeans', 'phenograph', 'leiden']
    )

    print(run_result)
    
    merge([None, args.output, '-o', args.output, '-d', '--csv'])

    # move pdf files to output / plots
    output_dir = pathlib.Path(args.output)
    pdfs = sorted(output_dir.rglob('*.pdf'))
    plots_dir = output_dir / 'plots'
    plots_dir.mkdir(exist_ok=True)

    if len(pdfs) > 0:
        print(textwrap.indent(
            'Moving pdf plots:\n',
            '    '
        ))
    for p in pdfs:
        if plots_dir not in p.parents:
            move_to = plots_dir / p.parent.name / p.name
            print(textwrap.indent(
                f'{str(p)}\n    -> {str(move_to)}',
                '    '
            ))
            move_to.parent.mkdir(exist_ok=True, parents=True)
            p.replace(move_to)
    print()
    return 0

def clustering(argv=sys.argv):

    parser = argparse.ArgumentParser()

    parser.add_argument(
        'csv',
        help='single-cell quantification table (CSV file) generated by mcmicro',
    )
    parser.add_argument(
        '-o', '--output', default='.',
        help='output directory'
    )
    clustering_methods = ['all', 'spatial', 'kmeans', 'phenograph', 'leiden', 'pass']
    parser.add_argument(
        '--clustering-method', default=['all'], choices=clustering_methods, nargs='+', 
        help='choice of clustering algorithms, "pass": do not run any clustering methods; "all": run all clustering methods; default: "all"'
    )

    args = parser.parse_args(argv[1:])

    _output_dir = pathlib.Path(args.output)
    mcmicro_csv_path = args.csv
    methods = set(args.clustering_method)

    assert '.csv' in mcmicro_csv_path, 'input file must be a csv file'
    
    if 'all' in methods:
        methods = clustering_methods[1:-1]

    if 'pass' in methods:
        pp.mcmicro_to_scimap(
            feature_table_path = mcmicro_csv_path, 
            output_dir = str(_output_dir)
        )
        return

    for method in methods:

        output_dir = _output_dir / method
        
        pp.mcmicro_to_scimap(
            feature_table_path = mcmicro_csv_path, 
            output_dir = str(output_dir)
        )

        adata_path = pathlib.Path(output_dir) / f'{pathlib.Path(mcmicro_csv_path).stem}.h5ad'

        if method == 'spatial':
            # Spatial clustering
            tl.spatial_expression(
                adata = str(adata_path), 
                output_dir = output_dir
            )
            tl.spatial_cluster(
                adata = str(adata_path), 
                df_name = "spatial_expression",
                output_dir = output_dir
            )
        
        else:
            # Expression clustering
            tl.cluster(
                adata = str(adata_path), 
                method = method,
                output_dir = output_dir
            )

            # Expression clustering plots
            pl.cluster_plots(
                adata = str(adata_path), 
                group_by  = method,
                output_dir = output_dir
            )
    return 0

def merge(argv=sys.argv):
    
    parser = argparse.ArgumentParser()

    parser.add_argument(
        'directory',
        help='recursively search for .h5ad files to merge',
    )
    parser.add_argument(
        '-o', '--output', default='.',
        help='output directory'
    )
    parser.add_argument(
        '-d', '--delete-merged', default=False, action='store_true',
        help='delete found input files after merging; default: False'
    )
    parser.add_argument(
        '--csv', default=False, action='store_true',
        help='output csv version of the merged h5ad file; default: False'
    )

    args = parser.parse_args(argv[1:])

    input_dir = pathlib.Path(args.directory)
    output_dir = pathlib.Path(args.output)
    delete_after = args.delete_merged
    output_csv = args.csv

    input_files = sorted(input_dir.rglob('*.h5ad'))
    output_file = output_dir / input_files[0].name
    
    if len(input_files) == 0:
        print(f'0 .h5ad found in {str(input_dir)}')
        return 1

    print_text = '''
        Merging:

        {}

        Writing:

        {}'''
    print_text = textwrap.dedent(print_text).format(
        "\n".join([str(p) for p in input_files]),
        str(output_file)
    )
    print(textwrap.indent(print_text, '    '))
    # Merge data
    hl.merge_adata_obs(
        adata = [
            str(p)
            for p in input_files
        ],
        output_dir = output_dir
    )
    if output_csv == True:
        hl.scimap_to_csv(
            adata = str(output_file),
            output_dir = output_dir
        )
        print(textwrap.indent(
            str(output_file.parent / f'{output_file.stem}.csv\n'),
            '    '
        ))
    

    if delete_after == True:
        if output_file in input_files:
            idx = input_files.index(output_file)
            input_files.pop(idx)
        print(
            textwrap.indent(
                'Deleting:\n\n' + "\n".join([str(f) for f in input_files]),
                '    '
            )
        )
        for f in input_files:
            f.unlink()
        print()
    return 0