from enum import Enum
from dataclasses import dataclass
from typing import Union, Literal, Optional
import ast
import numpy as np
import re


class Laterality(Enum):
    LEFT = "L"
    RIGHT = "R"
    UNKNOWN = "UNKNOWN"

    @classmethod
    def _missing_(cls, value) -> "Laterality":
        # if the lookup failed, check if a lower-case char was input
        # otherwise return UNKNOWN
        match value:
            case "l":
                return cls.LEFT
            case "r":
                return cls.RIGHT
            case _:
                return cls.UNKNOWN


class ViewPosition(Enum):
    CC = "CC"
    MLO = "MLO"
    ML = "ML"
    LM = "LM"
    XCCL = "XCCL"
    UNKNOWN = "UNKNOWN"

    @classmethod
    def _missing_(cls, value) -> "ViewPosition":
        if isinstance(value, str):
            # extract view position patterns (CC, MLO, or ML) after uppercasing
            match = re.search(r"\b(CC|MLO|ML)\b", value.upper())

            # if a match was found, return the Enum member by passing the extracted string back to cls()
            if match:
                return cls(match.group(1))

        # default to UNKNOWN if value was not a string or no matches were found
        return cls.UNKNOWN


class AlignmentDirection(Enum):
    RIGHT = "Right"
    LEFT = "Left"
    UP = "Up"
    DOWN = "Down"
    UNKNOWN = "Unknown"

    @property
    def opposite(self):
        opposites = {
            AlignmentDirection.RIGHT: AlignmentDirection.LEFT,
            AlignmentDirection.LEFT: AlignmentDirection.RIGHT,
            AlignmentDirection.UP: AlignmentDirection.DOWN,
            AlignmentDirection.DOWN: AlignmentDirection.UP,
        }
        return opposites.get(self, AlignmentDirection.UNKNOWN)


@dataclass
class Alignment:
    x: AlignmentDirection  # alignment looking across rows from left to right
    y: AlignmentDirection  # alignment looking down cols from top to bottom

    @property
    def _reference(self) -> "Alignment":
        return Alignment(AlignmentDirection.LEFT, AlignmentDirection.UP)

    def __hash__(self) -> int:
        return hash((self.x, self.y))

    def __repr__(self) -> str:
        return f"Alignment(X: {self.x.value}, Y: {self.y.value})"

    @staticmethod
    def _flip_required(source: AlignmentDirection, target: AlignmentDirection) -> bool:
        if target == AlignmentDirection.UNKNOWN:
            raise ValueError("Target alignment should not have an 'UNKNOWN' direction!")
        # we only need to flip if the source direction is the opposite of the target
        return source == target.opposite

    def _prepare_target(self, target: Optional["Alignment"]) -> "Alignment":
        target = target or self._reference

        # Check if axes are crossed (e.g. Target X is Up/Down), implying rotation
        if target.x in {AlignmentDirection.UP, AlignmentDirection.DOWN} or target.y in {
            AlignmentDirection.LEFT,
            AlignmentDirection.RIGHT,
        }:
            raise NotImplementedError("Rotation required! Implement method.")

        return target

    def realign_coords(
        self,
        coords: tuple[float, float, float, float],
        height: float,
        width: float,
        target_alignment: Optional["Alignment"] = None,
    ) -> tuple[float, float, float, float]:
        # provide the reference as default if no target_alignment was given
        target_alignment: Alignment = self._prepare_target(target_alignment)

        # unpack coords
        y_min, x_min, y_max, x_max = coords

        # check if left/right flipping is needed
        if self._flip_required(self.x, target_alignment.x):
            x_min, x_max = width - x_max, width - x_min

        # check if up/down flipping is needed
        if self._flip_required(self.y, target_alignment.y):
            y_min, y_max = height - y_max, height - y_min

        return (y_min, x_min, y_max, x_max)

    def realign_image(
        self,
        image: np.ndarray,
        target_alignment: Optional["Alignment"] = None,
    ) -> np.ndarray:
        # provide the reference as default if no target_alignment was given
        target_alignment: Alignment = self._prepare_target(target_alignment)

        # check if left/right flipping is needed
        if self._flip_required(self.x, target_alignment.x):
            image = image[..., ::-1]  # flip last axis

        # check if up/down flipping is needed
        if self._flip_required(self.y, target_alignment.y):
            image = image[..., ::-1, :]  # flip second-to-last axis

        return image

    def realign_coords_list(
        self,
        coords_list: list[list[float]],
        height: float,
        width: float,
        target_alignment: Optional["Alignment"] = None,
    ) -> list[list[float]]:
        # currently, this function still expects a list of lists and pylance will warn about
        # providing a tuple, TODO: fix this!
        # provide the reference as default if no target_alignment was given
        target_alignment: Alignment = self._prepare_target(target_alignment)

        # pre-determine if x/y flipping is required for all coords
        x_flip_required: bool = self._flip_required(self.x, target_alignment.x)
        y_flip_required: bool = self._flip_required(self.y, target_alignment.y)

        output_list: list[list[float]] = []
        for coords in coords_list:
            # unpack coords
            y_min, x_min, y_max, x_max = coords

            # check if left/right flipping is needed
            if x_flip_required:
                x_min, x_max = width - x_max, width - x_min

            # check if up/down flipping is needed
            if y_flip_required:
                y_min, y_max = height - y_max, height - y_min

            output_list.append([y_min, x_min, y_max, x_max])

        return output_list


