"""
ARC-AGI-2 Competition Solver
Dual-track approach: Program Synthesis (induction) + TTT (transduction)

Competition constraints:
- 4× NVIDIA L4 GPUs (24GB VRAM each)
- 12-hour wall-clock time for 240 tasks (~3 min/task avg)
- No internet access
- Pass@2 scoring (2 attempts per task, exact match)

Strategy:
Track A: SOAR-style program synthesis with Qwen-2.5-Coder-7B
Track B: TTT transduction with augmented inference + PoE
Ensemble: Combine both tracks with priority to verified solutions
"""
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

from arc_data import (
    load_arc_dataset_from_hf, grids_equal,
    D8_TRANSFORMS, augment_task, reverse_d8,
    create_color_permutation, reverse_color_permutation,
    grid_to_string, string_to_grid
)
from program_synthesis import (
    ProgramSynthesisEngine, evaluate_program_on_task,
    weighted_majority_vote, extract_python_code,
    format_examples_for_prompt, get_error_feedback,
    _execute_transform, ARCPrimitives
)


# ============================================================
# Configuration
# ============================================================

class Config:
    # Model settings
    PROGRAM_MODEL = "julien31/Soar-qwen-7b"  # SOAR fine-tuned Qwen for program synthesis
    TTT_MODEL = "nvidia/Mistral-NeMo-Minitron-8B-Base"  # For transduction TTT
    
    # Program synthesis settings
    N_PROGRAM_SAMPLES = 100  # Programs to sample per task
    N_PROGRAM_REFINEMENTS = 100  # Refinement budget per task
    PROGRAM_TEMPERATURE = 0.8
    PROGRAM_MAX_TOKENS = 1024
    
    # TTT settings
    TTT_LORA_RANK = 32
    TTT_STEPS = 64
    TTT_LR = 2e-4
    TTT_MAX_EXAMPLES = 250
    
    # Inference settings
    N_AUGMENTATIONS = 16
    DFS_THRESHOLD = 0.09  # 9% threshold for DFS
    
    # Time budget
    TOTAL_TIME_HOURS = 12
    TIME_PER_TASK_SECONDS = 150  # ~2.5 min per task (conservative)
    
    # Hardware
    N_GPUS = 4


# ============================================================
# Ensemble Logic
# ============================================================

