"""Tests for the arc_solver package.

Covers all three phases:
  grid.py       — perception (İlim)
  dsl.py        — DSL primitives (İrade vocabulary)
  synthesis.py  — program synthesis (İrade selection)
  solver.py     — search / execution (Kudret)
"""

from __future__ import annotations

import pytest

from arc_solver import (
    ArcSolver,
    BBox,
    Grid,
    Object,
    Pair,
    PairClues,
    Solution,
    Task,
    Transform,
)
from arc_solver import grid as G
from arc_solver import dsl
from arc_solver import synthesis as S


# =====================================================================
#   FIXTURES — reusable test grids and tasks
# =====================================================================

GRID_L = [
    [1, 0, 0],
    [1, 1, 0],
    [0, 0, 0],
]

GRID_L_ROT90 = [
    [0, 1, 1],
    [0, 1, 0],
    [0, 0, 0],
]

GRID_L_ROT180 = [
    [0, 0, 0],
    [0, 1, 1],
    [0, 0, 1],
]

GRID_L_FLIP_H = [
    [0, 0, 1],
    [0, 1, 1],
    [0, 0, 0],
]

GRID_L_FLIP_V = [
    [0, 0, 0],
    [1, 1, 0],
    [1, 0, 0],
]

GRID_SYMMETRIC = [
    [0, 1, 0],
    [1, 1, 1],
    [0, 1, 0],
]

GRID_4x4 = [
    [0, 0, 0, 0],
    [0, 1, 1, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 0],
]


def _make_task(pairs, test_pairs):
    """Helper to build a Task from (input, output) tuples."""
    return Task(
        task_id="test",
        train=[Pair(i, o) for i, o in pairs],
        test_inputs=[i for i, _ in test_pairs],
        test_outputs=[o for _, o in test_pairs],
    )


# =====================================================================
#   § 1  GRID UTILITIES
# =====================================================================

class TestGridUtilities:
    def test_grid_size(self):
        assert G.grid_size(GRID_L) == (3, 3)
        assert G.grid_size([[1]]) == (1, 1)
        assert G.grid_size([]) == (0, 0)

    def test_copy_grid(self):
        copy = G.copy_grid(GRID_L)
        assert copy == GRID_L
        copy[0][0] = 9
        assert GRID_L[0][0] == 1  # original untouched

    def test_make_grid(self):
        g = G.make_grid(2, 3, 5)
        assert g == [[5, 5, 5], [5, 5, 5]]

    def test_grids_equal(self):
        assert G.grids_equal(GRID_L, G.copy_grid(GRID_L))
        assert not G.grids_equal(GRID_L, GRID_L_ROT90)
        assert not G.grids_equal([[1]], [[1, 2]])


# =====================================================================
#   § 2  COLOUR ANALYSIS
# =====================================================================

class TestColourAnalysis:
    def test_grid_colors(self):
        assert G.grid_colors(GRID_L) == {0, 1}
        assert G.grid_colors([[0]]) == {0}

    def test_color_counts(self):
        counts = G.color_counts(GRID_L)
        assert counts[0] == 6
        assert counts[1] == 3

    def test_detect_background(self):
        assert G.detect_background(GRID_L) == 0
        # Grid where 5 is most frequent
        g = [[5, 5], [5, 1]]
        assert G.detect_background(g) == 5
        # Tie: 0 wins
        g = [[0, 1], [1, 0]]
        assert G.detect_background(g) == 0


# =====================================================================
#   § 3  OBJECT EXTRACTION
# =====================================================================

class TestObjectExtraction:
    def test_extract_objects_simple(self):
        objs = G.extract_objects(GRID_L, background=0)
        assert len(objs) == 1
        obj = objs[0]
        assert obj.color == 1
        assert obj.size == 3
        assert obj.bbox == BBox(0, 0, 1, 1)

    def test_extract_objects_multiple(self):
        g = [
            [1, 0, 2],
            [0, 0, 0],
            [3, 0, 0],
        ]
        objs = G.extract_objects(g, background=0)
        assert len(objs) == 3
        colors = {o.color for o in objs}
        assert colors == {1, 2, 3}

    def test_extract_multicolor_objects(self):
        g = [
            [1, 2, 0],
            [1, 0, 0],
            [0, 0, 0],
        ]
        objs = G.extract_multicolor_objects(g, background=0)
        # Color 1 and 2 are 4-connected → should be separate
        # Actually: (0,0)=1 connects to (1,0)=1 and (0,1)=2
        # (0,0)→(0,1) is non-bg, (0,0)→(1,0) is non-bg
        # All three form one connected region
        assert len(objs) == 1
        assert objs[0].size == 3


