#!/usr/bin/env python
import argparse
import glob
import io
import os
import subprocess
import shutil
import gzip
from typing import IO
from io import StringIO
import tempfile
from concurrent.futures import as_completed
import json
from collections import defaultdict

import pandas as pd
from loky import get_reusable_executor
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from tqdm.auto import tqdm

from masato.utils import read_table, write_table


def merge_pairs(
    fastq_dir: str, fastq_out: str, num_threads: int, backend: str = "vsearch"
) -> None:
    # merge paired end reads
    output_dir = os.path.dirname(fastq_out)
    os.makedirs(output_dir, exist_ok=True)
    if backend == "usearch":
        subprocess.run(
            "usearch11 -fastq_mergepairs "
            f"{fastq_dir}/*_R1.fastq "
            "-relabel @ "
            f"-fastqout {output_dir}/merged.fastq "
            "-fastq_maxdiffs 25 "
            "-fastq_minovlen 50 "
            f"-report {output_dir}/merge_report.txt "
            f"-tabbedout {output_dir}/merge.tsv",
            shell=True,
        )
    elif backend == "vsearch":
        # vsearch cannot merge multiple files or perform relabeling at once, so we first
        # relabel with seqtk and then merge samples, and finally merge pairs with
        # vsearch.
        # Besides, vsearch does not have tsv or report output
        os.makedirs(f"{output_dir}/temp", exist_ok=True)
        fs = []
        ps = []

        # `fastq_dir` could be a path regex so that we can easily do merging for a
        # subset of samples in a directory.
        if os.path.isdir(fastq_dir):
            files = sorted(glob.glob(f"{fastq_dir}/*.fastq"))
            if not files:
                raise ValueError(f"No fastq files found in {fastq_dir}")
        else:
            files = [i for i in sorted(glob.glob(fastq_dir)) if i.endswith(".fastq")]
            if not files:
                raise ValueError(f"No fastq files found matching {fastq_dir}")
        if not all(f.endswith("_R1.fastq") or f.endswith("_R2.fastq") for f in files):
            raise ValueError(
                "Make sure to rename all fastq files of interest to end with _R1.fastq or _R2.fastq"
            )
        non_pair_files = []
        for f in files:
            if f.endswith("_R1.fastq"):
                if f"{f[:-9]}_R2.fastq" not in files:
                    non_pair_files.append(f)
            elif f.endswith("_R2.fastq"):
                if f"{f[:-9]}_R1.fastq" not in files:
                    non_pair_files.append(f)
        if non_pair_files:
            raise ValueError(
                f"Make sure that each sample has both _R1.fastq and _R2.fastq files. "
                f"The following files are not paired: {non_pair_files}"
            )

        batch = 1000
        for i in range(0, len(files), batch):
            file_list = files[i : i + batch]
            for f in file_list:
                f_base = os.path.basename(f)
                sample_name = f_base.split("_")[0]
                # use Popen for non-blocking execution
                outfile = open(f"{output_dir}/temp/{f_base}", "w")
                process = subprocess.Popen(
                    ["seqtk", "rename", f, f"sample={sample_name} "],
                    stdout=outfile,
                )
                fs.append(outfile)
                ps.append(process)
            [p.wait() for p in tqdm(ps)]
            [f.close() for f in fs]

        subprocess.run(
            f"cat {output_dir}/temp/*_R1.fastq > {output_dir}/temp_R1.fastq", shell=True
        )
        subprocess.run(
            f"cat {output_dir}/temp/*_R2.fastq > {output_dir}/temp_R2.fastq", shell=True
        )
        subprocess.run(
            [
                "vsearch",
                "--fastq_mergepairs",
                f"{output_dir}/temp_R1.fastq",
                "--reverse",
                f"{output_dir}/temp_R2.fastq",
                "--fastqout",
                fastq_out,
                # "--fastaout",
                # f"{output_dir}/merged.fa",
                "--fastq_maxdiffs",
                "25",
                "--fastq_minovlen",
                "50",
                "--threads",
                str(num_threads),
            ]
        )
        os.remove(f"{output_dir}/temp_R1.fastq")
        os.remove(f"{output_dir}/temp_R2.fastq")
        shutil.rmtree(f"{output_dir}/temp")


