"""zipstrain.compare
========================
This module provides all comparison functions for zipstrain.

"""

from pathlib import Path
from typing import Optional, Union

import polars as pl
import duckdb

class PolarsANIExpressions:
    """ 
    Any kind of ANI calculation based on two profiles should be implemented as a method of this class.
    In defining this method, the following rules should be followed:
    
    -   The method returns a Polars expression (pl.Expr).
    
    -   When applied to a row, the method returns a zero if that position is a SNV. Otherwise it should return a number greater than zero.
    
    -   A, T, C, G columns in the first profile are named "A", "T", "C", "G" and in the second profile they are named "A_2", "T_2", "C_2", "G_2".
    
    1. popani: Population ANI based on the shared alleles between two profiles.
    2. conani: Consensus ANI based on the consensus alleles between two profiles.
    3. cosani_<threshold>: Generalized cosine similarity ANI where threshold is a float value between 0 and 1. Once the similarity is below the threshold, it is considered a SNV.
    """
    MPILE_1_BASES = ["A", "T", "C", "G"]
    MPILE_2_BASES = ["A_2", "T_2", "C_2", "G_2"]

    def popani(self):
        return pl.col("A")*pl.col("A_2") + pl.col("C")*pl.col("C_2") + pl.col("G")*pl.col("G_2") + pl.col("T")*pl.col("T_2")
    
    def conani(self):
        max_base_1=pl.max_horizontal(*[pl.col(base) for base in self.MPILE_1_BASES])
        max_base_2=pl.max_horizontal(*[pl.col(base) for base in self.MPILE_2_BASES])
        return pl.when((pl.col("A")==max_base_1) & (pl.col("A_2")==max_base_2) | 
                       (pl.col("T")==max_base_1) & (pl.col("T_2")==max_base_2) | 
                       (pl.col("C")==max_base_1) & (pl.col("C_2")==max_base_2) | 
                       (pl.col("G")==max_base_1) & (pl.col("G_2")==max_base_2)).then(1).otherwise(0)
    
    def generalized_cos_ani(self,threshold:float=0.4):
        dot_product = pl.col("A")*pl.col("A_2") + pl.col("C")*pl.col("C_2") + pl.col("G")*pl.col("G_2") + pl.col("T")*pl.col("T_2")
        magnitude_1 = (pl.col("A")**2 + pl.col("C")**2 + pl.col("G")**2 + pl.col("T")**2)**0.5
        magnitude_2 = (pl.col("A_2")**2 + pl.col("C_2")**2 + pl.col("G_2")**2 + pl.col("T_2")**2)**0.5
        cos_sim = dot_product / (magnitude_1 * magnitude_2)
        return pl.when(cos_sim >= threshold).then(1).otherwise(0)

    def __getattribute__(self, name):
        if name.startswith("cosani_"):
            try:
                threshold = float(name.split("_")[1])
            except ValueError:
                raise AttributeError(f"Invalid threshold in method name: {name}")
            return lambda: self.generalized_cos_ani(threshold)
        else:
            return super().__getattribute__(name)

def _duckdb_quote_sql_string(value: str) -> str:
    return value.replace("'", "''")


def _duckdb_from_source(conn: duckdb.DuckDBPyConnection, source: Union[str, Path, pl.LazyFrame], view_name: str) -> str:
    if isinstance(source, pl.LazyFrame):
        conn.register(view_name, source.collect().to_arrow())
        return view_name
    if isinstance(source, (str, Path)):
        path = _duckdb_quote_sql_string(str(source))
        return f"read_parquet('{path}')"
    raise TypeError(f"Unsupported source type for DuckDB input: {type(source)}")


def _duckdb_configure_connection(
    conn: duckdb.DuckDBPyConnection,
    memory_limit: Optional[str] = None,
    temp_directory: Optional[Union[str, Path]] = None,
    threads: Optional[int] = None,
) -> None:
    if memory_limit:
        conn.execute(f"SET memory_limit='{_duckdb_quote_sql_string(memory_limit)}'")
    if temp_directory:
        conn.execute(f"SET temp_directory='{_duckdb_quote_sql_string(str(temp_directory))}'")
    if threads is not None:
        if threads < 1:
            raise ValueError("threads must be >= 1")
        conn.execute(f"SET threads={int(threads)}")


