from typing import Optional, Any
import numpy as np
from utils import generate_rotation_matrices, compute_histogram
from histogram_transfer.pdf_transfer_1d import pdf_transfer_1d
from histogram_transfer.regrain import regrain


def apply_histogram(
    source_img: np.ndarray | Any,
    target_histogram: np.ndarray,
    *,
    iterations: int = 25,
    alpha: float = 0.3,
    source_histogram: Optional[np.ndarray] = None,
    should_regrain: bool = True,
):
    if not isinstance(source_img, np.ndarray):
        source_img = np.array(source_img)

    assert (
        target_histogram.shape[0]
        == target_histogram.shape[1]
        == target_histogram.shape[2]
    ), "Histograms must be 3D"

    bins = target_histogram.shape[0]
    assert 256 % bins == 0, "Bin size must be a factor of 256"

    if source_histogram is None:
        source_histogram = compute_histogram(source_img, bins=bins)
    else:
        assert (
            source_histogram.shape == target_histogram.shape
        ), "Source and target histograms must be the same shape"

    source_colours = np.array(
        [
            index for index, value in np.ndenumerate(source_histogram) if value > 0
        ]  # we ignore colours without occurences
    ).T  # unique rgb colours, 3xN
    source_colours_original = source_colours.copy()

    source_counts = np.array(
        [value for value in np.nditer(source_histogram) if value > 0]
    )  # cardinality of each rgb colour in source_colours, N

    target_histogram_colours = np.array(
        [index for index, value in np.ndenumerate(target_histogram) if value > 0]
    ).T
    target_histogram_counts = np.array(
        [value for value in np.nditer(target_histogram) if value > 0]
    )

    source_colour_adjustment = np.zeros_like(source_colours)
    for iteration, rotation in enumerate(generate_rotation_matrices(iterations)):
        rotated_source_colours = rotation @ source_colours
        rotated_target_histogram_colours = rotation @ target_histogram_colours

        rotated_source_colours.round(out=rotated_source_colours)
        rotated_target_histogram_colours.round(out=rotated_target_histogram_colours)

        assert rotation.shape[0] == 3
        for axis in range(rotation.shape[0]):  # for each axis (rgb)
            sorted_source_colours = sorted(
                enumerate(source_counts),
                key=lambda x: rotated_source_colours[axis, x[0]],
            )
            sorted_source_indices = [i for i, _ in sorted_source_colours]
            sorted_source_counts = [v for _, v in sorted_source_colours]

            sorted_target_histogram_counts = [
                v
                for _, v in sorted(
                    enumerate(target_histogram_counts),
                    key=lambda x: rotated_target_histogram_colours[axis, x[0]],
                )
            ]
            sorted_target_histogram_colours = [
                v for v in sorted(rotated_target_histogram_colours[axis, :])
            ]

            sorted_new_source_colours = pdf_transfer_1d(
                sorted_source_counts,
                sorted_target_histogram_counts,
                sorted_target_histogram_colours,
            )
            source_colour_adjustment[axis, :] = [
                v
                for _, v in sorted(
                    enumerate(sorted_new_source_colours),
                    key=lambda x: sorted_source_indices[x[0]],
                )
            ]
        source_colours = source_colours + (1 - (iteration / iterations) ** alpha) * (
            rotation.T @ (source_colour_adjustment - rotated_source_colours)
        )

    scale = 256 // bins
    source_colours = source_colours.T * scale
    source_colours.clip(0, 255, out=source_colours)
    lut = np.zeros((256, 256, 256, 3), dtype=np.uint8)

    for old, new in zip(source_colours_original.T, source_colours):
        for x in range(scale):
            for y in range(scale):
                for z in range(scale):
                    lut[old[0] * scale + x, old[1] * scale + y, old[2] * scale + z] = (
                        new
                    )

    result = lut[source_img[:, :, 0], source_img[:, :, 1], source_img[:, :, 2]]
    return regrain(source_img, result) if should_regrain else result