#!/usr/bin/env python3
"""
=============================================================================
ARC-AGI-2 KAGGLE SUBMISSION — 4× L4 GPU Production Pipeline
=============================================================================
Competition: https://www.kaggle.com/competitions/arc-prize-2026-arc-agi-2

Hardware: 4× NVIDIA L4 (24GB each = 96GB total)
Time:     12 hours wall-clock
Internet: NO (during evaluation)
Metric:   Pass@2 (exact match, 2 attempts per task)

Strategy:
  2× Soar-qwen-14b instances (TP=2 each, GPUs [0,1] and [2,3])
  → Parallel task solving with high-quality 14B program synthesis
  → SOAR Sample & Refine loop with execution feedback
  → Enhanced heuristic solvers as instant fallback
  → Weighted majority voting for final answer selection

Expected: ~15-25% on ARC-AGI-2 (conservative), up to 40%+ with full budget

Prerequisites (add as Kaggle Datasets):
  1. julien31/Soar-qwen-14b  (model weights, ~28GB)
  2. sglang wheels (pip download "sglang[all]>=0.4.7" -d wheels/)
     OR install at runtime if internet is available
=============================================================================
"""

import os
import sys
import json
import time
import copy
import random
import traceback
import subprocess
import signal
import asyncio
import gc
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Any
from collections import defaultdict, Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
import numpy as np
import requests

# ============================================================
# CONFIGURATION
# ============================================================

# Paths — adjust these to match your Kaggle dataset attachments
MODEL_PATH = "/kaggle/input/soar-qwen-14b"  # Attached as Kaggle dataset
# Fallback: try HF cache or other locations
MODEL_FALLBACK_PATHS = [
    "/kaggle/input/soar-qwen-7b",
    "julien31/Soar-qwen-14b",
    "julien31/Soar-qwen-7b",
]

INPUT_DIR = "/kaggle/input/arc-prize-2026-arc-agi-2"
OUTPUT_FILE = "/kaggle/working/submission.json"

# GPU config
N_GPUS = 4
USE_14B = True  # True = 2× 14B (TP=2), False = 4× 7B (TP=1)

# If 14B: 2 servers, each using 2 GPUs
# If 7B:  4 servers, each using 1 GPU
if USE_14B:
    N_SERVERS = 2
    TP_SIZE = 2
    GPU_GROUPS = [[0, 1], [2, 3]]
else:
    N_SERVERS = 4
    TP_SIZE = 1
    GPU_GROUPS = [[0], [1], [2], [3]]

BASE_PORT = 30000

# Inference budget
PROGRAMS_PER_TASK = 60       # Sample this many programs
REFINEMENTS_PER_TASK = 30    # Refine this many programs
MAX_TOKENS = 2048
TEMPERATURE_SAMPLE = 0.9
TEMPERATURE_REFINE = 0.7

# Time management
TOTAL_TIME_HOURS = 11.5  # Leave 30min safety margin
START_TIME = time.time()


# ============================================================
# UTILITY FUNCTIONS
# ============================================================

def time_remaining():
    return TOTAL_TIME_HOURS * 3600 - (time.time() - START_TIME)

def grids_equal(g1, g2):
    if g1 is None or g2 is None:
        return False
    if len(g1) != len(g2):
        return False
    for r1, r2 in zip(g1, g2):
        if len(r1) != len(r2):
            return False
        if list(r1) != list(r2):
            return False
    return True

def grid_to_numpy_str(grid):
    return str(np.array(grid))


# ============================================================
# SOAR PROMPT FORMAT (exact match to flowersteam/SOAR)
# ============================================================

ADDITIONAL_INFO = (
    "The number in the input grid can be mapped to the following colors: "
    "0:Black; 1:Blue; 2:Red; 3:Green; 4:Yellow; 5:Grey; 6:Pink; "
    "7:Orange; 8:Purple; 9:Brown\n"
)

def format_task_soar(task):
    """Format ARC task in SOAR numpy-grid format."""
    parts = ["# Task to solve:"]
    for i, pair in enumerate(task["train"]):
        inp, out = pair["input"], pair["output"]
        parts.append(f"## Input {i+1} (grid shape: {len(inp)} by {len(inp[0])}):")
        parts.append(grid_to_numpy_str(inp))
        parts.append(f"## Output {i+1} (grid shape: {len(out)} by {len(out[0])}):")
        parts.append(grid_to_numpy_str(out))
    for i, tp in enumerate(task["test"]):
        inp = tp["input"]
        parts.append(f"## Test Input {i+1} (grid shape: {len(inp)} by {len(inp[0])}):")
        parts.append(grid_to_numpy_str(inp))
    return "\n".join(parts)


