# Paper-scoped implementation; see SOURCE_PROVENANCE.json.
from __future__ import annotations
from dataclasses import asdict, dataclass, field
from typing import Any
import numpy as np
from wisp.search.grammar import CandidateSpec
from wisp.search.motifs import canonicalize_tags, family_motif_match

@dataclass(frozen=True)
class OperatorNode:
    name: str
    family: str
    information: float = 0.0
    invariance: float = 0.0
    dynamics: float = 0.0
    reconstructability: float = 0.0
    cost: float = 1.0
    tags: tuple[str, ...] = ()

    def to_dict(self) -> dict[str, Any]:
        d = asdict(self)
        d['tags'] = list(self.tags)
        return d
OPERATOR_LIBRARY: dict[str, OperatorNode] = {'identity': OperatorNode('identity', 'preprocess', information=1.0, reconstructability=1.0, cost=0.1, tags=('lossless',)), 'zscore': OperatorNode('zscore', 'preprocess', information=0.5, invariance=0.8, reconstructability=0.0, cost=0.2, tags=('shift_scale_invariant',)), 'magnitude': OperatorNode('magnitude', 'derive', information=0.5, invariance=0.8, cost=0.3, tags=('orientation_invariance',)), 'jerk': OperatorNode('jerk', 'derive', information=0.6, dynamics=0.8, cost=0.3, tags=('micro_burst',)), 'gram_eig': OperatorNode('gram_eig', 'derive', information=0.45, invariance=0.8, cost=0.4, tags=('device_shift', 'orientation_invariance')), 'intrinsic_timescale': OperatorNode('intrinsic_timescale', 'derive', information=0.45, dynamics=0.7, cost=0.4, tags=('state_persistence',)), 'multiscale_stats': OperatorNode('multiscale_stats', 'aggregate', information=0.65, dynamics=0.55, cost=1.0, tags=('local_shape',)), 'multiscale_interval_stats': OperatorNode('multiscale_interval_stats', 'aggregate', information=0.65, dynamics=0.55, cost=1.0, tags=('local_shape',)), 'dyadic_intervals': OperatorNode('dyadic_intervals', 'segment', information=0.6, dynamics=0.5, cost=0.8, tags=('interval_distribution',)), 'quantile_distribution': OperatorNode('quantile_distribution', 'aggregate', information=0.55, invariance=0.3, dynamics=0.4, cost=0.8, tags=('distribution_shape',)), 'random_shapelet': OperatorNode('random_shapelet', 'aggregate', information=0.7, invariance=0.2, dynamics=0.75, cost=1.4, tags=('local_shape', 'micro_burst')), 'aggregate_stats': OperatorNode('aggregate_stats', 'aggregate', information=0.55, cost=0.7, tags=('summary',)), 'segment': OperatorNode('segment', 'segment', information=0.7, cost=0.4, tags=('localization',)), 'fft_bands': OperatorNode('fft_bands', 'spectral', information=0.55, invariance=0.35, dynamics=0.55, cost=0.5, tags=('periodic_locomotion',)), 'autocorr': OperatorNode('autocorr', 'spectral', information=0.55, invariance=0.35, dynamics=0.75, cost=0.5, tags=('periodic_locomotion', 'state_persistence')), 'random_conv': OperatorNode('random_conv', 'random_conv', information=0.7, dynamics=0.65, cost=1.3, tags=('phase_tolerant_local_shape',)), 'ppv_pool': OperatorNode('ppv_pool', 'random_conv', information=0.45, invariance=0.5, cost=0.3, tags=('phase_tolerant_local_shape',)), 'ppv_max_mean_std': OperatorNode('ppv_max_mean_std', 'random_conv', information=0.55, invariance=0.5, cost=0.4, tags=('phase_tolerant_local_shape',)), 'symbolic_transition': OperatorNode('symbolic_transition', 'symbolic', information=0.35, invariance=0.5, dynamics=0.85, cost=0.9, tags=('state_dynamics', 'micro_pattern')), 'transition_histogram': OperatorNode('transition_histogram', 'symbolic', information=0.35, invariance=0.5, dynamics=0.8, cost=0.8, tags=('state_dynamics',)), 'quantize': OperatorNode('quantize', 'symbolic', information=0.25, invariance=0.4, cost=0.3, tags=('symbolic',)), 'symbolic': OperatorNode('symbolic', 'symbolic', information=0.35, dynamics=0.7, cost=0.8, tags=('state_dynamics',)), 'extratrees': OperatorNode('extratrees', 'classifier_cpu', information=0.4, cost=1.1, tags=('nonlinear_cpu',)), 'rf': OperatorNode('rf', 'classifier_cpu', information=0.4, cost=1.1, tags=('nonlinear_cpu',)), 'ridge': OperatorNode('ridge', 'classifier_cpu', information=0.3, cost=0.4, tags=('fast_cpu',)), 'logreg': OperatorNode('logreg', 'classifier_cpu', information=0.3, cost=0.6, tags=('calibrated_cpu',)), 'probability_ensemble': OperatorNode('probability_ensemble', 'classifier_cpu', information=0.55, cost=1.3, tags=('multi_view_cpu',)), 'posterior_weighting': OperatorNode('posterior_weighting', 'posterior', information=0.45, cost=0.5, tags=('posterior_over_motifs',)), 'invariant_physics': OperatorNode('invariant_physics', 'derive', information=0.55, invariance=0.8, cost=1.0, tags=('device_shift',)), 'feature_union': OperatorNode('feature_union', 'fusion', information=0.65, cost=1.2, tags=('multi_view_cpu',)), 'hmm': OperatorNode('hmm', 'postprocess', information=0.25, dynamics=0.9, cost=0.4, tags=('transition_smoothing', 'free_living_noise')), 'argmax': OperatorNode('argmax', 'postprocess', information=0.1, cost=0.1, tags=('window_only',))}

