from functools import reduce
from loguru import logger

import torch
import torch.nn.functional as F

from rubik.action import build_actions_tensor, parse_actions_str, sample_actions_str
from rubik.display import stringify
from rubik.tensor_utils import build_cube_tensor


class Cube:
    """
    A 4D tensor filled with colors. Dimensions have the following interpretation:
        - Face (from 0 to 5, with 0 = "Up", 1 = "Left", 2 = "Front", 3 = "Right", 4 = "Back", 5 = "Down").
        - X coordinate (from 0 to self.size - 1, from Left to Right).
        - Y coordinate (from 0 to self.size - 1, from Back to Front).
        - Z coordinate (from 0 to self.size - 1, from Down to Up).

    Colors filling each tensor cell are from 0 to 6, 0 being the "dark" color,
    the rest according to order given in "colors" attribute.
    """

    def __init__(self, colors: list[str], size: int):
        """
        Create Cube from a given list of 6 colors and size.
        Example:
            cube = Cube(['U', 'L', 'C', 'R', 'B', 'D'], size = 3)
        """
        tensor = build_cube_tensor(colors, size)
        self.coordinates = tensor.indices().transpose(0, 1).to(torch.int16)
        self.state = F.one_hot(tensor.values().long(), num_classes=7).to(torch.int16)
        self.actions = build_actions_tensor(size)
        self.history: list[list[int]] = []
        self.colors = colors
        self.size = size

    def to(self, device: str | torch.device) -> "Cube":
        device = torch.device(device)
        dtype = torch.int16 if device == torch.device("cpu") else torch.float32
        self.coordinates = self.coordinates.to(device=device, dtype=dtype)
        self.state = self.state.to(device=device, dtype=dtype)
        self.actions = self.actions.to(device=device, dtype=dtype)
        logger.info(f"Using device '{self.state.device}' and dtype '{dtype}'")
        return self

    def reset_history(self) -> None:
        """
        Reset internal history of moves.
        """
        self.history = []
        return

    def shuffle(self, num_moves: int, seed: int = 0) -> None:
        """
        Randomly shuffle the cube by the supplied number of steps, and reset history of moves.
        """
        moves = sample_actions_str(num_moves, self.size, seed=seed)
        self.rotate(moves)
        self.reset_history()
        return

    def rotate(self, moves: str) -> None:
        """
        Apply a sequence of moves (defined as plain string) to the cube.
        """
        actions = parse_actions_str(moves)
        for action in actions:
            self.rotate_once(*action)
        return

    def rotate_once(self, axis: int, slice: int, inverse: int) -> None:
        """
        Apply a move (defined as 3 coordinates) to the cube.
        """
        action = self.actions[axis, slice, inverse]
        self.state = action @ self.state
        self.history.append([axis, slice, inverse])
        return

    def compute_changes(self, moves: str) -> dict[int, int]:
        """
        combine a sequence of moves and return the resulting changes.
        """
        actions = parse_actions_str(moves)
        tensors = [self.actions[*action].to(torch.float32) for action in actions]
        result = reduce(lambda A, B: A @ B, tensors).to(torch.int16)
        return dict(result.indices().transpose(0, 1).tolist())

    def solve(self, policy: str) -> None:
        """
        Apply the specified solving policy to the cube.
        """
        raise NotImplementedError

    def __str__(self):
        """
        Compute a string representation of a cube.
        """
        state = self.state.argmax(dim=-1).to(device="cpu", dtype=torch.int16)
        return stringify(state, self.colors, self.size)