from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from repixelizer.analysis import analyze_tile_graph_source
from repixelizer.baselines import naive_resize_baseline
from repixelizer.io import save_rgba
from repixelizer.metrics import source_lattice_consistency_breakdown
from repixelizer.params import SolverHyperParams
from repixelizer.pipeline import run_pipeline
from repixelizer.synthetic import fake_pixelize, make_emblem
from repixelizer.tile_graph import (
    _extract_source_region_tiles,
    _segment_atomic_source_regions,
    build_tile_graph_model,
    clear_tile_graph_model_cache,
    optimize_tile_graph,
)
from repixelizer.types import InferenceResult


def test_tile_graph_model_extracts_candidates_and_adjacency_from_lattice_proposals() -> None:
    lowres = np.zeros((4, 4, 4), dtype=np.float32)
    lowres[..., 3] = 1.0
    lowres[:, 1] = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32)
    source = np.repeat(np.repeat(lowres, 8, axis=0), 8, axis=1)
    inference = InferenceResult(
        target_width=4,
        target_height=4,
        phase_x=0.0,
        phase_y=0.0,
        confidence=1.0,
        top_candidates=[],
    )

    model, stats = build_tile_graph_model(
        source,
        inference=inference,
        analysis=analyze_tile_graph_source(source, device="cpu"),
        device="cpu",
    )

    assert model.candidate_rgba.shape[0] > 4
    assert model.candidate_neighbor_rgba.shape[1:] == (4, 4)
    assert model.candidate_neighbor_mask.shape[1:] == (4,)
    assert bool(np.any(model.candidate_neighbor_mask))
    choice_counts = np.diff(model.cell_candidate_offsets)
    assert choice_counts.shape[0] == inference.target_width * inference.target_height
    assert np.all(choice_counts >= 1)
    assert stats.average_choices >= 1.0
    assert stats.edge_density > 0.0
    assert stats.model_device == "cpu"


def test_tile_graph_model_works_with_edge_only_analysis() -> None:
    lowres = np.zeros((4, 4, 4), dtype=np.float32)
    lowres[..., 3] = 1.0
    lowres[1:3, 1:3] = np.asarray([1.0, 0.7, 0.2, 1.0], dtype=np.float32)
    source = np.repeat(np.repeat(lowres, 8, axis=0), 8, axis=1)
    inference = InferenceResult(
        target_width=4,
        target_height=4,
        phase_x=0.0,
        phase_y=0.0,
        confidence=1.0,
        top_candidates=[],
    )

    model, stats = build_tile_graph_model(
        source,
        inference=inference,
        analysis=analyze_tile_graph_source(source, device="cpu"),
        device="cpu",
    )

    assert model.candidate_rgba.shape[0] > 0
    assert stats.model_device == "cpu"


def test_tile_graph_candidates_use_literal_source_pixel_colors() -> None:
    source = np.zeros((8, 8, 4), dtype=np.float32)
    source[..., 3] = 1.0
    palette = (
        np.asarray([1.0, 0.1, 0.1, 1.0], dtype=np.float32),
        np.asarray([0.1, 1.0, 0.1, 1.0], dtype=np.float32),
        np.asarray([0.1, 0.1, 1.0, 1.0], dtype=np.float32),
        np.asarray([1.0, 1.0, 0.1, 1.0], dtype=np.float32),
    )
    for y in range(source.shape[0]):
        for x in range(source.shape[1]):
            source[y, x] = palette[(2 * (y % 2)) + (x % 2)]

    inference = InferenceResult(
        target_width=4,
        target_height=4,
        phase_x=0.0,
        phase_y=0.0,
        confidence=1.0,
        top_candidates=[],
    )
    model, _stats = build_tile_graph_model(
        source,
        inference=inference,
        analysis=analyze_tile_graph_source(source),
        solver_params=SolverHyperParams(tile_graph_max_candidates=32, tile_graph_max_candidates_per_coord=2),
    )

    source_colors = {tuple(np.round(color, 4)) for color in source.reshape(-1, 4)}
    candidate_colors = {tuple(np.round(color, 4)) for color in model.candidate_rgba}

    assert candidate_colors.issubset(source_colors | {tuple(np.zeros(4, dtype=np.float32))})