# PatientOrientation -------------------------------------------------------------------------------------------------
# PatientOrientation -------------------------------------------------------------------------------------------------


class Direction(Enum):
    P = "P"
    A = "A"
    L = "L"
    R = "R"
    HL = "HL"
    HR = "HR"
    FL = "FL"
    FR = "FR"
    H = "H"
    F = "F"
    UNKNOWN = "UNKNOWN"

    @classmethod
    def _missing_(cls, _) -> "Direction":
        # if the lookup failed, return UNKNOWN
        return cls.UNKNOWN


class PatientOrientation:
    def __init__(self, orientation: Union[str, tuple[str, str]]) -> None:
        # parse input
        row, col = self._parse(orientation)
        self.row: Direction = row
        self.col: Direction = col

    @property
    def exact(self) -> bool:
        # returns true if the orientation elements were matched exactly
        return (self.row != Direction.UNKNOWN) and (self.col != Direction.UNKNOWN)

    def alignment(
        self,
        laterality: Union[str, Laterality],
        view_position: Union[str, ViewPosition],
    ) -> Alignment:
        if isinstance(laterality, str):
            laterality: Laterality = Laterality(laterality)
        if isinstance(view_position, str):
            view_position: ViewPosition = ViewPosition(view_position)

        def _check_alignment(
            laterality: Laterality,
            view_position: ViewPosition,
            direction: Direction,
        ) -> AlignmentDirection:
            # internal function to convert a Direction to an AlignmentDirection given the image laterality/view
            match (laterality, view_position, direction):
                # all views -----------------------------------------------\
                case (_, _, Direction.P):
                    return AlignmentDirection.RIGHT
                case (_, _, Direction.A):
                    return AlignmentDirection.LEFT
                # LCC views -----------------------------------------------\
                case (Laterality.LEFT, ViewPosition.CC, Direction.L):
                    return AlignmentDirection.DOWN
                case (Laterality.LEFT, ViewPosition.CC, Direction.R):
                    return AlignmentDirection.UP
                # RCC views -----------------------------------------------\
                case (Laterality.RIGHT, ViewPosition.CC, Direction.L):
                    return AlignmentDirection.UP
                case (Laterality.RIGHT, ViewPosition.CC, Direction.R):
                    return AlignmentDirection.DOWN
                # LMLO views ----------------------------------------------\
                case (Laterality.LEFT, ViewPosition.MLO, Direction.FR):
                    return AlignmentDirection.UP
                case (Laterality.LEFT, ViewPosition.MLO, Direction.HL):
                    return AlignmentDirection.DOWN
                # RMLO views ----------------------------------------------\
                case (Laterality.RIGHT, ViewPosition.MLO, Direction.FL):
                    return AlignmentDirection.UP
                case (Laterality.RIGHT, ViewPosition.MLO, Direction.HR):
                    return AlignmentDirection.DOWN
                # ML views ------------------------------------------------\
                case (_, ViewPosition.ML, Direction.F):
                    return AlignmentDirection.UP
                case (_, ViewPosition.ML, Direction.H):
                    return AlignmentDirection.DOWN
                # LM views ------------------------------------------------\
                case (_, ViewPosition.LM, Direction.F):
                    return AlignmentDirection.UP
                case (_, ViewPosition.LM, Direction.H):
                    return AlignmentDirection.DOWN
                # LXCCL views ------------------------------------------------\
                case (Laterality.LEFT, ViewPosition.XCCL, Direction.R):
                    return AlignmentDirection.UP
                case (Laterality.LEFT, ViewPosition.XCCL, Direction.L):
                    return AlignmentDirection.DOWN
                # RXCCL views ------------------------------------------------\
                case (Laterality.RIGHT, ViewPosition.XCCL, Direction.L):
                    return AlignmentDirection.UP
                case (Laterality.RIGHT, ViewPosition.XCCL, Direction.R):
                    return AlignmentDirection.DOWN
                # unknown combinations ------------------------------------\
                case _:
                    return AlignmentDirection.UNKNOWN

        return Alignment(
            _check_alignment(laterality, view_position, self.row),
            _check_alignment(laterality, view_position, self.col),
        )

    def _parse(
        self, orientation: Union[str, tuple[str, str]]
    ) -> tuple[Direction, Direction]:
        # if orientation is a str, e.g. "('P', 'R')", parse it into a tuple
        if isinstance(orientation, str):
            orientation: tuple[str, str] = ast.literal_eval(orientation)

        # if orientation is not a length 2 tuple at this stage, error
        if (not isinstance(orientation, tuple)) | (len(orientation) != 2):
            raise ValueError(f"Invalid orientation provided '{orientation}'!")

        # parse the row/col orientation strings into Direction enums
        return (
            Direction(orientation[0]),
            Direction(orientation[1]),
        )

    def __repr__(self) -> str:
        return f"PatientOrientation({self.row}, {self.col})"

    def __hash__(self) -> int:
        return hash((self.row, self.col))


