"""
zipstrain.utils
========================
This module contains the command-line interface (CLI) implementation for the zipstrain application.
"""
import click as click
import zipstrain.utils as ut
import zipstrain.compare as cp
import zipstrain.profile as pf
import zipstrain.task_manager as tm
import zipstrain.database as db
import zipstrain.build_db as bdb
import polars as pl
import pathlib


@click.group()
def cli():
    """ZipStrain CLI"""
    pass

@cli.group()
def utilities():
    """The commands in this group are related to various utility functions that mainly prepare input files for profiling and comparison."""
    pass

@utilities.command("build-null-model")
@click.option('--error-rate', '-e', default=0.001, help="Error rate for the sequencing technology.")
@click.option('--max-total-reads', '-m', default=10000, help="Maximum coverage to consider for a base")
@click.option('--p-threshold', '-p', default=0.05, help="Significance threshold for the Poisson distribution.")
@click.option('--output-file', '-o', required=True, help="Path to save the output Parquet file.")
@click.option('--model-type', '-t', default="poisson", type=click.Choice(['poisson']), help="Type of null model to build.")
def build_null_model(error_rate, max_total_reads, p_threshold, output_file, model_type):
    """
    Build a null model for sequencing errors based on the Poisson distribution.

    Args:
    error_rate (float): Error rate for the sequencing technology.
    max_total_reads (int): Maximum total reads to consider.
    p_threshold (float): Significance threshold for the Poisson distribution.
    """
    if model_type == "poisson":
        df_thresh = ut.build_null_poisson(error_rate, max_total_reads, p_threshold)
    else:
        raise ValueError(f"Unsupported model type: {model_type}")
    df_thresh = pl.DataFrame(df_thresh, schema=["cov", "max_error_count"])
    df_thresh.write_parquet(output_file)

@utilities.command("merge_parquet")
@click.option('--input-dir', '-i', required=True, help="Directory containing Parquet files to merge.") 
@click.option('--output-file', '-o', required=True, help="Path to save the merged Parquet file.")
def merge_parquet(input_dir, output_file):
    """
    Merge multiple Parquet files in a directory into a single Parquet file, adding gene information.

    Args:
    input_dir (str): Directory containing Parquet files to merge.
    output_file (str): Path to save the merged Parquet file.
    """
    input_path = pathlib.Path(input_dir)
    parquet_files = list(input_path.glob("*.parquet"))
    if not parquet_files:
        raise ValueError(f"No Parquet files found in directory: {input_dir}")
    
    mpileup_df = pl.concat([pl.scan_parquet(pf) for pf in parquet_files])
    mpileup_df.sink_parquet(pathlib.Path(output_file), compression='zstd')


@utilities.command("process_mpileup")
@click.option('--batch-size', '-s', default=10000, help="Buffer size for processing stdin from samtools.")
@click.option('--output-file', '-o', required=True, help="Location to save the output Parquet file.")
def process_mpileup(batch_size, output_file):
    """
    Process mpileup files and save the results in a Parquet file.

    Args:
    gene_range_table_loc (str): Path to the gene range table in TSV format.
    batch_bed (str): Path to the batch BED file.
    output_file (str): Path to save the output Parquet file.
    """
    ut.process_mpileup_function(batch_size, output_file)
    
@utilities.command("make_bed")
@click.option('--db-fasta-dir', '-d', required=True, help="Path to the database in fasta format.")
@click.option('--max-scaffold-length', '-m', default=500000, help="Maximum scaffold length to split into multiple entries.")
@click.option('--output-file', '-o', required=True, help="Path to save the output BED file.")
def make_bed(db_fasta_dir, max_scaffold_length, output_file):
    """
    Create a BED file from the database in fasta format.

    Args:
    db_fasta_dir (str): Path to the fasta file.
    max_scaffold_length (int): Splits scaffolds longer than this into multiple entries of length <= max_scaffold_length.
    output_file (str): Path to save the output BED file.
    """
    bed_df = ut.make_the_bed(db_fasta_dir, max_scaffold_length)
    bed_df.write_csv(output_file, separator='\t', include_header=False)