def test_tile_graph_candidate_buckets_are_contiguous_per_output_coord() -> None:
    source = np.zeros((8, 8, 4), dtype=np.float32)
    source[..., 3] = 1.0
    source[:, 2:4] = np.asarray([1.0, 0.9, 0.2, 1.0], dtype=np.float32)
    source[2:6, 5:7] = np.asarray([0.2, 0.8, 1.0, 1.0], dtype=np.float32)
    inference = InferenceResult(
        target_width=4,
        target_height=4,
        phase_x=0.0,
        phase_y=0.0,
        confidence=1.0,
        top_candidates=[],
    )

    model, _stats = build_tile_graph_model(
        source,
        inference=inference,
        analysis=analyze_tile_graph_source(source),
    )

    total_candidates = int(model.candidate_rgba.shape[0])
    assert int(model.cell_candidate_offsets[0]) == 0
    assert int(model.cell_candidate_offsets[-1]) == total_candidates
    assert np.array_equal(model.cell_candidate_indices, np.arange(total_candidates, dtype=np.int32))


def test_tile_graph_region_extraction_covers_every_occupied_output_cell() -> None:
    source = np.zeros((64, 64, 4), dtype=np.float32)
    white = np.asarray([1.0, 1.0, 1.0, 1.0], dtype=np.float32)
    source[6:10, 6:58] = white
    source[14:18, 6:50] = white
    source[22:26, 14:58] = white
    source[30:34, 6:50] = white
    source[38:42, 14:58] = white
    source[46:50, 6:58] = white
    source[6:50, 6:10] = white
    source[14:58, 54:58] = white

    inference = InferenceResult(
        target_width=8,
        target_height=8,
        phase_x=0.0,
        phase_y=-0.2,
        confidence=1.0,
        top_candidates=[],
    )
    params = SolverHyperParams()
    analysis = analyze_tile_graph_source(source, device="cpu")
    cell_h = source.shape[0] / inference.target_height
    cell_w = source.shape[1] / inference.target_width
    full_y_grid, full_x_grid = np.meshgrid(
        np.arange(source.shape[0], dtype=np.int32),
        np.arange(source.shape[1], dtype=np.int32),
        indexing="ij",
    )
    flat_rgba = source.reshape(-1, source.shape[-1]).astype(np.float32)
    flat_edge = analysis.edge_map.reshape(-1).astype(np.float32)
    flat_y = full_y_grid.reshape(-1).astype(np.int32)
    flat_x = full_x_grid.reshape(-1).astype(np.int32)

    components = _segment_atomic_source_regions(
        source_rgba=source,
        edge_map=analysis.edge_map,
        alpha_floor=params.alpha_transparent_threshold,
        color_threshold=params.tile_graph_component_color_threshold,
        alpha_threshold=params.tile_graph_component_alpha_threshold,
        device="cpu",
    )
    region_buckets = _extract_source_region_tiles(
        labeling=components,
        flat_rgba=flat_rgba,
        flat_edge=flat_edge,
        flat_x=flat_x,
        flat_y=flat_y,
        cell_w=cell_w,
        cell_h=cell_h,
        target_width=inference.target_width,
        target_height=inference.target_height,
        phase_x=inference.phase_x,
        phase_y=inference.phase_y,
        min_region_area_ratio=params.tile_graph_source_region_min_area_ratio,
        min_window_coverage=params.tile_graph_source_region_window_coverage,
    )

    opaque = flat_rgba[:, 3] >= params.alpha_transparent_threshold
    coord_x = np.clip(
        np.floor((flat_x.astype(np.float32) + 0.5) / cell_w - inference.phase_x).astype(np.int32),
        0,
        inference.target_width - 1,
    )
    coord_y = np.clip(
        np.floor((flat_y.astype(np.float32) + 0.5) / cell_h - inference.phase_y).astype(np.int32),
        0,
        inference.target_height - 1,
    )
    direct_indices = coord_y * inference.target_width + coord_x
    occupied_direct = np.bincount(
        direct_indices[opaque],
        minlength=inference.target_width * inference.target_height,
    )
    occupied_bucket = np.asarray([len(bucket) for bucket in region_buckets], dtype=np.int32)

    assert np.all(occupied_bucket[occupied_direct > 0] > 0)


