"""
SOAR Program Synthesis Evaluation for ARC-AGI-2
Uses julien31/Soar-qwen-7b to solve ARC tasks via program synthesis.

Exact prompt format from SOAR repository (flowersteam/SOAR).
"""
import os
import sys
import json
import time
import copy
import random
import traceback
from typing import List, Dict, Tuple, Optional, Any
from collections import defaultdict, Counter
import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

# ============================================================
# SOAR Prompt Format (exact replication from soar/prompt.py)
# ============================================================

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 grid_to_numpy_str(grid: List[List[int]]) -> str:
    """Format grid in numpy repr mode (default SOAR format)."""
    return str(np.array(grid))

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


def get_sampling_prompt(task: Dict) -> str:
    """Build the full sampling prompt for SOAR."""
    task_str = format_task_soar(task)
    
    user_msg = (
        "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"
        "Now, solve the following ARC-AGI task:\n\n"
        f"{task_str}"
    )
    
    return user_msg


def get_refinement_prompt(task: Dict, previous_code: str, 
                          execution_results: List[Dict]) -> str:
    """Build refinement prompt for SOAR."""
    task_str = format_task_soar(task)
    
    # Format previous implementation results
    n_correct = sum(1 for r in execution_results if r.get("correct", False))
    n_total = len(execution_results)
    
    prev_impl_parts = [
        f"```python\n{previous_code}\n```",
        f"This implementation of transform function correctly worked on {n_correct}/{n_total} train input-output pairs.",
        "Detailed results:"
    ]
    
    incorrect_outputs = []
    for i, result in enumerate(execution_results):
        if result.get("is_test", False):
            out_str = grid_to_numpy_str(result["output"]) if result.get("output") else "EXECUTION ERROR"
            prev_impl_parts.append(
                f"## Output Test {i+1} computed by `transform` (we don't know if it is correct or not)\n"
                f"The execution gave the following results:\n{out_str}"
            )
        elif result.get("correct", False):
            prev_impl_parts.append(f"## Output {i+1} computed by `transform` is correct.")
        else:
            out_str = grid_to_numpy_str(result["output"]) if result.get("output") else "EXECUTION ERROR"
            h = len(result["output"]) if result.get("output") else "?"
            w = len(result["output"][0]) if result.get("output") and result["output"] else "?"
            prev_impl_parts.append(
                f"## Output {i+1} computed by `transform` is incorrect.\n"
                f"The execution gave the following results (grid shape: {h} by {w}):\n{out_str}"
            )
            incorrect_outputs.append(f"Output {i+1}")
    
    if incorrect_outputs:
        prev_impl_parts.append(
            f"\nThe previous code give incorrect output for: {', '.join(incorrect_outputs)} "
            "Now, you need to fix the code to produce correct output for all inputs."
        )
    
    previous_implementation = "\n".join(prev_impl_parts)
    
    user_msg = (
        "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 so they give "
        "correct outputs as provided (it should also work for all test 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{previous_implementation}"
    )
    
    return user_msg


# ============================================================
# Code extraction
# ============================================================

def extract_transform_code(text: str) -> Optional[str]:
    """Extract Python transform function from LLM output (SOAR style)."""
    # Try ```python blocks
    if "```python" in text:
        parts = text.split("```python")
        for part in parts[1:]:
            end = part.find("```")
            if end != -1:
                code = part[:end].strip()
            else:
                code = 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
    
    # Try direct extraction
    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 not line.startswith("#"):
                if line.startswith("def ") or line.startswith("class ") or line.startswith("```"):
                    break
            func_lines.append(line)
        return "\n".join(func_lines).rstrip()
    
    return None


# ============================================================
# Safe execution
# ============================================================

def safe_execute(code: str, input_grid: List[List[int]], timeout_sec: float = 5.0) -> Optional[List[List[int]]]:
    """Execute transform function safely."""
    try:
        # Add common imports
        full_code = "import numpy as np\nfrom collections import Counter, defaultdict\nimport copy\n" + code
        namespace = {}
        exec(full_code, namespace)
        
        if "transform" not in namespace:
            return None
        
        result = namespace["transform"](copy.deepcopy(input_grid))
        
        if not isinstance(result, (list, np.ndarray)):
            return None
        
        if isinstance(result, np.ndarray):
            result = result.tolist()
        
        if len(result) == 0:
            return None
        
        # Validate
        for row in result:
            if not isinstance(row, (list, np.ndarray)):
                return None
            if isinstance(row, np.ndarray):
                row = row.tolist()
            for cell in row:
                if not isinstance(cell, (int, float, np.integer)):
                    return None
        
        result = [[int(c) for c in (r.tolist() if isinstance(r, np.ndarray) else r)] for r in result]
        return result
        
    except Exception:
        return None