class ARC_Solver:
    """
    Main solver that combines both tracks.
    Track A (program synthesis) provides verified solutions.
    Track B (TTT transduction) fills gaps.
    """
    
    def __init__(self, config: Config = None):
        self.config = config or Config()
        self.program_engine = None
        self.ttt_engine = None
        self.results = {}
        self.start_time = time.time()
    
    def load_models(self):
        """Load models for both tracks."""
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        
        print("Loading models...")
        
        # Track A: Program synthesis model (SOAR-Qwen-7B)
        print(f"  Loading {self.config.PROGRAM_MODEL}...")
        try:
            prog_tokenizer = AutoTokenizer.from_pretrained(
                self.config.PROGRAM_MODEL,
                trust_remote_code=True
            )
            prog_model = AutoModelForCausalLM.from_pretrained(
                self.config.PROGRAM_MODEL,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                trust_remote_code=True,
            )
            prog_model.eval()
            
            self.program_engine = ProgramSynthesisEngine(
                model=prog_model,
                tokenizer=prog_tokenizer,
                max_samples=self.config.N_PROGRAM_SAMPLES,
                max_refinements=self.config.N_PROGRAM_REFINEMENTS,
                temperature=self.config.PROGRAM_TEMPERATURE,
            )
            print("  ✓ Program synthesis model loaded")
        except Exception as e:
            print(f"  ✗ Failed to load program model: {e}")
        
        print("Models loaded.")
    
    def time_remaining(self) -> float:
        """Get remaining time in seconds."""
        elapsed = time.time() - self.start_time
        return self.config.TOTAL_TIME_HOURS * 3600 - elapsed
    
    def solve_with_programs(self, task: Dict, task_id: str, 
                            time_budget: float = 120) -> List[List[List[int]]]:
        """
        Track A: Solve using program synthesis.
        Programs that pass ALL training examples are verified solutions.
        """
        if self.program_engine is None:
            return []
        
        start = time.time()
        
        # Sample programs
        programs = self.program_engine.sample_programs(task)
        
        # Check time
        if time.time() - start > time_budget * 0.6:
            pass  # Skip refinement if running low on time
        else:
            # Refine
            programs = self.program_engine.refine_programs(task, programs)
        
        # Separate verified (100% accuracy) from partial
        verified = [(c, a, o) for c, a, o in programs if a == 1.0 and o is not None]
        partial = [(c, a, o) for c, a, o in programs if 0 < a < 1.0 and o is not None]
        
        if verified:
            # Use weighted vote among verified programs
            preds = weighted_majority_vote(verified, top_k=2)
            print(f"  [PROG] Task {task_id}: {len(verified)} verified programs → {len(preds)} predictions")
            return preds
        
        if partial:
            preds = weighted_majority_vote(partial, top_k=2)
            print(f"  [PROG] Task {task_id}: {len(partial)} partial programs → {len(preds)} predictions (unverified)")
            return preds
        
        print(f"  [PROG] Task {task_id}: No valid programs found")
        return []
    
    def solve_with_augmented_programs(self, task: Dict, task_id: str,
                                      time_budget: float = 120) -> List[List[List[int]]]:
        """
        Enhanced program synthesis using D8 augmentations.
        Try solving augmented versions of the task.
        """
        if self.program_engine is None:
            return []
        
        all_programs = []
        
        for t_name, _ in D8_TRANSFORMS[:4]:  # Use 4 augmentations to save time
            aug_task = augment_task(task, transform_name=t_name)
            programs = self.program_engine.sample_programs(aug_task, n_samples=10)
            
            # Map test outputs back to original space
            for code, acc, test_out in programs:
                if test_out is not None:
                    original_out = reverse_d8(test_out, t_name)
                    all_programs.append((code, acc, original_out))
        
        # Vote
        verified = [(c, a, o) for c, a, o in all_programs if a == 1.0 and o is not None]
        if verified:
            return weighted_majority_vote(verified, top_k=2)
        
        partial = [(c, a, o) for c, a, o in all_programs if a > 0 and o is not None]
        if partial:
            return weighted_majority_vote(partial, top_k=2)
        
        return []
    
    def solve_task(self, task: Dict, task_id: str = "unknown") -> List[List[List[int]]]:
        """
        Solve a single ARC task using the ensemble strategy.
        Priority:
        1. Verified program solutions (100% on training examples)
        2. TTT transduction with voting
        3. Partial program solutions (>50% on training examples)
        4. Any available prediction
        """
        time_budget = min(self.config.TIME_PER_TASK_SECONDS, self.time_remaining() - 60)
        
        if time_budget <= 0:
            print(f"  Task {task_id}: OUT OF TIME")
            return []
        
        predictions = []
        
        # Track A: Program synthesis
        prog_preds = self.solve_with_programs(task, task_id, time_budget=time_budget * 0.7)
        
        # If we have verified programs, use them directly
        if prog_preds:
            predictions = prog_preds
        
        # If program synthesis didn't find verified solutions, try augmented approach
        if not predictions and time_budget > 30:
            aug_preds = self.solve_with_augmented_programs(task, task_id, time_budget=time_budget * 0.3)
            if aug_preds:
                predictions = aug_preds
        
        # Ensure we have exactly 2 predictions
        if len(predictions) == 0:
            # Return empty grids as fallback
            return []
        elif len(predictions) == 1:
            # Duplicate single prediction
            predictions.append(predictions[0])
        
        return predictions[:2]
    
    def solve_all(self, tasks: List[Dict], task_ids: Optional[List[str]] = None) -> Dict:
        """
        Solve all tasks. Returns dict of task_id -> [pred1, pred2].
        """
        if task_ids is None:
            task_ids = [str(i) for i in range(len(tasks))]
        
        results = {}
        solved = 0
        
        for i, (task, task_id) in enumerate(zip(tasks, task_ids)):
            elapsed = time.time() - self.start_time
            remaining = self.config.TOTAL_TIME_HOURS * 3600 - elapsed
            
            print(f"\n[{i+1}/{len(tasks)}] Task {task_id} "
                  f"(elapsed: {elapsed/3600:.2f}h, remaining: {remaining/3600:.2f}h)")
            
            if remaining <= 60:
                print("  TIME'S UP - stopping")
                break
            
            try:
                preds = self.solve_task(task, task_id)
                if preds:
                    results[task_id] = preds
                    solved += 1
                else:
                    results[task_id] = []
            except Exception as e:
                print(f"  ERROR: {e}")
                traceback.print_exc()
                results[task_id] = []
        
        total_time = time.time() - self.start_time
        print(f"\n{'='*60}")
        print(f"RESULTS: {solved}/{len(tasks)} tasks produced predictions")
        print(f"Total time: {total_time/3600:.2f}h")
        print(f"{'='*60}")
        
        return results
    
    def evaluate(self, tasks: List[Dict], results: Dict, 
                 task_ids: Optional[List[str]] = None) -> Dict:
        """Evaluate results against ground truth (for tasks with known outputs)."""
        if task_ids is None:
            task_ids = [str(i) for i in range(len(tasks))]
        
        correct = 0
        evaluated = 0
        
        for task, task_id in zip(tasks, task_ids):
            gt = task["test"][0].get("output")
            if gt is None:
                continue
            
            evaluated += 1
            preds = results.get(task_id, [])
            
            for pred in preds:
                if grids_equal(pred, gt):
                    correct += 1
                    break
        
        accuracy = correct / evaluated if evaluated > 0 else 0
        
        print(f"\nEvaluation:")
        print(f"  Correct: {correct}/{evaluated}")
        print(f"  Pass@2 Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)")
        
        return {
            "correct": correct,
            "evaluated": evaluated,
            "accuracy": accuracy,
        }