# =====================================================================
#   § 4  GEOMETRIC TRANSFORMS
# =====================================================================

class TestGeometricTransforms:
    def test_rotate_90(self):
        assert G.rotate_90(GRID_L) == GRID_L_ROT90

    def test_rotate_180(self):
        assert G.rotate_180(GRID_L) == GRID_L_ROT180

    def test_rotate_270(self):
        r = G.rotate_270(GRID_L)
        # 270 CW = 90 CCW = rotate_90 three times
        expected = G.rotate_90(G.rotate_90(G.rotate_90(GRID_L)))
        assert r == expected

    def test_flip_h(self):
        assert G.flip_h(GRID_L) == GRID_L_FLIP_H

    def test_flip_v(self):
        assert G.flip_v(GRID_L) == GRID_L_FLIP_V

    def test_transpose(self):
        g = [[1, 2, 3], [4, 5, 6]]
        expected = [[1, 4], [2, 5], [3, 6]]
        assert G.transpose(g) == expected

    def test_identity_roundtrip(self):
        """Four 90° rotations return to original."""
        g = GRID_L
        for _ in range(4):
            g = G.rotate_90(g)
        assert g == GRID_L


# =====================================================================
#   § 5  MANIPULATION
# =====================================================================

class TestManipulation:
    def test_crop_to_content(self):
        cropped = G.crop_to_content(GRID_4x4, background=0)
        assert cropped == [[1, 1], [1, 0]]

    def test_extract_subgrid(self):
        sub = G.extract_subgrid(GRID_4x4, BBox(1, 1, 2, 2))
        assert sub == [[1, 1], [1, 0]]

    def test_overlay(self):
        base = G.make_grid(3, 3, 0)
        top = [[1, 1], [1, 0]]
        result = G.overlay(base, top, 0, 0)
        assert result[0] == [1, 1, 0]
        assert result[1] == [1, 0, 0]

    def test_overlay_transparent(self):
        base = [[5, 5], [5, 5]]
        top = [[1, 0], [0, 1]]
        result = G.overlay(base, top, 0, 0, transparent=0)
        assert result == [[1, 5], [5, 1]]

    def test_scale_grid(self):
        g = [[1, 2], [3, 4]]
        scaled = G.scale_grid(g, 2)
        assert G.grid_size(scaled) == (4, 4)
        assert scaled[0] == [1, 1, 2, 2]
        assert scaled[1] == [1, 1, 2, 2]

    def test_tile_grid(self):
        g = [[1]]
        tiled = G.tile_grid(g, 2, 3)
        assert tiled == [[1, 1, 1], [1, 1, 1]]

    def test_recolor(self):
        g = [[1, 0], [0, 1]]
        result = G.recolor(g, {1: 2})
        assert result == [[2, 0], [0, 2]]


# =====================================================================
#   § 6  GRAVITY
# =====================================================================

class TestGravity:
    def test_gravity_down(self):
        g = [[1, 0], [0, 1], [0, 0]]
        result = G.gravity_down(g)
        assert result == [[0, 0], [0, 0], [1, 1]]

    def test_gravity_up(self):
        g = [[0, 0], [0, 1], [1, 0]]
        result = G.gravity_up(g)
        assert result == [[1, 1], [0, 0], [0, 0]]

    def test_gravity_left(self):
        g = [[0, 1, 0], [0, 0, 2]]
        result = G.gravity_left(g)
        assert result == [[1, 0, 0], [2, 0, 0]]

    def test_gravity_right(self):
        g = [[1, 0, 0], [0, 2, 0]]
        result = G.gravity_right(g)
        assert result == [[0, 0, 1], [0, 0, 2]]


# =====================================================================
#   § 7  FILL ENCLOSED
# =====================================================================

class TestFillEnclosed:
    def test_fill_enclosed_basic(self):
        # 0 is bg (17 cells), 1 forms a ring enclosing one bg cell at (2,2)
        g = [
            [0, 0, 0, 0, 0],
            [0, 1, 1, 1, 0],
            [0, 1, 0, 1, 0],
            [0, 1, 1, 1, 0],
            [0, 0, 0, 0, 0],
        ]
        result = G.fill_enclosed(g, fill_color=2)
        assert result[2][2] == 2
        # Exterior bg cells untouched
        assert result[0][0] == 0

    def test_fill_enclosed_open(self):
        """Open region on the border should NOT be filled."""
        g = [
            [1, 0, 1],
            [1, 0, 1],
            [1, 1, 1],
        ]
        result = G.fill_enclosed(g, fill_color=2)
        # Column 1 is open at top → exterior → not filled
        assert result[0][1] == 0
        assert result[1][1] == 0


# =====================================================================
#   § 8  SYMMETRY DETECTION
# =====================================================================