def evaluate_code_on_task(code: str, task: Dict) -> Tuple[float, List[Dict], Optional[List[List[int]]]]:
    """
    Evaluate code on all training pairs and test input.
    Returns (accuracy, execution_results, test_output).
    """
    results = []
    n_correct = 0
    
    for pair in task["train"]:
        pred = safe_execute(code, pair["input"])
        correct = pred is not None and grids_equal(pred, pair["output"])
        if correct:
            n_correct += 1
        results.append({
            "output": pred,
            "correct": correct,
            "is_test": False,
        })
    
    accuracy = n_correct / len(task["train"]) if task["train"] else 0
    
    # Test output
    test_output = None
    if task.get("test"):
        test_output = safe_execute(code, task["test"][0]["input"])
        results.append({
            "output": test_output,
            "correct": None,  # We don't know ground truth
            "is_test": True,
        })
    
    return accuracy, results, test_output


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


# ============================================================
# Main SOAR evaluation
# ============================================================

def solve_task_soar(task: Dict, model, tokenizer, 
                    n_samples: int = 30, n_refinements: int = 20,
                    temperature: float = 0.8) -> List[List[List[int]]]:
    """
    Solve a single ARC task using SOAR-style program synthesis.
    
    1. Sample n_samples programs
    2. Evaluate each on training pairs
    3. Refine promising ones
    4. Weighted majority vote → top-2
    """
    
    # Phase 1: Sample programs
    sampling_prompt = get_sampling_prompt(task)
    messages = [{"role": "user", "content": sampling_prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    programs = []  # List of (code, accuracy, test_output)
    
    for i in range(n_samples):
        try:
            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=2048,
                    temperature=temperature,
                    top_p=0.95,
                    min_p=0.05,
                    do_sample=True,
                    pad_token_id=tokenizer.eos_token_id,
                    repetition_penalty=1.05,
                )
            
            response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
            code = extract_transform_code(response)
            
            if code is None:
                continue
            
            accuracy, exec_results, test_output = evaluate_code_on_task(code, task)
            programs.append((code, accuracy, test_output, exec_results))
            
            # Early stop if we find a perfect program
            if accuracy == 1.0:
                print(f"    Found perfect program at sample {i+1}")
                break
                
        except Exception as e:
            continue
    
    # Phase 2: Refine promising programs
    # Select programs to refine (REX-style: mix best + random)
    candidates_to_refine = []
    if programs:
        # Sort by accuracy
        sorted_progs = sorted(programs, key=lambda x: -x[1])
        # Top programs + random selection
        candidates_to_refine = sorted_progs[:5]
        if len(sorted_progs) > 5:
            candidates_to_refine += random.sample(sorted_progs[5:], min(3, len(sorted_progs)-5))
    
    for j, (code, acc, _, exec_results) in enumerate(candidates_to_refine):
        if acc == 1.0:
            continue  # Already perfect
        
        for r in range(min(3, n_refinements)):  # 3 refinement attempts per program
            try:
                repair_prompt = get_refinement_prompt(task, code, exec_results)
                messages = [{"role": "user", "content": repair_prompt}]
                repair_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
                
                inputs = tokenizer(repair_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=2048,
                        temperature=0.7,
                        top_p=0.95,
                        do_sample=True,
                        pad_token_id=tokenizer.eos_token_id,
                    )
                
                response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
                new_code = extract_transform_code(response)
                
                if new_code is None:
                    continue
                
                new_acc, new_exec, new_test = evaluate_code_on_task(new_code, task)
                programs.append((new_code, new_acc, new_test, new_exec))
                
                if new_acc == 1.0:
                    print(f"    Found perfect program via refinement")
                    break
                    
            except Exception:
                continue
    
    # Phase 3: Weighted majority vote
    vote_scores = defaultdict(float)
    
    for code, accuracy, test_output, _ in programs:
        if test_output is None:
            continue
        key = tuple(tuple(row) for row in test_output)
        # SOAR scoring: count + 1000 × accuracy
        vote_scores[key] += 1 + 1000 * accuracy
    
    if not vote_scores:
        return []
    
    sorted_votes = sorted(vote_scores.items(), key=lambda x: -x[1])
    results = []
    for key, score in sorted_votes[:2]:
        results.append([list(row) for row in key])
    
    return results


def main():
    print("=" * 60)
    print("SOAR Program Synthesis for ARC-AGI-2")
    print("=" * 60)
    
    # Load model
    model_id = "julien31/Soar-qwen-7b"
    print(f"\nLoading model: {model_id}")
    
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.eval()
    print("Model loaded!")
    
    # Load ARC-AGI-2 data
    print("\nLoading ARC-AGI-2 training data...")
    ds = load_dataset("arc-agi-community/arc-agi-2", split="train")
    tasks = []
    for row in ds:
        task = {"train": row["fewshots"], "test": row["question"]}
        tasks.append(task)
    print(f"Loaded {len(tasks)} tasks")
    
    # Also load ARC-AGI-1 eval for benchmarking
    print("Loading ARC-AGI-1 evaluation data...")
    ds_v1 = load_dataset("lordspline/arc-agi", split="evaluation")
    tasks_v1 = []
    for row in ds_v1:
        task = {"train": row["train"], "test": row["test"]}
        tasks_v1.append(task)
    print(f"Loaded {len(tasks_v1)} ARC-AGI-1 eval tasks")
    
    # Configuration
    N_SAMPLES = 30  # Programs per task
    N_REFINEMENTS = 15  # Refinement budget per task
    EVAL_SUBSET = 100  # Evaluate on first N tasks for speed
    
    # Evaluate on ARC-AGI-2 (first EVAL_SUBSET tasks)
    print(f"\n--- Evaluating on ARC-AGI-2 (first {EVAL_SUBSET} tasks) ---")
    correct_v2 = 0
    attempted_v2 = 0
    start_time = time.time()
    
    for i in range(min(EVAL_SUBSET, len(tasks))):
        task = tasks[i]
        gt = task["test"][0].get("output")
        
        elapsed = time.time() - start_time
        print(f"\n[{i+1}/{EVAL_SUBSET}] Task {i} (elapsed: {elapsed:.0f}s)")
        
        try:
            predictions = solve_task_soar(
                task, model, tokenizer, 
                n_samples=N_SAMPLES, 
                n_refinements=N_REFINEMENTS,
                temperature=0.8
            )
            
            if gt is not None:
                attempted_v2 += 1
                for pred in predictions:
                    if grids_equal(pred, gt):
                        correct_v2 += 1
                        print(f"  ✓ CORRECT!")
                        break
                else:
                    if predictions:
                        print(f"  ✗ Wrong ({len(predictions)} candidates)")
                    else:
                        print(f"  ✗ No predictions")
            
            # Print running accuracy
            if attempted_v2 > 0:
                print(f"  Running: {correct_v2}/{attempted_v2} = {correct_v2/attempted_v2*100:.1f}%")
                
        except Exception as e:
            print(f"  ERROR: {e}")
            traceback.print_exc()
    
    total_time = time.time() - start_time
    
    print(f"\n{'='*60}")
    print(f"ARC-AGI-2 Results:")
    print(f"  Correct: {correct_v2}/{attempted_v2}")
    print(f"  Pass@2: {correct_v2/attempted_v2*100:.2f}%" if attempted_v2 > 0 else "  N/A")
    print(f"  Time: {total_time:.0f}s ({total_time/EVAL_SUBSET:.1f}s/task)")
    print(f"{'='*60}")
    
    # Also quick eval on ARC-AGI-1
    print(f"\n--- Evaluating on ARC-AGI-1 (first {min(50, len(tasks_v1))} tasks) ---")
    correct_v1 = 0
    attempted_v1 = 0
    
    for i in range(min(50, len(tasks_v1))):
        task = tasks_v1[i]
        gt = task["test"][0].get("output")
        
        print(f"[{i+1}/50] Task {i}", end="")
        
        try:
            predictions = solve_task_soar(
                task, model, tokenizer,
                n_samples=20,
                n_refinements=10,
                temperature=0.8
            )
            
            if gt is not None:
                attempted_v1 += 1
                for pred in predictions:
                    if grids_equal(pred, gt):
                        correct_v1 += 1
                        print(f" ✓", end="")
                        break
                else:
                    print(f" ✗", end="")
            
            print(f" ({correct_v1}/{attempted_v1})" if attempted_v1 > 0 else "")
                
        except Exception as e:
            print(f" ERROR: {e}")
    
    print(f"\nARC-AGI-1 Results: {correct_v1}/{attempted_v1} = {correct_v1/attempted_v1*100:.2f}%" if attempted_v1 > 0 else "N/A")
    
    # Push results
    results = {
        "arc_agi_2": {"correct": correct_v2, "total": attempted_v2},
        "arc_agi_1": {"correct": correct_v1, "total": attempted_v1},
        "config": {
            "n_samples": N_SAMPLES,
            "n_refinements": N_REFINEMENTS,
            "model": model_id,
        }
    }
    
    with open("/app/results.json", "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to /app/results.json")


if __name__ == "__main__":
    main()