@utilities.command("get_genome_lengths")
@click.option('--stb-file', '-s', required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--bed-file', '-b', required=True, help="Path to the BED file.")
@click.option('--output-file', '-o', required=True, help="Path to save the output Parquet file.")
def get_genome_lengths(stb_file, bed_file, output_file):
    """
    Extract the genome length information from the scaffold-to-genome mapping table.

    Args:
    stb_file (str): Path to the scaffold-to-genome mapping file.
    bed_file (str): Path to the BED file containing genomic regions.
    output_file (str): Path to save the output Parquet file.
    """
    stb = pl.scan_csv(stb_file, separator='\t',has_header=False).with_columns(
        pl.col("column_1").alias("scaffold"),
        pl.col("column_2").alias("genome")
    )
    
    bed_table = pl.scan_csv(bed_file, separator='\t',has_header=False).with_columns(
        pl.col("column_1").alias("scaffold"),
        pl.col("column_2").cast(pl.Int64).alias("start"),
        pl.col("column_3").cast(pl.Int64).alias("end")
    ).select(["scaffold", "start", "end"])
    genome_length = ut.extract_genome_length(stb, bed_table)
    genome_length.sink_parquet(output_file, compression='zstd')

@utilities.command("genome_breadth_matrix")
@click.option('--profile', '-p', type=str, required=True, help="Path to the profile Parquet file.")
@click.option('--genome-length', '-g', type=str, required=True, help="Path to the genome length Parquet file.")
@click.option('--stb', '-s', type=str, required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--min-cov', '-c', default=1, help="Minimum coverage to consider a position.")
@click.option('--output-file', '-o', required=True, help="Path to save the output Parquet file.")
def genome_breadth_matrix(profile, genome_length, stb, min_cov, output_file):
    """
    Generate a genome breadth matrix from the given profiles and scaffold-to-genome mapping.

    Args:
    profiles (list): List of profiles in the format 'name:path_to_profile'.
    genome_length (str): Path to the genome length Parquet file.
    stb_file (str): Path to the scaffold-to-genome mapping file.
    min_cov (int): Minimum coverage to consider a position.
    output_file (str): Path to save the output Parquet file.
    """
    genome_length = pl.scan_parquet(genome_length)
    stb = pl.scan_csv(stb, separator='\t', has_header=False).select(
        pl.col("column_1").alias("scaffold"),
        pl.col("column_2").alias("genome")
    )
    profile_dir= pathlib.Path(profile)
    profile = pl.scan_parquet(profile)
    lf=ut.get_genome_breadth_matrix(profile,profile_dir.name, genome_length,stb, min_cov)
    lf.sink_parquet(output_file, compression='zstd')

@utilities.command("collect_breadth_tables")
@click.option('--breadth-tables-dir', '-d', required=True, help="Directory containing breadth tables in Parquet format.")
@click.option('--extension', '-e', default='parquet', help="File extension of the breadth tables.")
@click.option('--output-file', '-o', required=True, help="Path to save the collected breadth tables.")
def collect_breadth(breadth_tables_dir, extension, output_file):
    """
    Collect multiple genome breadth tables into a single LazyFrame.

    Args:
    breadth_tables_dir (str): Directory containing breadth tables in Parquet format.
    extension (str): File extension of the breadth tables.
    output_file (str): Path to save the collected breadth tables.
    """
    breadth_tables = list(pathlib.Path(breadth_tables_dir).glob(f"*.{extension}"))
    if not breadth_tables:
        raise ValueError(f"No breadth tables found in directory: {breadth_tables_dir}")

    lazy_frames = [pl.scan_parquet(str(pf)) for pf in breadth_tables]
    combined_lf = ut.collect_breadth_tables(lazy_frames)
    combined_lf.sink_parquet(output_file, compression='zstd')
    
@utilities.command("strain_heterogeneity")
@click.option('--profile-file', '-p', required=True, help="Path to the profile Parquet file.")
@click.option('--stb-file', '-s', required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--min-cov', '-c', default=5, help="Minimum coverage to consider a position.")
@click.option('--freq-threshold', '-f', default=0.8, help="Frequency threshold to define dominant nucleotide.")
@click.option('--output-file', '-o', required=True, help="Path to save the output Parquet file.")
def strain_heterogeneity(profile_file, stb_file, min_cov, freq_threshold, output_file):
    """
    Calculate strain heterogeneity for each genome based on nucleotide frequencies.

    Args:
    profile_file (str): Path to the profile Parquet file.
    stb_file (str): Path to the scaffold-to-genome mapping file.
    min_cov (int): Minimum coverage to consider a position.
    freq_threshold (float): Frequency threshold to define dominant nucleotide.
    output_file (str): Path to save the output Parquet file.
    """
    profile = pl.scan_parquet(profile_file)
    stb = pl.scan_csv(stb_file, separator="\t", has_header=False).with_columns(
        pl.col("column_1").alias("scaffold"),
        pl.col("column_2").alias("genome")
    ).select(["scaffold", "genome"])
    
    het_profile = pf.get_strain_hetrogeneity(profile, stb, min_cov=min_cov, freq_threshold=freq_threshold)
    het_profile.sink_parquet(output_file, compression='zstd')

@utilities.command("build-profile-db")
@click.option('--profile-db-csv', '-p', required=True, help="Path to the profile database CSV file.")
@click.option('--output-file', '-o', required=True, help="Path to save the output Parquet file.")
def build_profile_db(profile_db_csv, output_file):
    """
    Build a profile database from the given CSV file.

    Args:
    profile_db_csv (str): Path to the profile database CSV file.
    """
    profile_db = db.ProfileDatabase.from_csv(pathlib.Path(profile_db_csv))
    profile_db.save_as_new_database(pathlib.Path(output_file))


@utilities.command("build-genome-db")
@click.option('--tool', '-t', required=True, type=click.Choice(sorted(bdb.ADAPTERS.keys())), help="Abundance tool format to parse (for example, sylph).")
@click.option('--abundance-table', '-a', required=True, help="Path to abundance table (csv/tsv/parquet).")
@click.option('--db-file', '-d', required=True, help="Path to local genome DB parquet file.")
@click.option('--genomes-dir', '-g', required=True, help="Directory to store downloaded genomes.")
@click.option('--download/--no-download', default=True, help="Whether to download missing genomes.")
@click.option('--overwrite', is_flag=True, default=False, help="Redownload genomes even if existing files are found.")
@click.option('--report-file', '-r', default=None, help="Optional path to save a download report csv.")
def build_genome_db(tool, abundance_table, db_file, genomes_dir, download, overwrite, report_file):
    """
    Build or update a local genome database from an abundance table.

    The parser is selected by `--tool` and currently supports Sylph.
    """
    local_db, extracted, report = bdb.build_local_genome_db(
        tool_name=tool,
        abundance_table=pathlib.Path(abundance_table),
        db_path=pathlib.Path(db_file),
        genomes_dir=pathlib.Path(genomes_dir),
        download=download,
        overwrite=overwrite,
    )
    _ = local_db
    if report_file is not None:
        bdb.write_report_csv(report, pathlib.Path(report_file))
    click.echo(f"Extracted {extracted.height} genome accessions.")
    if download:
        downloaded = report.filter(pl.col("status") == "downloaded").height if report.height else 0
        failed = report.filter(pl.col("status") == "failed").height if report.height else 0
        click.echo(f"Downloaded: {downloaded}; Failed: {failed}.")


@utilities.command("build-genome-comparison-config")
@click.option('--profile-db', '-p', required=True, help="Path to the profile database Parquet file.")
@click.option('--gene-db-id', '-g', required=True, help="Gene database ID.")
@click.option('--reference-genome-id', '-r', required=True, help="Reference fasta ID.")
@click.option('--scope', '-s', default="all", help="Genome scope for comparison.")
@click.option('--min-cov', '-c', default=5, help="Minimum coverage to consider a position.")
@click.option('--min-gene-compare-len', '-l', default=200, help="Minimum gene length to consider for comparison.")
@click.option('--stb-file-loc', '-t', required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--current-comp-table', '-a', default=None, help="Path to the current comparison table in Parquet format.")
@click.option('--output-file', '-o', required=True, help="Path to save the output configuration JSON file.")
def build_genome_comparison_config(profile_db, gene_db_id, reference_genome_id, scope, min_cov, min_gene_compare_len, stb_file_loc, current_comp_table, output_file):
    """
    Build a comparison configuration JSON file from the given parameters.

    Args:
    profile_db (str): Path to the profile database Parquet file.
    gene_db_id (str): Gene database ID.
    reference_genome_id (str): Reference genome ID.
    scope (str): Genome scope for comparison.
    min_cov (int): Minimum coverage to consider a position.
    min_gene_compare_len (int): Minimum gene length to consider for comparison.
    stb_file_loc (str): Path to the scaffold-to-genome mapping file.
    current_comp_table (str): Path to the current comparison table in Parquet format.
    output_file (str): Path to save the output configuration JSON file.
    """
    conf_obj=db.GenomeComparisonConfig(
        gene_db_id=gene_db_id,
        reference_id=reference_genome_id,
        scope=scope,
        min_cov=min_cov,
        min_gene_compare_len=min_gene_compare_len,
        stb_file_loc=stb_file_loc,
    )
    profile_db=db.ProfileDatabase(profile_db)
    comp_obj=db.GenomeComparisonDatabase(
        profile_db=profile_db,
        config=conf_obj,
        comp_db_loc=current_comp_table
    )
    comp_obj.dump_obj(pathlib.Path(output_file))


@utilities.command("build-gene-comparison-config")
@click.option('--profile-db', '-p', required=True, help="Path to the profile database Parquet file.")
@click.option('--gene-db-id', '-g', required=True, help="Gene database ID.")
@click.option('--reference-genome-id', '-r', required=True, help="Reference fasta ID.")
@click.option('--scope', '-s', default="all:all", help="Genome scope for comparison.")
@click.option('--min-cov', '-c', default=5, help="Minimum coverage to consider a position.")
@click.option('--min-gene-compare-len', '-l', default=200, help="Minimum gene length to consider for comparison.")
@click.option('--stb-file-loc', '-t', required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--current-comp-table', '-a', default=None, help="Path to the current comparison table in Parquet format.")
@click.option('--output-file', '-o', required=True, help="Path to save the output configuration JSON file.")
def build_gene_comparison_config(profile_db, gene_db_id, reference_genome_id, scope, min_cov, min_gene_compare_len, stb_file_loc, current_comp_table, output_file):
    """
    Build a gene comparison configuration JSON file from the given parameters.

    Args:
    profile_db (str): Path to the profile database Parquet file.
    gene_db_id (str): Gene database ID.
    reference_genome_id (str): Reference genome ID.
    scope (str): Genome scope for comparison.
    min_cov (int): Minimum coverage to consider a position.
    min_gene_compare_len (int): Minimum gene length to consider for comparison.
    stb_file_loc (str): Path to the scaffold-to-genome mapping file.
    current_comp_table (str): Path to the current comparison table in Parquet format.
    output_file (str): Path to save the output configuration JSON file.
    """
    conf_obj=db.GeneComparisonConfig(
        gene_db_id=gene_db_id,
        reference_genome_id=reference_genome_id,
        scope=scope,
        min_cov=min_cov,
        min_gene_compare_len=min_gene_compare_len,
        stb_file_loc=stb_file_loc,
    )
    profile_db_obj=db.ProfileDatabase(profile_db)
    
    comp_obj=db.GeneComparisonDatabase(
        profile_db=profile_db_obj,
        config=conf_obj,
        comp_db_loc=current_comp_table
    )
    comp_obj.dump_obj(pathlib.Path(output_file))


@utilities.command("to-complete-table")
@click.option("--genome-comparison-object", "-g", required=True, help="Path to the genome comparison object in json format.")
@click.option("--output-file", "-o", required=True, help="Path to save the completed pairs csv file.")
def to_complete_table(genome_comparison_object, output_file):
    """
    Generate a table of completed genome comparison pairs and save it to a csv file.

    Parameters:
    genome_comparison_object (str): Path to the genome comparison object in json format.
    output_file (str): Path to save the completed pairs Parquet file.
    """
    genome_comp_db=db.GenomeComparisonDatabase.load_obj(pathlib.Path(genome_comparison_object))
    completed_pairs=genome_comp_db.to_complete_input_table()
    completed_pairs.sink_csv(pathlib.Path(output_file), engine="streaming")

@utilities.command("presence-profile")
@click.option('--profile-file', '-p', required=True, help="Path to the profile Parquet file.")
@click.option('--stb-file', '-s', required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--bed-file', '-b', required=True, help="Path to the BED file.")
@click.option('--read-loc-file', '-r', required=True, help="Path to the read location table.")
@click.option('--min-cov-fug', '-c', default=0.1, help="Minimum coverage to use fug.")
@click.option('--fug-threshold', '-f', default=2, help="FUG threshold.")
@click.option('--ber', '-e', default=0.5, help="Minimum ratio of breadth over expected breadth to consider presence.")
@click.option('--output-file', '-o', required=True, help="Path to save the output Parquet file.")
def presence_profile(profile_file, stb_file, bed_file, read_loc_file, min_cov_fug, fug_threshold, ber, output_file):
    """
    Generate a presence profile for genomes based on the given profile and read location data.

    Args:
    profile_file (str): Path to the profile Parquet file.
    stb_file (str): Path to the scaffold-to-genome mapping file.
    bed_file (str): Path to the BED file.
    read_loc_file (str): Path to the read location table.
    min_cov_fug (float): Minimum coverage to use fug.
    fug_threshold (float): FUG threshold.
    ber (float): Minimum ratio of breadth over expected breadth to consider presence.
    output_file (str): Path to save the output Parquet file.
    """
    profile = pl.scan_parquet(profile_file)
    stb = pl.scan_csv(stb_file, separator="\t", has_header=False).with_columns(
        pl.col("column_1").alias("scaffold"),
        pl.col("column_2").alias("genome")
    ).select(["scaffold", "genome"])
    bed = pl.scan_csv(bed_file, separator="\t", has_header=False).with_columns(
        pl.col("column_1").alias("scaffold"),
        pl.col("column_2").cast(pl.Int64).alias("start"),
        pl.col("column_3").cast(pl.Int64).alias("end")
    ).select(["scaffold", "start", "end"])
    read_loc_table = pl.scan_parquet(read_loc_file).rename({
        "chrom":"scaffold",
        "pos":"loc"
    })
    presence_df = ut.get_genome_stats(
        profile=profile,
        stb=stb,
        bed=bed,
        read_loc_table=read_loc_table,
        min_cov_use_fug=min_cov_fug,
        fug=fug_threshold,
        ber=ber
    )
    presence_df.sink_parquet(output_file, compression='zstd')

@utilities.command("process-read-locs")
@click.option("--output-file", "-o", required=True, help="Path to save the processed read locations Parquet file.")
def process_read_locs(output_file):
    """
    Process read locations and save them to a Parquet file.

    Args:
    output_file (str): Path to save the output Parquet file.
    """
    ut.process_read_location(output_file=pathlib.Path(output_file))

@cli.group()
def gene_tools():
    """Holds anything related to gene analysis."""
    pass

@utilities.command("generate_stb")
@click.option('--genomes-dir-file', '-g', required=True, help="Path to the genomes directory file. A text file with each line containing a genome fasta file path.")
@click.option('--output-file', '-o', required=True, help="Path to save the output scaffold-to-genome mapping file.")
@click.option('--extension', '-e', default=".fasta", help="File extension of the genome fasta files.")
def generate_stb(genomes_dir_file, output_file, extension):
    """
    Generate a scaffold-to-genome mapping file from the given genomes directory file.

    Args:
    genomes_dir_file (str): Path to the genomes directory file.
    output_file (str): Path to save the output scaffold-to-genome mapping file.
    extension (str): File extension of the genome fasta files.
    """
    with open(output_file, 'w') as out_f:
        for genome in pathlib.Path(genomes_dir_file).glob(f"*{extension}"):
            genome_name = genome.stem
            with open(genome, 'r') as gf:
                for line in gf:
                    if line.startswith('>'):
                        scaffold_name = line[1:].strip().split()[0]
                        out_f.write(f"{scaffold_name}\t{genome_name}\n")
    
        
    


@gene_tools.command("gene-range-table")
@click.option('--gene-file', '-g', required=True, help="location of gene file. Prodigal's nucleotide fasta output")
@click.option('--output-file', '-o', required=True, help="location to save output tsv file")
def get_gene_range_table(gene_file, output_file):
    """
    Main function to build and save the gene location table.

    Args:
    gene_file (str): Path to the gene FASTA file.
    output_file (str): Path to save the output TSV file.
    """
    gene_locs=pf.build_gene_range_table(pathlib.Path(gene_file))
    gene_locs.sink_csv(pathlib.Path(output_file), separator="\t", include_header=False)


@gene_tools.command("gene-loc-table")
@click.option('--gene-file', '-g', required=True, help="location of gene file. Prodigal's nucleotide fasta output")
@click.option('--scaffold-list', '-s', required=True, help="location of scaffold list. A text file with each line containing a scaffold name.")
@click.option('--output-file', '-o', required=True, help="location to save output parquet file")
def get_gene_loc_table(gene_file, scaffold_list, output_file):
    """
    Main function to build and save the gene location table.

    Args:
    gene_file (str): Path to the gene FASTA file.
    scaffold_list (str): Path to the scaffold list file.
    output_file (str): Path to save the output Parquet file.
    """
    scaffolds=set(pl.read_csv(pathlib.Path(scaffold_list), has_header=False,separator="\t").select(pl.col("column_1")).to_series().to_list())
    gene_locs=pf.build_gene_loc_table(pathlib.Path(gene_file), scaffolds)
    gene_locs.write_parquet(pathlib.Path(output_file))



@cli.group()
def compare():
    """The commands in this group are related to comparing profiled samples."""
    pass

@compare.command("single_compare_genome")
@click.option('--mpileup-contig-1', '-m1', required=True, help="Path to the first mpileup file.")
@click.option('--mpileup-contig-2', '-m2', required=True, help="Path to the second mpileup file.")
@click.option('--stb-file', '-s', required=True, help="Path to the scaffold to genome mapping file.")
@click.option('--min-cov', '-c', default=5, help="Minimum coverage to consider a position.")
@click.option('--min-gene-compare-len', '-l', default=100, help="Minimum gene length to consider for comparison.")
@click.option('--output-file', '-o', required=True, help="Path to save the parquet file.")
@click.option('--genome', '-g', default="all", help="If provided, do the comparison only for the specified genome.")
@click.option('--ani-method', '-a', default="popani", help="ANI calculation method to use (e.g., 'popani', 'conani', 'cosani_0.4').")
@click.option('--engine', type=click.Choice(["polars", "duckdb"]), default="polars", show_default=True, help="Execution engine for compare.")
@click.option('--duckdb-memory-limit', default=None, help="DuckDB memory limit (e.g., 2GB, 1024MB).")
@click.option('--duckdb-temp-directory', default=None, help="Directory DuckDB can use for spill files.")
@click.option('--duckdb-threads', type=int, default=None, help="Number of DuckDB worker threads.")
def single_compare_genome(mpileup_contig_1, mpileup_contig_2, stb_file, min_cov, min_gene_compare_len, output_file, genome, ani_method, engine, duckdb_memory_limit, duckdb_temp_directory, duckdb_threads):
    """
    Main function to compare two mpileup files and calculate genome and gene statistics.
    
    Args:
    mpileup_contig_1 (str): Path to the first mpileup file.
    mpileup_contig_2 (str): Path to the second mpileup file.
    min_cov (int): Minimum coverage to consider a position.
    min_gene_compare_len (int): Minimum gene length to consider for comparison.
    output_file (str): Path to save the parquet file.
    genome (str): If provided, do the comparison only for the specified genome.
    stb_file (str): Path to the scaffold to genome mapping file.
    """
    mpile_contig_1_name = pathlib.Path(mpileup_contig_1).name
    mpile_contig_2_name = pathlib.Path(mpileup_contig_2).name

    mpile_1_for_compare = mpileup_contig_1
    mpile_2_for_compare = mpileup_contig_2
    if genome != "all":
        mpile_1_for_compare, mpile_2_for_compare = cp.duckdb_prefilter_by_scope(
            mpile1=mpileup_contig_1,
            mpile2=mpileup_contig_2,
            genome_scope=genome,
            memory_limit=duckdb_memory_limit,
            temp_directory=duckdb_temp_directory,
            threads=duckdb_threads,
        )

    if engine == "duckdb":
        cp.duckdb_compare_genomes_to_parquet(
            mpile1=mpile_1_for_compare,
            mpile2=mpile_2_for_compare,
            output_file=output_file,
            stb_file=stb_file,
            sample_1_name=mpile_contig_1_name,
            sample_2_name=mpile_contig_2_name,
            min_cov=min_cov,
            min_gene_compare_len=min_gene_compare_len,
            genome_scope=genome,
            ani_method=ani_method,
            memory_limit=duckdb_memory_limit,
            temp_directory=duckdb_temp_directory,
            threads=duckdb_threads,
        )
        return

    comp = cp.compare_genomes(
        mpile_contig_1=mpile_1_for_compare,
        mpile_contig_2=mpile_2_for_compare,
        min_cov=min_cov,
        min_gene_compare_len=min_gene_compare_len,
        genome_scope=genome,
        ani_method=ani_method,
        duckdb_memory_limit=duckdb_memory_limit,
        duckdb_temp_directory=duckdb_temp_directory,
        duckdb_threads=duckdb_threads,
        engine="polars",
        stb_file=stb_file,
    ).with_columns(
        sample_1=pl.lit(mpile_contig_1_name),
        sample_2=pl.lit(mpile_contig_2_name),
    ).select(
        "genome",
        "total_positions",
        "share_allele_pos",
        "genome_pop_ani",
        "max_consecutive_length",
        "shared_genes_count",
        "identical_gene_count",
        "perc_id_genes",
        "sample_1",
        "sample_2",
    )
    comp.sink_parquet(output_file, compression='zstd')

@compare.command("single_compare_gene")
@click.option('--mpileup-contig-1', '-m1', required=True, help="Path to the first mpileup file.")
@click.option('--mpileup-contig-2', '-m2', required=True, help="Path to the second mpileup file.")
@click.option('--stb-file', '-s', required=True, help="Path to the scaffold to genome mapping file.")
@click.option('--min-cov', '-c', default=5, help="Minimum coverage to consider a position.")
@click.option('--min-gene-compare-len', '-l', default=100, help="Minimum gene length to consider for comparison.")
@click.option('--output-file', '-o', required=True, help="Path to save the parquet file.")
@click.option('--scope', '-g', default="all:all", help="If provided, do the comparison only for the specified genome-gene pair.")
@click.option('--ani-method', '-a', default="popani", help="ANI calculation method to use (e.g., 'popani', 'conani', 'cosani_0.4').")
@click.option('--engine', type=click.Choice(["polars", "duckdb"]), default="polars", show_default=True, help="Execution engine for compare.")
@click.option('--duckdb-memory-limit', default=None, help="DuckDB memory limit (e.g., 2GB, 1024MB).")
@click.option('--duckdb-temp-directory', default=None, help="Directory DuckDB can use for spill files.")
@click.option('--duckdb-threads', type=int, default=None, help="Number of DuckDB worker threads.")
def single_compare_gene(mpileup_contig_1, mpileup_contig_2, stb_file, min_cov, min_gene_compare_len, output_file, scope, ani_method, engine, duckdb_memory_limit, duckdb_temp_directory, duckdb_threads):
    """
    Compare two mpileup files and calculate gene-level comparison statistics.
    
    Args:
    mpileup_contig_1 (str): Path to the first mpileup file.
    mpileup_contig_2 (str): Path to the second mpileup file.
    stb_file (str): Path to the scaffold to genome mapping file.
    min_cov (int): Minimum coverage to consider a position.
    min_gene_compare_len (int): Minimum gene length to consider for comparison.
    output_file (str): Path to save the parquet file.
    scope (str): If provided, do the comparison only for the specified genome-gene pair.
    ani_method (str): ANI calculation method to use.
    """

    mpile_contig_1_name = pathlib.Path(mpileup_contig_1).name
    mpile_contig_2_name = pathlib.Path(mpileup_contig_2).name

    _ = stb_file
    if ":" not in scope:
        raise ValueError("scope must be in 'GENOME:GENE' format (e.g., all:all).")
    genome_scope, gene_scope = scope.split(":", 1)

    mpile_1_for_compare = mpileup_contig_1
    mpile_2_for_compare = mpileup_contig_2
    if genome_scope != "all" or gene_scope != "all":
        mpile_1_for_compare, mpile_2_for_compare = cp.duckdb_prefilter_by_scope(
            mpile1=mpileup_contig_1,
            mpile2=mpileup_contig_2,
            genome_scope=genome_scope,
            gene_scope=gene_scope,
            memory_limit=duckdb_memory_limit,
            temp_directory=duckdb_temp_directory,
            threads=duckdb_threads,
        )

    if engine == "duckdb":
        cp.duckdb_compare_genes_to_parquet(
            mpile1=mpile_1_for_compare,
            mpile2=mpile_2_for_compare,
            output_file=output_file,
            sample_1_name=mpile_contig_1_name,
            sample_2_name=mpile_contig_2_name,
            min_cov=min_cov,
            min_gene_compare_len=min_gene_compare_len,
            genome_scope=genome_scope,
            gene_scope=gene_scope,
            ani_method=ani_method,
            memory_limit=duckdb_memory_limit,
            temp_directory=duckdb_temp_directory,
            threads=duckdb_threads,
        )
        return

    gene_comp = cp.compare_genes(
        mpile_contig_1=mpile_1_for_compare,
        mpile_contig_2=mpile_2_for_compare,
        min_cov=min_cov,
        min_gene_compare_len=min_gene_compare_len,
        genome_scope=genome_scope,
        gene_scope=gene_scope,
        ani_method=ani_method,
        duckdb_memory_limit=duckdb_memory_limit,
        duckdb_temp_directory=duckdb_temp_directory,
        duckdb_threads=duckdb_threads,
        engine="polars",
    ).with_columns(
        sample_1=pl.lit(mpile_contig_1_name),
        sample_2=pl.lit(mpile_contig_2_name),
    ).select(
        "genome",
        "gene",
        "total_positions",
        "share_allele_pos",
        "ani",
        "sample_1",
        "sample_2",
    )
    gene_comp.sink_parquet(output_file, compression='zstd')

@cli.group()
def profile():
    """The commands in this group are related to profiling bam files."""
    pass


@profile.command("prepare_profiling",help="Prepare the files needed for profiling bam files and save them in the specified output directory.")
@click.option('--reference-fasta', '-r', required=True, help="Path to the reference genome in FASTA format.")
@click.option('--gene-fasta', '-g', required=True, help="Path to the gene annotations in FASTA format.")
@click.option('--stb-file', '-s', required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--output-dir', '-o', required=True, help="Directory to save the profiling database.")
def prepare_profiling(reference_fasta, gene_fasta, stb_file, output_dir):
    """
    Prepare the files needed for profiling bam files and save them in the specified output directory.
    """
    output_dir=pathlib.Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    bed_df = ut.make_the_bed(reference_fasta)
    bed_df.write_csv(output_dir / "genomes_bed_file.bed", separator='\t', include_header=False)
    gene_range_table = pf.build_gene_range_table(pathlib.Path(gene_fasta))
    gene_range_table.write_csv(output_dir / "gene_range_table.tsv", separator='\t', include_header=False)
    
    stb = pl.scan_csv(stb_file, separator='\t',has_header=False).with_columns(
        pl.col("column_1").alias("scaffold"),
        pl.col("column_2").alias("genome")
    )

    bed_df = bed_df.lazy()
    genome_length = ut.extract_genome_length(stb, bed_df)
    genome_length.sink_parquet(output_dir / "genome_lengths.parquet", compression='zstd')


@profile.command("profile-single")
@click.option('--bed-file', '-b', required=True, help="Path to the BED file describing regions to be profiled.")
@click.option('--bam-file', '-a', required=True, help="Path to the BAM file to be profiled.")
@click.option('--stb-file', '-s', required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--null-model', '-m', required=True, help="Path to the null model file. If not provided, the null model will be generated from the data.") 
@click.option('--gene-range-table', '-g', required=True, help="Path to the gene range table.")
@click.option('--num-workers', '-n', default=1, help="Number of workers to use for profiling.")
@click.option('--output-dir', '-o', required=True, help="Directory to save the profiling output.")
def profile_single(bed_file, bam_file, stb_file, null_model, gene_range_table, num_workers, output_dir):
    """
    Profile a single BAM file using the provided BED file and gene range table.
    
    """
    output_dir=pathlib.Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    stb= pl.scan_csv(stb_file, separator='\t',has_header=False).with_columns(
        pl.col("column_1").alias("scaffold"),
        pl.col("column_2").alias("genome")
    )
    null_model=pl.scan_parquet(null_model)
    pf.profile_bam(
        bed_file=bed_file,
        bam_file=bam_file,
        gene_range_table=gene_range_table,
        stb=stb,
        null_model=null_model,
        output_dir=output_dir,
        num_workers=num_workers,
    )

@cli.group()
def run():
    """The commands in this group are related to running zipstrain workflows."""
    pass


@run.command("profile")
@click.option('--input-table', '-i', required=True, help="Path to the input table in TSV format containing sample names and paths to bam files.")
@click.option('--stb-file', '-s', required=True, help="Path to the scaffold-to-genome mapping file.")
@click.option('--gene-range-table', '-g', required=True, help="Path to the gene range table file.")
@click.option('--bed-file', '-b', required=True, help="Path to the BED file for profiling regions.")
@click.option('--genome-length-file', '-l', required=True, help="Path to the genome length file.")
@click.option('--run-dir', '-r', required=True, help="Directory to save the run data.")
@click.option('--num-procs', '-n', default=8, help="Number of processors to use for each profiling task.")
@click.option('--max-concurrent-batches', '-m', default=5, help="Maximum number of concurrent batches to run.")
@click.option('--poll-interval', '-p', default=1, help="Polling interval in seconds to check the status of batches.")
@click.option('--execution-mode', '-e', default="local", help="Execution mode: 'local' or 'slurm'.")
@click.option('--slurm-config', '-c', default=None, help="Path to the SLURM configuration file in json format. Required if execution mode is 'slurm'.")
@click.option('--container-engine', '-o', default="local", help="Container engine to use: 'local', 'docker' or 'apptainer'.")
@click.option('--task-per-batch', '-t', default=10, help="Number of tasks to include in each batch.")
def profile(input_table, stb_file, gene_range_table, bed_file, genome_length_file, run_dir, num_procs, max_concurrent_batches, poll_interval, execution_mode, slurm_config, container_engine, task_per_batch):
    """
    Run BAM file profiling in batches using the specified execution mode and container engine.

    Args:
    input_table (str): Path to the input table in TSV format containing sample names and BAM file paths.
    stb_file (str): Path to the scaffold-to-genome mapping file.
    gene_range_table (str): Path to the gene range table file.
    bed_file (str): Path to the BED file for profiling regions.
    genome_length_file (str): Path to the genome length file.
    run_dir (str): Directory to save the run data.
    num_procs (int): Number of processors to use for each profiling task.
    max_concurrent_batches (int): Maximum number of concurrent batches to run.
    poll_interval (int): Polling interval in seconds to check the status of batches.
    execution_mode (str): Execution mode: 'local' or 'slurm'.
    slurm_config (str): Path to the SLURM configuration file in json format. Required if execution mode is 'slurm'.
    container_engine (str): Container engine to use: 'local', 'docker' or 'apptainer'.
    task_per_batch (int): Number of tasks to include in each batch.
    """
    # Load the BAM files table
    bams_lf = pl.scan_csv(input_table)
    
    # Validate required columns exist
    required_columns = {"sample_name", "bamfile"}
    actual_columns = set(bams_lf.collect_schema().names())
    if not required_columns.issubset(actual_columns):
        missing = required_columns - actual_columns
        raise ValueError(f"Input table missing required columns: {missing}")
    
    run_dir = pathlib.Path(run_dir)
    slurm_conf = None
    if execution_mode == "slurm":
        if slurm_config is None:
            raise ValueError("SLURM configuration file must be provided when execution mode is 'slurm'.")
        slurm_conf = tm.SlurmConfig.from_json(slurm_config)
    
    if container_engine == "local":
        container_engine_obj = tm.LocalEngine(address="")
    elif container_engine == "docker":
        container_engine_obj = tm.DockerEngine(address="parsaghadermazi/zipstrain:amd64")
    elif container_engine == "apptainer":
        container_engine_obj = tm.ApptainerEngine(address="docker://parsaghadermazi/zipstrain:amd64")
    else:
        raise ValueError("Invalid container engine. Choose from 'local', 'docker', or 'apptainer'.")
    
    tm.lazy_run_profile(
        run_dir=run_dir,
        container_engine=container_engine_obj,
        bams_lf=bams_lf,
        stb_file=pathlib.Path(stb_file),
        gene_range_table=pathlib.Path(gene_range_table),
        bed_file=pathlib.Path(bed_file),
        genome_length_file=pathlib.Path(genome_length_file),
        num_procs=num_procs,
        tasks_per_batch=task_per_batch,
        max_concurrent_batches=max_concurrent_batches,
        poll_interval=poll_interval,
        execution_mode=execution_mode,
        slurm_config=slurm_conf,
    )


@run.command("compare_genomes")
@click.option("--genome-comparison-object", "-g", required=True, help="Path to the genome comparison object in json format.")
@click.option("--run-dir", "-r", required=True, help="Directory to save the run data.")
@click.option("--max-concurrent-batches", "-m", default=5, help="Maximum number of concurrent batches to run.")
@click.option("--poll-interval", "-p", default=1, help="Polling interval in seconds to check the status of batches.")
@click.option("--execution-mode", "-e", default="local", help="Execution mode: 'local' or 'slurm'.")
@click.option("--slurm-config", "-s", default=None, help="Path to the SLURM configuration file in json format. Required if execution mode is 'slurm'.")
@click.option("--container-engine", "-c", default="local", help="Container engine to use: 'local', 'docker' or 'apptainer'.")
@click.option("--task-per-batch", "-t", default=10, help="Number of tasks to include in each batch.")
@click.option("--duckdb-memory-limit", "-d", default=None, help="DuckDB memory limit for compare tasks (e.g., 2GB).")
def compare_genomes(genome_comparison_object, run_dir, max_concurrent_batches, poll_interval, execution_mode, slurm_config, container_engine, task_per_batch, duckdb_memory_limit):
    """
    Run genome comparisons in batches using the specified execution mode and container engine.

    Args:
    genome_comparison_object (str): Path to the genome comparison object in json format.
    run_dir (str): Directory to save the run data.
    max_concurrent_batches (int): Maximum number of concurrent batches to run.
    poll_interval (int): Polling interval in seconds to check the status of batches.
    execution_mode (str): Execution mode: 'local' or 'slurm'.
    slurm_config (str): Path to the SLURM configuration file in json format. Required if execution mode is 'slurm'.
    container_engine (str): Container engine to use: 'local', 'docker' or 'apptainer'.
    task_per_batch (int): Number of tasks to include in each batch.
    """
    genome_comp_db=db.GenomeComparisonDatabase.load_obj(pathlib.Path(genome_comparison_object))
    run_dir=pathlib.Path(run_dir)
    slurm_conf=None
    if execution_mode == "slurm":
        if slurm_config is None:
            raise ValueError("SLURM configuration file must be provided when execution mode is 'slurm'.")
        slurm_conf = tm.SlurmConfig.from_json(slurm_config)
    
    if container_engine == "local":
        container_engine_obj = tm.LocalEngine(address="")
    elif container_engine == "docker":
        container_engine_obj = tm.DockerEngine(address="parsaghadermazi/zipstrain:amd64") #could go to a toml or json config file
    elif container_engine == "apptainer":
        container_engine_obj = tm.ApptainerEngine(address="docker://parsaghadermazi/zipstrain:amd64") #could go to a toml or json config file
    else:
        raise ValueError("Invalid container engine. Choose from 'local', 'docker', or 'apptainer'.")
    tm.lazy_run_compares(
        comps_db=genome_comp_db,
        container_engine=container_engine_obj,
        run_dir=run_dir,
        max_concurrent_batches=max_concurrent_batches,
        execution_mode=execution_mode,
        slurm_config=slurm_conf,
        duckdb_memory_limit=duckdb_memory_limit,
        tasks_per_batch=task_per_batch,
        poll_interval=poll_interval,
    )



@run.command("build-comp-database")
@click.option("--profile-db-dir", "-p", required=True, help="Directory containing profile either in parquet format.")
@click.option("--config-file", "-c", required=True, help="Path to the  genome comparsion database config file in json format.")
@click.option("--output-dir", "-o", required=True, help="Directory to genome comparison database object.")
@click.option("--comp-db-file", "-f", required=False, help="The initial database file. If provided only additional comparisons will be added to this database.")
def build_comp_database(profile_db_dir, config_file, output_dir, comp_db_file):
    """
    Build a genome comparison database from the given profiles and configuration.

    Parameters:
    profile_db_dir (str): Directory containing profile either in parquet format.
    config_file (str): Path to the genome comparison database config file in json format.
    """
    profile_db_dir=pathlib.Path(profile_db_dir)
    profile_db=db.ProfileDatabase(
        db_loc=profile_db_dir,
    )
    existing_db_loc=pathlib.Path(comp_db_file) if comp_db_file is not None else None
    if existing_db_loc is not None and not existing_db_loc.exists():
        raise FileNotFoundError(f"{existing_db_loc} does not exist.")
    obj=db.GenomeComparisonDatabase(
        profile_db=profile_db,
        config=db.GenomeComparisonConfig.from_json(pathlib.Path(config_file)),
        comp_db_loc=existing_db_loc,
    )
    obj.dump_obj(pathlib.Path(output_dir))


@run.command("compare_genes")
@click.option("--gene-comparison-object", "-g", required=True, help="Path to the gene comparison object in json format.")
@click.option("--run-dir", "-r", required=True, help="Directory to save the run data.")
@click.option("--max-concurrent-batches", "-m", default=5, help="Maximum number of concurrent batches to run.")
@click.option("--poll-interval", "-p", default=1, help="Polling interval in seconds to check the status of batches.")
@click.option("--execution-mode", "-e", default="local", help="Execution mode: 'local' or 'slurm'.")
@click.option("--slurm-config", "-s", default=None, help="Path to the SLURM configuration file in json format. Required if execution mode is 'slurm'.")
@click.option("--container-engine", "-c", default="local", help="Container engine to use: 'local', 'docker' or 'apptainer'.")
@click.option("--task-per-batch", "-t", default=10, help="Number of tasks to include in each batch.")
@click.option("--ani-method", "-n", default="popani", help="ANI calculation method to use (e.g., 'popani', 'conani', 'cosani_0.4').")
@click.option("--duckdb-memory-limit", "-d", default=None, help="DuckDB memory limit for compare tasks (e.g., 2GB).")
def compare_genes(gene_comparison_object, run_dir, max_concurrent_batches, poll_interval, execution_mode, slurm_config, container_engine, task_per_batch, ani_method, duckdb_memory_limit):
    """
    Run gene comparisons in batches using the specified execution mode and container engine.

    Args:
    genome_comparison_object (str): Path to the genome comparison object in json format.
    run_dir (str): Directory to save the run data.
    max_concurrent_batches (int): Maximum number of concurrent batches to run.
    poll_interval (int): Polling interval in seconds to check the status of batches.
    execution_mode (str): Execution mode: 'local' or 'slurm'.
    slurm_config (str): Path to the SLURM configuration file in json format. Required if execution mode is 'slurm'.
    container_engine (str): Container engine to use: 'local', 'docker' or 'apptainer'.
    task_per_batch (int): Number of tasks to include in each batch.
    ani_method (str): ANI calculation method to use.
    """
    genome_comp_db=db.GeneComparisonDatabase.load_obj(pathlib.Path(gene_comparison_object))
    run_dir=pathlib.Path(run_dir)
    slurm_conf=None
    if execution_mode == "slurm":
        if slurm_config is None:
            raise ValueError("SLURM configuration file must be provided when execution mode is 'slurm'.")
        slurm_conf = tm.SlurmConfig.from_json(slurm_config)
    
    if container_engine == "local":
        container_engine_obj = tm.LocalEngine(address="")
    elif container_engine == "docker":
        container_engine_obj = tm.DockerEngine(address="parsaghadermazi/zipstrain:amd64")
    elif container_engine == "apptainer":
        container_engine_obj = tm.ApptainerEngine(address="docker://parsaghadermazi/zipstrain:amd64")
    else:
        raise ValueError("Invalid container engine. Choose from 'local', 'docker', or 'apptainer'.")
    
    tm.lazy_run_gene_compares(
        comps_db=genome_comp_db,
        container_engine=container_engine_obj,
        run_dir=run_dir,
        max_concurrent_batches=max_concurrent_batches,
        execution_mode=execution_mode,
        slurm_config=slurm_conf,
        tasks_per_batch=task_per_batch,
        poll_interval=poll_interval,
        ani_method=ani_method,
        duckdb_memory_limit=duckdb_memory_limit,
    )
        
@cli.command("test")
def test():
    """Run basic tests to ensure ZipStrain is setup correctly."""
    ### Check samtools installation
    if all([ut.check_samtools()]):
        click.echo("ZipStrain setup looks good!")
    else:
        click.echo("There are issues with the ZipStrain setup. Please check the above messages.")

if __name__ == "__main__":
    cli()