def _duckdb_ani_expression(ani_method: str) -> str:
    if ani_method == "popani":
        return "p1.A*p2.A + p1.C*p2.C + p1.G*p2.G + p1.T*p2.T"
    if ani_method == "conani":
        return """
        CASE WHEN (
          (p1.A=GREATEST(p1.A,p1.T,p1.C,p1.G) AND p2.A=GREATEST(p2.A,p2.T,p2.C,p2.G)) OR
          (p1.T=GREATEST(p1.A,p1.T,p1.C,p1.G) AND p2.T=GREATEST(p2.A,p2.T,p2.C,p2.G)) OR
          (p1.C=GREATEST(p1.A,p1.T,p1.C,p1.G) AND p2.C=GREATEST(p2.A,p2.T,p2.C,p2.G)) OR
          (p1.G=GREATEST(p1.A,p1.T,p1.C,p1.G) AND p2.G=GREATEST(p2.A,p2.T,p2.C,p2.G))
        ) THEN 1 ELSE 0 END
        """
    if ani_method.startswith("cosani_"):
        thr = float(ani_method.split("_")[1])
        return f"""
        CASE WHEN (
          (
            p1.A*p2.A + p1.C*p2.C + p1.G*p2.G + p1.T*p2.T
          ) /
          NULLIF(
            SQRT(POWER(p1.A,2)+POWER(p1.C,2)+POWER(p1.G,2)+POWER(p1.T,2)) *
            SQRT(POWER(p2.A,2)+POWER(p2.C,2)+POWER(p2.G,2)+POWER(p2.T,2)),
            0
          )
        ) >= {thr} THEN 1 ELSE 0 END
        """
    raise ValueError(f"Unsupported ani_method: {ani_method}")


def _duckdb_shared_query(
    p1_source: str,
    p2_source: str,
    min_cov: int,
    genome_scope: str,
    ani_method: str,
    gene_scope: str = "all",
) -> str:
    ani_expr = _duckdb_ani_expression(ani_method)
    genome_scope_sql = _duckdb_quote_sql_string(genome_scope)
    gene_scope_sql = _duckdb_quote_sql_string(gene_scope)
    return f"""
    WITH p1f AS (
      SELECT
        chrom,
        pos,
        gene,
        genome,
        CAST(A AS INT32) AS A,
        CAST(T AS INT32) AS T,
        CAST(C AS INT32) AS C,
        CAST(G AS INT32) AS G
      FROM {p1_source} p1
      WHERE (A + T + C + G) >= {min_cov}
        AND ('{genome_scope_sql}' = 'all' OR genome = '{genome_scope_sql}')
        AND ('{gene_scope_sql}' = 'all' OR gene = '{gene_scope_sql}')
    ),
    p2f AS (
      SELECT
        chrom,
        pos,
        gene,
        genome,
        CAST(A AS INT32) AS A,
        CAST(T AS INT32) AS T,
        CAST(C AS INT32) AS C,
        CAST(G AS INT32) AS G
      FROM {p2_source} p2
      WHERE (A + T + C + G) >= {min_cov}
        AND ('{genome_scope_sql}' = 'all' OR genome = '{genome_scope_sql}')
        AND ('{gene_scope_sql}' = 'all' OR gene = '{gene_scope_sql}')
    )
    SELECT
      {ani_expr} AS surr,
      p1.chrom AS scaffold,
      p1.pos,
      p1.gene,
      p1.genome
    FROM p1f p1
    INNER JOIN p2f p2
      ON p1.genome = p2.genome AND p1.chrom = p2.chrom AND p1.pos = p2.pos
    """


def _duckdb_copy_query_to_parquet(conn: duckdb.DuckDBPyConnection, query: str, output_file: Union[str, Path]) -> None:
    output_sql = _duckdb_quote_sql_string(str(output_file))
    conn.execute(f"COPY ({query}) TO '{output_sql}' (FORMAT PARQUET, COMPRESSION ZSTD)")


