import sys
sys.path.append("..")
from faigen.data.sequence import Dna2VecList,regex_filter
import pandas as pd
import numpy as np
import os
from functools import partial
import configargparse
from pathlib import Path
from Bio.SeqRecord import SeqRecord
import yaml

def filter_by_count(df:pd.DataFrame, min=1)->pd.DataFrame:
    res=df.copy()
    drop = res.index[res.index.values[np.asarray(res.seq_count.values) < min]]
    return res.drop(drop, axis=0)


def filter_by_label(df:pd.DataFrame, word:str)->pd.DataFrame:
    res,mask=df.copy(),[]
    for x in df.label.values: mask.append(False if word in x else True)
    drop = res.index[mask]
    return res.drop(drop, axis=0)


def main():
    argp = configargparse.get_argument_parser()
    argp.add_argument('-i', help='input folder with Fasta files', type=str, default='/data/genomes/GenSeq_fastas')
    argp.add_argument('-o', help='output file name', type=str)
    argp.add_argument('-g', choices=['folder', 'file'], help='granularity', type=str, default="file")
    argp.add_argument('-l', choices=['description', 'id', 'file'], help='source of labels', type=str, default="description")
    argp.add_argument('-lsi', help='label selector (comma delimited numbers)', type=str)
    argp.add_argument('-lsr', help='regular expression for labeling', type=str)
    argp.add_argument('-rxkeep', help='keep if regular expression found', type=str)
    argp.add_argument('-rxdrop', help='drop if regular expression found', type=str)
    argp.add_argument('-d', help='label delimiter', type=str, default=" ")
    argp.add_argument('-split', help='split by folders, coma delimited string', type=str, default="train,valid,test")
    argp.add_argument('-portions', help='split by folders, coma delimited string', type=str, default="0.7,0.2,0.1")


    args = {k:v for k,v in vars(argp.parse_args()).items()}
    input= Path(args["i"]) if args["i"] is not None else Path(".")
    filters=[]
    if args["rxkeep"] is not None: filters.append(partial(regex_filter, rx=args["rxkeep"]))
    if args["rxdrop"] is not None: filters.append(partial(regex_filter, rx=args["rxdrop"], keep=False))

    all_fastas = Dna2VecList.from_folder(input, filters=filters if len(filters) > 0 else None ).items

    output = input / "inventory"
    print("Creating Inventory in", str(output))

    if not os.path.exists(output):
        os.makedirs(output)

    fn = "sequences" if args["o"] is None else args["o"]

    with open(output / f'{fn}_inventory.yml', 'w') as outfile:
        yaml.dump(args, outfile, default_flow_style=False)


    inventory = pd.DataFrame(data=all_fastas)
    if args['lsi'] is not None:
        lsi = [int(x) for x in args["lsi"].split(",")]
        label_source = inventory.loc[:,args['l']].values
        tokens = [np.asarray(x.split(args["d"])) for x in list(label_source)]
        inventory["label"] = [" ".join(t[lsi]) for t in tokens]

    inventory.to_csv(Path(output / f"{fn}.csv"))

    if "label" not in inventory.columns.values: return

    if args["g"] == "file":
        files_df = inventory.groupby(["file"]).agg({"id": ['count', list],"label":set,
                                                    "len": [list, min, max, np.mean, np.std], "description": list})
        files_df.columns=["seq_count", "id", "label", 'len', 'min' , 'max', 'mean', 'std', "description"]
        files_df["label"] = [list(x)[0] for x in list(files_df.label.values)]

    files_df.to_csv( output / f"{fn}_by_file.csv")


    label_df = inventory.groupby("label").agg({ "id": ["count", list], "len": [list, min, max, np.median],"file": list})
    label_df.columns = [ "seq_count", 'id', "lengths", "min", "max", "median","files"]
    label_df.files = [list(set(x)) for x in list(label_df.files.values)]
    label_df["file_count"] = [len(x) for x in list(label_df.files.values)]

    label_df.to_csv(output / f"{fn}_by_label.csv")
    label_df.to_pickle( output / f"{fn}_by_label.pkl")



if __name__ == '__main__':
    main()