from abc import ABC, abstractmethod

import numpy as np
from numpy.typing import NDArray

from Heimdall.fc import Fc


class Order(ABC):
    def __init__(
        self,
        fc: Fc,
    ):
        self.fc = fc

    @abstractmethod
    def __call__(
        self,
        cell_index: int,
        identity_indices: NDArray,
        expression_inputs: NDArray,
    ) -> NDArray:
        """Order cell tokens using metadata.

        Gene tokens can be reordered based on e.g. expression level, chromosome position, etc.

        Args:
            cell_tokenization: the stacked gene identity- and gene expression-based tokenization
                of a cell.

        """


class ExpressionOrder(Order):
    def __call__(self, cell_index: int, identity_indices: NDArray, expression_inputs: NDArray) -> NDArray:
        """Order cell tokens using metadata.

        Gene tokens are reordered based on expression level.
        Args:
            cell_tokenization: the stacked gene identity- and gene expression-based tokenization
                of a cell.

        """
        # cols = np.asarray(identity_indices, dtype=int)

        adata = self.fc.adata
        cell_slice = adata[cell_index].copy()
        raw_expression_values = np.squeeze(cell_slice.X.toarray())
        raw_expression_values = raw_expression_values[identity_indices]

        if "medians" in adata.var:
            medians = adata.var["medians"][identity_indices].astype(self.fc.dtype)
            raw_expression_values = raw_expression_values - medians

        # Sort non-zero values in descending order
        raw_expression_values = np.where(np.isnan(raw_expression_values), -np.inf, raw_expression_values)
        gene_order = np.argsort(raw_expression_values)[::-1]  # Indices for sorting descending

        return gene_order


class RandomOrder(Order):
    def __call__(self, cell_index: int, identity_indices: NDArray, expression_inputs: NDArray) -> NDArray:
        # TODO: consider cleaning up sampling (just sample all nonzero and all zero, then concat
        (nonzero_indices,) = np.where(expression_inputs != 0)
        (zero_indices,) = np.where(expression_inputs == 0)

        # First: sample/reorder nonzero expression tokens
        # num_nonzero_to_sample = min(len(nonzero_indices), self.fc.max_input_length)
        num_nonzero = len(nonzero_indices)
        num_zero = len(zero_indices)

        # selected_nonzero = self.fc.rng.choice(nonzero_indices, num_nonzero_to_sample, replace=False)
        selected_nonzero = self.fc.rng.choice(nonzero_indices, num_nonzero, replace=False)

        # If needed: sample zero-expression tokens to fill up
        # num_remaining = self.fc.max_input_length - num_nonzero_to_sample
        # if num_remaining > 0:
        if num_zero > 0:
            selected_zero = self.fc.rng.choice(zero_indices, num_zero, replace=False)
            gene_order = np.concatenate([selected_nonzero, selected_zero])
        else:
            gene_order = selected_nonzero

        # Optionally shuffle to avoid position bias, but we dont need to because the gene ids are the position
        # self.rng.shuffle(final_indices)

        return gene_order


class ChromosomeOrder(Order):
    def __call__(self, cell_index: int, identity_indices: NDArray, expression_inputs: NDArray) -> NDArray:
        """Order cell tokens using metadata.

        Gene tokens are reordered based on chromosome location.

        Args:
            cell_tokenization: the stacked gene identity- and gene expression-based tokenization
                of a cell.

        """

        choosen_chrom = self.fc.chroms.iloc[identity_indices]

        unique_chromosomes = np.unique(choosen_chrom)
        self.fc.shuffled_chromosomes = self.fc.rng.permutation(unique_chromosomes)

        gene_order = np.zeros(len(identity_indices), dtype=np.int32)
        for chromosome in self.fc.shuffled_chromosomes:
            (chromosome_index,) = np.where(choosen_chrom == chromosome)
            sort_by_start = np.argsort(
                self.fc.starts[chromosome_index],
            )  # start chromosome_indexations for this chromsome

            gene_order[chromosome_index] = chromosome_index[sort_by_start]

        return gene_order
