"""
Utilities for aligning model predictions between reference and variant sequences.

This module provides functions to handle the alignment of ML model predictions
when reference and variant sequences have position offsets due to structural variants.

The alignment logic properly handles:
- 1D predictions (chromatin accessibility, TF binding, etc.)
- 2D contact maps (Hi-C, Micro-C predictions)
- All variant types: SNV, INS, DEL, DUP, INV, BND

Key principle: Users must specify all model-specific parameters (bin_size, diag_offset)
as these vary across different prediction models.
"""

import numpy as np
import warnings
from typing import Tuple, List, Optional, Union
from dataclasses import dataclass

try:
    import torch

    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False


@dataclass
class VariantPosition:
    """
    Container for variant position information in both REF and ALT sequences.

    This class encapsulates the essential positional information needed to align
    predictions across reference and alternate sequences that may differ in length.
    """

    ref_pos: int  # Position in reference sequence (base pairs, 0-based)
    alt_pos: int  # Position in alternate sequence (base pairs, 0-based)
    svlen: int  # Length of structural variant (base pairs, signed for DEL/INS)
    variant_type: str  # Type of variant ('SNV', 'INS', 'DEL', 'DUP', 'INV', 'BND')

    def get_bin_positions(self, bin_size: int) -> Tuple[int, int, int]:
        """
        Convert base pair positions to bin indices.

        Args:
            bin_size: Number of base pairs per prediction bin

        Returns:
            Tuple of (ref_bin, alt_start_bin, alt_end_bin)
        """
        ref_bin = int(np.ceil(self.ref_pos / bin_size))
        alt_start_bin = int(np.ceil(self.alt_pos / bin_size))
        alt_end_bin = int(np.ceil((self.alt_pos + abs(self.svlen)) / bin_size))
        return ref_bin, alt_start_bin, alt_end_bin


