#!/usr/bin/env python3
"""
Safe Best-of-N merger using official neurogolf_utils.py scoring.

Uses the EXACT same code Kaggle uses:
- sanitize_model() for tensor name normalization
- score_network() with ORT Profiler for cost (memory + params)
- verify_subset() for validation

Override 6043 ONLY if:
1. Alternative passes ALL train + test + arc-gen examples
2. Alternative has strictly lower cost (memory + params) per official scorer

Usage:
  python merge_safe.py \
    --sub_6043 submission-6043.zip \
    --sub_5743 submission-5743.zip \
    --data_dir /kaggle/input/competitions/neurogolf-2026 \
    --output_zip /kaggle/working/merged.zip
"""

import json
import math
import os
import sys
import time
import traceback
import zipfile
import tempfile

import numpy as np
import onnx
import onnxruntime

# Import from official neurogolf_utils (must be in same directory or PYTHONPATH)
from neurogolf_utils import (
    sanitize_model, score_network, convert_to_numpy, run_network,
    _GRID_SHAPE, _BATCH_SIZE, _CHANNELS, _HEIGHT, _WIDTH
)


def encode_grid(grid):
    """Encode color grid to one-hot [1,10,30,30] tensor."""
    arr = np.array(grid, dtype=np.int32)
    h, w = arr.shape
    t = np.zeros((1, 10, 30, 30), dtype=np.float32)
    for r in range(h):
        for c in range(w):
            v = int(arr[r, c])
            if 0 <= v < 10:
                t[0, v, r, c] = 1.0
    return t


def validate_and_score(model_bytes, task_data):
    """Validate model and compute score using official neurogolf_utils.

    Returns (valid, memory, params, n_tested) or (False, None, None, n_tested) on failure.
    valid = True means model passes ALL examples.
    memory, params = official score_network() output.
    """
    # Save to temp file (neurogolf_utils works with files)
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as f:
            f.write(model_bytes)
            tmp_path = f.name

        # Load and sanitize (same as Kaggle does)
        model = onnx.load(tmp_path)
        sanitized = sanitize_model(model)
        if sanitized is None:
            return False, None, None, 0

        # Set up ORT session with profiling (same as Kaggle)
        options = onnxruntime.SessionOptions()
        options.enable_profiling = True
        options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
        options.profile_file_prefix = "merge_tmp"

        try:
            session = onnxruntime.InferenceSession(
                sanitized.SerializeToString(), options,
                providers=['CPUExecutionProvider'])
        except onnxruntime.ONNXRuntimeError as e:
            return False, None, None, 0

        # Validate against ALL examples
        all_examples = (task_data.get('train', []) +
                        task_data.get('test', []) +
                        task_data.get('arc-gen', []))

        n_tested = 0
        for example in all_examples:
            benchmark = convert_to_numpy(example)
            if benchmark is None:
                continue
            try:
                user_output = run_network(session, benchmark["input"])
                n_tested += 1
                if not np.array_equal(user_output, benchmark["output"]):
                    # End profiling to get trace path (cleanup)
                    try:
                        trace_path = session.end_profiling()
                        os.unlink(trace_path)
                    except:
                        pass
                    return False, None, None, n_tested
            except onnxruntime.ONNXRuntimeError:
                try:
                    trace_path = session.end_profiling()
                    os.unlink(trace_path)
                except:
                    pass
                return False, None, None, n_tested

        # All examples passed — now get the score
        trace_path = session.end_profiling()
        memory, params = score_network(sanitized, trace_path)

        # Cleanup trace file
        try:
            os.unlink(trace_path)
        except:
            pass

        if memory is None or params is None:
            return False, None, None, n_tested

        return True, memory, params, n_tested

    except Exception as e:
        print(f"    ERROR in validate_and_score: {e}")
        traceback.print_exc()
        return False, None, None, 0
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