class Rotation(Enum):
    NONE = 0.0
    FULL = 180.0
    UNKNOWN = float("nan")

    @classmethod
    def _missing_(cls, _) -> "Rotation":
        return cls.UNKNOWN


class HFlip(Enum):
    YES = "YES"
    NO = "NO"
    UNKNOWN = "UNKNOWN"

    @classmethod
    def _missing_(cls, _) -> "HFlip":
        return cls.UNKNOWN


class FieldOfView:
    def __init__(self, rotation: Union[str, float], h_flip: str) -> None:
        # parse inputs
        self.rotation: Rotation = Rotation(float(rotation))
        self.h_flip: HFlip = HFlip(h_flip)

    @property
    def exact(self) -> bool:
        # returns true if the fov elements were matched exactly
        return (self.h_flip != HFlip.UNKNOWN) and (self.rotation != Rotation.UNKNOWN)

    def alignment(self, laterality: Laterality) -> Alignment:
        match (laterality, self.rotation, self.h_flip):
            # laterality: L
            case (Laterality.LEFT, Rotation.FULL, HFlip.YES):
                return Alignment(AlignmentDirection.RIGHT, AlignmentDirection.UP)
            case (Laterality.LEFT, Rotation.FULL, HFlip.NO):
                return Alignment(AlignmentDirection.LEFT, AlignmentDirection.UP)
            case (Laterality.LEFT, Rotation.NONE, HFlip.YES):
                return Alignment(AlignmentDirection.LEFT, AlignmentDirection.DOWN)
            case (Laterality.LEFT, Rotation.NONE, HFlip.NO):
                return Alignment(AlignmentDirection.RIGHT, AlignmentDirection.DOWN)
            # laterality: R
            case (Laterality.RIGHT, Rotation.FULL, HFlip.YES):
                return Alignment(AlignmentDirection.RIGHT, AlignmentDirection.DOWN)
            case (Laterality.RIGHT, Rotation.FULL, HFlip.NO):
                return Alignment(AlignmentDirection.LEFT, AlignmentDirection.DOWN)
            case (Laterality.RIGHT, Rotation.NONE, HFlip.YES):
                return Alignment(AlignmentDirection.LEFT, AlignmentDirection.UP)
            case (Laterality.RIGHT, Rotation.NONE, HFlip.NO):
                return Alignment(AlignmentDirection.RIGHT, AlignmentDirection.UP)
            case (_, _, _):
                return Alignment(AlignmentDirection.UNKNOWN, AlignmentDirection.UNKNOWN)

    def __repr__(self) -> str:
        return f"FieldOfView({self.rotation}, {self.h_flip})"

    def __hash__(self) -> int:
        return hash((self.rotation, self.h_flip))