class PredictionAligner1D:
    """
    Aligns reference and alternate 1D prediction vectors for variant comparison.

    Handles alignment of 1D genomic predictions (e.g., chromatin accessibility,
    transcription factor binding, epigenetic marks) between reference and variant
    sequences that may differ in length due to structural variants.

    The aligner uses a masking strategy where positions that exist in one sequence
    but not the other are marked with NaN values, enabling direct comparison of
    corresponding genomic positions.

    Args:
        target_size: Expected number of bins in the prediction output
        bin_size: Number of base pairs per prediction bin (model-specific)

    Example:
        >>> aligner = PredictionAligner1D(target_size=896, bin_size=128)
        >>> ref_aligned, alt_aligned = aligner.align_predictions(
        ...     ref_pred, alt_pred, 'INS', variant_position
        ... )
    """

    def __init__(self, target_size: int, bin_size: int):
        """
        Initialize the 1D prediction aligner.

        Args:
            target_size: Expected number of bins in prediction (e.g., 896 for Enformer)
            bin_size: Base pairs per bin (e.g., 128 for Enformer)
        """
        self.target_size = target_size
        self.bin_size = bin_size

    def align_predictions(
        self,
        ref_pred: Union[np.ndarray, "torch.Tensor"],
        alt_pred: Union[np.ndarray, "torch.Tensor"],
        svtype: str,
        var_pos: VariantPosition,
    ) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
        """
        Main entry point for 1D prediction alignment.

        Args:
            ref_pred: Reference prediction vector (length N)
            alt_pred: Alternate prediction vector (length N)
            svtype: Variant type ('DEL', 'DUP', 'INS', 'INV', 'SNV')
            var_pos: Variant position information

        Returns:
            Tuple of (aligned_ref, aligned_alt) vectors with NaN masking applied

        Raises:
            ValueError: For unsupported variant types or if using BND (use align_bnd_predictions)
        """
        if svtype == "BND" or svtype == "SV_BND":
            raise ValueError("Use align_bnd_predictions() for breakends")

        # Normalize variant type names
        svtype_normalized = svtype.replace("SV_", "")

        if svtype_normalized in ["DEL", "DUP", "INS"]:
            return self._align_indel_predictions(
                ref_pred, alt_pred, svtype_normalized, var_pos
            )
        elif svtype_normalized == "INV":
            return self._align_inversion_predictions(ref_pred, alt_pred, var_pos)
        elif svtype_normalized in ["SNV", "MNV"]:
            # SNVs don't change coordinates, direct alignment
            is_torch = TORCH_AVAILABLE and torch.is_tensor(ref_pred)
            if is_torch:
                return ref_pred.clone(), alt_pred.clone()
            else:
                return ref_pred.copy(), alt_pred.copy()
        else:
            raise ValueError(f"Unknown variant type: {svtype}")

    def _align_indel_predictions(
        self,
        ref_pred: Union[np.ndarray, "torch.Tensor"],
        alt_pred: Union[np.ndarray, "torch.Tensor"],
        svtype: str,
        var_pos: VariantPosition,
    ) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
        """
        Align predictions for insertions, deletions, and duplications.

        Strategy:
        1. For DEL: Swap REF/ALT (deletion removes from REF)
        2. Insert NaN bins in shorter sequence
        3. Crop edges to maintain target size
        4. For DEL: Swap back

        This ensures that positions present in one sequence but not the other
        are marked with NaN, enabling fair comparison of overlapping regions.
        """
        is_torch = TORCH_AVAILABLE and torch.is_tensor(ref_pred)

        # Convert to numpy for manipulation
        if is_torch:
            ref_np = ref_pred.detach().cpu().numpy()
            alt_np = alt_pred.detach().cpu().numpy()
        else:
            ref_np = ref_pred
            alt_np = alt_pred

        # Swap for deletions (treat as insertion in reverse)
        if svtype == "DEL":
            ref_np, alt_np = alt_np, ref_np
            var_pos = VariantPosition(
                var_pos.alt_pos, var_pos.ref_pos, var_pos.svlen, svtype
            )

        # Get bin positions
        ref_bin, alt_start_bin, alt_end_bin = var_pos.get_bin_positions(self.bin_size)
        bins_to_add = alt_end_bin - alt_start_bin

        # Insert NaN bins in REF where variant exists in ALT
        ref_masked = self._insert_nan_bins(ref_np, ref_bin, bins_to_add)

        # Crop to maintain target size
        ref_masked = self._crop_vector(ref_masked, ref_bin, alt_start_bin)
        alt_masked = alt_np.copy()

        # Swap back for deletions
        if svtype == "DEL":
            ref_masked, alt_masked = alt_masked, ref_masked

        self._validate_size(ref_masked, alt_masked)

        # Convert back to torch if needed
        if is_torch:
            ref_masked = (
                torch.from_numpy(ref_masked).to(ref_pred.device).type(ref_pred.dtype)
            )
            alt_masked = (
                torch.from_numpy(alt_masked).to(alt_pred.device).type(alt_pred.dtype)
            )

        return ref_masked, alt_masked

    def _insert_nan_bins(
        self, vector: np.ndarray, position: int, num_bins: int
    ) -> np.ndarray:
        """
        Insert NaN values at specified position in vector.

        Args:
            vector: Input prediction vector
            position: Position to insert NaN values (bin index)
            num_bins: Number of NaN values to insert

        Returns:
            Vector with NaN values inserted
        """
        result = vector.copy()
        for offset in range(num_bins):
            insert_pos = position + offset
            result = np.insert(result, insert_pos, np.nan)
        return result

    def _crop_vector(
        self, vector: np.ndarray, ref_bin: int, alt_bin: int
    ) -> np.ndarray:
        """
        Crop vector edges to maintain target size.

        After inserting NaN bins, the vector is longer than expected.
        This function crops from edges proportionally to center the variant.

        Args:
            vector: Vector to crop
            ref_bin: Reference bin position
            alt_bin: Alternate bin position

        Returns:
            Cropped vector of target_size length
        """
        remove_left = ref_bin - alt_bin
        remove_right = len(vector) - self.target_size - remove_left

        # Apply cropping
        start = max(0, remove_left)
        end = len(vector) - max(0, remove_right)
        return vector[start:end]

    def _align_inversion_predictions(
        self,
        ref_pred: Union[np.ndarray, "torch.Tensor"],
        alt_pred: Union[np.ndarray, "torch.Tensor"],
        var_pos: VariantPosition,
    ) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
        """
        Align predictions for inversions.

        Strategy:
        1. Mask the inverted region in both vectors with NaN
        2. This allows comparison of only the flanking (unaffected) regions

        For strand-aware models, inversions can significantly affect predictions
        because regulatory elements now appear on the opposite strand. We mask
        the inverted region to focus comparison on unaffected flanking sequences.
        """
        is_torch = TORCH_AVAILABLE and torch.is_tensor(ref_pred)

        # Convert to numpy for manipulation
        if is_torch:
            ref_np = ref_pred.detach().cpu().numpy()
            alt_np = alt_pred.detach().cpu().numpy()
        else:
            ref_np = ref_pred.copy()
            alt_np = alt_pred.copy()

        var_start, _, var_end = var_pos.get_bin_positions(self.bin_size)

        # Mask inverted region in both REF and ALT
        ref_np[var_start : var_end + 1] = np.nan
        alt_np[var_start : var_end + 1] = np.nan

        self._validate_size(ref_np, alt_np)

        # Convert back to torch if needed
        if is_torch:
            ref_np = torch.from_numpy(ref_np).to(ref_pred.device).type(ref_pred.dtype)
            alt_np = torch.from_numpy(alt_np).to(alt_pred.device).type(alt_pred.dtype)

        return ref_np, alt_np

    def align_bnd_predictions(
        self,
        left_ref: Union[np.ndarray, "torch.Tensor"],
        right_ref: Union[np.ndarray, "torch.Tensor"],
        bnd_alt: Union[np.ndarray, "torch.Tensor"],
        breakpoint_bin: int,
    ) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
        """
        Align predictions for breakends (chromosomal rearrangements).

        BNDs join two distant loci, so we create a chimeric reference
        prediction from the two separate loci for comparison with the fusion ALT.

        Args:
            left_ref: Prediction from left locus
            right_ref: Prediction from right locus
            bnd_alt: Prediction from joined (alternate) sequence
            breakpoint_bin: Bin position of breakpoint

        Returns:
            Tuple of (chimeric_ref, alt) vectors
        """
        is_torch = TORCH_AVAILABLE and torch.is_tensor(left_ref)

        # Convert to numpy for manipulation
        if is_torch:
            left_np = left_ref.detach().cpu().numpy()
            right_np = right_ref.detach().cpu().numpy()
            alt_np = bnd_alt.detach().cpu().numpy()
        else:
            left_np = left_ref
            right_np = right_ref
            alt_np = bnd_alt

        # Extract the relevant portions from each reference
        left_portion = left_np[:breakpoint_bin]
        right_portion = right_np[-(self.target_size - breakpoint_bin) :]

        # Assemble chimeric reference (continuous, no masking)
        ref_chimeric = np.concatenate([left_portion, right_portion])

        self._validate_size(ref_chimeric, alt_np)

        # Convert back to torch if needed
        if is_torch:
            ref_chimeric = (
                torch.from_numpy(ref_chimeric).to(left_ref.device).type(left_ref.dtype)
            )
            alt_np = torch.from_numpy(alt_np).to(bnd_alt.device).type(bnd_alt.dtype)

        return ref_chimeric, alt_np

    def _validate_size(self, ref_vector: np.ndarray, alt_vector: np.ndarray):
        """
        Validate that vectors are the correct size.

        Args:
            ref_vector: Reference prediction vector
            alt_vector: Alternate prediction vector

        Raises:
            ValueError: If either vector has incorrect size
        """
        if len(ref_vector) != self.target_size:
            raise ValueError(
                f"Reference vector wrong size: {len(ref_vector)} vs {self.target_size}"
            )
        if len(alt_vector) != self.target_size:
            raise ValueError(
                f"Alternate vector wrong size: {len(alt_vector)} vs {self.target_size}"
            )