class TestSymmetry:
    def test_symmetric_grid(self):
        sym = G.detect_symmetry(GRID_SYMMETRIC)
        assert sym["h_mirror"] is True
        assert sym["v_mirror"] is True
        assert sym["rot180"] is True

    def test_asymmetric_grid(self):
        sym = G.detect_symmetry(GRID_L)
        assert sym["h_mirror"] is False
        assert sym["v_mirror"] is False


# =====================================================================
#   § 9  PERCEIVE (HIGH-LEVEL)
# =====================================================================

class TestPerceive:
    def test_perceive_basic(self):
        info = G.perceive(GRID_L)
        assert info.height == 3
        assert info.width == 3
        assert info.background == 0
        assert 1 in info.colors
        assert len(info.objects) == 1


# =====================================================================
#   § 10  DSL PRIMITIVES
# =====================================================================

class TestDSLPrimitives:
    def test_registry_not_empty(self):
        assert len(dsl.PRIMITIVES) > 20

    def test_all_in_registry(self):
        for t in dsl.PRIMITIVES:
            assert t.name in dsl.REGISTRY

    def test_identity(self):
        t = dsl.REGISTRY["identity"]
        assert t(GRID_L) == GRID_L

    def test_rot90_via_dsl(self):
        t = dsl.REGISTRY["rot90"]
        assert t(GRID_L) == GRID_L_ROT90

    def test_crop_via_dsl(self):
        t = dsl.REGISTRY["crop"]
        result = t(GRID_4x4)
        assert result == [[1, 1], [1, 0]]

    def test_compose(self):
        rot = dsl.REGISTRY["rot90"]
        ct = dsl.compose(rot, rot)
        assert ct.name == "rot90 >> rot90"
        assert ct(GRID_L) == G.rotate_180(GRID_L)

    def test_make_recolor(self):
        t = dsl.make_recolor({1: 3})
        result = t(GRID_L)
        assert result[0][0] == 3
        assert result[2][2] == 0

    def test_make_scale(self):
        t = dsl.make_scale(2)
        result = t([[1]])
        assert result == [[1, 1], [1, 1]]

    def test_make_keep_color(self):
        g = [[1, 2], [3, 0]]
        t = dsl.make_keep_color(1)
        result = t(g)
        assert result == [[1, 0], [0, 0]]

    def test_make_remove_color(self):
        g = [[1, 2], [3, 0]]
        t = dsl.make_remove_color(2)
        result = t(g)
        assert result == [[1, 0], [3, 0]]


# =====================================================================
#   § 11  SYNTHESIS — PAIR ANALYSIS
# =====================================================================

class TestSynthesisAnalysis:
    def test_analyze_pair_same_size(self):
        pair = Pair(GRID_L, GRID_L_ROT90)
        clues = S.analyze_pair(pair)
        assert clues.same_size is True

    def test_analyze_pair_diff_size(self):
        pair = Pair(GRID_4x4, [[1, 1], [1, 0]])
        clues = S.analyze_pair(pair)
        assert clues.same_size is False

    def test_analyze_pair_color_mapping(self):
        g_in = [[1, 0], [0, 1]]
        g_out = [[2, 0], [0, 2]]
        clues = S.analyze_pair(Pair(g_in, g_out))
        assert clues.color_mapping == {1: 2, 0: 0}


# =====================================================================
#   § 12  SYNTHESIS — PROGRAM FINDING
# =====================================================================