def duckdb_filter_join(
    mpile1: Union[str, Path, pl.LazyFrame],
    mpile2: Union[str, Path, pl.LazyFrame],
    min_cov: int,
    genome_scope: str = "all",   # or specific genome
    ani_method: str = "popani",  # "popani" / "conani" / "cosani_0.4"
    gene_scope: str = "all",
    memory_limit: Optional[str] = None,
    temp_directory: Optional[Union[str, Path]] = None,
    threads: Optional[int] = None,
) -> pl.LazyFrame:
    """Filter two profile sources and inner-join shared loci in DuckDB.

    Inputs can be parquet paths or Polars LazyFrames. Coverage, genome scope,
    and optional gene scope are pushed down into DuckDB. The returned lazy frame
    contains: `surr`, `scaffold`, `pos`, `gene`, and `genome`.
    """
    con = duckdb.connect()

    try:
        _duckdb_configure_connection(
            con,
            memory_limit=memory_limit,
            temp_directory=temp_directory,
            threads=threads,
        )
        p1_source = _duckdb_from_source(con, mpile1, "mpile1_source")
        p2_source = _duckdb_from_source(con, mpile2, "mpile2_source")
        query = _duckdb_shared_query(
            p1_source=p1_source,
            p2_source=p2_source,
            min_cov=min_cov,
            genome_scope=genome_scope,
            ani_method=ani_method,
            gene_scope=gene_scope,
        )
        table = con.execute(query).fetch_arrow_table()
        return pl.from_arrow(table).lazy()
    finally:
        con.close()


def duckdb_compare_genomes_to_parquet(
    mpile1: Union[str, Path, pl.LazyFrame],
    mpile2: Union[str, Path, pl.LazyFrame],
    output_file: Union[str, Path],
    stb_file: Union[str, Path],
    sample_1_name: str,
    sample_2_name: str,
    min_cov: int = 5,
    min_gene_compare_len: int = 100,
    genome_scope: str = "all",
    ani_method: str = "popani",
    memory_limit: Optional[str] = None,
    temp_directory: Optional[Union[str, Path]] = None,
    threads: Optional[int] = None,
) -> None:
    """Run genome comparison in DuckDB and write final output directly to parquet.

    This path avoids materializing large intermediate tables in Python memory.
    """
    con = duckdb.connect()
    try:
        _duckdb_configure_connection(
            con,
            memory_limit=memory_limit,
            temp_directory=temp_directory,
            threads=threads,
        )
        p1_source = _duckdb_from_source(con, mpile1, "mpile1_source")
        p2_source = _duckdb_from_source(con, mpile2, "mpile2_source")
        shared_query = _duckdb_shared_query(
            p1_source=p1_source,
            p2_source=p2_source,
            min_cov=min_cov,
            genome_scope=genome_scope,
            ani_method=ani_method,
        )
        stb_sql = _duckdb_quote_sql_string(str(stb_file))
        genome_scope_sql = _duckdb_quote_sql_string(genome_scope)
        sample_1_sql = _duckdb_quote_sql_string(sample_1_name)
        sample_2_sql = _duckdb_quote_sql_string(sample_2_name)
        query = f"""
        WITH shared AS (
          {shared_query}
        ),
        genomes AS (
          SELECT DISTINCT CAST(column1 AS VARCHAR) AS genome
          FROM read_csv('{stb_sql}', delim='\\t', header=false)
          WHERE ('{genome_scope_sql}' = 'all' OR CAST(column1 AS VARCHAR) = '{genome_scope_sql}')
        ),
        contig_base AS (
          SELECT
            s.*,
            LAG(scaffold) OVER (PARTITION BY genome ORDER BY scaffold, pos) AS prev_scaffold
          FROM shared s
        ),
        contig AS (
          SELECT
            *,
            SUM(
              CASE
                WHEN prev_scaffold IS NULL OR scaffold != prev_scaffold OR surr = 0 THEN 1
                ELSE 0
              END
            ) OVER (PARTITION BY genome ORDER BY scaffold, pos ROWS UNBOUNDED PRECEDING) AS group_id
          FROM contig_base
        ),
        pop AS (
          SELECT
            genome,
            COUNT(*)::BIGINT AS total_positions,
            SUM(CASE WHEN surr > 0 THEN 1 ELSE 0 END)::BIGINT AS share_allele_pos,
            SUM(CASE WHEN surr > 0 THEN 1 ELSE 0 END) * 100.0 / NULLIF(COUNT(*), 0) AS genome_pop_ani
          FROM contig
          GROUP BY genome
        ),
        max_blocks AS (
          SELECT
            genome,
            MAX(length)::BIGINT AS max_consecutive_length
          FROM (
            SELECT genome, scaffold, group_id, COUNT(*)::BIGINT AS length
            FROM contig
            GROUP BY genome, scaffold, group_id
          ) blocks
          GROUP BY genome
        ),
        gene_base AS (
          SELECT
            genome,
            gene,
            COUNT(*)::BIGINT AS total_positions,
            SUM(CASE WHEN surr > 0 THEN 1 ELSE 0 END)::BIGINT AS share_allele_pos
          FROM contig
          WHERE gene != 'NA'
          GROUP BY genome, gene
          HAVING COUNT(*) >= {min_gene_compare_len}
        ),
        gene_stats AS (
          SELECT
            genome,
            COUNT(*)::BIGINT AS shared_genes_count,
            SUM(CASE WHEN share_allele_pos = total_positions THEN 1 ELSE 0 END)::BIGINT AS identical_gene_count,
            SUM(CASE WHEN share_allele_pos = total_positions THEN 1 ELSE 0 END) * 100.0 / NULLIF(COUNT(*), 0) AS perc_id_genes
          FROM gene_base
          GROUP BY genome
        ),
        merged AS (
          SELECT
            g.genome,
            COALESCE(p.total_positions, 0)::BIGINT AS total_positions,
            COALESCE(p.share_allele_pos, 0)::BIGINT AS share_allele_pos,
            COALESCE(p.genome_pop_ani, 0.0)::DOUBLE AS genome_pop_ani,
            COALESCE(m.max_consecutive_length, 0)::BIGINT AS max_consecutive_length,
            COALESCE(gs.shared_genes_count, 0)::BIGINT AS shared_genes_count,
            COALESCE(gs.identical_gene_count, 0)::BIGINT AS identical_gene_count,
            COALESCE(gs.perc_id_genes, 0.0)::DOUBLE AS perc_id_genes
          FROM genomes g
          LEFT JOIN pop p ON g.genome = p.genome
          LEFT JOIN max_blocks m ON g.genome = m.genome
          LEFT JOIN gene_stats gs ON g.genome = gs.genome
        )
        SELECT
          genome,
          total_positions,
          share_allele_pos,
          genome_pop_ani,
          max_consecutive_length,
          shared_genes_count,
          identical_gene_count,
          perc_id_genes,
          '{sample_1_sql}' AS sample_1,
          '{sample_2_sql}' AS sample_2
        FROM merged
        ORDER BY genome
        """
        _duckdb_copy_query_to_parquet(con, query, output_file)
    finally:
        con.close()