class PredictionAligner2D:
    """
    Aligns reference and alternate prediction matrices for variant comparison.

    Handles alignment of 2D genomic predictions (e.g., Hi-C contact maps,
    Micro-C predictions) between reference and variant sequences that may
    differ in length due to structural variants.

    The aligner uses a masking strategy where matrix rows and columns that
    exist in one sequence but not the other are marked with NaN values.

    Args:
        target_size: Expected matrix dimension (NxN)
        bin_size: Number of base pairs per matrix bin (model-specific)
        diag_offset: Number of diagonal bins to mask (model-specific)

    Example:
        >>> aligner = PredictionAligner2D(
        ...     target_size=448,
        ...     bin_size=2048,
        ...     diag_offset=2
        ... )
        >>> ref_aligned, alt_aligned = aligner.align_predictions(
        ...     ref_matrix, alt_matrix, 'DEL', variant_position
        ... )
    """

    def __init__(self, target_size: int, bin_size: int, diag_offset: int):
        """
        Initialize the 2D prediction aligner.

        Args:
            target_size: Matrix dimension (e.g., 448 for Akita)
            bin_size: Base pairs per bin (e.g., 2048 for Akita)
            diag_offset: Diagonal masking offset (e.g., 2 for Akita)
        """
        self.target_size = target_size
        self.bin_size = bin_size
        self.diag_offset = diag_offset

    def align_predictions(
        self,
        ref_pred: Union[np.ndarray, "torch.Tensor"],
        alt_pred: Union[np.ndarray, "torch.Tensor"],
        svtype: str,
        var_pos: VariantPosition,
    ) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
        """
        Main entry point for 2D matrix alignment.

        Args:
            ref_pred: Reference prediction matrix (NxN)
            alt_pred: Alternate prediction matrix (NxN)
            svtype: Variant type ('DEL', 'DUP', 'INS', 'INV', 'SNV')
            var_pos: Variant position information

        Returns:
            Tuple of (aligned_ref, aligned_alt) matrices with NaN masking applied

        Raises:
            ValueError: For unsupported variant types or if using BND (use align_bnd_matrices)
        """
        if svtype == "BND" or svtype == "SV_BND":
            raise ValueError("Use align_bnd_matrices() for breakends")

        # Normalize variant type names
        svtype_normalized = svtype.replace("SV_", "")

        if svtype_normalized in ["DEL", "DUP", "INS"]:
            return self._align_indel_matrices(
                ref_pred, alt_pred, svtype_normalized, var_pos
            )
        elif svtype_normalized == "INV":
            return self._align_inversion_matrices(ref_pred, alt_pred, var_pos)
        elif svtype_normalized in ["SNV", "MNV"]:
            # SNVs don't change coordinates, direct alignment
            is_torch = TORCH_AVAILABLE and torch.is_tensor(ref_pred)
            if is_torch:
                return ref_pred.clone(), alt_pred.clone()
            else:
                return ref_pred.copy(), alt_pred.copy()
        else:
            raise ValueError(f"Unknown variant type: {svtype}")

    def _align_indel_matrices(
        self,
        ref_pred: Union[np.ndarray, "torch.Tensor"],
        alt_pred: Union[np.ndarray, "torch.Tensor"],
        svtype: str,
        var_pos: VariantPosition,
    ) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
        """
        Align matrices for insertions, deletions, and duplications.

        Strategy:
        1. For DEL: Swap REF/ALT (deletion removes from REF)
        2. Insert NaN bins (rows AND columns) in shorter matrix
        3. Crop edges to maintain target size
        4. For DEL: Swap back
        """
        is_torch = TORCH_AVAILABLE and torch.is_tensor(ref_pred)

        # Convert to numpy for manipulation
        if is_torch:
            ref_np = ref_pred.detach().cpu().numpy()
            alt_np = alt_pred.detach().cpu().numpy()
        else:
            ref_np = ref_pred
            alt_np = alt_pred

        # Swap for deletions (treat as insertion in reverse)
        if svtype == "DEL":
            ref_np, alt_np = alt_np, ref_np
            var_pos = VariantPosition(
                var_pos.alt_pos, var_pos.ref_pos, var_pos.svlen, svtype
            )

        # Get bin positions
        ref_bin, alt_start_bin, alt_end_bin = var_pos.get_bin_positions(self.bin_size)
        bins_to_add = alt_end_bin - alt_start_bin

        # Insert NaN bins in REF where variant exists in ALT
        ref_masked = self._insert_nan_bins(ref_np, ref_bin, bins_to_add)

        # Crop to maintain target size
        ref_masked = self._crop_matrix(ref_masked, ref_bin, alt_start_bin)
        alt_masked = alt_np.copy()

        # Swap back for deletions
        if svtype == "DEL":
            ref_masked, alt_masked = alt_masked, ref_masked

        self._validate_size(ref_masked, alt_masked)

        # Convert back to torch if needed
        if is_torch:
            ref_masked = (
                torch.from_numpy(ref_masked).to(ref_pred.device).type(ref_pred.dtype)
            )
            alt_masked = (
                torch.from_numpy(alt_masked).to(alt_pred.device).type(alt_pred.dtype)
            )

        return ref_masked, alt_masked

    def _insert_nan_bins(
        self, matrix: np.ndarray, position: int, num_bins: int
    ) -> np.ndarray:
        """
        Insert NaN bins (rows and columns) at specified position.

        For 2D matrices, we must insert both rows AND columns to maintain
        the square matrix structure and properly mask interactions.
        """
        result = matrix.copy()
        for offset in range(num_bins):
            insert_pos = position + offset
            result = np.insert(result, insert_pos, np.nan, axis=0)
            result = np.insert(result, insert_pos, np.nan, axis=1)
        return result

    def _crop_matrix(
        self, matrix: np.ndarray, ref_bin: int, alt_bin: int
    ) -> np.ndarray:
        """
        Crop matrix edges to maintain target size.

        After inserting NaN bins, the matrix is larger than expected.
        This function crops from edges proportionally to center the variant.
        """
        remove_left = ref_bin - alt_bin
        remove_right = len(matrix) - self.target_size - remove_left

        # Apply cropping
        start = max(0, remove_left)
        end = len(matrix) - max(0, remove_right)
        return matrix[start:end, start:end]

    def _align_inversion_matrices(
        self,
        ref_pred: Union[np.ndarray, "torch.Tensor"],
        alt_pred: Union[np.ndarray, "torch.Tensor"],
        var_pos: VariantPosition,
    ) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
        """
        Align matrices for inversions.

        Strategy: Mask the inverted region in both REF and ALT matrices using
        a cross-pattern (mask entire rows AND columns at the inversion position).

        Why cross-masking?
        - Inversions reverse the sequence, creating geometric rotation in contact maps
        - Masking rows removes interactions with the inverted region (one dimension)
        - Masking columns removes interactions from the inverted region (other dimension)
        - This ensures only flanking regions (unaffected by inversion) are compared

        The same NaN pattern is mirrored to ALT so both matrices have identical
        masked regions, enabling fair comparison of the unaffected areas.
        """
        is_torch = TORCH_AVAILABLE and torch.is_tensor(ref_pred)

        # Convert to numpy for manipulation
        if is_torch:
            ref_np = ref_pred.detach().cpu().numpy()
            alt_np = alt_pred.detach().cpu().numpy()
        else:
            ref_np = ref_pred.copy()
            alt_np = alt_pred.copy()

        var_start, _, var_end = var_pos.get_bin_positions(self.bin_size)

        # Mask inverted region in REF (cross-pattern: rows + columns)
        ref_np[var_start : var_end + 1, :] = np.nan
        ref_np[:, var_start : var_end + 1] = np.nan

        # Mirror NaN pattern to ALT (correct approach)
        nan_mask = ref_np.copy()
        nan_mask[np.invert(np.isnan(nan_mask))] = 0  # Non-NaN → 0, NaN stays NaN
        alt_np = alt_np + nan_mask  # Adding NaN propagates NaN to ALT

        self._validate_size(ref_np, alt_np)

        # Convert back to torch if needed
        if is_torch:
            ref_np = torch.from_numpy(ref_np).to(ref_pred.device).type(ref_pred.dtype)
            alt_np = torch.from_numpy(alt_np).to(alt_pred.device).type(alt_pred.dtype)

        return ref_np, alt_np

    def align_bnd_matrices(
        self,
        left_ref: Union[np.ndarray, "torch.Tensor"],
        right_ref: Union[np.ndarray, "torch.Tensor"],
        bnd_alt: Union[np.ndarray, "torch.Tensor"],
        breakpoint_bin: int,
    ) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
        """
        Align matrices for breakends (chromosomal rearrangements).

        BNDs join two distant loci, so we create a chimeric reference
        matrix from the two separate loci.

        Args:
            left_ref: Prediction from left locus
            right_ref: Prediction from right locus
            bnd_alt: Prediction from joined (alternate) sequence
            breakpoint_bin: Bin position of breakpoint

        Returns:
            Tuple of (chimeric_ref, alt) matrices
        """
        is_torch = TORCH_AVAILABLE and torch.is_tensor(left_ref)

        # Convert to numpy for manipulation
        if is_torch:
            left_np = left_ref.detach().cpu().numpy()
            right_np = right_ref.detach().cpu().numpy()
            alt_np = bnd_alt.detach().cpu().numpy()
        else:
            left_np = left_ref
            right_np = right_ref
            alt_np = bnd_alt

        # Assemble chimeric matrix from two loci
        ref_chimeric = self._assemble_chimeric_matrix(left_np, right_np, breakpoint_bin)

        self._validate_size(ref_chimeric, alt_np)

        # Convert back to torch if needed
        if is_torch:
            ref_chimeric = (
                torch.from_numpy(ref_chimeric).to(left_ref.device).type(left_ref.dtype)
            )
            alt_np = torch.from_numpy(alt_np).to(bnd_alt.device).type(bnd_alt.dtype)

        return ref_chimeric, alt_np

    def _assemble_chimeric_matrix(
        self, left_matrix: np.ndarray, right_matrix: np.ndarray, breakpoint: int
    ) -> np.ndarray:
        """
        Assemble chimeric matrix from two loci.

        Structure:
        - Upper left quadrant: left locus
        - Lower right quadrant: right locus
        - Upper right/lower left quadrants: NaN (no trans prediction)
        """
        matrix = np.zeros((self.target_size, self.target_size))

        # Fill upper left quadrant (left locus)
        matrix[:breakpoint, :breakpoint] = left_matrix[:breakpoint, :breakpoint]

        # Fill lower right quadrant (right locus)
        matrix[breakpoint:, breakpoint:] = right_matrix[breakpoint:, breakpoint:]

        # Fill transition quadrants with NaN
        matrix[:breakpoint, breakpoint:] = np.nan
        matrix[breakpoint:, :breakpoint] = np.nan

        # Mask diagonals as specified by model
        for offset in range(-self.diag_offset + 1, self.diag_offset):
            if offset < 0:
                np.fill_diagonal(matrix[abs(offset) :, :], np.nan)
            else:
                np.fill_diagonal(matrix[:, offset:], np.nan)

        return matrix

    def _validate_size(self, ref_matrix: np.ndarray, alt_matrix: np.ndarray):
        """
        Validate that matrices are the correct size.

        Args:
            ref_matrix: Reference prediction matrix
            alt_matrix: Alternate prediction matrix

        Raises:
            ValueError: If either matrix has incorrect size
        """
        if ref_matrix.shape[0] != self.target_size:
            raise ValueError(
                f"Reference matrix wrong size: {ref_matrix.shape[0]} vs {self.target_size}"
            )
        if alt_matrix.shape[0] != self.target_size:
            raise ValueError(
                f"Alternate matrix wrong size: {alt_matrix.shape[0]} vs {self.target_size}"
            )


