"""Task 219 Solver — Template pattern extension with periodicity handling (255/265 PASS).

Rule:
1. 15x10 grid with groups of rows containing 8s
2. Template = group with most 8s
3. Each group is extended to match the template pattern using color 1
4. Alignment (top vs bottom) determined by overlap heuristic
5. Per-row offset handles shifted periodic patterns

Status: 255/265 — 10 remaining failures are 'merged row' cases where
group rows contain columns from MULTIPLE template rows. These need a
per-row alignment approach that conflicts with the per-row offset approach.

ONNX: NOT VIABLE until solver reaches 265/265.
"""
import numpy as np
import json


def find_groups(inp):
    """Find groups of consecutive rows with 8s."""
    H, W = inp.shape
    groups = []
    i = 0
    while i < H:
        if np.any(inp[i] == 8):
            start = i
            while i < H and np.any(inp[i] == 8):
                i += 1
            groups.append((start, i - 1))
        else:
            i += 1
    return groups


def alignment_score(g_patterns, t_patterns, t_height, offset):
    """Score alignment: fraction of group cells in aligned template rows."""
    g_height = len(g_patterns)
    matched = 0
    total = 0
    for i in range(g_height):
        t_idx = offset + i
        if t_idx < 0 or t_idx >= t_height:
            total += len(g_patterns[i])
            continue
        t_cols = set(int(c) for c in t_patterns[t_idx])
        for c in g_patterns[i]:
            total += 1
            if c in t_cols:
                matched += 1
    return matched / max(total, 1)


def compute_extension(inp, actual_row, t_row_cols, max_col, W, all_t_patterns=None):
    """Compute extension columns for a row using per-row offset."""
    t_row_set = set(int(c) for c in t_row_cols)
    
    row_8_cols = set(np.where(inp[actual_row] == 8)[0].tolist()) if 0 <= actual_row < inp.shape[0] else set()
    
    if row_8_cols:
        group_rightmost = max(row_8_cols)
        template_cols_le = {c for c in t_row_set if c <= group_rightmost}
        
        if template_cols_le:
            t_rightmost = max(template_cols_le)
            offset = group_rightmost - t_rightmost
            ext_cols = []
            for c in t_row_cols:
                shifted = int(c) + offset
                if int(c) > t_rightmost and shifted < W and shifted > group_rightmost:
                    ext_cols.append(shifted)
            return ext_cols
        else:
            return [int(c) for c in t_row_cols if int(c) > max_col and int(c) < W]
    else:
        return [int(c) for c in t_row_cols if int(c) > max_col and int(c) < W]


def solve_219(inp_grid):
    inp = np.array(inp_grid)
    H, W = inp.shape
    out = inp.copy()
    
    groups = find_groups(inp)
    if len(groups) < 2:
        return out.tolist()
    
    group_8_counts = [int(np.sum(inp[s:e+1] == 8)) for s, e in groups]
    template_idx = int(np.argmax(group_8_counts))
    t_start, t_end = groups[template_idx]
    t_height = t_end - t_start + 1
    
    t_patterns = []
    for r in range(t_start, t_end + 1):
        t_patterns.append(list(np.where(inp[r] == 8)[0]))
    
    for g_idx, (g_start, g_end) in enumerate(groups):
        if g_idx == template_idx:
            continue
        
        g_height = g_end - g_start + 1
        
        max_col = -1
        for r in range(g_start, g_end + 1):
            cols_8 = np.where(inp[r] == 8)[0]
            if len(cols_8) > 0:
                max_col = max(max_col, int(cols_8.max()))
        
        g_patterns = []
        for r in range(g_start, g_end + 1):
            g_patterns.append(set(np.where(inp[r] == 8)[0].tolist()))
        
        top_offset = 0
        bot_offset = t_height - g_height
        
        score_top = alignment_score(g_patterns, t_patterns, t_height, top_offset)
        score_bot = alignment_score(g_patterns, t_patterns, t_height, bot_offset)
        
        if score_bot > score_top:
            base_row = g_end - t_height + 1
        else:
            base_row = g_start
        
        for t_row_idx in range(t_height):
            actual_row = base_row + t_row_idx
            if actual_row < 0 or actual_row >= H:
                continue
            
            ext_cols = compute_extension(inp, actual_row, t_patterns[t_row_idx], max_col, W, all_t_patterns=t_patterns)
            for c in ext_cols:
                if out[actual_row, c] == 0:
                    out[actual_row, c] = 1
    
    return out.tolist()


if __name__ == "__main__":
    import sys, os
    task_data = 'task-data'
    if not os.path.exists(task_data):
        task_data = '../task-data'
    with open(f"{task_data}/task219.json") as f:
        data = json.load(f)
    all_ex = data["train"] + data["test"] + data.get("arc-gen", [])
    right = sum(1 for ex in all_ex if solve_219(ex["input"]) == ex["output"])
    print(f"task219: {right}/{len(all_ex)} PASS")