"""
test_miniproject_b.py — run all tests for Mini-project B.

Usage:
    python3 test_miniproject_b.py

Run this file after implementing each function in miniproject_b.py.
It imports your code, tests each function in order, and reports which
sections are passing.  No need to modify this file.

Tests run in dependency order: if a section fails, any section that
depends on it is skipped and shown as (not run) in the summary.

Dependency graph:
  Step 0 (reconstruct)          — must pass; all others depend on it
  Step 0 (prepare_map)          — must pass; provides img_gray
  1A pattern_difference         — independent
  1B find_similar_pattern_pos.  — depends on 1A (calls pattern_difference)
  1C highlight_rectangle        — independent
  1D highlight_similar_pattern  — depends on 1B and 1C
  2A downsample                 — independent
  2B find_most_likely_positions — depends on 1A and 2A
  2C faster_search              — depends on 1A, 2A, and 2B
  2D faster_highlight           — depends on 1C and 2C
"""

import os
import time
import numpy as np
from miniprojectutils import *
import miniproject_b as student


# ── Helpers ───────────────────────────────────────────────────────────────────

def _run(label: str, fn) -> bool:
    """Run fn(), print a labelled pass/fail block, and return True if passed."""
    print(f"\n── {label} {'─' * max(0, 52 - len(label))}")
    try:
        fn()
        return True
    except NotImplementedError as e:
        print(f"   ✗ Not implemented: {e}")
        print(f"   → Open miniproject_b.py and complete this function.")
        return False
    except AssertionError as e:
        print(f"   ✗ Test failed: {e}")
        return False
    except Exception as e:
        import traceback
        print(f"   ✗ Unexpected error: {type(e).__name__}: {e}")
        traceback.print_exc()
        return False


def _skip(label: str, reason: str) -> bool:
    """Mark a section as skipped due to a failing dependency."""
    print(f"\n── {label} {'─' * max(0, 52 - len(label))}")
    print(f"   - Skipped: {reason}")
    return False


# ── Step 0: reconstruction quality + prepare_map ──────────────────────────────

def test_reconstruction() -> None:
    """
    Run reconstruct() from miniproject_a independently and verify every tile
    is placed in its correct position.  Must pass before pattern search runs —
    a wrong map makes all subsequent tests meaningless.
    """
    from miniproject_a import (
        reconstruct, GIVEN_POS, N_MISSING, GRID, SOLUTION,
    )

    os.makedirs("imgs/out", exist_ok=True)
    given_tiles   = load_given_tiles("tiles", GIVEN_POS)
    missing_tiles = load_missing_tiles("tiles", N_MISSING)

    save_grid(given_tiles, missing_tiles, {}, GRID,
              "imgs/out/cartoon_map_partial_grid.jpg")
    print("   ✓ Saved imgs/out/cartoon_map_partial_grid.jpg")

    arrangement = reconstruct(given_tiles, missing_tiles, GRID)

    assert isinstance(arrangement, dict), \
        f"reconstruct must return a dict, got {type(arrangement)}"
    assert len(arrangement) == N_MISSING, \
        f"arrangement has {len(arrangement)} entries, expected {N_MISSING}"
    assert sorted(arrangement.keys()) == list(range(N_MISSING)), \
        "arrangement must contain each file index exactly once"

    valid_slots = {(r, c) for r in range(GRID) for c in range(GRID)
                   if (r + c) % 2 == 1}
    for fi, pos in arrangement.items():
        assert pos in valid_slots, \
            f"tile {fi} placed at {pos} which is not a valid empty slot"

    save_grid(given_tiles, missing_tiles, arrangement, GRID,
              "imgs/out/cartoon_map_reconstructed.jpg")
    print("   ✓ Saved imgs/out/cartoon_map_reconstructed.jpg")

    correct = sum(1 for fi, pos in arrangement.items() if pos == SOLUTION[fi])
    total   = N_MISSING
    assert correct == total, (
        f"{correct}/{total} tiles placed correctly. "
        f"Wrong file indices: "
        f"{[fi for fi, pos in arrangement.items() if pos != SOLUTION[fi]][:10]}"
        f"{'...' if (total - correct) > 10 else ''}"
    )
    print(f"   ✓ Correct! All {total} missing tiles placed correctly.")


