import click


@click.group()
def tasks():
    pass
    

@tasks.command(help="Parse the tsv file generated by command 'search'. "
                    "If you want to output the result into a specified file, you may "
                    "use the redirection symbol \"> path\".")
@click.option("-i", "--input", type=str, help="tsv file path.", default=None)
@click.option("--id", "--id", type=str, help="output information of specified query id in tsv file.", default=None)
def parse(input, id):
    from Results import format as fmt
    
    if not input:
        click.echo("Error: No tsv file input, please input a tsv file.")
        return
        
    fmt.parse_tsv(input, id)


@tasks.command(help="Convert tsv files to MSAs.")
@click.option("-i", "--input", type=str, help="tsv file path.", default=None)
@click.option("-o", "--output", type=str, help="the directory for storing results.", default=None)
@click.option("-t", "--filter_threshold", type=float, help="sequence similarity threshold. default: 0.9", default=0.9)
@click.option("--align", type=bool, help="whether align sequences or not. default: 1(True)", default=True)
@click.option("--outfmt", type=click.Choice(['fasta', 'a3m']), help="output result in specified format. "
                                                                    "default: a3m", default='a3m')
@click.option("--id", "--id", type=str, help="output result of specified query id in tsv file.", default=None)
def tsv2msa(input, output, filter_threshold, align, outfmt, id):
    from Results import format as fmt
    if not input:
        click.echo("Error: No tsv file input, please input a tsv file.")
        return
        
    fmt.tsv_to_msa(input, output, filter_threshold, align, outfmt, id)


@tasks.command(help="Search query sequences in standard UniRef30 database. The database is divided into several parts "
                    "based on sequence length because of memory limitation. The search will be executed on every "
                    "sub-database and all results will finally be aggregated into one.")
@click.option("--q_vec_path", type=str, help="npy path of query sequences.", default=None)
@click.option("--q_info_path", type=str, help="information of query sequences.", default=None)
@click.option("--db_path", type=str, help="database path", default=None)
@click.option("-o", "--output", type=str, help="output result to specified file. input related path will be used "
                                               "by default.", default=None)
@click.option("-t", "--threshold", type=float, help="threshold that the model use to detect homologies. default: 1",
              default=1)
@click.option("--nprobe", type=int, help="search the first n clusters. f you want to search all clusters, "
                                         "set this to 0. default: 1", default=1)
@click.option("--min_num", type=int, help="min num limitation of homologies for one query sequence. default: 0", default=0)
@click.option("--max_num", type=int, help="max num limitation of homologies for one query sequence. if you want no "
                                          "limitation, set this to 0. default: 500", default=500)