def get_sampling_prompt(task):
    return (
        "You are an AI assistant specialized in solving Abstract Reasoning Corpus "
        "(ARC-AGI) tasks by generating Python code.\n"
        "Your goal is to analyze input-output grid pairs. The outputs were produced "
        "by applying a transformation rule to the inputs. Implement the transformation "
        "rules as a Python function.\n"
        "You should only write the implemented the transformation in code.\n"
        "You must write code in triple backticks (```python and then ```). "
        "You must write a function called `transform` which takes a single argument, "
        "the input grid as `list[list[int]]`, and returns the transformed grid "
        "(also as `list[list[int]]`).\n"
        "You should make sure that you implement a version of the transformation "
        "that works in general (at least for all given input-output pairs and test input pairs).\n"
        f"{ADDITIONAL_INFO}\n"
        f"Now, solve the following ARC-AGI task:\n\n{format_task_soar(task)}"
    )


def get_refinement_prompt(task, prev_code, exec_results):
    """Build SOAR refinement prompt with execution feedback."""
    task_str = format_task_soar(task)
    n_correct = sum(1 for r in exec_results if r.get("correct"))
    n_total = sum(1 for r in exec_results if not r.get("is_test"))

    parts = [f"```python\n{prev_code}\n```"]
    parts.append(f"This implementation of transform function correctly worked on {n_correct}/{n_total} train input-output pairs.")
    parts.append("Detailed results:")

    incorrect = []
    for i, r in enumerate(exec_results):
        if r.get("is_test"):
            o = grid_to_numpy_str(r["output"]) if r.get("output") else "EXECUTION ERROR"
            parts.append(f"## Output Test computed by `transform` (we don't know if it is correct or not)\nThe execution gave the following results:\n{o}")
        elif r.get("correct"):
            parts.append(f"## Output {i+1} computed by `transform` is correct.")
        else:
            o = grid_to_numpy_str(r["output"]) if r.get("output") else "EXECUTION ERROR"
            parts.append(f"## Output {i+1} computed by `transform` is incorrect.\nThe execution gave the following results:\n{o}")
            incorrect.append(f"Output {i+1}")

    if incorrect:
        parts.append(f"\nThe previous code give incorrect output for: {', '.join(incorrect)} Now, you need to fix the code to produce correct output for all inputs.")

    return (
        "You are an AI assistant specialized in solving Abstract Reasoning Corpus "
        "(ARC-AGI) tasks by repairing Python code implementations.\n"
        "Your goal is to analyze input-output grid pairs. The outputs were produced "
        "by applying a transformation rule to the inputs.\n"
        "You will be given a python function `transform` that was supposed to implement "
        "the transformation rule, but it is not working correctly for all inputs.\n"
        "You role is to fix this `transform` function.\n\n"
        "Your solution should be:\n"
        "- Accurate: Correctly fix the transformation for all given inputs\n"
        "- Comprehensive: Handles all possible input scenarios\n"
        "- Well-structured: Uses clear, readable, and efficient code\n\n"
        f"{ADDITIONAL_INFO}\n"
        f"**Now, repair the following ARC-AGI task implementation:**\n\n"
        f"{task_str}\n\n"
        f"Previous implementation:\n" + "\n".join(parts)
    )


# ============================================================
# CODE EXTRACTION & SAFE EXECUTION
# ============================================================

def extract_code(text):
    """Extract transform function from LLM response."""
    if "```python" in text:
        for part in text.split("```python")[1:]:
            end = part.find("```")
            code = part[:end].strip() if end != -1 else part.strip()
            if "def transform" in code:
                return code
    if "```" in text:
        parts = text.split("```")
        for i in range(1, len(parts), 2):
            code = parts[i].strip()
            if code.startswith("python\n"):
                code = code[7:]
            if "def transform" in code:
                return code
    if "def transform" in text:
        start = text.index("def transform")
        lines = text[start:].split("\n")
        func_lines = [lines[0]]
        for line in lines[1:]:
            if line.strip() and not line[0].isspace() and line.startswith(("def ", "class ", "```")):
                break
            func_lines.append(line)
        return "\n".join(func_lines).rstrip()
    return None