def test_prepare_map() -> Image:
    """
    Run prepare_map, verify the output files were saved, and return img_gray
    so all subsequent tests can use the same grayscale image.
    """
    img_gray = student.prepare_map("tiles")

    assert img_gray is not None, "prepare_map returned None"
    assert img_gray.ndim == 2, \
        f"img_gray should be 2D (H, W), got shape {img_gray.shape}"
    h, w = img_gray.shape
    assert h > 0 and w > 0, "img_gray has zero dimension"
    assert os.path.exists("imgs/out/cartoon_map_reconstructed.jpg"), \
        "imgs/out/cartoon_map_reconstructed.jpg was not saved"
    assert os.path.exists("imgs/out/cartoon_map_gray.jpg"), \
        "imgs/out/cartoon_map_gray.jpg was not saved"

    print(f"   ✓ prepare_map: grayscale image shape = {img_gray.shape}")
    print("   ✓ imgs/out/cartoon_map_reconstructed.jpg and cartoon_map_gray.jpg saved.")
    return img_gray


# ── Part 1A: pattern_difference ───────────────────────────────────────────────

def test_pattern_difference(img_gray: Image) -> None:
    # Out-of-bounds positions must return exactly 1.0
    assert student.pattern_difference(img_gray, img_gray, (-1, 0)) == 1.0, \
        "negative row should return 1.0"
    assert student.pattern_difference(img_gray, img_gray, (0, -1)) == 1.0, \
        "negative col should return 1.0"
    assert student.pattern_difference(img_gray, img_gray, (1, 0)) == 1.0, \
        "pattern larger than remaining image should return 1.0"
    print("   ✓ pattern_difference: out-of-bounds positions return 1.0.")

    # Self-comparison must score exactly 0.0
    ph, pw    = 20, 30
    score_self = student.pattern_difference(img_gray, img_gray[:ph, :pw], (0, 0))
    assert score_self == 0.0, \
        f"self-comparison should be 0.0, got {score_self}"
    print(f"   ✓ pattern_difference: self-comparison = {score_self:.4f}.")

    # Score must be in [0, 1]
    score = student.pattern_difference(img_gray, img_gray[:ph, :pw], (10, 10))
    assert 0.0 <= score <= 1.0, \
        f"score must be in [0, 1], got {score}"
    print(f"   ✓ pattern_difference: score at (10,10) = {score:.4f}  (in [0,1]).")

    # All-zero image vs all-255 pattern → score must be 1.0
    zeros   = np.zeros((10, 10), dtype=np.int16)
    maxvals = np.full((5, 5), 255, dtype=np.int16)
    score_max = student.pattern_difference(zeros, maxvals, (0, 0))
    assert abs(score_max - 1.0) < 1e-6, \
        f"max difference should be 1.0, got {score_max}"
    print(f"   ✓ pattern_difference: max difference = {score_max:.4f}  (expected 1.0).")


# ── Part 1B: find_similar_pattern_position ────────────────────────────────────