# ============================================================
# Quick evaluation without models (using heuristic solvers)
# ============================================================

class HeuristicSolver:
    """
    Simple heuristic solver for common ARC patterns.
    No ML model needed — uses pattern matching.
    """
    
    @staticmethod
    def try_identity(task: Dict) -> Optional[List[List[int]]]:
        """Check if output == input."""
        for pair in task["train"]:
            if not grids_equal(pair["input"], pair["output"]):
                return None
        return copy.deepcopy(task["test"][0]["input"])
    
    @staticmethod
    def try_color_replacement(task: Dict) -> Optional[List[List[int]]]:
        """Check if output is input with colors replaced."""
        # Infer color mapping from first example
        inp0 = task["train"][0]["input"]
        out0 = task["train"][0]["output"]
        
        if len(inp0) != len(out0) or len(inp0[0]) != len(out0[0]):
            return None
        
        color_map = {}
        for r in range(len(inp0)):
            for c in range(len(inp0[0])):
                ic = inp0[r][c]
                oc = out0[r][c]
                if ic in color_map:
                    if color_map[ic] != oc:
                        return None
                color_map[ic] = oc
        
        # Verify on other examples
        for pair in task["train"][1:]:
            inp, out = pair["input"], pair["output"]
            if len(inp) != len(out) or len(inp[0]) != len(out[0]):
                return None
            for r in range(len(inp)):
                for c in range(len(inp[0])):
                    expected = color_map.get(inp[r][c])
                    if expected is None or expected != out[r][c]:
                        return None
        
        # Apply to test
        test_inp = task["test"][0]["input"]
        result = []
        for row in test_inp:
            result.append([color_map.get(c, c) for c in row])
        return result
    
    @staticmethod
    def try_rotation(task: Dict) -> Optional[List[List[int]]]:
        """Check if output is a rotation of input."""
        for k in [1, 2, 3]:  # 90, 180, 270 degrees
            all_match = True
            for pair in task["train"]:
                rotated = np.rot90(np.array(pair["input"]), k=-k).tolist()
                if not grids_equal(rotated, pair["output"]):
                    all_match = False
                    break
            if all_match:
                return np.rot90(np.array(task["test"][0]["input"]), k=-k).tolist()
        return None
    
    @staticmethod
    def try_flip(task: Dict) -> Optional[List[List[int]]]:
        """Check if output is a flip of input."""
        for flip_fn in [np.fliplr, np.flipud]:
            all_match = True
            for pair in task["train"]:
                flipped = flip_fn(np.array(pair["input"])).tolist()
                if not grids_equal(flipped, pair["output"]):
                    all_match = False
                    break
            if all_match:
                return flip_fn(np.array(task["test"][0]["input"])).tolist()
        return None
    
    @staticmethod
    def try_transpose(task: Dict) -> Optional[List[List[int]]]:
        """Check if output is transpose of input."""
        all_match = True
        for pair in task["train"]:
            transposed = np.array(pair["input"]).T.tolist()
            if not grids_equal(transposed, pair["output"]):
                all_match = False
                break
        if all_match:
            return np.array(task["test"][0]["input"]).T.tolist()
        return None
    
    @staticmethod
    def try_crop_to_nonzero(task: Dict) -> Optional[List[List[int]]]:
        """Check if output is the input cropped to non-zero bounding box."""
        for pair in task["train"]:
            arr = np.array(pair["input"])
            nonzero = np.argwhere(arr != 0)
            if len(nonzero) == 0:
                return None
            r1, c1 = nonzero.min(axis=0)
            r2, c2 = nonzero.max(axis=0)
            cropped = arr[r1:r2+1, c1:c2+1].tolist()
            if not grids_equal(cropped, pair["output"]):
                return None
        
        test_arr = np.array(task["test"][0]["input"])
        nonzero = np.argwhere(test_arr != 0)
        if len(nonzero) == 0:
            return None
        r1, c1 = nonzero.min(axis=0)
        r2, c2 = nonzero.max(axis=0)
        return test_arr[r1:r2+1, c1:c2+1].tolist()
    
    @staticmethod
    def try_scale(task: Dict) -> Optional[List[List[int]]]:
        """Check if output is input scaled by some factor."""
        for factor in [2, 3, 4, 5]:
            all_match = True
            for pair in task["train"]:
                inp = pair["input"]
                out = pair["output"]
                if len(out) != len(inp) * factor or len(out[0]) != len(inp[0]) * factor:
                    all_match = False
                    break
                # Check each cell
                for r in range(len(inp)):
                    for c in range(len(inp[0])):
                        for dr in range(factor):
                            for dc in range(factor):
                                if out[r*factor+dr][c*factor+dc] != inp[r][c]:
                                    all_match = False
                                    break
                            if not all_match:
                                break
                        if not all_match:
                            break
                    if not all_match:
                        break
                if not all_match:
                    break
            if all_match:
                inp = task["test"][0]["input"]
                result = []
                for row in inp:
                    for dr in range(factor):
                        result.append([c for c in row for dc in range(factor)])
                return result
        return None
    
    def solve(self, task: Dict) -> Optional[List[List[int]]]:
        """Try all heuristic solvers."""
        solvers = [
            self.try_identity,
            self.try_color_replacement,
            self.try_rotation,
            self.try_flip,
            self.try_transpose,
            self.try_crop_to_nonzero,
            self.try_scale,
        ]
        
        for solver in solvers:
            try:
                result = solver(task)
                if result is not None:
                    return result
            except Exception:
                continue
        
        return None