# =============================================================================
# Utility Functions for Contact Maps
# =============================================================================


def vector_to_contact_matrix(
    vector: Union[np.ndarray, "torch.Tensor"], matrix_size: int, diag_offset: int = 0
) -> Union[np.ndarray, "torch.Tensor"]:
    """
    Convert flattened upper triangular vector to full contact matrix.

    This function reconstructs a full symmetric contact matrix from its upper
    triangular representation, following the pattern used in genomic contact map models.
    Supports diagonal masking where near-diagonal elements are excluded.

    Args:
        vector: Flattened upper triangular matrix
                Expected length depends on diag_offset:
                - diag_offset=0: matrix_size * (matrix_size + 1) / 2 (includes diagonal)
                - diag_offset=k: (matrix_size - k) * (matrix_size - k + 1) / 2
        matrix_size: Dimension of the output square matrix
        diag_offset: Diagonal offset for masking (default=0, no masking)
                     diag_offset=2 means skip main diagonal and first off-diagonal

    Returns:
        Full symmetric contact matrix of shape (matrix_size, matrix_size)
        Elements within diag_offset of the diagonal are set to NaN

    Example:
        >>> # Full upper triangle (diag_offset=0, default)
        >>> vector = np.array([1, 2, 3, 4, 5, 6])
        >>> matrix = vector_to_contact_matrix(vector, 3)
        >>> # Result: [[1, 2, 3], [2, 4, 5], [3, 5, 6]]
    """
    # Handle both PyTorch tensors and NumPy arrays
    is_torch = TORCH_AVAILABLE and torch.is_tensor(vector)

    if is_torch:
        # Initialize matrix with NaN
        matrix = torch.full(
            (matrix_size, matrix_size),
            float("nan"),
            dtype=vector.dtype,
            device=vector.device,
        )
        # Get upper triangle indices with diagonal offset
        triu_indices = torch.triu_indices(matrix_size, matrix_size, offset=diag_offset)
        matrix[triu_indices[0], triu_indices[1]] = vector
        # Make symmetric by copying upper triangle to lower (preserving NaN)
        for i in range(matrix_size):
            for j in range(i + diag_offset, matrix_size):
                if not torch.isnan(matrix[i, j]):
                    matrix[j, i] = matrix[i, j]
    else:
        # Initialize matrix with NaN
        matrix = np.full((matrix_size, matrix_size), np.nan, dtype=vector.dtype)
        # Get upper triangle indices with diagonal offset
        triu_indices = np.triu_indices(matrix_size, k=diag_offset)
        matrix[triu_indices] = vector
        # Make symmetric by copying upper triangle to lower (preserving NaN)
        for i in range(matrix_size):
            for j in range(i + diag_offset, matrix_size):
                if not np.isnan(matrix[i, j]):
                    matrix[j, i] = matrix[i, j]

    return matrix


