#!/usr/bin/env python3
"""
fix_scoring_bugs.py — Fix ONNX models with scoring bugs.

Bug 1: dims=[0] initializers (scalar_shape)
  - The scorer rejects any initializer with dims containing 0
  - Fix: Replace Reshape(x, scalar_shape) with Squeeze(x)
  
Bug 2: Missing value_info (only fixable if shapes are static)
  - Run shape inference to populate value_info
  - Only works if model has no dynamic-shape ops (NonZero, Where with data-dependent shapes)

Usage:
    python fix_scoring_bugs.py --input task291.onnx --output fixed/task291.onnx --task-num 291
"""
import argparse
import os
import sys
import onnx
import onnxruntime as ort
import numpy as np
import json
from onnx import helper, TensorProto, numpy_helper, shape_inference


def fix_scalar_shape(model):
    """Replace Reshape(x, scalar_shape) with Squeeze(x) where scalar_shape has dims=[0]."""
    # Find problematic initializers
    bad_inits = {init.name for init in model.graph.initializer if any(d <= 0 for d in init.dims)}
    
    if not bad_inits:
        return model, False
    
    # Remove bad initializers
    new_inits = [init for init in model.graph.initializer if init.name not in bad_inits]
    
    # Replace Reshape nodes that use bad initializers with Squeeze
    new_nodes = []
    fixed = False
    for node in model.graph.node:
        if node.op_type == 'Reshape' and any(inp in bad_inits for inp in node.input):
            squeeze_node = helper.make_node(
                'Squeeze',
                inputs=[node.input[0]],
                outputs=list(node.output),
            )
            new_nodes.append(squeeze_node)
            fixed = True
        else:
            new_nodes.append(node)
    
    del model.graph.initializer[:]
    model.graph.initializer.extend(new_inits)
    del model.graph.node[:]
    model.graph.node.extend(new_nodes)
    
    return model, fixed


def fix_missing_value_info(model):
    """Run shape inference to add value_info for models that lack it."""
    if len(model.graph.value_info) > 0:
        return model, False
    
    try:
        inferred = shape_inference.infer_shapes(model, strict_mode=True)
        if len(inferred.graph.value_info) > 0:
            return inferred, True
    except:
        pass
    
    try:
        inferred = shape_inference.infer_shapes(model, strict_mode=False)
        if len(inferred.graph.value_info) > 0:
            return inferred, True
    except:
        pass
    
    return model, False


def validate_model(model_path, task_num, task_data_dir):
    """Validate model produces correct outputs for all examples."""
    task_json = os.path.join(task_data_dir, f'task{task_num:03d}.json')
    if not os.path.exists(task_json):
        print(f"  Warning: task data not found at {task_json}")
        return True  # Can't validate without data
    
    sess = ort.InferenceSession(model_path)
    with open(task_json) as f:
        data = json.load(f)
    
    right, wrong = 0, 0
    for ex in data['train'] + data['test'] + data.get('arc-gen', []):
        inp = np.zeros((1, 10, 30, 30), dtype=np.float32)
        for r, row in enumerate(ex['input']):
            for c, v in enumerate(row):
                if r < 30 and c < 30:
                    inp[0][v][r][c] = 1.0
        result = sess.run(['output'], {'input': inp})
        out = (result[0] > 0.0).astype(float)
        exp = np.zeros((1, 10, 30, 30), dtype=np.float32)
        for r, row in enumerate(ex['output']):
            for c, v in enumerate(row):
                if r < 30 and c < 30:
                    exp[0][v][r][c] = 1.0
        if np.array_equal(out, exp):
            right += 1
        else:
            wrong += 1
    
    print(f"  Validation: {right} pass, {wrong} fail")
    return wrong == 0


def main():
    parser = argparse.ArgumentParser(description="Fix ONNX scoring bugs")
    parser.add_argument('--input', required=True, help='Input ONNX model')
    parser.add_argument('--output', required=True, help='Output ONNX model')
    parser.add_argument('--task-num', type=int, required=True, help='Task number')
    parser.add_argument('--task-data-dir', default='/app/task-data', help='Task data directory')
    args = parser.parse_args()
    
    print(f"Loading {args.input}...")
    model = onnx.load(args.input)
    print(f"  Nodes: {len(model.graph.node)}")
    print(f"  Initializers: {len(model.graph.initializer)}")
    print(f"  Value_info: {len(model.graph.value_info)}")
    
    # Apply fixes
    model, fixed_scalar = fix_scalar_shape(model)
    if fixed_scalar:
        print("  ✓ Fixed scalar_shape (dims=[0]) bug")
    
    model, fixed_vi = fix_missing_value_info(model)
    if fixed_vi:
        print(f"  ✓ Added value_info via shape inference ({len(model.graph.value_info)} entries)")
    
    if not fixed_scalar and not fixed_vi:
        print("  No fixes needed")
        return
    
    # Save
    os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True)
    onnx.save(model, args.output)
    print(f"  Saved to {args.output}")
    
    # Validate
    if validate_model(args.output, args.task_num, args.task_data_dir):
        print("  ✓ Model passes all examples!")
    else:
        print("  ✗ Model FAILS validation — do NOT use!")
        sys.exit(1)


if __name__ == '__main__':
    main()