def duckdb_compare_genes_to_parquet(
    mpile1: Union[str, Path, pl.LazyFrame],
    mpile2: Union[str, Path, pl.LazyFrame],
    output_file: Union[str, Path],
    sample_1_name: str,
    sample_2_name: str,
    min_cov: int = 5,
    min_gene_compare_len: int = 100,
    genome_scope: str = "all",
    gene_scope: str = "all",
    ani_method: str = "popani",
    memory_limit: Optional[str] = None,
    temp_directory: Optional[Union[str, Path]] = None,
    threads: Optional[int] = None,
) -> None:
    """Run gene comparison in DuckDB and write final output directly to parquet."""
    con = duckdb.connect()
    try:
        _duckdb_configure_connection(
            con,
            memory_limit=memory_limit,
            temp_directory=temp_directory,
            threads=threads,
        )
        p1_source = _duckdb_from_source(con, mpile1, "mpile1_source")
        p2_source = _duckdb_from_source(con, mpile2, "mpile2_source")
        shared_query = _duckdb_shared_query(
            p1_source=p1_source,
            p2_source=p2_source,
            min_cov=min_cov,
            genome_scope=genome_scope,
            gene_scope=gene_scope,
            ani_method=ani_method,
        )
        sample_1_sql = _duckdb_quote_sql_string(sample_1_name)
        sample_2_sql = _duckdb_quote_sql_string(sample_2_name)
        query = f"""
        WITH shared AS (
          {shared_query}
        )
        SELECT
          genome,
          gene,
          COUNT(*)::BIGINT AS total_positions,
          SUM(CASE WHEN surr > 0 THEN 1 ELSE 0 END)::BIGINT AS share_allele_pos,
          SUM(CASE WHEN surr > 0 THEN 1 ELSE 0 END) * 100.0 / NULLIF(COUNT(*), 0) AS ani,
          '{sample_1_sql}' AS sample_1,
          '{sample_2_sql}' AS sample_2
        FROM shared
        GROUP BY genome, gene
        HAVING COUNT(*) >= {min_gene_compare_len}
        ORDER BY genome, gene
        """
        _duckdb_copy_query_to_parquet(con, query, output_file)
    finally:
        con.close()