def contact_matrix_to_vector(
    matrix: Union[np.ndarray, "torch.Tensor"], diag_offset: int = 0
) -> Union[np.ndarray, "torch.Tensor"]:
    """
    Convert full contact matrix to flattened upper triangular vector.

    This function extracts the upper triangular portion of a contact matrix,
    which is the standard representation for genomic contact maps. Supports
    diagonal masking to exclude near-diagonal elements.

    Args:
        matrix: Full symmetric contact matrix of shape (N, N)
        diag_offset: Diagonal offset for extraction (default=0, includes diagonal)
                     diag_offset=2 means skip main diagonal and first off-diagonal

    Returns:
        Flattened upper triangular vector
        Length depends on diag_offset:
        - diag_offset=0: N * (N + 1) / 2 (includes diagonal)
        - diag_offset=k: (N - k) * (N - k + 1) / 2

    Example:
        >>> # Full upper triangle (diag_offset=0, default)
        >>> matrix = np.array([[1, 2, 3], [2, 4, 5], [3, 5, 6]])
        >>> vector = contact_matrix_to_vector(matrix)
        >>> # Result: [1, 2, 3, 4, 5, 6]
    """
    # Handle both PyTorch tensors and NumPy arrays
    is_torch = TORCH_AVAILABLE and torch.is_tensor(matrix)

    if is_torch:
        triu_indices = torch.triu_indices(
            matrix.shape[0], matrix.shape[1], offset=diag_offset
        )
        return matrix[triu_indices[0], triu_indices[1]]
    else:
        triu_indices = np.triu_indices(matrix.shape[0], k=diag_offset)
        return matrix[triu_indices]