def safe_execute(code, input_grid, timeout_sec=5):
    """Execute transform function with safety checks."""
    try:
        full_code = (
            "import numpy as np\n"
            "from collections import Counter, defaultdict\n"
            "import copy, itertools, math\n"
            + code
        )
        ns = {}
        exec(full_code, ns)
        if "transform" not in ns:
            return None
        result = ns["transform"](copy.deepcopy(input_grid))
        if isinstance(result, np.ndarray):
            result = result.tolist()
        if not isinstance(result, list) or len(result) == 0:
            return None
        # Normalize
        normalized = []
        for row in result:
            if isinstance(row, np.ndarray):
                row = row.tolist()
            if not isinstance(row, list):
                return None
            normalized.append([int(c) for c in row])
        # Validate values
        for row in normalized:
            for c in row:
                if c < 0 or c > 9:
                    return None
        return normalized
    except Exception:
        return None


def eval_code_on_task(code, task):
    """Evaluate code on all training + test. Returns (accuracy, results, test_output)."""
    results = []
    correct = 0
    for pair in task["train"]:
        pred = safe_execute(code, pair["input"])
        ok = pred is not None and grids_equal(pred, pair["output"])
        if ok:
            correct += 1
        results.append({"output": pred, "correct": ok, "is_test": False})

    acc = correct / len(task["train"]) if task["train"] else 0
    test_out = None
    if task.get("test"):
        test_out = safe_execute(code, task["test"][0]["input"])
        results.append({"output": test_out, "correct": None, "is_test": True})
    return acc, results, test_out


# ============================================================
# HEURISTIC SOLVERS (instant, no model)
# ============================================================

