"""ZeroGPU neural-Q-learning demo for solving 2x2x2 Rubik's Cube states."""

from __future__ import annotations

import random
from typing import Iterable, List, Tuple

import gradio as gr


Vec = Tuple[int, int, int]
State = Tuple[int, ...]

FACE_NAMES = ("U", "R", "F", "D", "L", "B")
FACE_COLORS = ("#f5f5f5", "#dc2626", "#16a34a", "#facc15", "#f97316", "#2563eb")
MOVES = tuple(move for face in FACE_NAMES for move in (face, f"{face}'"))
SOLVED: State = tuple(color for color in range(6) for _ in range(4))

# Each face defines its outward normal plus its screen-right and screen-down axes.
FACE_GEOMETRY: Tuple[Tuple[Vec, Vec, Vec], ...] = (
    ((0, 1, 0), (1, 0, 0), (0, 0, 1)),    # U
    ((1, 0, 0), (0, 0, -1), (0, -1, 0)),  # R
    ((0, 0, 1), (1, 0, 0), (0, -1, 0)),   # F
    ((0, -1, 0), (1, 0, 0), (0, 0, -1)),  # D
    ((-1, 0, 0), (0, 0, 1), (0, -1, 0)),  # L
    ((0, 0, -1), (-1, 0, 0), (0, -1, 0)), # B
)


def add(*vectors: Vec) -> Vec:
    return tuple(sum(vector[i] for vector in vectors) for i in range(3))  # type: ignore[return-value]


def scale(vector: Vec, amount: int) -> Vec:
    return tuple(amount * part for part in vector)  # type: ignore[return-value]


def dot(left: Vec, right: Vec) -> int:
    return sum(a * b for a, b in zip(left, right))


def cross(left: Vec, right: Vec) -> Vec:
    return (
        left[1] * right[2] - left[2] * right[1],
        left[2] * right[0] - left[0] * right[2],
        left[0] * right[1] - left[1] * right[0],
    )


def rotate_quarter(vector: Vec, axis: Vec, direction: int) -> Vec:
    """Rotate a grid vector 90 degrees about `axis` using right-hand direction."""
    parallel = scale(axis, dot(axis, vector))
    perpendicular = cross(axis, vector) if direction == 1 else cross(vector, axis)
    return add(parallel, perpendicular)


def sticker_keys() -> List[Tuple[Vec, Vec]]:
    keys = []
    for normal, right, down in FACE_GEOMETRY:
        for row in range(2):
            for col in range(2):
                position = add(
                    normal,
                    scale(right, -1 if col == 0 else 1),
                    scale(down, -1 if row == 0 else 1),
                )
                keys.append((normal, position))
    return keys


STICKER_KEYS = sticker_keys()
KEY_TO_INDEX = {key: index for index, key in enumerate(STICKER_KEYS)}


def build_permutation(face_index: int, direction: int) -> Tuple[int, ...]:
    axis = FACE_GEOMETRY[face_index][0]
    destination = list(range(24))
    for source, (normal, position) in enumerate(STICKER_KEYS):
        if dot(position, axis) == 1:
            rotated = (rotate_quarter(normal, axis, direction), rotate_quarter(position, axis, direction))
            destination[source] = KEY_TO_INDEX[rotated]
    return tuple(destination)


PERMUTATIONS = {face: build_permutation(index, 1) for index, face in enumerate(FACE_NAMES)}
PERMUTATIONS.update({f"{face}'": build_permutation(index, -1) for index, face in enumerate(FACE_NAMES)})


def apply_move(state: State, move: str) -> State:
    """Apply one legal face turn to a sticker-state representation."""
    output = [0] * 24
    for source, destination in enumerate(PERMUTATIONS[move]):
        output[destination] = state[source]
    return tuple(output)


def apply_moves(state: State, moves: Iterable[str]) -> State:
    for move in moves:
        state = apply_move(state, move)
    return state


def random_scramble(length: int, rng: random.Random) -> List[str]:
    scramble: List[str] = []
    last_face = ""
    for _ in range(length):
        move = rng.choice([candidate for candidate in MOVES if candidate[0] != last_face])
        scramble.append(move)
        last_face = move[0]
    return scramble