def get_shared_locs(mpile_contig_1:pl.LazyFrame, mpile_contig_2:pl.LazyFrame,ani_method:str="popani") -> pl.LazyFrame:
    """
    Returns a lazyframe with ATCG information for shared scaffolds and positions between two mpileup files.

    Args:
        mpile_contig_1 (pl.LazyFrame): The first mpileup LazyFrame.
        mpile_contig_2 (pl.LazyFrame): The second mpileup LazyFrame.
        ani_method (str): The ANI calculation method to use. Default is "popani".
    
    Returns:
        pl.LazyFrame: Merged LazyFrame containing shared scaffolds and positions with ATCG information.
    """
    ani_expr=getattr(PolarsANIExpressions(), ani_method)()

    mpile_contig= mpile_contig_1.join(
        mpile_contig_2,
        on=["chrom", "pos"],
        how="inner",
        suffix="_2"  # To distinguish lf2 columns
    ).with_columns(
        ani_expr.alias("surr")
    ).select(
        pl.col("surr"),
        scaffold=pl.col("chrom"),
        pos=pl.col("pos"),
        gene=pl.col("gene")
    )
    return mpile_contig

def add_contiguity_info(mpile_contig:pl.LazyFrame) -> pl.LazyFrame:
    """ Adds group id information to the lazy frame. If on the same scaffold and not popANI, then they are in the same group.
    
    Args:
        mpile_contig (pl.LazyFrame): The input LazyFrame containing mpileup data.
    
    Returns:
        pl.LazyFrame: Updated LazyFrame with group id information added.
    """

    mpile_contig= mpile_contig.sort(["scaffold", "pos"])
    mpile_contig = mpile_contig.with_columns([
        (pl.col("scaffold").shift(1).fill_null(pl.col("scaffold")).alias("prev_scaffold")),
    ])
    mpile_contig = mpile_contig.with_columns([
        (((pl.col("scaffold") != pl.col("prev_scaffold")) | (pl.col("surr") == 0))).cum_sum().alias("group_id")
    ])
    return mpile_contig

def add_genome_info(mpile_contig:pl.LazyFrame, scaffold_to_genome:pl.LazyFrame) -> pl.LazyFrame:
    """
    Adds genome information to the mpileup LazyFrame based on scaffold to genome mapping.
    
    Args:
        mpile_contig (pl.LazyFrame): The input LazyFrame containing mpileup data.
        scaffold_to_genome (pl.LazyFrame): The LazyFrame mapping scaffolds to genomes.
    
    Returns:
        pl.LazyFrame: Updated LazyFrame with genome information added.
    """
    return mpile_contig.join(
        scaffold_to_genome, on="scaffold", how="left"
    ).fill_null("NA")

def calculate_pop_ani(mpile_contig:pl.LazyFrame) -> pl.LazyFrame:
    """
    Calculates the population ANI (Average Nucleotide Identity) for the given mpileup LazyFrame.
    NOTE: Remember that this function should be applied to the merged mpileup using get_shared_locs.

    Args:
        mpile_contig (pl.LazyFrame): The input LazyFrame containing mpileup data.
    
    Returns:
        pl.LazyFrame: Updated LazyFrame with population ANI information added.
    """
    return mpile_contig.group_by("genome").agg(
            total_positions=pl.len(),
            share_allele_pos=(pl.col("surr") > 0 ).sum()
        ).with_columns(
            genome_pop_ani=pl.col("share_allele_pos")/pl.col("total_positions")*100,
        )

def get_longest_consecutive_blocks(mpile_contig:pl.LazyFrame) -> pl.LazyFrame:
    """
    Calculates the longest consecutive blocks for each genome in the mpileup LazyFrame for any genome.
    
    Args:
        mpile_contig (pl.LazyFrame): The input LazyFrame containing mpileup data.
    
    Returns:
        pl.LazyFrame: Updated LazyFrame with longest consecutive blocks information added.
    """
    block_lengths = (
        mpile_contig.group_by(["genome", "scaffold", "group_id"])
        .agg(pl.len().alias("length"))
    ) 
    return block_lengths.group_by("genome").agg(pl.col("length").max().alias("max_consecutive_length"))