@dataclass
class TemporalProgram:
    name: str
    route: str
    nodes: list[OperatorNode]
    params: dict[str, Any] = field(default_factory=dict)
    parent_id: str | None = None
    mutation: str | None = None

    @classmethod
    def from_candidate(cls, spec: CandidateSpec) -> 'TemporalProgram':
        nodes = [OPERATOR_LIBRARY.get(op, OperatorNode(op, 'unknown', cost=1.0)) for op in spec.operators]
        return cls(name=spec.name, route=spec.route, nodes=nodes, params=dict(spec.params), parent_id=spec.parent_id, mutation=spec.mutation)

    def to_candidate(self) -> CandidateSpec:
        return CandidateSpec(name=self.name, route=self.route, params=dict(self.params), operators=[node.name for node in self.nodes], parent_id=self.parent_id, mutation=self.mutation)

    @property
    def operators(self) -> list[str]:
        return [node.name for node in self.nodes]

    @property
    def families(self) -> list[str]:
        return [node.family for node in self.nodes]

    def certificate(self) -> dict[str, Any]:
        if not self.nodes:
            return {}
        arr = np.asarray([[n.information, n.invariance, n.dynamics, n.reconstructability, n.cost] for n in self.nodes], dtype=np.float64)
        tags: set[str] = set()
        for node in self.nodes:
            tags.update(node.tags)
        canonical = canonicalize_tags(tags)
        return {'name': self.name, 'route': self.route, 'n_operators': len(self.nodes), 'operator_trace': self.operators, 'family_trace': self.families, 'mean_information': float(np.mean(arr[:, 0])), 'mean_invariance': float(np.mean(arr[:, 1])), 'mean_dynamics': float(np.mean(arr[:, 2])), 'mean_reconstructability': float(np.mean(arr[:, 3])), 'estimated_cost': float(np.sum(arr[:, 4])), 'tags': sorted(tags), 'canonical_motifs': sorted(canonical), 'task_representation_congruence': self.task_congruence()}

    def motif_match(self, motif_weights: dict[str, float]) -> dict[str, Any]:
        match = family_motif_match(self.name, motif_weights)
        return {'family': match.family, 'score': match.score, 'contributions': match.contributions}

    def task_congruence(self) -> dict[str, float]:
        tags = set()
        for node in self.nodes:
            tags.update(node.tags)
        score = {'Capture24_free_living': 0.0, 'WISDM_micro_actions': 0.0, 'HHAR_device_shift': 0.0, 'PAMAP2_multisensor': 0.0, 'MHEALTH_small_scripted': 0.0}
        score['Capture24_free_living'] += 0.35 if 'free_living_noise' in tags else 0.0
        score['Capture24_free_living'] += 0.25 if 'transition_smoothing' in tags else 0.0
        score['Capture24_free_living'] += 0.2 if 'periodic_locomotion' in tags else 0.0
        score['Capture24_free_living'] += 0.2 if 'orientation_invariance' in tags else 0.0
        score['WISDM_micro_actions'] += 0.35 if 'micro_burst' in tags or 'micro_pattern' in tags else 0.0
        score['WISDM_micro_actions'] += 0.25 if 'local_shape' in tags else 0.0
        score['WISDM_micro_actions'] += 0.2 if 'state_dynamics' in tags else 0.0
        score['WISDM_micro_actions'] += 0.2 if 'phase_tolerant_local_shape' in tags else 0.0
        score['HHAR_device_shift'] += 0.45 if 'device_shift' in tags else 0.0
        score['HHAR_device_shift'] += 0.35 if 'orientation_invariance' in tags else 0.0
        score['HHAR_device_shift'] += 0.2 if 'fast_cpu' in tags or 'phase_tolerant_local_shape' in tags else 0.0
        score['PAMAP2_multisensor'] += 0.3 if 'multi_view_cpu' in tags else 0.0
        score['PAMAP2_multisensor'] += 0.25 if 'local_shape' in tags else 0.0
        score['PAMAP2_multisensor'] += 0.25 if 'periodic_locomotion' in tags else 0.0
        score['PAMAP2_multisensor'] += 0.2 if 'state_dynamics' in tags else 0.0
        score['MHEALTH_small_scripted'] += 0.35 if 'fast_cpu' in tags else 0.0
        score['MHEALTH_small_scripted'] += 0.25 if 'local_shape' in tags else 0.0
        score['MHEALTH_small_scripted'] += 0.2 if 'interval_distribution' in tags else 0.0
        score['MHEALTH_small_scripted'] += 0.2 if 'multi_view_cpu' in tags else 0.0
        return {k: float(min(v, 1.0)) for k, v in score.items()}

def jaccard_novelty(a: CandidateSpec, b: CandidateSpec) -> float:
    A = set(a.operators)
    B = set(b.operators)
    if not A and (not B):
        return 0.0
    return 1.0 - len(A & B) / max(len(A | B), 1)