def test_find_similar_pattern_position(img_gray: Image) -> None:
    # Self-verifying crop test: crop a known region and ask the function to find it
    corner  = student.PATTERN_CORNER
    size    = student.PATTERN_SIZE
    pattern = student.crop_image(img_gray, corner, size)

    found = student.find_similar_pattern_position(img_gray, pattern)

    assert isinstance(found, tuple) and len(found) == 2, \
        f"must return a (row, col) tuple, got {found}"
    assert found == corner, (
        f"expected to find the crop at {corner}, found {found} "
        f"(off by {abs(found[0]-corner[0])} rows, {abs(found[1]-corner[1])} cols)"
    )
    print(f"   ✓ find_similar_pattern_position: correct position {found}.")

    # Pattern larger than image → must return (0, 0)
    h, w  = img_gray.shape
    big   = np.zeros((h + 1, w + 1), dtype=np.int16)
    guard = student.find_similar_pattern_position(img_gray, big)
    assert guard == (0, 0), \
        f"pattern larger than image should return (0,0), got {guard}"
    print("   ✓ find_similar_pattern_position: returns (0,0) for oversized pattern.")


# ── Part 1C: highlight_rectangle ─────────────────────────────────────────────

def test_highlight_rectangle(img_gray: Image) -> None:
    corner = student.PATTERN_CORNER
    size   = student.PATTERN_SIZE
    linewidth = 10
    value     = 255

    result = student.highlight_rectangle(img_gray, corner, size, value, linewidth)

    # Must return a new array — input must be unchanged
    assert result is not img_gray, \
        "highlight_rectangle must return a NEW image, not the input"
    assert result.shape == img_gray.shape, \
        f"output shape {result.shape} does not match input {img_gray.shape}"
    print("   ✓ highlight_rectangle: returns a new image with correct shape.")

    # Interior pixels must be unchanged
    r0, c0 = corner
    rh, cw = size
    assert np.array_equal(result[r0 : r0 + rh, c0 : c0 + cw],
                          img_gray[r0 : r0 + rh, c0 : c0 + cw]), \
        "interior pixels were modified — they must stay unchanged"
    print("   ✓ highlight_rectangle: interior pixels unchanged.")

    # Top border band must equal value
    top_band = result[r0 - linewidth : r0, c0 - linewidth : c0 + cw + linewidth]
    assert np.all(top_band == value), \
        f"top border band should be {value}, found other values"
    print(f"   ✓ highlight_rectangle: top border band = {value}.")

    os.makedirs("imgs/out", exist_ok=True)
    save_image(result, "imgs/out/cartoon_map_original_pattern_highlighted.jpg")
    print("   ✓ Saved imgs/out/cartoon_map_original_pattern_highlighted.jpg")


# ── Part 1D: highlight_similar_pattern ───────────────────────────────────────

def test_highlight_similar_pattern(img_gray: Image) -> None:
    corner  = student.PATTERN_CORNER
    size    = student.PATTERN_SIZE
    pattern = student.crop_image(img_gray, corner, size)
    linewidth = 10
    value     = 0


    result = student.highlight_similar_pattern(img_gray, pattern, value, linewidth)

    assert result is not img_gray, \
        "highlight_similar_pattern must return a new image"
    assert result.shape == img_gray.shape, \
        f"output shape {result.shape} does not match input {img_gray.shape}"

    # Border at the correct position should be 0 (black)
    r0, c0 = corner
    check_row = max(0, r0 - 1)
    assert result[check_row, c0] == 0, \
        "expected black border above the found region"

    os.makedirs("imgs/out", exist_ok=True)
    save_image(pattern, "imgs/out/cartoon_map_original_pattern_crop.jpg")
    save_image(result, "imgs/out/cartoon_map_found_pattern_highlighted.jpg")
    print("   ✓ highlight_similar_pattern: correct position highlighted.")
    print("   ✓ Saved imgs/out/cartoon_map_found_pattern_highlighted.jpg")


# ── Part 2A: downsample ───────────────────────────────────────────────────────