def get_gene_ani(mpile_contig:pl.LazyFrame, min_gene_compare_len:int) -> pl.LazyFrame:
    """
    Calculates gene ANI (Average Nucleotide Identity) for each gene in each genome.
    
    Args:
        mpile_contig (pl.LazyFrame): The input LazyFrame containing mpileup data.
        min_gene_compare_len (int): Minimum length of the gene to consider for comparison.
    
    Returns:
        pl.LazyFrame: Updated LazyFrame with gene ANI information added.
    """
    return mpile_contig.group_by(["genome", "gene"]).agg(
        total_positions=pl.len(),
        share_allele_pos=(pl.col("surr") > 0).sum()
    ).filter(pl.col("total_positions") >= min_gene_compare_len).with_columns(
        identical=(pl.col("share_allele_pos") == pl.col("total_positions")),
    ).filter(pl.col("gene") != "NA").group_by("genome").agg(
        shared_genes_count=pl.len(),
        identical_gene_count=pl.col("identical").sum()
    ).with_columns(perc_id_genes=pl.col("identical_gene_count") / pl.col("shared_genes_count") * 100)

def get_unique_scaffolds(mpile_contig:pl.LazyFrame,batch_size:int=10000) -> set:
    """
    Retrieves unique scaffolds from the mpileup LazyFrame.

    Args:
        mpile_contig (pl.LazyFrame): The input LazyFrame containing mpileup data.
        batch_size (int): The number of rows to process in each batch. Default is 10000.
    Returns:
        set: A set of unique scaffold names.
    """
    scaffolds = set()
    start_index = 0
    while True:
        batch = mpile_contig.slice(start_index, batch_size).select("chrom").collect()
        if batch.height == 0:
            break
        scaffolds.update(batch["chrom"].to_list())
        start_index += batch_size
    return scaffolds 


def compare_genomes(mpile_contig_1:Union[str, Path, pl.LazyFrame],
              mpile_contig_2:Union[str, Path, pl.LazyFrame],
              min_cov:int=5,
              min_gene_compare_len:int=100,
              genome_scope:str="all",
              ani_method:str="popani",
              duckdb_memory_limit: Optional[str] = None,
              duckdb_temp_directory: Optional[Union[str, Path]] = None,
              duckdb_threads: Optional[int] = None,
            )-> pl.LazyFrame:
    """
    Compare two profiles and return genome-level statistics.

    DuckDB is used for join/filter and Polars for final aggregation.
    Optional DuckDB settings allow memory and spill tuning.
    """
    lf= duckdb_filter_join(
        mpile1=mpile_contig_1,
        mpile2=mpile_contig_2,
        min_cov=min_cov,
        ani_method=ani_method,
        genome_scope=genome_scope,
        memory_limit=duckdb_memory_limit,
        temp_directory=duckdb_temp_directory,
        threads=duckdb_threads,
    )
    lf= add_contiguity_info(lf)
    genome_comp= calculate_pop_ani(lf)
    ## Calculate longest consecutive blocks
    max_consecutive_per_genome = get_longest_consecutive_blocks(lf)
    ## Calculate gene ani for each gene in each genome
    gene= get_gene_ani(lf, min_gene_compare_len)
    genome_comp=genome_comp.join(max_consecutive_per_genome, on="genome", how="left")
    genome_comp=genome_comp.join(gene, on="genome", how="left")
    return genome_comp



def compare_genes(mpile_contig_1:Union[str, Path, pl.LazyFrame],
              mpile_contig_2:Union[str, Path, pl.LazyFrame],
              min_cov:int=5,
              min_gene_compare_len:int=100,
              genome_scope:str="all",
              gene_scope:str="all",
              ani_method:str="popani",
              duckdb_memory_limit: Optional[str] = None,
              duckdb_temp_directory: Optional[Union[str, Path]] = None,
              duckdb_threads: Optional[int] = None,
            )-> pl.LazyFrame:
    """
    Compare two profiles and return gene-level ANI statistics.
    """
 
    lf = duckdb_filter_join(
        mpile1=mpile_contig_1,
        mpile2=mpile_contig_2,
        min_cov=min_cov,
        ani_method=ani_method,
        genome_scope=genome_scope,
        gene_scope=gene_scope,
        memory_limit=duckdb_memory_limit,
        temp_directory=duckdb_temp_directory,
        threads=duckdb_threads,
    )
    gene_comp = lf.group_by(["genome", "gene"]).agg(
        total_positions=pl.len(),
        share_allele_pos=(pl.col("surr") > 0).sum()
    ).filter(pl.col("total_positions") >= min_gene_compare_len).with_columns(
        ani=pl.col("share_allele_pos") / pl.col("total_positions") * 100,
    )
    return gene_comp