def test_tile_graph_projects_component_overlaps_directly_onto_output_cells() -> None:
    source = np.zeros((6, 12, 4), dtype=np.float32)
    source[..., 3] = 1.0
    red = np.asarray([1.0, 0.0, 0.0, 1.0], dtype=np.float32)
    blue = np.asarray([0.0, 0.0, 1.0, 1.0], dtype=np.float32)
    source[:, :6] = red
    source[:, 6:] = blue
    inference = InferenceResult(
        target_width=2,
        target_height=1,
        phase_x=0.5,
        phase_y=0.0,
        confidence=1.0,
        top_candidates=[],
    )

    model, _stats = build_tile_graph_model(
        source,
        inference=inference,
        analysis=analyze_tile_graph_source(source),
        device="cpu",
    )

    left_start = int(model.cell_candidate_offsets[0])
    left_end = int(model.cell_candidate_offsets[1])
    right_start = int(model.cell_candidate_offsets[1])
    right_end = int(model.cell_candidate_offsets[2])
    left_candidates = {
        tuple(np.round(color, 4))
        for color in model.candidate_rgba[model.cell_candidate_indices[left_start:left_end]]
    }
    right_candidates = {
        tuple(np.round(color, 4))
        for color in model.candidate_rgba[model.cell_candidate_indices[right_start:right_end]]
    }

    assert tuple(np.round(red, 4)) in left_candidates
    assert tuple(np.round(blue, 4)) in right_candidates
    assert left_candidates.issubset({tuple(np.round(red, 4)), tuple(np.round(blue, 4))})
    assert right_candidates.issubset({tuple(np.round(red, 4)), tuple(np.round(blue, 4))})


def test_tile_graph_edge_cells_include_multiple_same_cell_contour_colors() -> None:
    source = np.zeros((8, 8, 4), dtype=np.float32)
    source[..., 3] = 1.0
    source[:4, 0] = np.asarray([1.0, 0.9, 0.2, 1.0], dtype=np.float32)
    source[:4, 1] = np.asarray([0.85, 0.6, 0.1, 1.0], dtype=np.float32)
    source[:4, 2] = np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32)
    source[:4, 3] = np.asarray([0.9, 0.2, 0.1, 1.0], dtype=np.float32)
    source[:4, 4:] = np.asarray([0.2, 0.8, 0.2, 1.0], dtype=np.float32)
    source[4:, :] = np.asarray([0.2, 0.2, 0.8, 1.0], dtype=np.float32)

    inference = InferenceResult(
        target_width=2,
        target_height=2,
        phase_x=0.0,
        phase_y=0.0,
        confidence=1.0,
        top_candidates=[],
    )
    model, _stats = build_tile_graph_model(
        source,
        inference=inference,
        analysis=analyze_tile_graph_source(source),
        device="cpu",
    )

    start = int(model.cell_candidate_offsets[0])
    end = int(model.cell_candidate_offsets[1])
    candidates = {
        tuple(np.round(color, 4))
        for color in model.candidate_rgba[model.cell_candidate_indices[start:end]]
    }
    non_black_candidates = {color for color in candidates if color[:3] != (np.float32(0.0), np.float32(0.0), np.float32(0.0))}

    assert len(candidates) >= 3
    assert tuple(np.round(np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32), 4)) in candidates
    assert len(non_black_candidates) >= 2


