"""
ARC-AGI Data Loading and Augmentation Pipeline
Handles both ARC-AGI-1 and ARC-AGI-2 data formats.
"""
import json
import copy
import random
import itertools
from typing import List, Tuple, Dict, Any, Optional
import numpy as np


# ============================================================
# Data loading
# ============================================================

def load_arc_dataset_from_hf(dataset_name: str = "arc-agi-community/arc-agi-2", split: str = "train"):
    """Load ARC dataset from HuggingFace Hub."""
    from datasets import load_dataset
    ds = load_dataset(dataset_name, split=split)
    tasks = []
    for row in ds:
        if "fewshots" in row:
            # ARC-AGI-2 format
            task = {
                "train": row["fewshots"],
                "test": row["question"]
            }
        else:
            # ARC-AGI-1 format
            task = {
                "train": row["train"],
                "test": row["test"]
            }
        tasks.append(task)
    return tasks


def load_arc_dataset_from_json(path: str) -> Dict[str, Any]:
    """Load a single ARC task from JSON file."""
    with open(path, 'r') as f:
        return json.load(f)


# ============================================================
# Grid operations
# ============================================================

def grid_to_numpy(grid: List[List[int]]) -> np.ndarray:
    return np.array(grid, dtype=np.int32)


def numpy_to_grid(arr: np.ndarray) -> List[List[int]]:
    return arr.tolist()


def grids_equal(g1: List[List[int]], g2: List[List[int]]) -> bool:
    """Check if two grids are exactly equal."""
    if len(g1) != len(g2):
        return False
    for r1, r2 in zip(g1, g2):
        if len(r1) != len(r2):
            return False
        if r1 != r2:
            return False
    return True


# ============================================================
# D8 Symmetry Group Augmentations (rotations + reflections)
# ============================================================

def rotate_90(grid: np.ndarray) -> np.ndarray:
    """Rotate 90 degrees clockwise."""
    return np.rot90(grid, k=-1)

def rotate_180(grid: np.ndarray) -> np.ndarray:
    return np.rot90(grid, k=-2)

def rotate_270(grid: np.ndarray) -> np.ndarray:
    return np.rot90(grid, k=-3)

def flip_horizontal(grid: np.ndarray) -> np.ndarray:
    return np.fliplr(grid)

def flip_vertical(grid: np.ndarray) -> np.ndarray:
    return np.flipud(grid)

def transpose(grid: np.ndarray) -> np.ndarray:
    return grid.T

def anti_transpose(grid: np.ndarray) -> np.ndarray:
    return np.rot90(grid.T, k=2)

# Identity + 7 non-trivial = 8 D8 symmetry operations
D8_TRANSFORMS = [
    ("identity", lambda g: g.copy()),
    ("rot90", rotate_90),
    ("rot180", rotate_180),
    ("rot270", rotate_270),
    ("flip_h", flip_horizontal),
    ("flip_v", flip_vertical),
    ("transpose", transpose),
    ("anti_transpose", anti_transpose),
]

# Inverse operations (to reverse augmentations)
D8_INVERSES = {
    "identity": "identity",
    "rot90": "rot270",
    "rot180": "rot180",
    "rot270": "rot90",
    "flip_h": "flip_h",
    "flip_v": "flip_v",
    "transpose": "transpose",
    "anti_transpose": "anti_transpose",
}

def get_d8_transform(name: str):
    for n, fn in D8_TRANSFORMS:
        if n == name:
            return fn
    raise ValueError(f"Unknown transform: {name}")


def apply_d8_to_pair(inp: List[List[int]], out: List[List[int]], transform_name: str):
    """Apply a D8 transform to an input-output pair."""
    fn = get_d8_transform(transform_name)
    new_inp = numpy_to_grid(fn(grid_to_numpy(inp)))
    new_out = numpy_to_grid(fn(grid_to_numpy(out)))
    return new_inp, new_out


def reverse_d8(grid: List[List[int]], transform_name: str) -> List[List[int]]:
    """Reverse a D8 transform."""
    inv_name = D8_INVERSES[transform_name]
    fn = get_d8_transform(inv_name)
    return numpy_to_grid(fn(grid_to_numpy(grid)))


# ============================================================
# Color Permutation Augmentations
# ============================================================

def create_color_permutation(seed: Optional[int] = None) -> Dict[int, int]:
    """Create a random permutation of colors 0-9."""
    rng = random.Random(seed)
    colors = list(range(10))
    shuffled = colors.copy()
    rng.shuffle(shuffled)
    return dict(zip(colors, shuffled))