def test_downsample(img_gray: Image) -> None:
    # Shape check after blur+downsample
    factor      = 4
    blurred     = student.simple_blur(img_gray, ksize=factor)
    downsampled = student.downsample(blurred, factor)

    bh, bw = blurred.shape
    dh, dw = downsampled.shape
    assert abs(dh - bh // factor) <= 1 and abs(dw - bw // factor) <= 1, \
        f"downsample(factor=4) on {bh}×{bw} → expected ~{bh//factor}×{bw//factor}, got {dh}×{dw}"
    print(f"   ✓ downsample: {bh}×{bw} → {dh}×{dw} with factor={factor}.")

    # Pixel value check: must equal arr[::factor, ::factor]
    tiny   = np.arange(20, dtype=np.int16).reshape(4, 5)
    result = student.downsample(tiny, 2)
    assert np.array_equal(result, tiny[::2, ::2]), \
        f"downsample(factor=2) should equal arr[::2, ::2], got {result}"
    print("   ✓ downsample: pixel values match arr[::factor, ::factor].")

    os.makedirs("imgs/out", exist_ok=True)
    save_image(downsampled, "imgs/out/cartoon_map_downsampled.jpg")
    print("   ✓ Saved imgs/out/cartoon_map_downsampled.jpg")


# ── Part 2B: find_most_likely_positions ───────────────────────────────────────

def test_find_most_likely_positions(img_gray: Image) -> None:
    factor  = 4
    top_n   = 5
    blurred = student.simple_blur(img_gray, ksize=factor)
    small   = student.downsample(blurred, factor)

    corner = student.PATTERN_CORNER
    size   = student.PATTERN_SIZE
    sc_r, sc_c = corner[0] // factor, corner[1] // factor
    pat_small  = student.crop_image(
        small, (sc_r, sc_c), (size[0] // factor, size[1] // factor)
    )

    positions = student.find_most_likely_positions(small, pat_small, top_n)

    assert isinstance(positions, list), \
        f"must return a list, got {type(positions)}"
    assert len(positions) == top_n, \
        f"must return exactly {top_n} positions, got {len(positions)}"
    assert all(isinstance(p, tuple) and len(p) == 2 for p in positions), \
        "all entries must be (row, col) tuples"
    print(f"   ✓ find_most_likely_positions: returned {len(positions)} positions.")

    # Positions must be sorted best (lowest score) first
    scores = [student.pattern_difference(small, pat_small, p) for p in positions]
    for i in range(len(scores) - 1):
        assert scores[i] <= scores[i + 1] + 1e-9, (
            f"positions not sorted: score[{i}]={scores[i]:.4f} "
            f"> score[{i+1}]={scores[i+1]:.4f}"
        )
    print("   ✓ find_most_likely_positions: positions sorted by score (best first).")

    # The true pattern position at low resolution must be in the top_n
    assert (sc_r, sc_c) in positions, (
        f"true position {(sc_r, sc_c)} not found in top {top_n} candidates: {positions}"
    )
    print(f"   ✓ find_most_likely_positions: true position {(sc_r, sc_c)} in top {top_n}.")


# ── Part 2C: faster_search ────────────────────────────────────────────────────

def test_faster_search(img_gray: Image) -> None:
    corner  = student.PATTERN_CORNER
    size    = student.PATTERN_SIZE
    pattern = student.crop_image(img_gray, corner, size)
    linewidth = 10
    value     = 0

    t0        = time.time()
    found_pos = student.faster_search(img_gray, pattern, factor=4, top_n=10)
    elapsed   = time.time() - t0

    assert isinstance(found_pos, tuple) and len(found_pos) == 2, \
        f"faster_search must return a (row, col) tuple, got {found_pos}"
    assert found_pos == corner, (
        f"expected {corner}, found {found_pos} "
        f"(off by {abs(found_pos[0]-corner[0])} rows, "
        f"{abs(found_pos[1]-corner[1])} cols)"
    )
    print(f"   ✓ faster_search: correct position {found_pos} in {elapsed:.2f}s.")

    os.makedirs("imgs/out", exist_ok=True)
    result = student.highlight_rectangle(img_gray, found_pos, size, value, linewidth)
    save_image(result, "imgs/out/cartoon_map_faster_search_pattern_highlighted.jpg")
    print("   ✓ Saved imgs/out/cartoon_map_faster_search_pattern_highlighted.jpg")


# ── Part 2D: faster_highlight ─────────────────────────────────────────────────

def test_faster_highlight(img_gray: Image) -> None:
    corner  = student.PATTERN_CORNER
    size    = student.PATTERN_SIZE
    pattern = student.crop_image(img_gray, corner, size)

    result = student.faster_highlight(img_gray, pattern, 0, 3, factor=4, top_n=10)

    assert result is not img_gray, \
        "faster_highlight must return a new image"
    assert result.shape == img_gray.shape, \
        f"output shape {result.shape} does not match input {img_gray.shape}"

    # os.makedirs("imgs/out", exist_ok=True)
    # save_image(result, "imgs/out/cartoon_map_faster_highlight.jpg")
    print("   ✓ faster_highlight: returned correctly shaped image.")
    print("   ✓ Saved imgs/out/cartoon_map_faster_highlight.jpg")


# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 56)
    print("  Mini-project B — Tests  (9 sections)")
    print("=" * 56)

    results = {}

    # ── Step 0: reconstruction + prepare_map ──────────────────────────────────
    # Reconstruction must pass first — everything else depends on a correct map.
    results["Step 0 - reconstruct (miniproject_a)"] = _run(
        "Step 0 - reconstruct (miniproject_a)",
        test_reconstruction
    )

    if not results["Step 0 - reconstruct (miniproject_a)"]:
        print("\n   Reconstruction is not correct.")
        print("   Fix find_best_tile in miniproject_a.py and re-run this file.")
        print("   The remaining tests require a correct map and cannot run yet.")
        # Print the full summary with all remaining sections marked as not run
        remaining = [
            "Step 0 - prepare_map",
            "1A - pattern_difference", "1B - find_similar_pattern_position",
            "1C - highlight_rectangle", "1D - highlight_similar_pattern",
            "2A - downsample", "2B - find_most_likely_positions",
            "2C - faster_search", "2D - faster_highlight",
        ]
        print("\n" + "=" * 56)
        print(f"  Summary: 0/9 sections passing")
        print("=" * 56)
        print(f"  ✗  Step 0 - reconstruct (miniproject_a)")
        for label in remaining:
            print(f"  -  {label}  (not run)")
        print("=" * 56)
        raise SystemExit(1)

    # prepare_map must succeed to get img_gray for all remaining tests
    print(f"\n── Step 0 - prepare_map {'─' * 30}")
    try:
        img_gray = test_prepare_map()
        results["Step 0 - prepare_map"] = True
        print("   ✓ prepare_map passed.")
    except Exception as e:
        import traceback
        print(f"   ✗ prepare_map failed: {type(e).__name__}: {e}")
        traceback.print_exc()
        results["Step 0 - prepare_map"] = False
        print("   Cannot continue without the grayscale image.")
        for label in ["1A - pattern_difference", "1B - find_similar_pattern_position",
                      "1C - highlight_rectangle", "1D - highlight_similar_pattern",
                      "2A - downsample", "2B - find_most_likely_positions",
                      "2C - faster_search", "2D - faster_highlight"]:
            results[label] = False
        # Fall through to print summary
        passed = sum(results.values())
        total  = len(results)
        print("\n" + "=" * 56)
        print(f"  Summary: {passed}/{total} sections passing")
        print("=" * 56)
        for label, ok in results.items():
            mark = "✓" if ok else "✗"
            print(f"  {mark}  {label}")
        print("=" * 56)
        raise SystemExit(1)

    # ── Part 1A ───────────────────────────────────────────────────────────────
    results["1A - pattern_difference"] = _run(
        "Part 1A - pattern_difference",
        lambda: test_pattern_difference(img_gray)
    )

    # ── Part 1B ───────────────────────────────────────────────────────────────
    # find_similar_pattern_position calls pattern_difference on every position
    if results["1A - pattern_difference"]:
        results["1B - find_similar_pattern_position"] = _run(
            "Part 1B - find_similar_pattern_position",
            lambda: test_find_similar_pattern_position(img_gray)
        )
    else:
        results["1B - find_similar_pattern_position"] = _skip(
            "Part 1B - find_similar_pattern_position",
            "pattern_difference must pass first"
        )

    # ── Part 1C ───────────────────────────────────────────────────────────────
    results["1C - highlight_rectangle"] = _run(
        "Part 1C - highlight_rectangle",
        lambda: test_highlight_rectangle(img_gray)
    )

    # ── Part 1D ───────────────────────────────────────────────────────────────
    # highlight_similar_pattern calls find_similar_pattern_position AND highlight_rectangle
    if results["1B - find_similar_pattern_position"] and results["1C - highlight_rectangle"]:
        results["1D - highlight_similar_pattern"] = _run(
            "Part 1D - highlight_similar_pattern",
            lambda: test_highlight_similar_pattern(img_gray)
        )
    else:
        missing = []
        if not results["1B - find_similar_pattern_position"]:
            missing.append("find_similar_pattern_position")
        if not results["1C - highlight_rectangle"]:
            missing.append("highlight_rectangle")
        results["1D - highlight_similar_pattern"] = _skip(
            "Part 1D - highlight_similar_pattern",
            f"{' and '.join(missing)} must pass first"
        )

    # ── Part 2A ───────────────────────────────────────────────────────────────
    results["2A - downsample"] = _run(
        "Part 2A - downsample",
        lambda: test_downsample(img_gray)
    )

    # ── Part 2B ───────────────────────────────────────────────────────────────
    # find_most_likely_positions calls pattern_difference and uses downsample output
    if results["1A - pattern_difference"] and results["2A - downsample"]:
        results["2B - find_most_likely_positions"] = _run(
            "Part 2B - find_most_likely_positions",
            lambda: test_find_most_likely_positions(img_gray)
        )
    else:
        missing = []
        if not results["1A - pattern_difference"]:
            missing.append("pattern_difference")
        if not results["2A - downsample"]:
            missing.append("downsample")
        results["2B - find_most_likely_positions"] = _skip(
            "Part 2B - find_most_likely_positions",
            f"{' and '.join(missing)} must pass first"
        )

    # ── Part 2C ───────────────────────────────────────────────────────────────
    # faster_search calls simple_blur, downsample, find_most_likely_positions,
    # and pattern_difference
    if results["2B - find_most_likely_positions"]:
        results["2C - faster_search"] = _run(
            "Part 2C - faster_search",
            lambda: test_faster_search(img_gray)
        )
    else:
        results["2C - faster_search"] = _skip(
            "Part 2C - faster_search",
            "find_most_likely_positions must pass first"
        )

    # ── Part 2D ───────────────────────────────────────────────────────────────
    # faster_highlight calls faster_search AND highlight_rectangle
    if results["2C - faster_search"] and results["1C - highlight_rectangle"]:
        results["2D - faster_highlight"] = _run(
            "Part 2D - faster_highlight",
            lambda: test_faster_highlight(img_gray)
        )
    else:
        missing = []
        if not results["2C - faster_search"]:
            missing.append("faster_search")
        if not results["1C - highlight_rectangle"]:
            missing.append("highlight_rectangle")
        results["2D - faster_highlight"] = _skip(
            "Part 2D - faster_highlight",
            f"{' and '.join(missing)} must pass first"
        )

    # ── Summary ───────────────────────────────────────────────────────────────
    passed = sum(results.values())
    total  = len(results)
    print("\n" + "=" * 56)
    print(f"  Summary: {passed}/{total} sections passing")
    print("=" * 56)
    for label, ok in results.items():
        mark = "✓" if ok else "✗"
        print(f"  {mark}  {label}")
    print("=" * 56)