@click.option("-v", "--verbose", type=bool, help="display details. default: 1(True)", default=True)
def standard_search(q_vec_path, q_info_path, db_path, output, threshold, nprobe, min_num, max_num,
                      verbose):
    import os
    import pandas as pd
    import click
    import numpy as np
    from utils import TimeCounter
    from search import pipeline
    
    assert min_num <= max_num
    
    var_dict = {'q_vec_path': q_vec_path,
                'q_info_path': q_info_path,
                'db_index_path': db_path}
    
    for k, v in var_dict.items():
        if not v:
            click.echo(f"Error: No {k} input, please input {k}.")
            return
    
    output = q_vec_path.replace(".npy", "_result.tsv") if not output else output + '_result.tsv'
    dir_name = db_path.split("/")[-1]
    db_name = [f"{i}_{i+100}" for i in range(0, 1100, 100)]
    aggregation_list = []
    for name in db_name:
        for i in range(10):
            db_index_path = f"{db_path}_{name}/{dir_name}_{name}_p{i}_IVF_Flat.index"
            db_pointer_path = f"{db_path}_{name}/{dir_name}_{name}_p{i}_pointer.tsv"
            db = f"{db_path}_{name}/{dir_name}_{name}_{i}.faa"
            
            temp_output = output + name + str(i)
            
            if not os.path.exists(db_index_path):
                continue
            
            pipeline(q_vec_path, db_index_path, q_info_path, db_pointer_path, db, temp_output,
                     t=threshold, nprobe=nprobe, max_num=max_num, min_num=min_num, verbose=verbose)
            
            aggregation_list.append(temp_output)
        click.echo(f"{dir_name}_{name} search finished.")
        
    query_dict = {}
    seq_dict = {}
    with TimeCounter("Aggregating results...", verbose):
        for temp_output in aggregation_list:
            if not os.path.exists(temp_output):
                continue
                
            res = pd.read_csv(temp_output, sep='\t')
            
            for key in res['query_id'].value_counts().keys():
                if key not in query_dict.keys():
                    data = res[res['query_id'] == key]
                    query_dict[key] = data
                    seq_dict[key] = data.iloc[0, 4]
                else:
                    query_dict[key] = pd.concat([query_dict[key], res[res['query_id'] == key]], axis=0).sort_values(by='distance')
                    query_dict[key].iloc[0, 4] = seq_dict[key]
                    if max_num != 0:
                        query_dict[key] = query_dict[key].iloc[:max_num]
            
            os.remove(temp_output)
            
        # output result
        res_list = []
        q_vec = np.load(q_vec_path)
        # t = q_vec[:, -1]
        for i, query_id in enumerate(pd.read_csv(q_info_path, sep='\t')['id']):
            if query_id in query_dict.keys():
                # min num limitation
                tail_num = max(min_num, sum(query_dict[query_id]['distance'] <= threshold))
                
                res_list.append(query_dict[query_id].iloc[:tail_num])
            
            else:
                print(f"No homology detected for {query_id}!")
                
        if res_list:
            pd.concat(res_list, axis=0).to_csv(output, sep='\t', index=False)


@tasks.command(help="Search query sequences in user customized database")
@click.option("--q_vec_path", type=str, help="npy path of query sequences.", default=None)
@click.option("--q_info_path", type=str, help="information of query sequences.", default=None)
@click.option("--db_index_path", type=str, help="faiss index of database", default=None)
@click.option("--db_pointer_path", type=str, help="pointer for database.", default=None)
@click.option("--db_path", type=str, help="database path.", default=None)
@click.option("-o", "--output", type=str, help="output result to specified file. input related path will be used "
                                               "by default.", default=None)
@click.option("-t", "--threshold", type=float, help="threshold that the model use to detect homologies. default: 1", default=1)
@click.option("--nprobe", type=int, help="search the first n clusters. f you want to search all clusters, "
                                         "set this to 0. default: 1", default=1)
@click.option("--max_num", type=int, help="max num limitation of homologies for one query sequence. if you want no "
                                          "limitation, set this to 0. default: 500", default=500)
@click.option("--min_num", type=int, help="min num limitation of homologies for one query sequence. default: 0", default=0)
@click.option("-v", "--verbose", type=bool, help="display details. default: 1(True)", default=True)
def customized_search(q_vec_path, q_info_path, db_index_path, db_pointer_path, db_path,
                      output, threshold, nprobe, max_num, min_num, verbose):
    from search import pipeline
    
    var_dict = {'q_vec_path': q_vec_path,
                'q_info_path': q_info_path,
                'db_index_path': db_index_path,
                'db_pointer_path': db_pointer_path,
                'db_path': db_path}
    
    for k, v in var_dict.items():
        if not v:
            click.echo(f"Error: No {k} input, please input {k}.")
            return
    
    output = q_vec_path.replace(".npy", "_result.tsv") if not output else output + '_result.tsv'
    
    pipeline(q_vec_path, db_index_path, q_info_path, db_pointer_path, db_path, output,
             t=threshold, nprobe=nprobe, max_num=max_num, min_num=min_num, verbose=verbose)


