#!/usr/bin/env python3
"""Static profiling for ONNX models.

The new neurogolf_utils.py (May 14 2026) changed score_network() to require
(sanitized_model, trace_path) instead of just a filename. This requires running
ORT with profiling enabled, which is too heavy for local model generation.

Strategy: Use static fallback for local scoring during model generation.
Real scoring happens on Kaggle at submission time via the official utils.
Models are NOT rejected locally — they're validated via inference correctness.

Returns (macs=0, memory, params) for backward compatibility with solver_registry.py.
"""

import math
import onnx
from onnx import numpy_helper
from .constants import BANNED_OPS, GH, GW


def score_network(path):
    """Score network locally. Returns (macs, memory, params) or (None, None, None).
    
    macs is always 0 (no longer used in Kaggle scoring since May 4 2026).
    memory and params are static estimates sufficient for local development.
    Real scoring uses ORT profiler on Kaggle.
    """
    result = _static_profile(path)
    if result is None:
        return None, None, None
    memory, params = result
    return 0, memory, params


def estimate_score(path):
    """Estimate score under new formula: 25 - ln(memory + params)."""
    result = _static_profile(path)
    if result is None:
        return None
    memory, params = result
    cost = memory + params
    if cost <= 0:
        return 25.0
    return max(1.0, 25.0 - math.log(cost))


def _static_profile(path):
    """Static profiling: estimate memory + params.
    
    memory = sum of all initializer bytes + estimated intermediate tensor bytes
    params = sum of all initializer element counts + Constant node values
    
    Returns (memory, params) or None if model is invalid.
    """
    try:
        model = onnx.load(path)
    except Exception:
        return None
    
    params = 0
    memory = 0  # bytes
    
    # Count initializers (weights)
    for init in model.graph.initializer:
        try:
            a = numpy_helper.to_array(init)
            params += a.size
            memory += a.nbytes
        except Exception:
            pass
    
    # Count Constant nodes
    for nd in model.graph.node:
        if nd.op_type == 'Constant':
            for attr in nd.attribute:
                if attr.name == 'value' and attr.t and attr.t.ByteSize() > 0:
                    try:
                        a = numpy_helper.to_array(attr.t)
                        params += a.size
                        memory += a.nbytes
                    except Exception:
                        pass
                elif attr.name == 'value_floats':
                    params += len(attr.floats)
                    memory += len(attr.floats) * 4
                elif attr.name == 'value_ints':
                    params += len(attr.ints)
                    memory += len(attr.ints) * 8
        
        # Banned op check
        if nd.op_type.upper() in {op.upper() for op in BANNED_OPS}:
            print(f"WARNING: Banned op '{nd.op_type}' found in {path}")
            return None
    
    # Estimate intermediate tensor memory (node outputs that aren't 'output')
    n_intermediates = 0
    for nd in model.graph.node:
        for out_name in nd.output:
            if out_name and out_name != 'output':
                n_intermediates += 1
    
    # Conservative estimate: average intermediate ~20KB (mix of small and large tensors)
    memory += n_intermediates * 20000
    
    return int(memory), int(params)