"""Conservative state dependencies for compiler-owned expressions.

This is structural bookkeeping, not C parsing or numerical Jacobian evaluation.
Native multi-output operations, projections and mode-dependent writes must be
registered explicitly by the lowering path. Unresolved reachable inputs disable
coloring for the complete model.
"""
from collections import deque
from dataclasses import dataclass
import hashlib
import json
import re


# Reviewed expression leaves. Unknown arrays/scalars remain unresolved instead
# of silently becoming constants; the flow schedule supplies its own explicit IR.
_ARRAY = re.compile(r'\b[A-Za-z_]\w*\[\d+\](?:\.[A-Za-z_]\w*)?|\bgas_\d+\.[A-Za-z_]\w*')
_IDENTIFIER = re.compile(r'\b[A-Za-z_]\w*\b')
_NUMBER = re.compile(r'(?<![\w.])(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?')
_FUNCTIONS = frozenset({
    'fmax', 'fmin', 'fabs', 'sqrt', 'copysign', 'pow',
    'native_signal', 'native_contact', 'native_limit_force',
    'native_pipe_flow_context', 'native_pipe_flow_cached_context',
    'native_temperature_ph_context', 'native_density_context',
})
_CONSTANTS = frozenset({'t', 'properties', 'NAN', 'INFINITY', 'NULL', 'true', 'false'})


def expression_inputs(expression):
    """Extract leaves from a reviewed expression; unfamiliar names fail closed."""
    inputs = set()
    def leaf(match):
        key = match.group()
        # Per-evaluation cache scratch is an implementation detail of reviewed
        # kernels; each value is determined by that kernel's explicit arguments.
        if not re.fullmatch(r'pipe_cache\[\d+\]', key):
            inputs.add(key)
        return ' '
    remainder = _ARRAY.sub(leaf, expression)
    remainder = _NUMBER.sub(' ', remainder)
    for name in _IDENTIFIER.findall(remainder):
        if name in _FUNCTIONS or name in _CONSTANTS or re.fullmatch(r'(?:medium|signal)_\d+', name):
            continue
        inputs.add(name)
    return frozenset(inputs)


@dataclass(frozen=True)
class JacobianStructure:
    state_count: int
    rows: tuple[tuple[int, ...], ...] = ()
    colors: tuple[int, ...] = ()
    reason: str | None = None

    @classmethod
    def dense(cls, state_count, reason):
        return cls(state_count, reason=reason)

    @property
    def enabled(self):
        return self.reason is None

    @property
    def color_count(self):
        return max(self.colors, default=-1) + 1 if self.enabled else self.state_count

    @property
    def nonzeros(self):
        return sum(map(len, self.rows)) if self.enabled else self.state_count ** 2

    def manifest(self, *, canonical_rhs=False):
        # Use proven coloring automatically when it reduces finite-difference
        # groups; unsupported or unhelpful patterns retain dense differences.
        runtime_eligible = self.enabled and 0 < self.color_count < self.state_count
        fallback_reason = None if runtime_eligible else (
            self.reason or 'Coloring does not reduce finite-difference groups')
        pattern = {'rows': self.rows, 'columnColors': self.colors}
        policy = ('CVODE colored forward differences; canonical-property-cache RHS' if canonical_rhs else 'CVODE colored forward differences') if runtime_eligible else 'CVODE default dense differences'
        return {
            'policyScope': 'default-runtime',
            'defaultRuntimePolicy': policy,
            'verification': '--verify-jacobian',
            'policy': policy,
            'rhsPolicy': 'canonical-property-cache finite differences' if canonical_rhs and runtime_eligible else 'ordinary model_eval',
            'runtimeEligible': runtime_eligible, 'runtimeFallbackReason': fallback_reason,
            'canonicalRhs': canonical_rhs, 'ordinaryRhsUnchanged': True,
            'enabled': self.enabled, 'reason': self.reason, 'stateCount': self.state_count,
            'nonzeros': self.nonzeros, 'density': self.nonzeros / self.state_count ** 2 if self.state_count else 0,
            'colorCount': self.color_count,
            'patternSha256': hashlib.sha256(json.dumps(pattern, separators=(',', ':')).encode()).hexdigest() if self.enabled else None,
            'columnColors': list(self.colors) if self.enabled else None,
        }

    def header_lines(self, *, canonical_rhs=False):
        lines = [f'#define MODEL_JACOBIAN_COLORED {int(self.enabled)}',
                 f'#define MODEL_JACOBIAN_CANONICAL_RHS {int(canonical_rhs)}',
                 f'#define MODEL_JACOBIAN_COLOR_COUNT {self.color_count}',
                 f'#define MODEL_JACOBIAN_NNZ {self.nonzeros}']
        if canonical_rhs:
            lines += ['int model_eval_jacobian(double t, const double *y, double *dy, double *w);']
        if self.enabled:
            lines += ['extern const int model_jacobian_column_color[NSTATES];',
                      'extern const int model_jacobian_col_ptr[NSTATES+1];',
                      'extern const int model_jacobian_row_index[MODEL_JACOBIAN_NNZ];']
        return lines

    def source_lines(self):
        if not self.enabled:
            return []
        columns = [[] for _ in range(self.state_count)]
        for row, entries in enumerate(self.rows):
            for column in entries:
                columns[column].append(row)
        pointers = [0]
        indices = []
        for column in columns:
            indices.extend(column)
            pointers.append(len(indices))
        return [f'const int {name}[{len(values)}] = {{'+','.join(map(str, values))+'};'
                for name, values in [('model_jacobian_column_color', self.colors),
                                     ('model_jacobian_col_ptr', pointers),
                                     ('model_jacobian_row_index', indices)]]