class TestSynthesisProgram:
    def test_finds_rot90(self):
        """Synthesise the transform for a rotate-90 task."""
        task = _make_task(
            pairs=[
                (GRID_L, GRID_L_ROT90),
                ([[0, 1], [0, 0]], [[0, 0], [0, 1]]),
            ],
            test_pairs=[
                ([[1, 1, 0], [0, 0, 0], [0, 0, 0]],
                 [[0, 0, 1], [0, 0, 1], [0, 0, 0]]),
            ],
        )
        t = S.synthesize(task)
        assert t is not None
        assert "rot90" in t.name

    def test_finds_flip_h(self):
        task = _make_task(
            pairs=[
                (GRID_L, GRID_L_FLIP_H),
                ([[1, 0], [0, 0]], [[0, 1], [0, 0]]),
            ],
            test_pairs=[
                ([[1, 1, 0], [0, 0, 0]], [[0, 1, 1], [0, 0, 0]]),
            ],
        )
        t = S.synthesize(task)
        assert t is not None
        assert "flip_h" in t.name

    def test_finds_recolor(self):
        """Synthesise a colour remapping."""
        task = _make_task(
            pairs=[
                ([[1, 0], [0, 1]], [[2, 0], [0, 2]]),
                ([[1, 1], [0, 0]], [[2, 2], [0, 0]]),
            ],
            test_pairs=[
                ([[0, 1], [1, 1]], [[0, 2], [2, 2]]),
            ],
        )
        t = S.synthesize(task)
        assert t is not None

    def test_finds_crop(self):
        task = _make_task(
            pairs=[
                (GRID_4x4, [[1, 1], [1, 0]]),
                ([[0, 0, 0], [0, 2, 0], [0, 0, 0]], [[2]]),
            ],
            test_pairs=[
                ([[0, 3, 0], [0, 0, 0]], [[3]]),
            ],
        )
        t = S.synthesize(task)
        assert t is not None
        assert "crop" in t.name

    def test_returns_none_when_impossible(self):
        """No DSL composition can turn every pixel to 9 from 0."""
        task = _make_task(
            pairs=[
                ([[0, 0], [0, 0]], [[9, 9], [9, 9]]),
            ],
            test_pairs=[
                ([[0]], [[9]]),
            ],
        )
        # Recolor {0: 9} should be found from color mapping
        t = S.synthesize(task)
        # Actually this IS solvable via make_recolor({0: 9})
        # Let's test a truly unsolvable case instead
        task2 = _make_task(
            pairs=[
                ([[1, 2], [3, 4]], [[5, 6], [7, 8]]),
                ([[1, 2], [3, 4]], [[9, 0], [1, 2]]),  # contradicts first
            ],
            test_pairs=[],
        )
        t2 = S.synthesize(task2)
        assert t2 is None


# =====================================================================
#   § 13  SOLVER END-TO-END
# =====================================================================

class TestSolverEndToEnd:
    def test_solve_rotate_task(self):
        task = _make_task(
            pairs=[
                (GRID_L, GRID_L_ROT90),
                ([[0, 1], [0, 0]], [[0, 0], [0, 1]]),
            ],
            test_pairs=[
                ([[1, 1, 0], [0, 0, 0], [0, 0, 0]],
                 [[0, 0, 1], [0, 0, 1], [0, 0, 0]]),
            ],
        )
        solver = ArcSolver()
        sol = solver.solve_task(task)
        assert len(sol.predictions) == 1
        assert G.grids_equal(sol.predictions[0], task.test_outputs[0])

    def test_solve_flip_task(self):
        task = _make_task(
            pairs=[
                (GRID_L, GRID_L_FLIP_H),
                ([[1, 0], [0, 0]], [[0, 1], [0, 0]]),
            ],
            test_pairs=[
                ([[1, 1, 0], [0, 0, 0]], [[0, 1, 1], [0, 0, 0]]),
            ],
        )
        solver = ArcSolver()
        sol = solver.solve_task(task)
        assert len(sol.predictions) == 1
        assert G.grids_equal(sol.predictions[0], task.test_outputs[0])

    def test_batch_solve(self):
        tasks = [
            _make_task(
                pairs=[(GRID_L, GRID_L_ROT90)],
                test_pairs=[(GRID_L, GRID_L_ROT90)],
            ),
            _make_task(
                pairs=[(GRID_L, GRID_L_FLIP_H)],
                test_pairs=[(GRID_L, GRID_L_FLIP_H)],
            ),
        ]
        solver = ArcSolver()
        result = solver.batch_solve(tasks)
        stats = result["stats"]
        assert stats["total"] == 2
        assert stats["correct"] == 2
        assert stats["accuracy"] == 1.0


# =====================================================================
#   § 14  TASK I/O
# =====================================================================

class TestTaskIO:
    def test_from_json(self):
        data = {
            "train": [
                {"input": [[1, 0], [0, 1]], "output": [[0, 1], [1, 0]]},
            ],
            "test": [
                {"input": [[1, 1], [0, 0]], "output": [[0, 0], [1, 1]]},
            ],
        }
        task = Task.from_json(data, task_id="demo")
        assert task.task_id == "demo"
        assert len(task.train) == 1
        assert len(task.test_inputs) == 1
        assert len(task.test_outputs) == 1

    def test_from_json_no_test_output(self):
        data = {
            "train": [{"input": [[1]], "output": [[2]]}],
            "test": [{"input": [[3]]}],
        }
        task = Task.from_json(data)
        assert len(task.test_inputs) == 1
        assert len(task.test_outputs) == 0


# =====================================================================
#   § 15  BBOX
# =====================================================================

class TestBBox:
    def test_dimensions(self):
        bb = BBox(1, 2, 4, 5)
        assert bb.height == 4
        assert bb.width == 4
        assert bb.area == 16

    def test_unit_bbox(self):
        bb = BBox(0, 0, 0, 0)
        assert bb.height == 1
        assert bb.width == 1
        assert bb.area == 1