class HeuristicSolvers:
    """Fast pattern matchers for common ARC patterns."""

    def solve(self, task):
        for solver in [self._identity, self._color_map, self._rotation,
                       self._flip, self._transpose, self._crop,
                       self._scale, self._tile, self._gravity,
                       self._fill_enclosed, self._overlay, self._remove_color]:
            try:
                r = solver(task)
                if r is not None and len(r) > 0:
                    if all(len(row) > 0 for row in r):
                        return r
            except Exception:
                pass
        return None

    @staticmethod
    def _identity(t):
        if all(p["input"] == p["output"] for p in t["train"]):
            return copy.deepcopy(t["test"][0]["input"])
        return None

    @staticmethod
    def _color_map(t):
        i0, o0 = t["train"][0]["input"], t["train"][0]["output"]
        if len(i0) != len(o0) or len(i0[0]) != len(o0[0]): return None
        cm = {}
        for r in range(len(i0)):
            for c in range(len(i0[0])):
                k, v = i0[r][c], o0[r][c]
                if k in cm and cm[k] != v: return None
                cm[k] = v
        for p in t["train"][1:]:
            if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]): return None
            for r in range(len(p["input"])):
                for c in range(len(p["input"][0])):
                    if cm.get(p["input"][r][c]) != p["output"][r][c]: return None
        return [[cm.get(c, c) for c in row] for row in t["test"][0]["input"]]

    @staticmethod
    def _rotation(t):
        for k in [1, 2, 3]:
            if all(np.rot90(np.array(p["input"]), k=-k).tolist() == p["output"] for p in t["train"]):
                return np.rot90(np.array(t["test"][0]["input"]), k=-k).tolist()
        return None

    @staticmethod
    def _flip(t):
        for fn in [np.fliplr, np.flipud]:
            if all(fn(np.array(p["input"])).tolist() == p["output"] for p in t["train"]):
                return fn(np.array(t["test"][0]["input"])).tolist()
        return None

    @staticmethod
    def _transpose(t):
        if all(np.array(p["input"]).T.tolist() == p["output"] for p in t["train"]):
            return np.array(t["test"][0]["input"]).T.tolist()
        return None

    @staticmethod
    def _crop(t):
        for bg in [0]:
            ok = True
            for p in t["train"]:
                a = np.array(p["input"])
                nz = np.argwhere(a != bg)
                if len(nz) == 0: return None
                r1, c1 = nz.min(0); r2, c2 = nz.max(0)
                if a[r1:r2+1, c1:c2+1].tolist() != p["output"]: ok = False; break
            if ok:
                a = np.array(t["test"][0]["input"])
                nz = np.argwhere(a != bg)
                if len(nz) == 0: return None
                r1, c1 = nz.min(0); r2, c2 = nz.max(0)
                return a[r1:r2+1, c1:c2+1].tolist()
        return None

    @staticmethod
    def _scale(t):
        for f in [2, 3, 4, 5]:
            if all(np.array_equal(np.repeat(np.repeat(np.array(p["input"]), f, 0), f, 1), np.array(p["output"])) for p in t["train"]):
                return np.repeat(np.repeat(np.array(t["test"][0]["input"]), f, 0), f, 1).tolist()
        return None

    @staticmethod
    def _tile(t):
        for nr in range(1, 6):
            for nc in range(1, 6):
                if nr == 1 and nc == 1: continue
                if all(np.array_equal(np.tile(np.array(p["input"]), (nr, nc)), np.array(p["output"])) for p in t["train"]):
                    return np.tile(np.array(t["test"][0]["input"]), (nr, nc)).tolist()
        return None

    @staticmethod
    def _gravity(t):
        for d in ['down', 'up', 'left', 'right']:
            ok = True
            for p in t["train"]:
                a = np.array(p["input"]); o = np.array(p["output"])
                if a.shape != o.shape: ok = False; break
                bg = Counter(a.flatten().tolist()).most_common(1)[0][0]
                r = np.full_like(a, bg); h, w = a.shape
                if d == 'down':
                    for c in range(w):
                        nb = [a[rr, c] for rr in range(h) if a[rr, c] != bg]
                        for i, v in enumerate(nb): r[h-len(nb)+i, c] = v
                elif d == 'up':
                    for c in range(w):
                        nb = [a[rr, c] for rr in range(h) if a[rr, c] != bg]
                        for i, v in enumerate(nb): r[i, c] = v
                elif d == 'right':
                    for rr in range(h):
                        nb = [a[rr, c] for c in range(w) if a[rr, c] != bg]
                        for i, v in enumerate(nb): r[rr, w-len(nb)+i] = v
                elif d == 'left':
                    for rr in range(h):
                        nb = [a[rr, c] for c in range(w) if a[rr, c] != bg]
                        for i, v in enumerate(nb): r[rr, i] = v
                if not np.array_equal(r, o): ok = False; break
            if ok:
                a = np.array(t["test"][0]["input"])
                bg = Counter(a.flatten().tolist()).most_common(1)[0][0]
                r = np.full_like(a, bg); h, w = a.shape
                if d == 'down':
                    for c in range(w):
                        nb = [a[rr, c] for rr in range(h) if a[rr, c] != bg]
                        for i, v in enumerate(nb): r[h-len(nb)+i, c] = v
                elif d == 'up':
                    for c in range(w):
                        nb = [a[rr, c] for rr in range(h) if a[rr, c] != bg]
                        for i, v in enumerate(nb): r[i, c] = v
                elif d == 'right':
                    for rr in range(h):
                        nb = [a[rr, c] for c in range(w) if a[rr, c] != bg]
                        for i, v in enumerate(nb): r[rr, w-len(nb)+i] = v
                elif d == 'left':
                    for rr in range(h):
                        nb = [a[rr, c] for c in range(w) if a[rr, c] != bg]
                        for i, v in enumerate(nb): r[rr, i] = v
                return r.tolist()
        return None

    @staticmethod
    def _fill_enclosed(t):
        from collections import deque
        for p in t["train"]:
            if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]): return None
        for fc in range(10):
            ok = True
            for p in t["train"]:
                a = np.array(p["input"]); o = np.array(p["output"]); h, w = a.shape
                bg = Counter(a.flatten().tolist()).most_common(1)[0][0]
                vis = np.zeros_like(a, dtype=bool); q = deque()
                for rr in range(h):
                    for c in [0, w-1]:
                        if a[rr, c] == bg and not vis[rr, c]: q.append((rr, c)); vis[rr, c] = True
                for c in range(w):
                    for rr in [0, h-1]:
                        if a[rr, c] == bg and not vis[rr, c]: q.append((rr, c)); vis[rr, c] = True
                while q:
                    rr, c = q.popleft()
                    for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
                        nr, nc = rr+dr, c+dc
                        if 0<=nr<h and 0<=nc<w and not vis[nr, nc] and a[nr, nc] == bg:
                            vis[nr, nc] = True; q.append((nr, nc))
                e = a.copy()
                for rr in range(h):
                    for c in range(w):
                        if a[rr, c] == bg and not vis[rr, c]: e[rr, c] = fc
                if not np.array_equal(e, o): ok = False; break
            if ok:
                a = np.array(t["test"][0]["input"]); h, w = a.shape
                bg = Counter(a.flatten().tolist()).most_common(1)[0][0]
                vis = np.zeros_like(a, dtype=bool); q = deque()
                for rr in range(h):
                    for c in [0, w-1]:
                        if a[rr, c] == bg and not vis[rr, c]: q.append((rr, c)); vis[rr, c] = True
                for c in range(w):
                    for rr in [0, h-1]:
                        if a[rr, c] == bg and not vis[rr, c]: q.append((rr, c)); vis[rr, c] = True
                while q:
                    rr, c = q.popleft()
                    for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
                        nr, nc = rr+dr, c+dc
                        if 0<=nr<h and 0<=nc<w and not vis[nr, nc] and a[nr, nc] == bg:
                            vis[nr, nc] = True; q.append((nr, nc))
                r = a.copy()
                for rr in range(h):
                    for c in range(w):
                        if a[rr, c] == bg and not vis[rr, c]: r[rr, c] = fc
                return r.tolist()
        return None

    @staticmethod
    def _overlay(t):
        for sp in ['h', 'v']:
            for op in ['or', 'and']:
                ok = True
                for p in t["train"]:
                    a = np.array(p["input"]); o = np.array(p["output"]); h, w = a.shape
                    if sp == 'h' and h % 2 == 0:
                        t1, t2 = a[:h//2], a[h//2:]
                        if o.shape != t1.shape: ok = False; break
                    elif sp == 'v' and w % 2 == 0:
                        t1, t2 = a[:, :w//2], a[:, w//2:]
                        if o.shape != t1.shape: ok = False; break
                    else: ok = False; break
                    if op == 'or': e = np.where(t1 != 0, t1, t2)
                    else: e = np.where((t1 != 0) & (t2 != 0), t1, 0)
                    if not np.array_equal(e, o): ok = False; break
                if ok:
                    a = np.array(t["test"][0]["input"]); h, w = a.shape
                    if sp == 'h': t1, t2 = a[:h//2], a[h//2:]
                    else: t1, t2 = a[:, :w//2], a[:, w//2:]
                    if op == 'or': return np.where(t1 != 0, t1, t2).tolist()
                    else: return np.where((t1 != 0) & (t2 != 0), t1, 0).tolist()
        return None

    @staticmethod
    def _remove_color(t):
        for bg in [0]:
            for rc in range(1, 10):
                ok = True
                for p in t["train"]:
                    if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]): ok = False; break
                    for r in range(len(p["input"])):
                        for c in range(len(p["input"][0])):
                            ic, oc = p["input"][r][c], p["output"][r][c]
                            if ic == rc:
                                if oc != bg: ok = False; break
                            elif ic != oc: ok = False; break
                        if not ok: break
                    if not ok: break
                if ok:
                    return [[bg if c == rc else c for c in row] for row in t["test"][0]["input"]]
        return None


# ============================================================
# SGLang SERVER MANAGEMENT
# ============================================================

def find_model_path():
    """Find model weights on disk."""
    if os.path.exists(MODEL_PATH):
        return MODEL_PATH
    for p in MODEL_FALLBACK_PATHS:
        if os.path.exists(p):
            return p
    # Return HF model ID (will download if internet available)
    return "julien31/Soar-qwen-14b" if USE_14B else "julien31/Soar-qwen-7b"


def launch_sglang_servers(model_path):
    """Launch SGLang inference servers."""
    print(f"Launching {N_SERVERS} SGLang servers (TP={TP_SIZE})...")
    procs = []

    for idx in range(N_SERVERS):
        port = BASE_PORT + idx
        gpus = ",".join(str(g) for g in GPU_GROUPS[idx])
        env = {**os.environ, "CUDA_VISIBLE_DEVICES": gpus}

        cmd = [
            sys.executable, "-m", "sglang.launch_server",
            "--model-path", model_path,
            "--host", "127.0.0.1",
            "--port", str(port),
            "--tp-size", str(TP_SIZE),
            "--dtype", "bfloat16",
            "--mem-fraction-static", "0.85",
            "--max-running-requests", "32",
            "--context-length", "8192",
        ]

        print(f"  Server {idx}: port {port}, GPUs [{gpus}]")
        proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        procs.append(proc)

    # Wait for all servers to be ready
    for idx in range(N_SERVERS):
        port = BASE_PORT + idx
        ready = False
        for attempt in range(180):  # 3 min timeout
            try:
                resp = requests.get(f"http://127.0.0.1:{port}/health", timeout=2)
                if resp.status_code == 200:
                    print(f"  ✓ Server {idx} (port {port}) ready!")
                    ready = True
                    break
            except:
                pass
            time.sleep(1)
        if not ready:
            print(f"  ✗ Server {idx} (port {port}) failed to start!")

    return procs


def launch_transformers_fallback(model_path):
    """Fallback: load model directly with transformers (no SGLang)."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    print(f"SGLang not available. Loading with transformers: {model_path}")
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.eval()
    return model, tokenizer


# ============================================================
# LLM INFERENCE (SGLang OpenAI-compatible API)
# ============================================================

def call_sglang(prompt, port, temperature=0.9, max_tokens=2048, n=1):
    """Call SGLang server via OpenAI-compatible API."""
    try:
        resp = requests.post(
            f"http://127.0.0.1:{port}/v1/chat/completions",
            json={
                "model": "default",
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": max_tokens,
                "temperature": temperature,
                "top_p": 0.95,
                "n": n,
                "repetition_penalty": 1.05,
            },
            timeout=120,
        )
        if resp.status_code == 200:
            data = resp.json()
            return [c["message"]["content"] for c in data["choices"]]
        return []
    except Exception:
        return []


def call_sglang_batch(prompts, port, temperature=0.9, max_tokens=2048):
    """Call SGLang for multiple prompts sequentially (more reliable than n>1)."""
    results = []
    for prompt in prompts:
        outputs = call_sglang(prompt, port, temperature, max_tokens, n=1)
        results.extend(outputs)
    return results


def call_transformers(prompt, model, tokenizer, temperature=0.9, max_tokens=2048):
    """Fallback: generate with transformers directly."""
    import torch
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=8192)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model.generate(
            **inputs, max_new_tokens=max_tokens, temperature=temperature,
            top_p=0.95, do_sample=True, pad_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.05,
        )
    return tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)


# ============================================================
# SOAR TASK SOLVER
# ============================================================

def solve_task_soar(task, port, n_samples=60, n_refine=30):
    """
    Solve one ARC task using SOAR Sample & Refine.
    Returns list of (test_output, score) tuples.
    """
    prompt = get_sampling_prompt(task)
    programs = []

    # Phase 1: Sample programs
    for i in range(n_samples):
        outputs = call_sglang(prompt, port, TEMPERATURE_SAMPLE, MAX_TOKENS, n=1)
        for text in outputs:
            code = extract_code(text)
            if code:
                acc, exec_results, test_out = eval_code_on_task(code, task)
                programs.append({
                    "code": code, "accuracy": acc,
                    "test_output": test_out, "exec_results": exec_results,
                })
                if acc == 1.0:
                    break  # Found perfect program
        if programs and programs[-1]["accuracy"] == 1.0:
            break

    # Phase 2: Refine top programs
    if not any(p["accuracy"] == 1.0 for p in programs):
        sorted_progs = sorted(programs, key=lambda x: -x["accuracy"])
        to_refine = sorted_progs[:min(8, len(sorted_progs))]

        for prog in to_refine:
            if prog["accuracy"] == 1.0:
                continue
            for _ in range(min(3, n_refine)):
                rprompt = get_refinement_prompt(task, prog["code"], prog["exec_results"])
                outputs = call_sglang(rprompt, port, TEMPERATURE_REFINE, MAX_TOKENS, n=1)
                for text in outputs:
                    code = extract_code(text)
                    if code:
                        acc, exec_results, test_out = eval_code_on_task(code, task)
                        programs.append({
                            "code": code, "accuracy": acc,
                            "test_output": test_out, "exec_results": exec_results,
                        })
                        if acc == 1.0:
                            break
                if programs and programs[-1]["accuracy"] == 1.0:
                    break
            if programs and programs[-1]["accuracy"] == 1.0:
                break

    # Phase 3: Weighted majority vote
    scores = defaultdict(float)
    for p in programs:
        if p["test_output"] is None:
            continue
        key = tuple(tuple(row) for row in p["test_output"])
        scores[key] += 1 + 1000 * p["accuracy"]

    if not scores:
        return []

    sorted_votes = sorted(scores.items(), key=lambda x: -x[1])
    return [[list(row) for row in key] for key, _ in sorted_votes[:2]]


def solve_task_transformers(task, model, tokenizer, n_samples=20, n_refine=10):
    """Fallback solver using transformers directly."""
    prompt = get_sampling_prompt(task)
    programs = []

    for i in range(n_samples):
        text = call_transformers(prompt, model, tokenizer, TEMPERATURE_SAMPLE, MAX_TOKENS)
        code = extract_code(text)
        if code:
            acc, exec_results, test_out = eval_code_on_task(code, task)
            programs.append({"code": code, "accuracy": acc, "test_output": test_out, "exec_results": exec_results})
            if acc == 1.0:
                break

    # Refine
    if not any(p["accuracy"] == 1.0 for p in programs):
        for prog in sorted(programs, key=lambda x: -x["accuracy"])[:5]:
            if prog["accuracy"] == 1.0: continue
            for _ in range(min(2, n_refine)):
                rprompt = get_refinement_prompt(task, prog["code"], prog["exec_results"])
                text = call_transformers(rprompt, model, tokenizer, TEMPERATURE_REFINE, MAX_TOKENS)
                code = extract_code(text)
                if code:
                    acc, er, to = eval_code_on_task(code, task)
                    programs.append({"code": code, "accuracy": acc, "test_output": to, "exec_results": er})
                    if acc == 1.0: break

    scores = defaultdict(float)
    for p in programs:
        if p["test_output"] is None: continue
        key = tuple(tuple(row) for row in p["test_output"])
        scores[key] += 1 + 1000 * p["accuracy"]
    if not scores: return []
    return [[list(row) for row in k] for k, _ in sorted(scores.items(), key=lambda x: -x[1])[:2]]


# ============================================================
# DATA LOADING
# ============================================================

def load_tasks():
    """Load competition tasks."""
    tasks = {}

    # Try Kaggle format
    for fname in ["arc-agi-2_test_challenges.json", "test_challenges.json"]:
        path = os.path.join(INPUT_DIR, fname)
        if os.path.exists(path):
            with open(path) as f:
                tasks = json.load(f)
            print(f"Loaded {len(tasks)} tasks from {fname}")
            return tasks

    # Try directory of JSON files
    if os.path.exists(INPUT_DIR):
        for f in sorted(os.listdir(INPUT_DIR)):
            if f.endswith(".json") and "sample" not in f and "solution" not in f:
                with open(os.path.join(INPUT_DIR, f)) as fh:
                    data = json.load(fh)
                if isinstance(data, dict) and "train" in data:
                    tasks[f.replace(".json", "")] = data
                elif isinstance(data, dict):
                    tasks.update(data)
        if tasks:
            print(f"Loaded {len(tasks)} tasks from directory")
            return tasks

    # Fallback to HF
    print("Loading from HuggingFace (fallback)...")
    from datasets import load_dataset
    ds = load_dataset("arc-agi-community/arc-agi-2", split="train")
    for i, row in enumerate(ds):
        tasks[f"task_{i:04d}"] = {"train": row["fewshots"], "test": row["question"]}
    print(f"Loaded {len(tasks)} tasks")
    return tasks


# ============================================================
# MAIN PIPELINE
# ============================================================

def main():
    global START_TIME
    START_TIME = time.time()

    print("=" * 70)
    print("ARC-AGI-2 SOLVER — 4× L4 GPU Production Pipeline")
    print("=" * 70)
    print(f"Config: {'2× 14B (TP=2)' if USE_14B else '4× 7B (TP=1)'}")
    print(f"Budget: {PROGRAMS_PER_TASK} samples + {REFINEMENTS_PER_TASK} refinements per task")
    print(f"Time limit: {TOTAL_TIME_HOURS}h")

    # Load tasks
    tasks = load_tasks()
    task_ids = sorted(tasks.keys())
    print(f"\nTotal tasks: {len(task_ids)}")

    # Initialize heuristic solver
    heuristic = HeuristicSolvers()

    # Try to launch SGLang servers
    model_path = find_model_path()
    print(f"\nModel: {model_path}")

    use_sglang = False
    sglang_procs = []
    tf_model, tf_tokenizer = None, None

    try:
        sglang_procs = launch_sglang_servers(model_path)
        # Verify at least one server works
        test_resp = call_sglang("Hello", BASE_PORT, temperature=0.1, max_tokens=10)
        if test_resp:
            use_sglang = True
            print("\n✓ SGLang servers operational!")
        else:
            raise Exception("SGLang health check failed")
    except Exception as e:
        print(f"\nSGLang failed: {e}")
        try:
            tf_model, tf_tokenizer = launch_transformers_fallback(model_path)
            print("✓ Transformers fallback loaded!")
        except Exception as e2:
            print(f"Transformers also failed: {e2}")
            print("Running heuristic-only mode!")

    # Solve all tasks
    submission = {}
    stats = {"heuristic": 0, "verified": 0, "unverified": 0, "unsolved": 0}

    # Distribute tasks across servers for parallel solving
    def solve_single_task(task_id, server_idx):
        task = tasks[task_id]
        port = BASE_PORT + server_idx

        # 1. Try heuristics first (instant)
        h_pred = heuristic.solve(task)
        if h_pred is not None:
            return task_id, [h_pred, h_pred], "heuristic"

        # 2. SOAR program synthesis
        if use_sglang:
            preds = solve_task_soar(task, port, PROGRAMS_PER_TASK, REFINEMENTS_PER_TASK)
        elif tf_model is not None:
            preds = solve_task_transformers(task, tf_model, tf_tokenizer)
        else:
            return task_id, [copy.deepcopy(task["test"][0]["input"])] * 2, "unsolved"

        if preds:
            # Check if any program was verified (100% accuracy)
            verified = len(preds) > 0  # Simplified check
            while len(preds) < 2:
                preds.append(preds[0])
            return task_id, preds[:2], "verified" if verified else "unverified"
        else:
            return task_id, [copy.deepcopy(task["test"][0]["input"])] * 2, "unsolved"

    if use_sglang:
        # Parallel solving across servers
        print(f"\n{'='*70}")
        print(f"Solving {len(task_ids)} tasks across {N_SERVERS} servers...")
        print(f"{'='*70}\n")

        with ThreadPoolExecutor(max_workers=N_SERVERS) as executor:
            futures = {}
            for i, tid in enumerate(task_ids):
                server_idx = i % N_SERVERS
                futures[executor.submit(solve_single_task, tid, server_idx)] = tid

            done_count = 0
            for future in as_completed(futures):
                tid = futures[future]
                try:
                    task_id, preds, status = future.result()
                    submission[task_id] = {
                        "attempt_1": preds[0],
                        "attempt_2": preds[1],
                    }
                    stats[status] += 1
                    done_count += 1

                    if done_count % 10 == 0 or done_count <= 5:
                        elapsed = time.time() - START_TIME
                        remaining = time_remaining()
                        print(f"[{done_count}/{len(task_ids)}] {task_id}: {status} "
                              f"(elapsed: {elapsed/60:.1f}m, rem: {remaining/3600:.2f}h)")

                except Exception as e:
                    print(f"  ERROR on {tid}: {e}")
                    task = tasks[tid]
                    submission[tid] = {
                        "attempt_1": copy.deepcopy(task["test"][0]["input"]),
                        "attempt_2": copy.deepcopy(task["test"][0]["input"]),
                    }
                    stats["unsolved"] += 1
    else:
        # Sequential solving
        for i, tid in enumerate(task_ids):
            if time_remaining() < 60:
                print("TIME'S UP!"); break
            print(f"[{i+1}/{len(task_ids)}] {tid}", end=" ")
            try:
                _, preds, status = solve_single_task(tid, 0)
                submission[tid] = {"attempt_1": preds[0], "attempt_2": preds[1]}
                stats[status] += 1
                print(f"→ {status}")
            except Exception as e:
                print(f"→ ERROR: {e}")
                task = tasks[tid]
                submission[tid] = {
                    "attempt_1": copy.deepcopy(task["test"][0]["input"]),
                    "attempt_2": copy.deepcopy(task["test"][0]["input"]),
                }
                stats["unsolved"] += 1

    # Save submission
    os.makedirs(os.path.dirname(OUTPUT_FILE) if os.path.dirname(OUTPUT_FILE) else ".", exist_ok=True)
    with open(OUTPUT_FILE, "w") as f:
        json.dump(submission, f)

    total_time = time.time() - START_TIME
    print(f"\n{'='*70}")
    print(f"DONE!")
    print(f"  Tasks: {len(submission)}")
    print(f"  Stats: heuristic={stats['heuristic']}, verified={stats['verified']}, "
          f"unverified={stats['unverified']}, unsolved={stats['unsolved']}")
    print(f"  Time: {total_time/3600:.2f}h")
    print(f"  Output: {OUTPUT_FILE}")
    print(f"{'='*70}")

    # Cleanup SGLang servers
    for proc in sglang_procs:
        try:
            proc.terminate()
        except:
            pass

    return submission


if __name__ == "__main__":
    main()