def apply_color_permutation(grid: List[List[int]], perm: Dict[int, int]) -> List[List[int]]:
    """Apply color permutation to a grid."""
    return [[perm.get(c, c) for c in row] for row in grid]


def reverse_color_permutation(grid: List[List[int]], perm: Dict[int, int]) -> List[List[int]]:
    """Reverse a color permutation."""
    inv_perm = {v: k for k, v in perm.items()}
    return apply_color_permutation(grid, inv_perm)


# ============================================================
# Augmented Task Creation (for TTT + training)
# ============================================================

def augment_task(task: Dict, transform_name: str = "identity", 
                 color_perm: Optional[Dict[int, int]] = None,
                 permute_examples: bool = False) -> Dict:
    """
    Apply augmentations to an ARC task:
    1. D8 geometric transform
    2. Color permutation
    3. Example order permutation
    """
    new_task = {"train": [], "test": []}
    
    # Apply to training pairs
    train_pairs = list(task["train"])
    if permute_examples:
        random.shuffle(train_pairs)
    
    for pair in train_pairs:
        inp, out = pair["input"], pair["output"]
        # D8 transform
        inp, out = apply_d8_to_pair(inp, out, transform_name)
        # Color permutation
        if color_perm:
            inp = apply_color_permutation(inp, color_perm)
            out = apply_color_permutation(out, color_perm)
        new_task["train"].append({"input": inp, "output": out})
    
    # Apply to test pairs
    for pair in task["test"]:
        inp = pair["input"]
        fn = get_d8_transform(transform_name)
        inp = numpy_to_grid(fn(grid_to_numpy(inp)))
        if color_perm:
            inp = apply_color_permutation(inp, color_perm)
        test_pair = {"input": inp}
        if "output" in pair and pair["output"] is not None:
            out = pair["output"]
            out = numpy_to_grid(fn(grid_to_numpy(out)))
            if color_perm:
                out = apply_color_permutation(out, color_perm)
            test_pair["output"] = out
        new_task["test"].append(test_pair)
    
    return new_task


def create_leave_one_out_tasks(task: Dict) -> List[Dict]:
    """
    Create leave-one-out ICL tasks for TTT (Akyürek et al.).
    For K training pairs, create K synthetic tasks where each pair
    plays "test" once while others serve as demonstrations.
    """
    train_pairs = task["train"]
    K = len(train_pairs)
    loo_tasks = []
    
    for j in range(K):
        # The j-th pair becomes the "test"
        demos = [train_pairs[i] for i in range(K) if i != j]
        test_pair = train_pairs[j]
        
        loo_task = {
            "train": demos,
            "test": [{"input": test_pair["input"], "output": test_pair["output"]}]
        }
        loo_tasks.append(loo_task)
    
    return loo_tasks


def create_ttt_dataset(task: Dict, n_augmentations: int = 16, max_examples: int = 250) -> List[Dict]:
    """
    Create full TTT dataset for a task:
    1. Generate leave-one-out tasks
    2. Apply D8 augmentations to each
    3. Apply color permutations
    4. Permute example orders
    Cap at max_examples per task.
    """
    loo_tasks = create_leave_one_out_tasks(task)
    ttt_dataset = []
    
    for loo_task in loo_tasks:
        for t_name, _ in D8_TRANSFORMS:
            # Apply D8 transform
            aug_task = augment_task(loo_task, transform_name=t_name)
            ttt_dataset.append(aug_task)
            
            # Also with random color permutation
            if len(ttt_dataset) < max_examples:
                color_perm = create_color_permutation()
                aug_task_c = augment_task(loo_task, transform_name=t_name, color_perm=color_perm)
                ttt_dataset.append(aug_task_c)
            
            if len(ttt_dataset) >= max_examples:
                break
        if len(ttt_dataset) >= max_examples:
            break
    
    random.shuffle(ttt_dataset)
    return ttt_dataset[:max_examples]


# ============================================================
# Grid serialization for LLM input
# ============================================================

def grid_to_string(grid: List[List[int]], separator: str = " ") -> str:
    """Convert a grid to string representation for LLM input."""
    return "\n".join(separator.join(str(c) for c in row) for row in grid)


def string_to_grid(s: str) -> List[List[int]]:
    """Parse grid string back to list of lists."""
    rows = s.strip().split("\n")
    grid = []
    for row in rows:
        cells = row.strip().split()
        grid.append([int(c) for c in cells])
    return grid