# ============================================================
# Main entry point
# ============================================================

def main():
    """Main evaluation pipeline."""
    print("=" * 60)
    print("ARC-AGI-2 Solver v1.0")
    print("Dual-track: Program Synthesis + Test-Time Training")
    print("=" * 60)
    
    # Load data
    print("\nLoading ARC-AGI-2 training data...")
    tasks = load_arc_dataset_from_hf("arc-agi-community/arc-agi-2", "train")
    print(f"Loaded {len(tasks)} tasks")
    
    # Quick heuristic evaluation
    print("\n--- Heuristic Solver Evaluation ---")
    heuristic = HeuristicSolver()
    h_correct = 0
    h_total = 0
    
    for i, task in enumerate(tasks):
        gt = task["test"][0].get("output")
        if gt is None:
            continue
        h_total += 1
        
        pred = heuristic.solve(task)
        if pred is not None and grids_equal(pred, gt):
            h_correct += 1
    
    print(f"Heuristic solver: {h_correct}/{h_total} = {h_correct/h_total*100:.1f}%")
    
    # Full solver (if models available)
    print("\n--- Full Solver ---")
    solver = ARC_Solver()
    
    try:
        solver.load_models()
        
        # Solve all tasks
        task_ids = [str(i) for i in range(len(tasks))]
        results = solver.solve_all(tasks, task_ids)
        
        # Evaluate
        eval_results = solver.evaluate(tasks, results, task_ids)
        
    except Exception as e:
        print(f"Model loading failed (expected on CPU): {e}")
        print("Running heuristic-only evaluation...")


if __name__ == "__main__":
    main()