"""Audit FORM's emitted C arithmetic, excluding compile-time constant arithmetic.

The timed/validated engine harness is preserved unchanged. This postprocessing
step corrects the original counter's treatment of literal fractions such as 3./4.
Those divisions are constants, as in the Symbolica/Mathematica instruction counts.
"""
import argparse
import ast
from collections import Counter
from copy import deepcopy
import hashlib
import json
from pathlib import Path
import re


def constant(node):
    if isinstance(node,ast.Constant):return isinstance(node.value,(int,float))
    if isinstance(node,ast.UnaryOp):return constant(node.operand)
    if isinstance(node,ast.BinOp):return constant(node.left) and constant(node.right)
    if isinstance(node,ast.Call) and isinstance(node.func,ast.Name) and node.func.id=='pow':
        return all(constant(arg) for arg in node.args)
    return False


def count_operations(source):
    arithmetic=Counter();calls=Counter();folded=Counter();scope=Counter()
    def add(name,amount,inside):
        arithmetic[name]+=amount
        scope['inside_function_arguments' if inside else 'outside_function_arguments']+=amount
    def visit(node,inside=False):
        if constant(node):
            if isinstance(node,ast.BinOp):folded[type(node.op).__name__]+=1
            return
        if isinstance(node,ast.BinOp):
            add({ast.Add:'add',ast.Sub:'add',ast.Mult:'mul',ast.Div:'div'}[type(node.op)],1,inside)
        elif isinstance(node,ast.UnaryOp) and isinstance(node.op,ast.USub):add('neg',1,inside)
        elif isinstance(node,ast.Call):
            assert isinstance(node.func,ast.Name)
            if node.func.id=='pow':
                exponent=ast.literal_eval(node.args[1]);assert isinstance(exponent,int) and exponent>=1
                add('mul',exponent.bit_length()+exponent.bit_count()-2,inside)
                visit(node.args[0],inside);return
            assert node.func.id in {'exp','sin','cos'}
            calls[node.func.id]+=1
            inside=True
        for child in ast.iter_child_nodes(node):visit(child,inside)
    for statement in source.split(';'):
        if '=' in statement:
            visit(ast.parse(re.sub(r'\s+',' ',statement.split('=',1)[1]).strip(),mode='eval').body)
    return {'arithmetic':dict(arithmetic),'arithmetic_operations':sum(arithmetic.values()),
            'function_calls':dict(calls),'excluded_constant_subtrees':dict(folded),'arithmetic_scope':dict(scope)}


def audit_results(record):
    data=deepcopy(record);digest=hashlib.sha256(Path(__file__).read_bytes()).hexdigest();changed=False
    for row in data['results']:
        if not row['engine'].startswith('form_') or row['status']!='completed':continue
        generated=Path(row['artifacts_dir'])/'generated.inc'
        source=generated.read_text();counts=count_operations(source)
        row.setdefault('original_operation_counter', {key:deepcopy(row[key]) for key in ['arithmetic','arithmetic_operations','function_calls']})
        row.update(counts);row['generated_source_sha256']=hashlib.sha256(source.encode()).hexdigest()
        row['operation_counter_source_sha256']=digest;changed=True
    if changed:
        data['operation_count_audit']={'source':'count-native-output.py','source_sha256':digest,
            'method':'Same arithmetic cost model as the timed harness, excluding constant-only arithmetic subtrees. Original counter values are retained per row. Build times, runtime samples and verification are unchanged.'}
    return data


if __name__=='__main__':
    p=argparse.ArgumentParser(description=__doc__);p.add_argument('input',type=Path);p.add_argument('--output',type=Path,required=True)
    a=p.parse_args();a.output.write_text(json.dumps(audit_results(json.loads(a.input.read_text())),indent=2)+'\n')