def test_tile_graph_reconstructs_synthetic_thin_feature_better_than_naive() -> None:
    source = np.zeros((16, 16, 4), dtype=np.float32)
    source[4:12, 8] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
    source[4, 7:10] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
    source[11, 7] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
    source[10, 6] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)
    source[9, 5] = np.asarray([0.95, 0.95, 0.95, 1.0], dtype=np.float32)

    fake = fake_pixelize(
        source,
        upscale=9,
        phase_x=0.18,
        phase_y=-0.12,
        blur_radius=0.45,
        warp_strength=0.2,
        warp_detail=5,
        seed=11,
    )
    inference = InferenceResult(
        target_width=16,
        target_height=16,
        phase_x=0.18,
        phase_y=-0.12,
        confidence=1.0,
        top_candidates=[],
    )
    artifacts, _ = optimize_tile_graph(
        fake,
        inference=inference,
        analysis=analyze_tile_graph_source(fake),
        steps=0,
        seed=7,
        device="cpu",
    )
    naive = naive_resize_baseline(fake, width=16, height=16)

    tile_score = source_lattice_consistency_breakdown(
        fake,
        artifacts.target_rgba,
        target_width=16,
        target_height=16,
        phase_x=0.18,
        phase_y=-0.12,
    )["score"]
    initial_tile_score = source_lattice_consistency_breakdown(
        fake,
        artifacts.initial_rgba,
        target_width=16,
        target_height=16,
        phase_x=0.18,
        phase_y=-0.12,
    )["score"]
    naive_score = source_lattice_consistency_breakdown(
        fake,
        naive,
        target_width=16,
        target_height=16,
        phase_x=0.18,
        phase_y=-0.12,
    )["score"]

    assert tile_score <= initial_tile_score + 1e-6
    assert tile_score <= naive_score


def test_tile_graph_keeps_shallow_connected_stroke_from_fanning_out() -> None:
    source = np.zeros((48, 48, 4), dtype=np.float32)
    for y in range(48):
        x = int(round(8 + y * 0.65))
        source[max(0, y - 1) : min(48, y + 2), max(0, x - 1) : min(48, x + 2)] = np.asarray(
            [1.0, 1.0, 1.0, 1.0],
            dtype=np.float32,
        )

    artifacts, _ = optimize_tile_graph(
        source,
        inference=InferenceResult(
            target_width=6,
            target_height=6,
            phase_x=0.0,
            phase_y=0.0,
            confidence=1.0,
            top_candidates=[],
        ),
        analysis=analyze_tile_graph_source(source),
        steps=0,
        seed=7,
        device="cpu",
    )

    occupied_per_row = [
        int(np.count_nonzero(artifacts.target_rgba[y, :, 3] > 0.5))
        for y in range(artifacts.target_rgba.shape[0])
    ]

    assert sum(count > 1 for count in occupied_per_row) <= 1
    assert sum(occupied_per_row) <= artifacts.target_rgba.shape[0] + 1


def test_tile_graph_preserves_transparent_background_on_sparse_detail() -> None:
    source = np.zeros((12, 12, 4), dtype=np.float32)
    source[5:7, 5:7] = np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32)

    fake = fake_pixelize(
        source,
        upscale=8,
        phase_x=0.12,
        phase_y=-0.08,
        blur_radius=0.35,
        warp_strength=0.1,
        warp_detail=4,
        seed=5,
    )
    artifacts, _ = optimize_tile_graph(
        fake,
        inference=InferenceResult(
            target_width=12,
            target_height=12,
            phase_x=0.12,
            phase_y=-0.08,
            confidence=1.0,
            top_candidates=[],
        ),
        analysis=analyze_tile_graph_source(fake),
        steps=0,
        seed=7,
        device="cpu",
    )

    alpha_zero_ratio = float((artifacts.target_rgba[..., 3] <= 0.05).mean())

    assert alpha_zero_ratio >= 0.75


