from typing import Dict

import numpy as np
from einops import rearrange
from monai.transforms.transform import Transform


class OrientationGuidanceMultipleLabelDeepEditd(Transform):
    def __init__(self, ref_image="image", label_names=None):
        """
        Convert the guidance to the RAS orientation
        """
        self.ref_image = ref_image
        self.label_names = label_names

    def transform_points(self, point, affine):
        """transform point to the coordinates of the transformed image
        point: numpy array [bs, N, 3]
        """
        bs, n = point.shape[:2]
        point = np.concatenate((point, np.ones((bs, n, 1))), axis=-1)
        point = rearrange(point, "b n d -> d (b n)")
        point = affine @ point
        point = rearrange(point, "d (b n)-> b n d", b=bs)[:, :, :3]
        return point

    def __call__(self, data):
        d: Dict = dict(data)
        for key_label in self.label_names.keys():
            points = d.get(key_label, [])
            if len(points) < 1:
                continue
            reoriented_points = self.transform_points(
                np.array(points)[None],
                np.linalg.inv(d[self.ref_image].meta["affine"].numpy()) @ d[self.ref_image].meta["original_affine"],
            )
            d[key_label] = reoriented_points[0]
        return d