import torch
from typing import Callable


def split_distance(
    a: torch.Tensor,
    b: torch.Tensor,
    dist_fn_1: Callable,
    dist_fn_2: Callable,
    split_idx: int,
) -> torch.Tensor:
    """
    Splits both `a` and `b` in two parts at `split_idx`, then applies two
    different distance functions for each part.

    This is useful on a mixed dataset containing embeddings and continous values.

    Parameters
    ----------
    a : torch.Tensor
        The first Tensor
    b : torch.Tensor
        The second Tensor
    dist_fn_1 : Callable
        The distance function for the first chunks
    dist_fn_2 : Callable
        The distance function for the second chunks
    split_idx : int
        The split point for `a` and `b`
    """
    a1, a2 = torch.split(a, [split_idx, a.size(-1) - split_idx], dim=-1)
    b1, b2 = torch.split(b, [split_idx, b.size(-1) - split_idx], dim=-1)
    dist_1 = dist_fn_1(a1, b1)
    dist_2 = dist_fn_2(a2, b2)

    def normalize(a: torch.Tensor) -> torch.Tensor:
        return (a - a.min()) / (a.max() - a.min())

    dist_1 = normalize(dist_1)
    dist_2 = normalize(dist_2)
    return (dist_1 * a1.shape[-1] / a.shape[-1]) + (dist_2 * a2.shape[-1] / a.shape[-1])