def test_tile_graph_cuda_matches_cpu_on_small_case() -> None:
    torch = pytest.importorskip("torch")
    if not torch.cuda.is_available():
        pytest.skip("CUDA is not available in this environment")

    source = make_emblem(16, 16)
    fake = fake_pixelize(source, upscale=8, phase_x=0.1, phase_y=-0.05, blur_radius=0.4, warp_strength=0.15, warp_detail=4)
    inference = InferenceResult(
        target_width=16,
        target_height=16,
        phase_x=0.1,
        phase_y=-0.05,
        confidence=1.0,
        top_candidates=[],
    )
    cpu_analysis = analyze_tile_graph_source(fake, device="cpu")
    cuda_analysis = analyze_tile_graph_source(fake, device="cuda")

    cpu_artifacts, cpu_diag = optimize_tile_graph(
        fake,
        inference=inference,
        analysis=cpu_analysis,
        steps=0,
        seed=7,
        device="cpu",
    )
    cuda_artifacts, cuda_diag = optimize_tile_graph(
        fake,
        inference=inference,
        analysis=cuda_analysis,
        steps=0,
        seed=7,
        device="cuda",
    )

    assert cpu_diag["tile_graph_model_device"] == "cpu"
    assert cpu_diag["tile_graph_solver_device"] == "cpu"
    assert cuda_diag["tile_graph_model_device"] == "cuda"
    assert cuda_diag["tile_graph_solver_device"] == "cuda"
    assert np.allclose(cpu_artifacts.target_rgba, cuda_artifacts.target_rgba, atol=1e-5)


def test_tile_graph_model_cache_reuses_fixed_lattice_build() -> None:
    clear_tile_graph_model_cache()
    source = make_emblem(16, 16)
    fake = fake_pixelize(source, upscale=8, phase_x=0.1, phase_y=-0.05, blur_radius=0.4, warp_strength=0.15, warp_detail=4)
    inference = InferenceResult(
        target_width=16,
        target_height=16,
        phase_x=0.1,
        phase_y=-0.05,
        confidence=1.0,
        top_candidates=[],
    )
    analysis = analyze_tile_graph_source(fake, device="cpu")

    first, first_stats = build_tile_graph_model(
        fake,
        inference=inference,
        analysis=analysis,
        device="cpu",
    )
    second, second_stats = build_tile_graph_model(
        fake,
        inference=inference,
        analysis=analysis,
        device="cpu",
    )

    assert first_stats.cache_hit is False
    assert second_stats.cache_hit is True
    assert np.array_equal(first.candidate_rgba, second.candidate_rgba)


def test_pipeline_tile_graph_mode_writes_reconstruction_diagnostics(tmp_path: Path) -> None:
    source = make_emblem(24, 24)
    fake = fake_pixelize(source, upscale=10, phase_x=0.15, phase_y=0.2, blur_radius=0.6, warp_strength=0.2, warp_detail=5)
    input_path = tmp_path / "input.png"
    output_path = tmp_path / "nested" / "output.png"
    diagnostics_dir = tmp_path / "diagnostics"

    save_rgba(input_path, fake)
    result = run_pipeline(
        input_path,
        output_path,
        diagnostics_dir=diagnostics_dir,
        steps=24,
        reconstruction_mode="tile-graph",
        device="cpu",
    )

    import json

    run_json = json.loads((diagnostics_dir / "run.json").read_text(encoding="utf-8"))
    assert output_path.exists()
    assert run_json["settings"]["reconstruction_mode"] == "tile-graph"
    assert run_json["reconstruction"]["mode"] == "tile-graph"
    assert run_json["reconstruction"]["tile_graph_model_device"] == "cpu"
    assert "tile_graph_model_cache_hit" in run_json["reconstruction"]
    assert run_json["reconstruction"]["tile_graph_candidate_count"] > 0
    assert set(run_json["source_fidelity"].keys()) == {"snap_initial", "solver_target", "final_output"}
    assert result.diagnostics["reconstruction"]["mode"] == "tile-graph"