@tasks.command(help="Convert sequences into vectors")
@click.option("-i", "--input", type=str, help="fasta path.", default=None)
@click.option("-o", "--output", type=str, help="output result to specified file. input related path will be used "
                                               "by default.", default=None)
@click.option("--devices", type=str, help="choose which GPUs will be used. e.g. '0,1,2,3'. if not specified, all GPUs "
                                          "will be used by default. if no GPU available, please set this to 'cpu'", default=None)
@click.option("--batch", type=int, help="batch size adopted. default: 128", default=128)
@click.option("--model", type=str, help="model parameter file. e.g. 'your_path/model_name.pt'", default=None)
@click.option("-v", "--verbose", type=bool, help="display details. default: 1(True)", default=True)
@click.option("--type", type=str, help="type of protein sequences: query or database. default: query", default='query')
def build(input, output, devices, model, batch, type, verbose):
    import click
    import torch
    from vector_construction import fasta2vec
    
    if not input:
        click.echo("Error: No fasta file input, please input a fasta file.")
        return
    
    if not model:
        click.echo("Error: No model path input, please input a model path.")
        return
        
    if not devices:
        cnt = torch.cuda.device_count()
        devices = list(range(cnt)) if cnt != 0 else ['cpu']
    elif devices == 'cpu':
        devices = ['cpu']
    else:
        devices = [int(i) for i in devices.split(',')]
    
    if not output:
        output = input
    
    fasta2vec(devices, model, input, output, batch_size=batch, type=type, cover=True, verbose=verbose)


@tasks.command(help="Construct indices of given vectors.")
@click.option("-i", "--input", type=str, help="npy path.", default=None)
@click.option("-db", "--database", type=str, help="fasta format database path.", default=None)
@click.option("-o", "--output", type=str, help="specified index and pointer path. input related path will be used "
                                               "by default.", default=None)
@click.option("--param", type=str, help="faiss index parameters. for more information, please check "
                                                   "https://github.com/facebookresearch/faiss/wiki/Faiss-indexes."
                                                   "default: IVFx, Flat", default=None)
@click.option("-v", "--verbose", type=bool, help="display details. default: 1(True)", default=True)
def vec2index(input, database, output, param, verbose):
    import click
    import faiss
    import numpy as np
    from math import ceil
    from utils import TimeCounter
    from LoadingData.data_construction import construct_pointer
    from Faiss.index_construction import contruct_faiss_index
    
    if not input:
        click.echo("--input Error: No npy file input, please input a npy format file.")
        return
    
    if not database:
        click.echo("--database Error: No fasta file input, please input a fasta file.")
        return

    pointer_out = input.replace(".npy", "_pointer.tsv") if not output else output + "_pointer.tsv"
    output = input.replace(".npy", ".index") if not output else output + '.index'
    
    vectors = np.load(input)
    dim, measure = 1280, faiss.METRIC_L2
    if not param:
        n = ceil(vectors.shape[0] / 256)
        param = f'IVF{n}, Flat'

    with TimeCounter("Constructing index...", verbose):
        contruct_faiss_index(vectors, output, dim, measure, param)
    
    with TimeCounter("Constructing pointer...", verbose):
        construct_pointer(database, pointer_out)


@tasks.command(help="Align a fasta file.")
@click.option("-i", "--input", type=str, help="tsv file path.", default=None)
@click.option("-o", "--output", type=str, help="output file path.", default=None)
@click.option("--outfmt", type=click.Choice(['fasta', 'a3m']), help="output result in specified format. "
                                                                    "default: a3m", default='a3m')
def align(input, output, outfmt):
    from Results import format as fmt
    from utils import TimeCounter
    
    if not input:
        click.echo("Error: No tsv file input, please input a tsv file.")
        return
    
    if not output:
        click.echo("Error: Missing output path!")
        return
    
    with TimeCounter("Aligning..."):
        fmt.align(input, output, outfmt)
    
    
if __name__ == '__main__':
    tasks()
    
