from __future__ import annotations

import random
from typing import Any

from data.schema import COLORS, OBJECT_IDS, RELATIONS, SHAPES, Fact, SceneObject

HELD_OUT_SHAPES = {
    "red": "triangle",
    "blue": "square",
    "green": "star",
    "yellow": "circle",
}


def _relation(first: SceneObject, second: SceneObject, rng: random.Random) -> Fact:
    relation = rng.choice(RELATIONS)
    if relation == "left_of":
        first = SceneObject(first.object_id, first.color, first.shape, 55, 112)
        second = SceneObject(second.object_id, second.color, second.shape, 169, 112)
    elif relation == "right_of":
        first = SceneObject(first.object_id, first.color, first.shape, 169, 112)
        second = SceneObject(second.object_id, second.color, second.shape, 55, 112)
    elif relation == "above":
        first = SceneObject(first.object_id, first.color, first.shape, 112, 55)
        second = SceneObject(second.object_id, second.color, second.shape, 112, 169)
    else:
        first = SceneObject(first.object_id, first.color, first.shape, 112, 169)
        second = SceneObject(second.object_id, second.color, second.shape, 112, 55)
    return Fact(first.object_id, relation, second.object_id), first, second


def generate_scene(
    scene_id: int, seed: int, num_objects: int = 2, held_out_composition: bool = False
) -> dict[str, Any]:
    if num_objects != 2:
        raise ValueError("The stage-one MVP currently supports exactly two objects.")
    rng = random.Random((seed * 1_000_003) + scene_id)
    colors = rng.sample(COLORS, k=num_objects)
    if held_out_composition:
        first_shape = HELD_OUT_SHAPES[colors[0]]
        second_choices = [shape for shape in SHAPES if shape != first_shape]
        shapes = [first_shape, rng.choice(second_choices)]
    else:
        shapes = []
        for color in colors:
            choices = [shape for shape in SHAPES if shape != HELD_OUT_SHAPES[color]]
            choices = [shape for shape in choices if shape not in shapes]
            shapes.append(rng.choice(choices))
    objects = [
        SceneObject(OBJECT_IDS[index], colors[index], shapes[index], 0, 0)
        for index in range(num_objects)
    ]
    relation, first, second = _relation(objects[0], objects[1], rng)
    objects = [first, second]
    facts = []
    for item in objects:
        facts.extend(
            [
                Fact(item.object_id, "color", item.color),
                Fact(item.object_id, "shape", item.shape),
            ]
        )
    facts.append(relation)
    return {
        "scene_id": f"scene-{scene_id:07d}",
        "objects": [item.to_dict() for item in objects],
        "visual_facts": [fact.to_dict() for fact in facts],
    }