class StateDependencies:
    def __init__(self, state_count):
        self.state_count = state_count
        self.seeds = {f'y[{i}]': 1 << i for i in range(state_count)}
        self.inputs = {}

    def assign(self, target, inputs):
        # Union repeated writes, rather than dropping dependencies from another
        # branch or an earlier in-place value (stop motion is mode dependent).
        self.inputs.setdefault(target, set()).update(inputs)

    def expression(self, target, expression):
        self.assign(target, expression_inputs(expression))

    def project_states(self, offsets):
        refs = {f'y[{i}]' for i in offsets}
        for target in refs:
            self.assign(target, refs)

    def stop_motion(self, velocity_index, position_index):
        refs = {f'y[{velocity_index}]', f'y[{position_index}]',
                f'dy[{velocity_index}]', f'dy[{position_index}]'}
        for target in (f'dy[{velocity_index}]', f'dy[{position_index}]'):
            self.assign(target, refs)

    def computation(self, operation):
        # A schedule SCC reaches a fixed point in build(), so every member gains
        # every external state dependency even through pressure/stream loops.
        for output in operation.outputs:
            self.assign(output, operation.inputs)

    def build(self):
        if self.state_count == 0:
            return JacobianStructure.dense(1, 'Algebraic-only internal state uses default differences')
        required = {f'dy[{i}]' for i in range(self.state_count)}
        pending = list(required)
        reachable = set()
        missing = set()
        while pending:
            key = pending.pop()
            if key in reachable:
                continue
            reachable.add(key)
            if key not in self.inputs and key not in self.seeds:
                missing.add(key)
            pending.extend(self.inputs.get(key, ()))
        if missing:
            return JacobianStructure.dense(self.state_count,
                'Unresolved structural inputs: '+', '.join(sorted(missing)[:16]))
        consumers = {key: set() for key in reachable}
        for target in reachable:
            for source in self.inputs.get(target, ()):
                consumers[source].add(target)
        masks = dict(self.seeds)
        queue = deque(key for key in self.seeds if key in reachable)
        queued = set(queue)
        while queue:
            source = queue.popleft()
            queued.remove(source)
            for target in consumers[source]:
                merged = masks.get(target, 0) | masks[source]
                if merged != masks.get(target, 0):
                    masks[target] = merged
                    if target not in queued:
                        queue.append(target)
                        queued.add(target)
        # A diagonal overestimate is safe and keeps isolated/dummy rows valid
        # without introducing zero-length C arrays or uncolored state columns.
        rows = tuple(tuple(column for column in range(self.state_count)
                           if (masks.get(f'dy[{row}]', 0) | (1 << row)) >> column & 1)
                     for row in range(self.state_count))
        conflicts = [set() for _ in range(self.state_count)]
        for entries in rows:
            for column in entries:
                conflicts[column].update(set(entries) - {column})
        colors = {}
        while len(colors) < self.state_count:
            column = max((i for i in range(self.state_count) if i not in colors),
                         key=lambda i: (len({colors[j] for j in conflicts[i] if j in colors}),
                                        len(conflicts[i]), -i))
            used = {colors[j] for j in conflicts[column] if j in colors}
            color = 0
            while color in used:
                color += 1
            colors[column] = color
        ordered = tuple(colors[i] for i in range(self.state_count))
        for entries in rows:
            if len({ordered[column] for column in entries}) != len(entries):
                raise AssertionError('Jacobian coloring contains a row conflict')
        return JacobianStructure(self.state_count, rows, ordered)