def cube_preview(state: State) -> str:
    """Render a CSS-only, rotatable 3D cube preview for the supplied sequence."""
    faces = []
    for face_index, face in enumerate(FACE_NAMES):
        stickers = "".join(
            f"<span style='background:{FACE_COLORS[state[face_index * 4 + sticker]]}'></span>"
            for sticker in range(4)
        )
        faces.append(f"<div class='cube-face cube-{face.lower()}'>{stickers}</div>")
    return """
<style>
.cube-stage { min-height: 310px; display: grid; place-items: center; perspective: 720px; overflow: hidden; }
.cube-3d { width: 160px; height: 160px; position: relative; transform-style: preserve-3d; transform: rotateX(-25deg) rotateY(38deg); animation: cube-float 5s ease-in-out infinite; }
.cube-face { position: absolute; width: 160px; height: 160px; display: grid; grid-template-columns: repeat(2, 1fr); gap: 5px; padding: 5px; box-sizing: border-box; background: #0f172a; border: 2px solid #020617; backface-visibility: hidden; }
.cube-face span { border-radius: 5px; border: 1px solid rgba(15, 23, 42, .8); box-shadow: inset 0 0 10px rgba(255, 255, 255, .22); }
.cube-f { transform: translateZ(80px); } .cube-b { transform: rotateY(180deg) translateZ(80px); }
.cube-r { transform: rotateY(90deg) translateZ(80px); } .cube-l { transform: rotateY(-90deg) translateZ(80px); }
.cube-u { transform: rotateX(90deg) translateZ(80px); } .cube-d { transform: rotateX(-90deg) translateZ(80px); }
@keyframes cube-float { 50% { transform: rotateX(-19deg) rotateY(58deg) translateY(-7px); } }
</style>
<div class='cube-stage' role='img' aria-label='Interactive-style 3D 2 by 2 cube preview'>
  <div class='cube-3d'>%s</div>
</div>
""" % "".join(faces)


def parse_sequence(text: str) -> Tuple[List[str] | None, str]:
    """Normalize a quarter-turn sequence for preview and model inference."""
    sequence = text.upper().replace("’", "'").replace(",", " ").split()
    if not sequence:
        return None, "Enter a sequence such as `R U F' R'`, or generate one below."
    invalid = [move for move in sequence if move not in MOVES]
    if invalid:
        return None, "Use quarter turns U, R, F, D, L, B and optional primes, separated by spaces."
    return sequence, "Initial state ready for deployed-model inference."


def load_user_input(text: str) -> Tuple[str, str, str]:
    """Apply a supplied sequence and show its 3D initial-state preview."""
    sequence, message = parse_sequence(text)
    if sequence is None:
        return cube_preview(SOLVED), text, message
    normalized = " ".join(sequence)
    return cube_preview(apply_moves(SOLVED, sequence)), normalized, message


def make_random_problem(length: int, seed: int) -> Tuple[str, str, str]:
    """Generate a random legal turn sequence and show its 3D initial state."""
    scramble = random_scramble(int(length), random.Random(int(seed)))
    sequence = " ".join(scramble)
    return cube_preview(apply_moves(SOLVED, scramble)), sequence, "Random initial state ready for deployed-model inference."


def solve_with_deployed_model(initial_state: str) -> str:
    """Reserve the inference endpoint for the pretrained 2x2x2 cube model."""
    sequence, message = parse_sequence(initial_state)
    if sequence is None:
        return message
    return "Model checkpoint pending: this input is ready for deployed-model inference once the trained artifact is added."


with gr.Blocks(title="ZeroGPU Cube Q-Learning Lab") as demo:
    gr.Markdown(
        "# ZeroGPU Cube Q-Learning Lab\n"
        "Provide an initial 2x2x2 cube state and ask a deployed local model for a solution. "
        "Training happens outside this website; this Space is inference-only."
    )
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Initial state")
            state_input = gr.Textbox(
                label="Initial turn sequence",
                placeholder="R U F' R'",
                lines=2,
            )
            load_button = gr.Button("Save initial state")
            random_length = gr.Slider(1, 12, value=7, step=1, label="Random scramble length")
            random_seed = gr.Number(value=42, precision=0, label="Random scramble seed")
            random_button = gr.Button("Generate random initial state")
            input_report = gr.Markdown("Enter the initial state that the deployed model should solve.")
        with gr.Column():
            cube = gr.HTML(cube_preview(SOLVED), label="3D cube preview")
            gr.Markdown("### Model inference")
            model_status = gr.Markdown("**Model status:** checkpoint not yet included in this Space.")
            solve_button = gr.Button("Solve with deployed model", variant="primary")
            solution_report = gr.Markdown("The model solution will appear here.")
            solution_moves = gr.Textbox(label="Predicted solution moves", interactive=False)
    gr.Examples(
        examples=[["R U F' R'"], ["U R F U' L B'"], ["R' D L U' B F"]],
        inputs=[state_input],
        outputs=[cube, state_input, input_report],
        fn=load_user_input,
        cache_examples=True,
        cache_mode="lazy",
        label="Input examples",
    )

    load_button.click(load_user_input, inputs=state_input, outputs=[cube, state_input, input_report])
    random_button.click(make_random_problem, inputs=[random_length, random_seed], outputs=[cube, state_input, input_report])
    solve_button.click(solve_with_deployed_model, inputs=state_input, outputs=solution_report)


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(), mcp_server=True)