def qc(
    input_fastq: str,
    output_fasta: str | None = None,
    num_threads: int = 16,
    backend: str = "vsearch",
) -> str:
    """Output fasta is default to be modified from input fastq, but can be specified.
    For example, if input fastq file is /data/merged.fq.gz, then output fasta will be
    /data/merged.filtered.fq.gz, unless otherwise specified.
    """
    if output_fasta is None:
        output_dir = os.path.dirname(input_fastq)
        filename = os.path.basename(input_fastq).split(".")[0]
        output_fastx = os.path.join(output_dir, f"{filename}.filtered.fq")
    else:
        output_fastx = output_fasta.rstrip(".gz")

    if backend == "usearch":
        num_splits = None
        # Split the merged FASTQ file into multiple pieces
        total_lines = int(subprocess.check_output(["wc", "-l", input_fastq]).split()[0])
        lines_per_file = ((total_lines + num_splits - 1) // num_splits // 4) * 4
        base_dir = os.path.dirname(input_fastq)
        if num_splits > 1:
            subprocess.run(
                [
                    "split",
                    "-l",
                    str(lines_per_file),
                    "-d",
                    "--additional-suffix=.fastq",
                    input_fastq,
                    os.path.join(base_dir, "_split_"),
                ]
            )

            # Perform QC on each split FASTQ file
            with open(output_fastx, "w") as outfile:
                for split_file in glob.glob(os.path.join(base_dir, "_split_*.fastq")):
                    subprocess.run(
                        [
                            "usearch11",
                            "-fastq_filter",
                            split_file,
                            "-fastaout",
                            "/dev/stdout",
                            "-fastq_maxee",
                            "0.5",
                            "-fastq_minlen",
                            "200",
                            "-fastq_maxns",
                            "0",
                            "-relabel",
                            "filtered",
                            "-threads",
                            str(num_threads),
                        ],
                        stdout=outfile,
                    )
                    os.remove(split_file)
        else:
            subprocess.run(
                [
                    "usearch11",
                    "-fastq_filter",
                    input_fastq,
                    "-fastaout",
                    output_fastx,
                    "-fastq_maxee",
                    "0.5",
                    "-fastq_minlen",
                    "200",
                    "-fastq_maxns",
                    "0",
                    "-relabel",
                    "filtered",
                    "-threads",
                    str(num_threads),
                ]
            )
    elif backend == "vsearch":
        # we don't need to split anymore since there is no memory limit.
        subprocess.run(
            [
                "vsearch",
                "--fastq_filter",
                input_fastq,
                "--fastqout",
                output_fastx,
                "--fastq_maxee",
                "0.5",
                "--fastq_minlen",
                "100",
                "--fastq_maxns",
                "0",
                "--relabel",
                "filtered",
                # "--threads",
                # str(num_threads),
            ]
        )
        subprocess.run(["pigz", "-f", output_fastx])
        output_fastx += ".gz"
    return output_fastx


def split_fastq(args):
    # Calculate the number of lines in the FASTQ file (each read consists of 4 lines)
    total_lines = int(subprocess.check_output(["wc", "-l", args.input_file]).split()[0])
    # Calculate the number of lines for each split file, ensuring that the total number of lines is divisible by 4
    lines_per_file = ((total_lines + args.num_splits - 1) // args.num_splits // 4) * 4

    # Split the file using the `split` command
    subprocess.run(
        [
            "split",
            "-l",
            str(lines_per_file),
            "-d",
            "--additional-suffix=.fastq",
            args.input_file,
            os.path.join(args.output_dir, "split_"),
        ]
    )

    print(f"FASTQ file has been split into {args.num_splits} pieces.")


def add_depth_to_metadata(fastq_dir: str, metadata_path: str) -> None:
    # Generate metadata for the input FASTQ files

    # data = []
    # for f in os.listdir(fastq_dir):
    #     if f.endswith("_R1.fastq"):
    #         sample_name = f.replace("_R1.fastq", "")
    #         count = sum(1 for line in open(os.path.join(fastq_dir, f))) / 4
    #         data.append([sample_name, count])
    command = f"seqkit stats {fastq_dir}/*_R1.fastq -T"
    result = subprocess.run(command, capture_output=True, text=True, shell=True)

    # Check if the command was successful
    if result.returncode != 0:
        print("Error executing command:", result.stderr)
        exit()

    # Split the captured output into lines
    lines = result.stdout.splitlines()

    # Use pandas to read the tab-separated output into a DataFrame
    df_rc = pd.read_csv(io.StringIO("\n".join(lines)), sep="\t")
    df_rc["sample"] = df_rc.file.apply(
        lambda x: os.path.basename(x).replace("_R1.fastq", "")
    )
    df_rc = df_rc[["sample", "num_seqs"]].set_index("sample")
    df_rc.columns = ["read_count"]

    df = pd.read_table(metadata_path, index_col="sample")
    if "read_count" in df.columns:
        df.drop("read_count", axis=1, inplace=True)
    # raise error if there are samples in the metadata that are not in the fastq
    # directory, and raise warning if there are samples in the fastq directory that
    # are not in the metadata (these samples will be ignored)
    extra_samples = df_rc.index.difference(df.index)
    if len(extra_samples) > 0:
        print(
            f"WARNING: These samples in the fastq directory are not present in the metadata: {extra_samples}"
        )
    extra_samples = df.index.difference(df_rc.index)
    if len(extra_samples) > 0:
        raise ValueError(
            f"These samples in the metadata are not present in the fastq directory: {extra_samples}"
        )
    else:
        print("Great! All samples in the metadata are present in the fastq directory.")
    df_res = pd.merge(df, df_rc, left_index=True, right_index=True, how="inner")
    df_res.to_csv(metadata_path, sep="\t")


def subsample(
    fastq_dir: str,
    metadata_path: str,
    output_dir: str,
    output_metadata_path: str,
    num_subsamples: int,
    mode: str,
    seed: int = 100,
    metadata_only: bool = False,
) -> None:
    # Load the metadata CSV file into a pandas DataFrame
    df_meta = pd.read_table(metadata_path, index_col="sample")

    if mode == "reference":
        df_meta["read_count_subsample"] = df_meta["reference"].map(
            df_meta["read_count"]
        )
    elif mode == "min":
        df_meta["read_count_subsample"] = df_meta.read_count.min() - 1
    else:
        raise ValueError(f"Invalid mode: {mode}")

    os.makedirs(output_dir, exist_ok=True)

    # Begin subsampling
    subsample_metadata = []
    for sample_name, row in tqdm(df_meta.iterrows(), total=len(df_meta)):
        read_count = min(row["read_count_subsample"], row["read_count"])
        num_subs = (
            num_subsamples if row["read_count_subsample"] < row["read_count"] else 1
        )
        del row["read_count_subsample"], row["read_count"]
        ps = []
        for iteration in range(1, num_subs + 1):
            # Subsample R1 and R1 and output to the subsampled directory
            subsample_metadata.append(
                {
                    "sample": f"{sample_name}-subsampled-{iteration}",
                    "original_sample": sample_name,
                    "read_count": read_count,
                }
                | row.to_dict()
            )
            if metadata_only:
                continue
            # Increment the seed value
            seed += 1
            for suffix in ["_R1.fastq", "_R2.fastq"]:
                fastq_file = os.path.join(fastq_dir, f"{sample_name}{suffix}")
                if not os.path.isfile(fastq_file):
                    print(f"WARNING: {fastq_file} does not exist. Skipping.")
                    continue
                # can also use seqtk, but will need to redirect stdout to file
                p = subprocess.Popen(  # redirect stdout and stderr to /dev/null
                    [
                        "vsearch",
                        "--fastx_subsample",
                        fastq_file,
                        "--sample_size",
                        str(read_count),
                        "-randseed",
                        str(seed),
                        "--fastqout",
                        os.path.join(
                            output_dir, f"{sample_name}-subsampled-{iteration}{suffix}"
                        ),
                    ],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
                ps.append(p)
        [p.wait() for p in ps]
    pd.DataFrame(subsample_metadata).to_csv(output_metadata_path, sep="\t", index=False)


def db_construct(
    input_fastx: str, output_fasta: str, num_threads: int, backend: str = "vsearch"
) -> str:
    """Output fasta is default to be modified from input fastq, but can be specified.
    For example, if input fastq file is /data/merged.filtered.fq.gz, then output fasta
    will be /data/merged.uniq.fa.gz, unless otherwise specified.
    """
    if output_fasta is None:
        output_dir, base_name = os.path.split(input_fastx)
        base_name_noext = base_name.split(".")[0]
        output_path = f"{output_dir}/{base_name_noext}.uniq.fa"
    else:
        output_path = output_fasta.rstrip(".gz")

    if backend == "usearch":
        subprocess.run(
            [
                "usearch11",
                "-fastx_uniques",
                input_fastx,
                "-fastaout",
                f"{output_dir}/{base_name_noext}.uniq.fa",
                "-sizeout",
            ],
            check=True,
        )
        subprocess.run(
            [
                "usearch11",
                "-sortbysize",
                f"{output_dir}/{base_name_noext}.uniq.fa",
                "-fastaout",
                f"{output_dir}/{base_name_noext}.sorted.fa",
                "-minsize",
                "2",
            ],
            check=True,
        )
        subprocess.run(
            [
                "usearch11",
                "-search_exact",
                input_fastx,
                "-db",
                f"{output_dir}/{base_name_noext}.sorted.fa",
                "-strand",
                "both",
                "-notmatched",
                f"{output_dir}/{base_name_noext}.unmatched.fa",
                "-t",
                str(num_threads),
            ],
            check=True,
        )
    elif backend == "vsearch":
        subprocess.run(
            [
                "vsearch",
                "--fastx_uniques",
                input_fastx,
                "--fastaout",
                output_path,
                "--sizeout",
            ]
        )
        subprocess.run(["pigz", "-f", output_path])
        output_path += ".gz"
        return output_path


def cluster_uparse(
    input_fastq: str, db_fasta: str, output_dir: str, num_threads: int
) -> None:
    if output_dir is None:
        output_dir = os.path.dirname(input_fastq)
    os.makedirs(output_dir, exist_ok=True)
    subprocess.run(
        [
            "usearch11",
            "-cluster_otus",
            db_fasta,
            "-otus",
            f"{output_dir}/uparse_otu.fa",
            "-relabel",
            "OTU",
        ],
        check=True,
    )
    subprocess.run(
        [
            "usearch10",
            "-usearch_global",
            input_fastq,
            "-db",
            f"{output_dir}/uparse_otu.fa",
            "-strand",
            "both",
            "-id",
            "0.97",
            "-otutabout",
            f"{output_dir}/uparse_otu.tsv",
            "-biomout",
            f"{output_dir}/uparse_otu.biom",
            "-threads",
            str(num_threads),
            "-sample_delim",
            ".",
        ],
        check=True,
    )
    subprocess.run(
        [
            "usearch11",
            "-calc_distmx",
            f"{output_dir}/uparse_otu.fa",
            "-tabbedout",
            f"{output_dir}/uparse_otu.dist",
            "-maxdist",
            "0.2",
            "-termdist",
            "0.3",
            "-threads",
            str(num_threads),
        ],
        check=True,
    )
    subprocess.run(
        [
            "usearch11",
            "-cluster_aggd",
            f"{output_dir}/uparse_otu.dist",
            "-treeout",
            f"{output_dir}/uparse_otu.tree",
            "-clusterout",
            f"{output_dir}/uparse_otu.clust",
            "-id",
            "0.8",
            "-linkage",
            "min",
        ],
        check=True,
    )


def unoise3(
    input_fastq: str | IO[str] | gzip.GzipFile,
    output_fasta: str | IO[str] | gzip.GzipFile | None,
    min_size: int,
    min_length: int,
    relabel_prefix: str | None = None,
    size_out: bool = False,
    num_threads: int = 16,
    stderr: bool = True,
) -> None | str:
    stderr = subprocess.DEVNULL if not stderr else None
    input_arg, stdin = _decide_io_arg(input_fastq)
    output_arg, stdout = _decide_io_arg(output_fasta)

    args_qc = [
        "vsearch",
        "--fastq_filter",
        input_arg,
        "--fastqout",
        "-",
        "--fastq_maxee",
        "0.5",
        "--fastq_minlen",
        str(min_length),
        "--fastq_maxns",
        "0",
        "--relabel",
        "filtered",
        # "--threads",
        # str(num_threads),
    ]
    # print_command(args_qc)
    qc_proc = subprocess.Popen(
        args_qc,
        stdin=stdin,
        stdout=subprocess.PIPE,
        stderr=stderr,
    )

    args_uniq = [
        "vsearch",
        "--fastx_uniques",
        "-",
        "--fastaout",
        "-",
        "--sizeout",
    ]
    # print_command(args_uniq)
    uniq_proc = subprocess.Popen(
        args_uniq, stdin=qc_proc.stdout, stdout=subprocess.PIPE, stderr=stderr
    )

    args_unoise3 = [
        "vsearch",
        "--cluster_unoise",
        "-",
        "--minsize",
        str(min_size),
        "--sizeout",
        "--centroids",
        "-",
        "--threads",
        str(num_threads),
    ]
    # print_command(args_unoise3)
    unoise3_proc = subprocess.Popen(
        args_unoise3, stdin=uniq_proc.stdout, stdout=subprocess.PIPE, stderr=stderr
    )

    args_uchime3 = [
        "vsearch",
        "--uchime3_denovo",
        "-",
        "--nonchimeras",
        output_arg,
        "--relabel",
        relabel_prefix or "ZOTU",
    ]
    if size_out:
        args_uchime3.append("--sizeout")
    uchime3_proc = subprocess.Popen(
        args_uchime3, stdin=unoise3_proc.stdout, stdout=stdout, stderr=stderr
    )

    qc_proc.stdout.close()
    uniq_proc.stdout.close()
    unoise3_proc.stdout.close()
    if stdin is subprocess.PIPE:
        qc_proc.communicate(input_fastq.encode())
    uchime3_out, _ = uchime3_proc.communicate()
    if stdout is subprocess.PIPE:
        return uchime3_out.decode()


def search_global(
    input_fastq: str | gzip.GzipFile | IO[str] | None,
    zotu_fasta: str | gzip.GzipFile | IO[str] | None,
    output_tsv: str | gzip.GzipFile | IO[str] | None,
    not_matched_fasta: str | gzip.GzipFile | IO[str] | None = None,
    id_: float = 0.97,
    num_threads: int = 8,
    stderr: bool = True,
) -> None | str:
    """
    When `output_csv` is None, output is in captured in the stdout of returned process.
    When `not_matched_fasta` is None, unmatched query sequences are not saved as fasta.
    """
    fastq_arg, stdin = _decide_io_arg(input_fastq)
    db_arg, stdin_db = _decide_io_arg(zotu_fasta)
    tsv_arg, stdout = _decide_io_arg(output_tsv)
    if not_matched_fasta is None:
        not_matched_arg, stdout_not_matched = None, None
    else:
        not_matched_arg, stdout_not_matched = _decide_io_arg(not_matched_fasta)
    if stdin is not None and stdin_db is not None:
        raise ValueError(
            "`input_fastq` and `zotu_fasta` cannot both be file-like objects."
        )
    if stdout is not None and stdout_not_matched is not None:
        raise ValueError(
            "`output_tsv` and `not_matched_fasta` cannot both be file-like objects."
        )
    stderr = subprocess.DEVNULL if not stderr else None
    # https://github.com/torognes/vsearch/issues/552
    # https://github.com/torognes/vsearch/issues/467
    # https://github.com/torognes/vsearch/issues/392
    args_search = [
        "vsearch",
        "--usearch_global",
        fastq_arg,
        "--db",
        db_arg,
        "--strand",
        "plus",
        "--id",
        str(id_),
        "--maxaccepts",
        "10",
        "--maxhits",
        "1",
        "--threads",
        str(num_threads),
    ]
    if tsv_arg == "-":
        args_search.extend(["--otutabout", "-"])
    elif tsv_arg.endswith(".tsv") or tsv_arg.endswith(".tsv.gz"):
        args_search.extend(["--otutabout", tsv_arg])
    elif output_tsv.endswith(".biom"):
        args_search.extend(["--biomout", tsv_arg])
    else:
        raise ValueError(
            f"Invalid output file extension {output_tsv}. Must be one of .tsv, .biom"
        )
    if not_matched_arg is not None:
        args_search.extend(["--notmatched", not_matched_arg])
    sg_proc = subprocess.Popen(
        args_search,
        stdin=stdin,
        stdout=stdout,
        stderr=stderr,
    )
    sg_proc_out, sg_proc_err = sg_proc.communicate(
        input_fastq.encode() if stdin is subprocess.PIPE else None
    )
    if stdout is subprocess.PIPE:
        return sg_proc_out.decode()


def search_global_add_unknown(
    input_fastq: str,
    zotu_fasta: str,
    output_tsv: str,
    id_: float,
    unknown_name: str,
    num_threads: int,
) -> None:
    if not os.path.isfile(zotu_fasta):
        raise ValueError(f"ZOTU fasta file (database) {zotu_fasta} does not exist")
    if not os.path.isfile(input_fastq):
        raise ValueError(f"Query fastq file {input_fastq} does not exist")
    if output_tsv is None:
        output_dir = os.path.dirname(input_fastq)
        output_tsv = os.path.join(output_dir, "unoise3_zotu.biom")
    else:
        output_dir = os.path.dirname(output_tsv)
        os.makedirs(output_dir, exist_ok=True)
    output_not_matched = os.path.join(output_dir, "notmatched.fa")

    search_global(
        input_fastq, zotu_fasta, output_tsv, output_not_matched, id_, num_threads
    )

    # add a ZOTU_UNKNOWN to the zotu table by counting the number of reads that do not
    # matched to anything in database
    zotu_table = read_table(output_tsv)
    counter = {i: 0 for i in zotu_table.columns}
    for record in SeqIO.parse(output_not_matched, "fasta"):
        counter[record.id.split("=")[1]] += 1
    zotu_table.loc[unknown_name] = pd.Series(counter)
    write_table(zotu_table, output_tsv)


def _decide_io_arg(arg) -> tuple:
    if arg is None:  # PIPE as IO
        return "-", subprocess.PIPE
    elif isinstance(arg, str):
        if "\n" in arg:  # string literal corresponding to file content
            return "-", subprocess.PIPE
        else:  # file path
            arg = os.path.abspath(arg)
            if (
                not os.path.isfile(arg)
                and not os.path.isdir(arg)
                and not os.path.isdir(os.path.dirname(arg))
            ):
                raise ValueError(
                    f"{arg} does not appear to be stirng liternal and is also not a file path."
                )
            return arg, None
    elif isinstance(arg, StringIO):
        return "-", subprocess.PIPE
    else:
        if (
            not isinstance(arg, gzip.GzipFile)
            and not isinstance(arg, IO[str])
            and not isinstance(arg, subprocess.Popen)
        ):
            raise ValueError(f"Invalid argument {arg} of type {type(arg)}")
        return "-", arg


def _workflow_one_sample(
    seqs_sample: list[str],
    min_size: int,
    prefix: str | None = None,
    search: bool = True,
) -> tuple[list[str], list[str], list[int]]:
    num_qs = len(seqs_sample)
    input_fastq = "".join(seqs_sample)
    db_fasta = tempfile.NamedTemporaryFile()
    unoise3(
        input_fastq,
        db_fasta.name,
        min_size=min_size,
        relabel_prefix=prefix,
        size_out=False,
        num_threads=1,
        stderr=False,
    )
    try:
        zotus_name, zotus_seq = zip(
            *(
                (record.id, str(record.seq))
                for record in SeqIO.parse(db_fasta.name, "fasta")
            )
        )
    except ValueError:
        zotus_name, zotus_seq = [], []
    else:
        zotus_name, zotus_seq = list(zotus_name), list(zotus_seq)
    if zotus_name and search:
        count_tsv = search_global(
            input_fastq, db_fasta.name, None, num_threads=1, stderr=False
        )
        counts = pd.read_table(StringIO(count_tsv), index_col=0).iloc[:, 0].tolist()
        num_unknown = num_qs - sum(counts)
        if num_unknown > 0:
            zotus_name.append("#UNKNOWN")
            zotus_seq.append("#UNKNOWN")
            counts.append(num_unknown)
    else:
        counts = []
        if num_qs > 0:
            zotus_name = ["#UNKNOWN"]
            zotus_seq = ["#UNKNOWN"]
            counts = [num_qs]
    db_fasta.close()
    return zotus_name, zotus_seq, counts


def workflow_per_sample(
    input_fastq: str,
    output_path: str,
    min_size: int,
    prefix: str | None = None,
    num_threads: int = 8,
    search: bool = True,
) -> None:
    """Split input fastq file by sample and run a separate unoise3 workflow for each
        sample, so that each sample has one set of ZOTUs and counts.

    We assume sequences in the input fastq file are named like `@sample=<sample1>` and
        sequences from the same sample form consecutive blocks.

    If output path ends with json, save results in a json file like:
        `{<sample>: {"zotus": ["ATCG", "GCTA", ...], "counts": [10, 2, ...]}, ...}`.
    If output path ends with fasta, save results in a fasta file file like:
        ```
        >ATCG
        ATCG
        ```
        where sequence names are just sequences themselves. Identical sequences appear
        only once.
    """
    from masato.utils import smart_open

    current_sample = None
    fastq_for_current_sample = []
    future2sample = {}
    samples = set()
    # sample2future = {}
    results = {}

    executor = (
        get_reusable_executor(max_workers=num_threads, kill_workers=True, reuse=False)
        if num_threads > 1
        else None
    )
    f: IO[str] = smart_open(input_fastq, "r")
    for entry in zip(f, f, f, f):
        # Parse the header line to extract the sample name
        header: str = entry[0]
        sample_name = header.split()[0].split("=")[1].strip()
        if sample_name != current_sample and fastq_for_current_sample:
            if sample_name in samples:
                raise ValueError(
                    "Sequences in fastq file are not grouped by sample, found "
                    f"{sample_name} both before {current_sample} and after"
                )
            if executor:
                future = executor.submit(
                    _workflow_one_sample,
                    fastq_for_current_sample,
                    min_size,
                    prefix,
                    search,
                )
                future2sample[future] = current_sample
                # sample2future[current_sample] = future
            else:
                results[current_sample] = _workflow_one_sample(
                    fastq_for_current_sample, min_size, prefix, search
                )
            fastq_for_current_sample = []  # Reset for the next sample
            samples.add(current_sample)
        current_sample = sample_name
        fastq_for_current_sample.append("".join(entry))
    f.close()
    # Don't forget to process the last sample
    if fastq_for_current_sample:
        if executor:
            future = executor.submit(
                _workflow_one_sample, fastq_for_current_sample, min_size, prefix, search
            )
            future2sample[future] = current_sample
            # sample2future[current_sample] = future
            for future in tqdm(as_completed(future2sample), total=len(future2sample)):
                # start the task and save the future object
                sample = future2sample[future]
                results[sample] = future.result()
            results = {s: results[s] for s in samples}
        else:
            results[current_sample] = _workflow_one_sample(
                fastq_for_current_sample, min_size, prefix, search
            )
    # Assuming each task's result could be aggregated into a final result dictionary
    # This step depends on how you implement the process_queue and task results handling
    # You would collect results from each completed task here
    # This could involve collecting returned values or reading from a shared resource
    ext = os.path.splitext(output_path)[1]
    if ext == ".json":
        results = {s: {"zotus": z, "counts": c} for s, (_, z, c) in results.items()}
        with open(output_path, "w") as outfile:
            json.dump(results, outfile, indent=4)
    elif ext in set([".fa", ".fasta", ".fna"]):
        prefix = prefix or "ZOTU"
        zotus = {}
        for zns, zss, counts in results.values():
            for zn, zs, count in zip(zns, zss, counts):
                if zn == "#UNKNOWN":
                    continue
                zotus[zs] = zotus.get(zs, 0) + int(count.split("=")[1])
        # sort by frequency
        records = [
            SeqRecord(Seq(z), id=f"{prefix}{i}", description=str(zotus[z]))
            for i, z in enumerate(sorted(zotus, key=zotus.get, reverse=True), 1)
        ]
        with open(output_path, "w") as outfile:
            SeqIO.write(records, outfile, "fasta")
    else:
        raise ValueError(
            f"Invalid output file extension {ext}. Must be one of .json, .fa, .fasta, .fna"
        )


def aggregate_samples(
    input_json: str, output_fasta: str, output_count: str, prefix: str | None = None
) -> None:
    """Take the json file generated by workflow_per_sample and aggregate the results into
    fasta file of ZOTUs and count file.

    ZOTUs are labeled as <prefix><idx> where idx is determined by total counts across
    all samples, starting from <prefix>1 (e.g. 16S_ZOTU1).
    """
    if prefix is None:
        prefix = "ZOTU"
    # Load JSON data from file
    with open(input_json) as json_file:
        data = json.load(json_file)

    # Extract unique sequences and calculate their total counts, keeping track of "#UNKNOWN" separately
    sequence_counts = defaultdict(int)
    key_sequence_counts = defaultdict(lambda: defaultdict(int))
    unknown_counts = defaultdict(int)
    for key, value in data.items():
        for sequence, count in zip(value["zotus"], value["counts"]):
            if sequence == "#UNKNOWN":
                unknown_counts[key] = count
            else:
                sequence_counts[sequence] += count
                key_sequence_counts[key][sequence] = count

    # Sort sequences by total counts from high to low
    sorted_sequences = sorted(sequence_counts.items(), key=lambda x: x[1], reverse=True)

    # Create a list of SeqRecord objects with names ZOTU1, ZOTU2, ..., and total counts as description
    seq_records = []
    sequence_to_zotu = {}
    for i, (sequence, count) in enumerate(sorted_sequences, start=1):
        zotu_name = f"{prefix}{i}"
        seq_records.append(
            SeqRecord(Seq(sequence), id=zotu_name, description=f"total_counts={count}")
        )
        sequence_to_zotu[sequence] = zotu_name

    # Write to FASTA file
    with open(output_fasta, "w") as fasta_file:
        SeqIO.write(seq_records, fasta_file, "fasta")

    # Create a DataFrame with ZOTU names as index and counts as values
    df_data = defaultdict(list)
    index = []

    for sequence, zotu_name in sequence_to_zotu.items():
        index.append(zotu_name)
        for key in key_sequence_counts:
            df_data[key].append(key_sequence_counts[key].get(sequence, 0))

    # Add #UNKNOWN to the DataFrame as the last row
    # if "#UNKNOWN" in unknown_count:
    if unknown_counts:
        index.append(f"{prefix}_UNKNOWN")
        for key in key_sequence_counts:
            df_data[key].append(unknown_counts[key])
        # df_data[f"{prefix}_UNKNOWN"] = df_data.pop("#UNKNOWN")

    df = pd.DataFrame(df_data, index=index)
    df.index.name = "#OTU_ID"

    # Save the DataFrame to a CSV file (optional)
    write_table(df, output_count)


def cluster_unoise3(
    uniq_fasta: str, out_fasta: str, minsize: int, relabel: str = "ZOTU"
):
    if os.path.isdir(out_fasta):
        raise ValueError(f"Output fasta file {out_fasta} is a directory")
    output_dir = os.path.dirname(out_fasta)
    os.makedirs(output_dir, exist_ok=True)
    # zotu_init = os.path.join(output_dir, f"_{out_fasta}")
    # vsearch does not have `--otutab`, `--otutab_stats`, `--calc_distmx`, and
    # `--cluster_aggd`. As suggested on the github issue
    #  (https://github.com/torognes/vsearch/issues/392), we use `--search_global`
    # to replace `--otutab`, but the other commands are not implemented.
    # Besides, in usearch chimera is removed as part of the clustering step, but
    # in vsearch chimera removal is a separate step.
    unoise3_proc = subprocess.Popen(
        [
            "vsearch",
            "--cluster_unoise",
            uniq_fasta,
            "--minsize",
            str(minsize),
            "--centroids",
            # f"{output_dir}/temp.fa",
            "-",
            # "--strand",
            # "both",
        ],
        stdout=subprocess.PIPE,
    )
    uchime3_proc = subprocess.Popen(
        [
            "vsearch",
            "--uchime3_denovo",
            "-",
            "--nonchimeras",
            out_fasta,
            "--relabel",
            relabel,
        ],
        stdin=unoise3_proc.stdout,
    )
    uchime3_proc.wait()
    # os.remove(zotu_init)


def _fq2fa(input_fastq: str, output_fasta: str) -> None:
    """Convert fastq (could be gzipped) to output fasta (gzipped) using seqtk"""
    try:
        # Constructing the seqtk command
        cmd = f"seqtk seq -a {input_fastq} | pigz -f > {output_fasta}"

        # Executing the command
        subprocess.run(cmd, shell=True, check=True)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Error executing command: {e.stderr}")


def _cluster_unoise3(
    input_fastq: str,
    db_fasta: str,
    minsize: int,
    id_: float,
    output_dir: str | None,
    num_threads: int,
    backend: str = "vsearch",
) -> None:
    if output_dir is None:
        output_dir = os.path.dirname(input_fastq)
    os.makedirs(output_dir, exist_ok=True)
    if backend == "usearch":
        subprocess.run(
            [
                "usearch11",
                "-unoise3",
                db_fasta,
                "-zotus",
                f"{output_dir}/unoise3_zotu.fa",
                "-tabbedout",
                f"{output_dir}/unoise3_zotu.tsv",
            ],
            check=True,
        )
        subprocess.run(
            [
                "usearch11",
                "-fastx_relabel",
                f"{output_dir}/unoise3_zotu.fa",
                "-prefix",
                "ZOTU",
                "-fastaout",
                f"{output_dir}/unoise3_zotu_relabel.fa",
                "-keep_annots",
            ],
            check=True,
        )
        os.rename(
            f"{output_dir}/unoise3_zotu_relabel.fa", f"{output_dir}/unoise3_zotu.fa"
        )
        subprocess.run(
            [
                "usearch10",
                "-otutab",
                input_fastq,
                "-db",
                f"{output_dir}/unoise3_zotu.fa",
                "-otutabout",
                f"{output_dir}/unoise3_zotu.tsv",
                "-biomout",
                f"{output_dir}/unoise3_zotu.biom",
                "-mapout",
                f"{output_dir}/unoise3_zotu.map",
                "-notmatched",
                f"{output_dir}/unoise3_zotu_notmatched.fa",
                "-dbmatched",
                f"{output_dir}/unoise3_zotu_dbmatched.fa",
                "-sizeout",
                "-strand",
                "both",
                "-threads",
                str(num_threads),
                "-sample_delim",
                ".",
            ],
            check=True,
        )
        subprocess.run(
            [
                "usearch10",
                "-otutab_stats",
                f"{output_dir}/unoise3_zotu.tsv",
                "-output",
                f"{output_dir}/unoise3_zotu_report.txt",
            ],
            check=True,
        )
        subprocess.run(
            [
                "usearch11",
                "-calc_distmx",
                f"{output_dir}/unoise3_zotu.fa",
                "-tabbedout",
                f"{output_dir}/unoise3_zotu.dist",
                "-maxdist",
                "0.2",
                "-termdist",
                "0.3",
            ],
            check=True,
        )
        subprocess.run(
            [
                "usearch11",
                "-cluster_aggd",
                f"{output_dir}/unoise3_zotu.dist",
                "-treeout",
                f"{output_dir}/unoise3_zotu.tree",
                "-clusterout",
                f"{output_dir}/unoise3_zotu.clust",
                "-id",
                "0.8",
                "-linkage",
                "min",
            ],
            check=True,
        )
    elif backend == "vsearch":
        # vsearch does not have `--otutab`, `--otutab_stats`, `--calc_distmx`, and
        # `--cluster_aggd`. As suggested on the github issue
        #  (https://github.com/torognes/vsearch/issues/392), we use `--search_global`
        # to replace `--otutab`, but the other commands are not implemented.
        # Besides, in usearch chimera is removed as part of the clustering step, but
        # in vsearch chimera removal is a separate step.
        subprocess.run(
            [
                "vsearch",
                "--cluster_unoise",
                db_fasta,
                "--minsize",
                str(minsize),
                "--centroids",
                f"{output_dir}/temp.fa",
                "--strand",
                "both",
            ],
            check=True,
        )
        subprocess.run(
            [
                "vsearch",
                "--uchime3_denovo",
                f"{output_dir}/temp.fa",
                "--nonchimeras",
                f"{output_dir}/unoise3_zotu.fa",
                "--relabel",
                "ZOTU",
            ],
        )
        os.remove(f"{output_dir}/temp.fa")
        # convert input fastq to fasta
        with open(f"{os.path.splitext(input_fastq)[0]}.fa", "w") as f:
            subprocess.run(["seqtk", "seq", "-A", input_fastq], stdout=f, check=True)
        subprocess.run(
            [
                "vsearch",
                "--usearch_global",
                f"{os.path.splitext(input_fastq)[0]}.fa",
                "--db",
                f"{output_dir}/unoise3_zotu.fa",
                "--strand",
                "plus",
                "--id",
                str(id_),
                "--otutabout",
                f"{output_dir}/unoise3_zotu.tsv",
                "--biomout",
                f"{output_dir}/unoise3_zotu.biom",
                "--alnout",
                f"{output_dir}/unoise3_zotu.aln",
                "--dbmatched",
                f"{output_dir}/unoise3_zotu_dbmatched.fa",
                "--dbnotmatched",
                f"{output_dir}/unoise3_zotu_dbnotmatched.fa",
                "--sizeout",
                "--threads",
                str(num_threads),
            ]
        )


def tax_nbc(
    input_fasta: str,
    db_fasta: str | None,
    output_path: str,
    num_threads: str,
    backend: str = "rdp_classifier",
) -> None:
    if backend == "usearch":
        if db_fasta is None:
            raise ValueError("db_fasta is required when using usearch")
        _ = subprocess.run(
            [
                "usearch11",
                "-nbc_tax",
                input_fasta,
                "-db",
                db_fasta,
                "-strand",
                "both",
                "-tabbedout",
                output_path,
                "-threads",
                str(num_threads),
            ],
            check=True,
        )
    elif backend == "rdp_classifier":
        if db_fasta is not None:
            print("WARNING: db_fasta is ignored when using rdp_classifier")
        subprocess.run(["classifier", "-o", output_path, input_fasta], check=True)
    else:
        raise ValueError(f"Invalid backend: {backend}")


def tax_sintax(
    input_fastq: str,
    db_udb: str,
    output_path: str,
    num_threads: str,
    backend: str = "vsearch",
) -> None:
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    if backend == "usearch":
        subprocess.run(
            [
                "usearch10",
                "-sintax",
                input_fastq,
                "-db",
                db_udb,
                "-strand",
                "both",
                "-sintax_cutoff",
                "0.8",
                "-threads",
                str(num_threads),
                "-tabbedout",
                output_path,
            ],
            check=True,
        )
    elif backend == "vsearch":
        subprocess.run(
            [
                "vsearch",
                "--sintax",
                input_fastq,
                "--db",
                db_udb,
                "--strand",
                "both",
                "--sintax_cutoff",
                "0.8",
                "--threads",
                str(num_threads),
                "--tabbedout",
                output_path,
            ],
            check=True,
        )


def main():
    parser = argparse.ArgumentParser(description="Perform QC on FASTQ files")
    subparsers = parser.add_subparsers(dest="subcommand")

    merge_pairs_parser = subparsers.add_parser(
        "merge_pairs", help="Merge paired-end reads"
    )
    merge_pairs_parser.add_argument(
        "-i", "--fastq_dir", help="Input directory containing paired-end fastq files"
    )
    merge_pairs_parser.add_argument(
        "-o", "--fastq_out", help="Output path for merged fastq file"
    )
    merge_pairs_parser.add_argument(
        "-t", "--num_threads", type=int, default=16, help="Number of threads to use"
    )

    db_construct_parser = subparsers.add_parser(
        "db_construct", help="Construct a database from a FASTA file"
    )
    db_construct_parser.add_argument(
        "-i", "--input_fastq", help="Input FASTA file to construct the database from"
    )
    db_construct_parser.add_argument(
        "-o",
        "--output_fasta",
        help="Output path for database FASTA file",
        type=str,
        default=None,
    )
    db_construct_parser.add_argument(
        "-t", "--num_threads", type=int, default=16, help="Number of threads to use"
    )

    cluster_uparse_parser = subparsers.add_parser(
        "cluster_uparse", help="Cluster reads using UPARSE"
    )
    cluster_uparse_parser.add_argument(
        "-i", "--input_fastq", help="Input FASTQ file to cluster", type=str
    )
    cluster_uparse_parser.add_argument(
        "-d", "--db_fasta", help="Input FASTA file to cluster against", type=str
    )
    cluster_uparse_parser.add_argument(
        "-o",
        "--output_dir",
        help="Output directory for UPARSE results",
        type=str,
        default=None,
    )
    cluster_uparse_parser.add_argument(
        "-t", "--num_threads", type=int, default=16, help="Number of threads to use"
    )

    cluster_unoise3_parser = subparsers.add_parser(
        "cluster_unoise3", help="Cluster reads using UNOISE3"
    )
    cluster_unoise3_parser.add_argument(
        "-i", "--uniq_fasta", help="Input FASTQ file to cluster", type=str
    )
    cluster_unoise3_parser.add_argument(
        "-m", "--minsize", type=int, default=8, help="Minimum cluster size"
    )
    cluster_unoise3_parser.add_argument(
        "-o", "--out_fasta", help="Output path for ZOTU fasta file"
    )
    unoise3_parser = subparsers.add_parser(
        "unoise3", help="Cluster reads using UNOISE3"
    )
    unoise3_parser.add_argument(
        "-i",
        "--input_fastq",
        help="Input FASTQ file to cluster",
        type=str,
        required=True,
    )
    unoise3_parser.add_argument(
        "-o",
        "--output_fasta",
        help="Output path for ZOTU fasta file",
        type=str,
        required=True,
    )
    unoise3_parser.add_argument(
        "-m", "--minsize", type=int, default=8, help="Minimum cluster size"
    )
    unoise3_parser.add_argument(
        "-l",
        "--min_length",
        type=int,
        default=100,
        help="Minimum length of sequences to cluster",
    )
    unoise3_parser.add_argument(
        "-l", "--relabel_prefix", type=str, default=None, help="Prefix for ZOTU labels"
    )
    unoise3_parser.add_argument(
        "-t", "--num_threads", type=int, default=4, help="Number of threads to use"
    )
    search_global_parser = subparsers.add_parser(
        "search_global", help="Search reads against a database"
    )
    search_global_parser.add_argument(
        "-i", "--input_fastq", help="Input FASTQ file to search", type=str
    )
    search_global_parser.add_argument(
        "-d", "--zotu_fasta", help="Input FASTA file to search against", type=str
    )
    search_global_parser.add_argument(
        "-o",
        "--output_tsv",
        help="Output path for ZOTU tsv file",
        type=str,
        default=None,
    )
    search_global_parser.add_argument(
        "--id", type=float, default=0.97, help="Minimum cluster identity"
    )
    search_global_parser.add_argument(
        "--unknown_name",
        type=str,
        default="ZOTU_UNKNOWN",
        help="Name for unknown ZOTUs",
    )
    search_global_parser.add_argument(
        "-t", "--num_threads", type=int, default=8, help="Number of threads to use"
    )
    workflow_per_sample_parser = subparsers.add_parser(
        "workflow_per_sample", help="Run UNOISE3 workflow per sample"
    )
    workflow_per_sample_parser.add_argument(
        "-i", "--input_fastq", help="Input FASTQ file to cluster", type=str
    )
    workflow_per_sample_parser.add_argument(
        "-o", "--output_path", help="Output path", type=str
    )
    workflow_per_sample_parser.add_argument(
        "-m", "--min_size", type=int, default=8, help="Minimum cluster size"
    )
    workflow_per_sample_parser.add_argument(
        "-l",
        "--relabel_prefix",
        type=str,
        default="ZOTU",
        help="Prefix for ZOTU labels",
    )
    workflow_per_sample_parser.add_argument(
        "-t", "--num_threads", type=int, default=16, help="Number of threads to use"
    )
    workflow_per_sample_parser.add_argument(
        "--search",
        action="store_true",
        help="Search reads against ZOTU database",
    )
    aggregate_samples_parser = subparsers.add_parser(
        "aggregate_samples", help="Aggregate ZOTU counts across samples"
    )
    aggregate_samples_parser.add_argument(
        "-i", "--input_json", help="Input JSON file with ZOTU counts", type=str
    )
    aggregate_samples_parser.add_argument(
        "-of", "--output_fasta", help="Output path for aggregated ZOTU fasta", type=str
    )
    aggregate_samples_parser.add_argument(
        "-oc", "--output_count", help="Output path for aggregated ZOTU counts", type=str
    )
    aggregate_samples_parser.add_argument(
        "-p",
        "--prefix",
        help="Prefix for ZOTU labels",
        type=str,
        default=None,
    )

    tax_nbc_parser = subparsers.add_parser(
        "tax_nbc", help="Taxonomy classification using NBC"
    )
    tax_nbc_parser.add_argument(
        "-i", "--input_fasta", help="Input FASTA file to classify", type=str
    )
    tax_nbc_parser.add_argument(
        "-d", "--db_fasta", help="Input FASTA file to classify against", type=str
    )
    tax_nbc_parser.add_argument(
        "-o", "--output_path", help="Output path for NBC results", type=str
    )
    tax_nbc_parser.add_argument(
        "-t", "--num_threads", type=int, default=16, help="Number of threads to use"
    )

    tax_sintax_parser = subparsers.add_parser(
        "tax_sintax", help="Taxonomy classification using SINTAX"
    )
    tax_sintax_parser.add_argument(
        "-i", "--input_fastq", help="Input FASTQ file to classify", type=str
    )
    tax_sintax_parser.add_argument(
        "-d", "--db_udb", help="Input udb file to classify against", type=str
    )
    tax_sintax_parser.add_argument(
        "-o", "--output_path", help="Output path for SINTAX results", type=str
    )
    tax_sintax_parser.add_argument(
        "-t", "--num_threads", type=int, default=16, help="Number of threads to use"
    )

    args = parser.parse_args()

    if args.subcommand == "db_construct":
        fastx_post_qc = qc(args.input_fastq, None, args.num_threads)
        db_construct(fastx_post_qc, args.output_fasta, args.num_threads)
    elif args.subcommand == "cluster_uparse":
        cluster_uparse(
            args.input_fastq, args.db_fasta, args.output_dir, args.num_threads
        )
    elif args.subcommand == "cluster_unoise3":
        cluster_unoise3(args.uniq_fasta, args.minsize, args.out_fasta)
    elif args.subcommand == "unoise3":
        unoise3(
            args.input_fastq,
            args.output_fasta,
            min_size=args.minsize,
            min_length=args.min_length,
            relabel_prefix=args.relabel_prefix,
            num_threads=args.num_threads,
        )
    elif args.subcommand == "search_global":
        search_global_add_unknown(
            args.input_fastq,
            args.zotu_fasta,
            args.output_tsv,
            id_=args.id,
            unknown_name=args.unknown_name,
            num_threads=args.num_threads,
        )
    elif args.subcommand == "workflow_per_sample":
        workflow_per_sample(
            args.input_fastq,
            args.output_path,
            min_size=args.min_size,
            prefix=args.relabel_prefix,
            num_threads=args.num_threads,
            search=args.search,
        )
    elif args.subcommand == "aggregate_samples":
        aggregate_samples(
            args.input_json, args.output_fasta, args.output_count, prefix=args.prefix
        )
    elif args.subcommand == "tax_nbc":
        tax_nbc(args.input_fasta, args.db_fasta, args.output_path, args.num_threads)
    elif args.subcommand == "tax_sintax":
        tax_sintax(args.input_fastq, args.db_udb, args.output_path, args.num_threads)
    else:
        raise ValueError(f"Invalid subcommand: {args.subcommand}")


if __name__ == "__main__":
    main()
