#!/usr/bin/env python -u
"""
Task Scanner: Analyzes all unsolved tasks to identify pattern clusters.
Run on Kaggle to see which solver categories have the most potential hits.

Usage:
    python -u task_scanner.py --data_dir /kaggle/input/competitions/neurogolf-2026/

Output: Categorized list of tasks by pattern type.
"""
import os, sys, json, glob, numpy as np
from collections import defaultdict

os.environ["PYTHONUNBUFFERED"] = "1"

SOLVED = {15,16,53,56,60,73,78,82,83,87,95,116,120,127,129,135,136,139,140,142,
           152,164,171,172,179,194,210,211,214,220,223,230,241,256,258,261,266,
           276,282,294,299,305,309,311,317,322,331,337,352,373,380}


def load_task(path):
    with open(path) as f:
        return json.load(f)


def get_shapes(td):
    """Get (input_shapes, output_shapes) for train examples."""
    in_shapes = []
    out_shapes = []
    for ex in td.get("train", []):
        inp = np.array(ex["input"])
        out = np.array(ex["output"])
        in_shapes.append(inp.shape)
        out_shapes.append(out.shape)
    return in_shapes, out_shapes


def analyze_task(td):
    """Analyze a single task and return category tags."""
    tags = []
    in_shapes, out_shapes = get_shapes(td)
    if not in_shapes:
        return ["empty"]

    # Shape analysis
    same_shape = all(i == o for i, o in zip(in_shapes, out_shapes))
    fixed_in = len(set(in_shapes)) == 1
    fixed_out = len(set(out_shapes)) == 1

    if same_shape:
        tags.append("same_shape")
    else:
        tags.append("diff_shape")
        # Check if output is smaller (crop-like)
        if all(o[0] <= i[0] and o[1] <= i[1] for i, o in zip(in_shapes, out_shapes)):
            tags.append("output_smaller")
        # Check if output is larger (expand-like)
        elif all(o[0] >= i[0] and o[1] >= i[1] for i, o in zip(in_shapes, out_shapes)):
            tags.append("output_larger")
        else:
            tags.append("mixed_shape_change")

    if fixed_in:
        tags.append("fixed_input_shape")
    if fixed_out:
        tags.append("fixed_output_shape")

    # Size ratios
    if fixed_in and fixed_out:
        ih, iw = in_shapes[0]
        oh, ow = out_shapes[0]
        if oh == ih and ow == iw:
            pass  # same shape
        elif oh % ih == 0 and ow % iw == 0:
            tags.append(f"scale_{oh//ih}x{ow//iw}")
        elif ih % oh == 0 and iw % ow == 0:
            tags.append(f"downsample_{ih//oh}x{iw//ow}")

    # Color analysis
    train_exs = td.get("train", [])
    in_colors_all = set()
    out_colors_all = set()
    color_preserved = True
    new_colors_introduced = False

    for ex in train_exs:
        inp = np.array(ex["input"])
        out = np.array(ex["output"])
        ic = set(inp.flatten())
        oc = set(out.flatten())
        in_colors_all.update(ic)
        out_colors_all.update(oc)
        if oc - ic:
            new_colors_introduced = True
        if ic != oc:
            color_preserved = False

    if color_preserved:
        tags.append("colors_preserved")
    if new_colors_introduced:
        tags.append("new_colors_added")
    if out_colors_all < in_colors_all:
        tags.append("colors_removed")

    n_in_colors = len(in_colors_all)
    n_out_colors = len(out_colors_all)
    tags.append(f"in_colors={n_in_colors}")
    tags.append(f"out_colors={n_out_colors}")

    # Symmetry checks (on output)
    for ex in train_exs[:1]:
        out = np.array(ex["output"])
        h, w = out.shape
        if h > 1 and np.array_equal(out, out[::-1]):
            tags.append("output_v_symmetric")
        if w > 1 and np.array_equal(out, out[:, ::-1]):
            tags.append("output_h_symmetric")
        if h == w and np.array_equal(out, out.T):
            tags.append("output_transpose_symmetric")

    # Check if output is constant (all same value)
    for ex in train_exs:
        out = np.array(ex["output"])
        if len(np.unique(out)) == 1:
            tags.append("constant_output")
            break

    # Check overlay pattern: output = input + something
    if same_shape:
        overlay_consistent = True
        for ex in train_exs:
            inp = np.array(ex["input"])
            out = np.array(ex["output"])
            diff = (inp != out)
            # Check if changes are only additions (input unchanged where same)
            if not np.all(out[~diff] == inp[~diff]):
                overlay_consistent = False
                break
        if overlay_consistent:
            # Check if the diff pixels follow a pattern
            diffs = []
            for ex in train_exs:
                inp = np.array(ex["input"])
                out = np.array(ex["output"])
                d = (inp != out).astype(int)
                diffs.append(d.sum())
            if all(d > 0 for d in diffs):
                tags.append("overlay_pattern")
                avg_diff = np.mean(diffs)
                tags.append(f"avg_changed_pixels={avg_diff:.0f}")

    # Check if output contains input as subgrid
    if not same_shape and all(o[0] >= i[0] and o[1] >= i[1] for i, o in zip(in_shapes, out_shapes)):
        contains_input = True
        for ex in train_exs:
            inp = np.array(ex["input"])
            out = np.array(ex["output"])
            ih, iw = inp.shape
            oh, ow = out.shape
            found = False
            for r in range(oh - ih + 1):
                for c in range(ow - iw + 1):
                    if np.array_equal(out[r:r+ih, c:c+iw], inp):
                        found = True
                        break
                if found:
                    break
            if not found:
                contains_input = False
                break
        if contains_input:
            tags.append("output_contains_input")

    # Flood fill / propagation hints
    if same_shape:
        for ex in train_exs[:1]:
            inp = np.array(ex["input"])
            out = np.array(ex["output"])
            # Check if output has more non-zero pixels (spreading)
            if (out != 0).sum() > (inp != 0).sum():
                tags.append("pixels_spread")
            elif (out != 0).sum() < (inp != 0).sum():
                tags.append("pixels_removed")

    return tags


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", default="/kaggle/input/competitions/neurogolf-2026/")
    args = parser.parse_args()

    task_files = sorted(glob.glob(os.path.join(args.data_dir, "task*.json")))
    print(f"Found {len(task_files)} tasks", flush=True)
    print(f"Already solved: {len(SOLVED)}", flush=True)
    print(f"Scanning unsolved tasks...\n", flush=True)

    # Categorize
    category_tasks = defaultdict(list)
    task_tags = {}

    for f in task_files:
        task_id = os.path.basename(f).replace(".json", "")
        try:
            task_num = int(task_id.replace("task", ""))
        except ValueError:
            continue
        if task_num in SOLVED:
            continue

        td = load_task(f)
        tags = analyze_task(td)
        task_tags[task_num] = tags
        for tag in tags:
            category_tasks[tag].append(task_num)

    # Report
    print("=" * 70, flush=True)
    print("PATTERN CLUSTERS (sorted by count)", flush=True)
    print("=" * 70, flush=True)

    # Most actionable categories
    actionable = [
        "same_shape", "diff_shape", "output_smaller", "output_larger",
        "colors_preserved", "new_colors_added", "colors_removed",
        "overlay_pattern", "pixels_spread", "pixels_removed",
        "output_contains_input", "output_v_symmetric", "output_h_symmetric",
        "constant_output",
    ]

    for cat in actionable:
        tasks = category_tasks.get(cat, [])
        if tasks:
            print(f"\n{cat} ({len(tasks)} tasks):", flush=True)
            print(f"  Tasks: {tasks[:20]}{'...' if len(tasks)>20 else ''}", flush=True)

    # Scale/downsample patterns
    print(f"\n--- Scale patterns ---", flush=True)
    for cat in sorted(category_tasks.keys()):
        if cat.startswith("scale_") or cat.startswith("downsample_"):
            tasks = category_tasks[cat]
            print(f"  {cat}: {tasks}", flush=True)

    # Summary stats
    print(f"\n{'='*70}", flush=True)
    print("SOLVER OPPORTUNITY SUMMARY", flush=True)
    print(f"{'='*70}", flush=True)

    opportunities = [
        ("overlay_pattern (add pixels to input)", "overlay_pattern"),
        ("pixels_spread (flood fill / propagation)", "pixels_spread"),
        ("output_smaller (crop/extract)", "output_smaller"),
        ("output_larger (expand/border/tile)", "output_larger"),
        ("new_colors_added (computed colors)", "new_colors_added"),
        ("colors_removed (filter/mask)", "colors_removed"),
        ("output_contains_input (embed input in larger)", "output_contains_input"),
        ("symmetric output (completion)", "output_v_symmetric"),
    ]

    for desc, tag in opportunities:
        count = len(category_tasks.get(tag, []))
        print(f"  {desc}: {count} tasks", flush=True)

    print(f"\nTotal unsolved: {len(task_tags)}", flush=True)
    print("Done.", flush=True)


if __name__ == "__main__":
    main()