def task_to_prompt(task: Dict, include_test_output: bool = False) -> str:
    """
    Convert an ARC task to a text prompt for an LLM.
    Uses the format from Akyürek et al.
    """
    parts = []
    
    # Training examples
    for i, pair in enumerate(task["train"]):
        parts.append(f"Example {i+1}:")
        parts.append(f"Input:")
        parts.append(grid_to_string(pair["input"]))
        parts.append(f"Output:")
        parts.append(grid_to_string(pair["output"]))
        parts.append("")
    
    # Test input
    parts.append("Test:")
    parts.append("Input:")
    parts.append(grid_to_string(task["test"][0]["input"]))
    parts.append("Output:")
    
    if include_test_output and "output" in task["test"][0] and task["test"][0]["output"] is not None:
        parts.append(grid_to_string(task["test"][0]["output"]))
    
    return "\n".join(parts)


def task_to_program_prompt(task: Dict) -> str:
    """
    Convert an ARC task to a prompt for program synthesis (SOAR-style).
    The model should generate a Python transform() function.
    """
    parts = [
        "Given the following input-output grid transformation examples, "
        "write a Python function `transform(input_grid: list[list[int]]) -> list[list[int]]` "
        "that implements the transformation.\n"
    ]
    
    for i, pair in enumerate(task["train"]):
        parts.append(f"Example {i+1}:")
        parts.append(f"  Input:  {pair['input']}")
        parts.append(f"  Output: {pair['output']}")
    
    parts.append(f"\nTest Input: {task['test'][0]['input']}")
    parts.append("\nWrite the transform function:")
    
    return "\n".join(parts)


# ============================================================
# Evaluation utilities
# ============================================================

def evaluate_program(code: str, task: Dict) -> Tuple[bool, Optional[List[List[int]]]]:
    """
    Execute a program on the task's training examples and test input.
    Returns (all_train_correct, test_output_or_None).
    """
    try:
        namespace = {"__builtins__": __builtins__}
        exec(code, namespace)
        
        if "transform" not in namespace:
            return False, None
        
        transform_fn = namespace["transform"]
        
        # Check all training examples
        all_correct = True
        for pair in task["train"]:
            try:
                predicted = transform_fn(copy.deepcopy(pair["input"]))
                if not grids_equal(predicted, pair["output"]):
                    all_correct = False
                    break
            except Exception:
                all_correct = False
                break
        
        # Get test output
        test_output = None
        try:
            test_output = transform_fn(copy.deepcopy(task["test"][0]["input"]))
        except Exception:
            pass
        
        return all_correct, test_output
        
    except Exception:
        return False, None


def evaluate_predictions(tasks: List[Dict], predictions: List[List[List[List[int]]]]) -> Dict:
    """
    Evaluate pass@2 predictions against ground truth.
    predictions[i] = [attempt1, attempt2] for task i.
    """
    correct = 0
    total = len(tasks)
    
    for task, preds in zip(tasks, predictions):
        gt = task["test"][0].get("output")
        if gt is None:
            continue
        
        for pred in preds:
            if grids_equal(pred, gt):
                correct += 1
                break
    
    return {
        "correct": correct,
        "total": total,
        "accuracy": correct / total if total > 0 else 0,
        "pass_at_2": correct / total if total > 0 else 0,
    }


if __name__ == "__main__":
    # Test with sample data
    print("Loading ARC-AGI-2 dataset...")
    tasks = load_arc_dataset_from_hf("arc-agi-community/arc-agi-2", "train")
    print(f"Loaded {len(tasks)} tasks")
    
    # Test augmentations
    task = tasks[0]
    print(f"\nTask 0: {len(task['train'])} train pairs, {len(task['test'])} test pairs")
    print(f"Train input shape: {len(task['train'][0]['input'])}x{len(task['train'][0]['input'][0])}")
    
    # Test D8 augmentations
    for name, _ in D8_TRANSFORMS:
        aug = augment_task(task, transform_name=name)
        print(f"  {name}: input shape {len(aug['train'][0]['input'])}x{len(aug['train'][0]['input'][0])}")
    
    # Test leave-one-out
    loo_tasks = create_leave_one_out_tasks(task)
    print(f"\nLeave-one-out: {len(loo_tasks)} tasks created from {len(task['train'])} demos")
    
    # Test TTT dataset
    ttt_ds = create_ttt_dataset(task, max_examples=50)
    print(f"TTT dataset: {len(ttt_ds)} examples")
    
    # Test prompt generation
    prompt = task_to_prompt(task)
    print(f"\nPrompt length: {len(prompt)} chars")
    print(prompt[:500])
    
    print("\n✅ All tests passed!")