def main():
    import argparse
    parser = argparse.ArgumentParser(description='Safe Best-of-N Merger (official scorer)')
    parser.add_argument('--sub_6043', required=True, help='6043 submission (base)')
    parser.add_argument('--sub_5743', default='', help='5743 submission (override source)')
    parser.add_argument('--sub_solver', default='', help='Our solver submission (override source)')
    parser.add_argument('--data_dir', required=True, help='Task JSON directory')
    parser.add_argument('--output_zip', required=True, help='Output merged zip')
    parser.add_argument('--threshold', type=float, default=0.5,
                        help='Only override if cost < threshold * base_cost (default 0.5)')
    parser.add_argument('--tasks', type=str, default='',
                        help='Comma-separated task IDs to process (default: all)')
    args = parser.parse_args()

    t0 = time.time()

    print(f"merge_safe.py — using official neurogolf_utils.py scorer")
    print(f"Threshold: {args.threshold} (override only if cost < {args.threshold*100:.0f}% of base)")
    sys.stdout.flush()

    # Load base (6043)
    models_6043 = {}
    with zipfile.ZipFile(args.sub_6043, 'r') as zf:
        for tid in range(1, 401):
            fname = f'task{tid:03d}.onnx'
            if fname in zf.namelist():
                models_6043[tid] = zf.read(fname)
    print(f"Loaded {len(models_6043)} from 6043 (base)")
    sys.stdout.flush()

    # Load override sources
    override_sources = {}
    if args.sub_5743 and os.path.exists(args.sub_5743):
        with zipfile.ZipFile(args.sub_5743, 'r') as zf:
            for tid in range(1, 401):
                fname = f'task{tid:03d}.onnx'
                if fname in zf.namelist():
                    override_sources.setdefault(tid, []).append(
                        ('5743', zf.read(fname)))
        print(f"Loaded 5743 as override source")
        sys.stdout.flush()

    if args.sub_solver and os.path.exists(args.sub_solver):
        with zipfile.ZipFile(args.sub_solver, 'r') as zf:
            for tid in range(1, 401):
                fname = f'task{tid:03d}.onnx'
                if fname in zf.namelist():
                    override_sources.setdefault(tid, []).append(
                        ('solver', zf.read(fname)))
        print(f"Loaded solver as override source")
        sys.stdout.flush()

    # Determine which tasks to process
    if args.tasks:
        task_ids = [int(t) for t in args.tasks.split(',')]
    else:
        task_ids = sorted(override_sources.keys())

    # Process each task
    output = dict(models_6043)
    overrides = 0
    failures = 0
    kept = 0
    skipped = 0
    total_gain = 0.0

    print(f"\nProcessing {len(task_ids)} tasks...")
    print(f"{'='*70}")
    sys.stdout.flush()

    for tid in task_ids:
        if tid not in override_sources:
            kept += 1
            continue

        task_path = os.path.join(args.data_dir, f'task{tid:03d}.json')
        if not os.path.exists(task_path):
            print(f"  Task {tid:>3}: SKIP — task file not found")
            skipped += 1
            continue

        try:
            with open(task_path) as f:
                task_data = json.load(f)
        except Exception as e:
            print(f"  Task {tid:>3}: ERROR loading task: {e}")
            skipped += 1
            continue

        n_arcgen = len(task_data.get('arc-gen', []))
        n_total = (len(task_data.get('train', [])) +
                   len(task_data.get('test', [])) + n_arcgen)

        # Score base (6043) model
        if tid not in models_6043:
            kept += 1
            continue

        base_valid, base_mem, base_params, base_tested = validate_and_score(
            models_6043[tid], task_data)

        if not base_valid or base_mem is None:
            # 6043 model itself fails — unusual but possible
            print(f"  Task {tid:>3}: WARNING — 6043 model failed scoring (mem={base_mem}, params={base_params})")
            kept += 1
            continue

        base_cost = base_mem + base_params
        base_score = max(1.0, 25.0 - math.log(max(1, base_cost)))

        # Try each override source
        best_source = None
        best_cost = base_cost
        best_bytes = None
        best_tested = 0

        for source_name, source_bytes in override_sources[tid]:
            # Validate and score using official utils
            src_valid, src_mem, src_params, src_tested = validate_and_score(
                source_bytes, task_data)

            if not src_valid:
                failures += 1
                print(f"  Task {tid:>3}: {source_name} FAILED validation "
                      f"(at example {src_tested}/{n_total})")
                sys.stdout.flush()
                continue

            src_cost = src_mem + src_params

            # Only override if significantly cheaper
            if src_cost < base_cost * args.threshold and src_cost < best_cost:
                best_source = source_name
                best_cost = src_cost
                best_bytes = source_bytes
                best_tested = src_tested

        if best_source:
            output[tid] = best_bytes
            overrides += 1
            new_score = max(1.0, 25.0 - math.log(max(1, best_cost)))
            gain = new_score - base_score
            total_gain += gain
            print(f"  Task {tid:>3}: OVERRIDE with {best_source} "
                  f"(cost {base_cost:,} → {best_cost:,}, "
                  f"score {base_score:.2f} → {new_score:.2f}, "
                  f"+{gain:.2f} pts, "
                  f"validated {best_tested}/{n_total} examples)")
            sys.stdout.flush()
        else:
            kept += 1

    elapsed = time.time() - t0
    print(f"\n{'='*70}")
    print(f"Results ({elapsed:.1f}s):")
    print(f"  Overrides: {overrides} (total estimated gain: +{total_gain:.2f} pts)")
    print(f"  Kept 6043: {kept}")
    print(f"  Failed validation: {failures}")
    print(f"  Skipped: {skipped}")
    print(f"  Expected LB: ~6043 + {total_gain:.1f} = ~{6043 + total_gain:.0f}")
    sys.stdout.flush()

    # Write output
    with zipfile.ZipFile(args.output_zip, 'w', zipfile.ZIP_DEFLATED) as zf:
        for tid in range(1, 401):
            if tid in output:
                zf.writestr(f'task{tid:03d}.onnx', output[tid])

    total_size = sum(len(v) for v in output.values())
    print(f"\nOutput: {args.output_zip} ({total_size:,} bytes, {len(output)} tasks)")
    sys.stdout.flush()


if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(f"\nFATAL ERROR: {e}", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        print(f"\nFATAL ERROR: {e}")
        traceback.print_exc()
        sys.exit(1)