def align_predictions_by_coordinate(
    ref_preds: Union[np.ndarray, "torch.Tensor"],
    alt_preds: Union[np.ndarray, "torch.Tensor"],
    metadata_row: dict,
    bin_size: int,
    prediction_type: str,
    matrix_size: Optional[int] = None,
    diag_offset: int = 0,
) -> Tuple[Union[np.ndarray, "torch.Tensor"], Union[np.ndarray, "torch.Tensor"]]:
    """
    Align reference and alt predictions using coordinate transformation and variant type awareness.

    This is the main public API for prediction alignment. It handles both 1D prediction
    vectors (e.g., chromatin accessibility, TF binding) and 2D matrices (e.g., Hi-C contact maps),
    routing to the appropriate alignment strategy based on variant type.

    IMPORTANT: Model-specific parameters (bin_size, matrix_size) must be explicitly
    provided by the user. There are no defaults because these vary across different models.

    Args:
        ref_preds: Reference predictions array (from model with edge cropping)
        alt_preds: Alt predictions array (same shape as ref_preds)
        metadata_row: Dictionary with variant information containing:
            - 'variant_type': Type of variant (SNV, INS, DEL, DUP, INV, BND)
            - 'window_start': Start position of window (0-based)
            - 'variant_pos0': Variant position (0-based, absolute genomic coordinate)
            - 'svlen': Length of structural variant (optional, for symbolic alleles)
        bin_size: Number of base pairs per prediction bin (REQUIRED, model-specific)
            Examples: 2048 for Akita
        prediction_type: Type of predictions ("1D" or "2D")
            - "1D": Vector predictions (chromatin accessibility, TF binding, etc.)
            - "2D": Matrix predictions (Hi-C contact maps, Micro-C, etc.)
        matrix_size: Size of contact matrix (REQUIRED for 2D type)
            Examples: 448 for Akita
        diag_offset: Number of diagonal bins to mask (default: 0 for no masking)
            Set to 0 if your model doesn't mask diagonals, or to model-specific value
            Examples: 2 for Akita, 0 for models without diagonal masking

    Returns:
        Tuple of (aligned_ref_preds, aligned_alt_preds) with NaN masking applied
        at positions that differ between reference and alternate sequences

    Raises:
        ValueError: If prediction_type is invalid, required parameters are missing,
            or variant type is unsupported

    Example (1D predictions):
        >>> ref_aligned, alt_aligned = align_predictions_by_coordinate(
        ...     ref_preds=ref_scores,
        ...     alt_preds=alt_scores,
        ...     metadata_row={'variant_type': 'INS', 'window_start': 0,
        ...                   'variant_pos0': 500, 'svlen': 100},
        ...     bin_size=128,
        ...     prediction_type="1D"
        ... )

    Example (2D contact maps with diagonal masking):
        >>> ref_aligned, alt_aligned = align_predictions_by_coordinate(
        ...     ref_preds=ref_contact_map,
        ...     alt_preds=alt_contact_map,
        ...     metadata_row={'variant_type': 'DEL', 'window_start': 0,
        ...                   'variant_pos0': 50000, 'svlen': -2048},
        ...     bin_size=2048,
        ...     prediction_type="2D",
        ...     matrix_size=448,
        ...     diag_offset=2  # Optional: use 0 if no diagonal masking
        ... )

    Example (2D contact maps without diagonal masking):
        >>> ref_aligned, alt_aligned = align_predictions_by_coordinate(
        ...     ref_preds=ref_contact_map,
        ...     alt_preds=alt_contact_map,
        ...     metadata_row={'variant_type': 'INS', 'window_start': 0,
        ...                   'variant_pos0': 1000, 'svlen': 500},
        ...     bin_size=1000,
        ...     prediction_type="2D",
        ...     matrix_size=512
        ...     # diag_offset defaults to 0 (no masking)
        ... )
    """
    # Validate prediction type and parameters
    if prediction_type not in ["1D", "2D"]:
        raise ValueError(
            f"prediction_type must be '1D' or '2D', got '{prediction_type}'"
        )

    if prediction_type == "2D":
        if matrix_size is None:
            raise ValueError("matrix_size must be provided for 2D prediction type")

    # Extract variant information from metadata
    variant_type = metadata_row.get("variant_type", "unknown")
    window_start = metadata_row.get("window_start", 0)
    variant_pos0 = metadata_row.get("variant_pos0")

    # For backward compatibility, check for effective_variant_start (deprecated)
    if variant_pos0 is not None:
        abs_variant_pos = variant_pos0
    else:
        # Fall back to old field name if present (for backward compatibility)
        effective_variant_start = metadata_row.get("effective_variant_start", 0)
        abs_variant_pos = window_start + effective_variant_start

    svlen = metadata_row.get("svlen", None)

    # Calculate svlen from alleles if not present (for non-symbolic DEL/INS variants)
    # Symbolic variants (SV_DEL, SV_INS, SV_INV, SV_DUP) have svlen in INFO field
    # Regular variants (DEL, INS) need svlen calculated from allele lengths
    if svlen is None or svlen == 0:
        ref_allele = metadata_row.get("ref", "")
        alt_allele = metadata_row.get("alt", "")
        if ref_allele and alt_allele:
            # svlen = len(alt) - len(ref)
            # For DEL: negative (e.g., 1 - 13 = -12)
            # For INS: positive (e.g., 7 - 1 = 6)
            svlen = len(alt_allele) - len(ref_allele)

    # Create VariantPosition object
    var_pos = VariantPosition(
        ref_pos=abs_variant_pos,
        alt_pos=abs_variant_pos,
        svlen=svlen if svlen is not None else 0,
        variant_type=variant_type,
    )

    # Determine target size from predictions
    if prediction_type == "1D":
        # Check if predictions are multi-dimensional (multiple targets)
        is_torch = TORCH_AVAILABLE and torch.is_tensor(ref_preds)
        is_numpy = isinstance(ref_preds, np.ndarray)

        if is_torch:
            ndim = len(ref_preds.shape)
        elif is_numpy:
            ndim = ref_preds.ndim
        else:
            ndim = 1

        # Handle multi-target predictions [n_targets, n_bins]
        if ndim > 1:
            target_size = ref_preds.shape[-1]  # Number of bins
            aligner = PredictionAligner1D(target_size=target_size, bin_size=bin_size)

            # Align each target separately
            n_targets = ref_preds.shape[0]
            ref_aligned_list = []
            alt_aligned_list = []

            for target_idx in range(n_targets):
                ref_target = ref_preds[target_idx]
                alt_target = alt_preds[target_idx]
                ref_aligned, alt_aligned = aligner.align_predictions(
                    ref_target, alt_target, variant_type, var_pos
                )
                ref_aligned_list.append(ref_aligned)
                alt_aligned_list.append(alt_aligned)

            # Stack results back into multi-target format
            if is_torch:
                ref_result = torch.stack(ref_aligned_list)
                alt_result = torch.stack(alt_aligned_list)
            else:
                ref_result = np.stack(ref_aligned_list)
                alt_result = np.stack(alt_aligned_list)

            return ref_result, alt_result
        else:
            # Single target prediction [n_bins]
            target_size = len(ref_preds)
            aligner = PredictionAligner1D(target_size=target_size, bin_size=bin_size)
            return aligner.align_predictions(
                ref_preds, alt_preds, variant_type, var_pos
            )
    else:  # 2D
        # Check if predictions are 1D (flattened upper triangular) or 2D (full matrix)
        is_torch = TORCH_AVAILABLE and torch.is_tensor(ref_preds)

        if is_torch:
            ndim = len(ref_preds.shape)
        else:
            ndim = ref_preds.ndim

        # If 1D, convert to 2D matrix
        if ndim == 1:
            ref_matrix = vector_to_contact_matrix(
                ref_preds, matrix_size, diag_offset=diag_offset
            )
            alt_matrix = vector_to_contact_matrix(
                alt_preds, matrix_size, diag_offset=diag_offset
            )

            # Align matrices
            aligner = PredictionAligner2D(
                target_size=matrix_size, bin_size=bin_size, diag_offset=diag_offset
            )
            aligned_ref_matrix, aligned_alt_matrix = aligner.align_predictions(
                ref_matrix, alt_matrix, variant_type, var_pos
            )

            # Convert back to flattened format
            aligned_ref_vector = contact_matrix_to_vector(
                aligned_ref_matrix, diag_offset=diag_offset
            )
            aligned_alt_vector = contact_matrix_to_vector(
                aligned_alt_matrix, diag_offset=diag_offset
            )

            return aligned_ref_vector, aligned_alt_vector
        else:
            # Already 2D matrices
            aligner = PredictionAligner2D(
                target_size=matrix_size, bin_size=bin_size, diag_offset=diag_offset
            )
            return aligner.align_predictions(
                ref_preds, alt_preds, variant_type, var_pos
            )
