compact_sets/src/path.py

#!/usr/bin/env python
"""
Path and PathStep classes for storing path state from PathFinder.
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from .pitch import Pitch
from .chord import Chord

if TYPE_CHECKING:
    from .graph import Candidate


@dataclass
class PathStep:
    """Stores data for a single step (edge) in the path."""

    source_node: Chord
    destination_node: Chord
    source_chord: Chord
    destination_chord: Chord
    transposition: Pitch | None = None
    movements: dict[int, int] = field(default_factory=dict)
    scores: dict[str, float] = field(default_factory=dict)
    candidates: list["Candidate"] = field(default_factory=list)
    last_visited_count_before: dict | None = None
    last_visited_count_after: dict | None = None
    sustain_count_before: tuple[int, ...] | None = None
    sustain_count_after: tuple[int, ...] | None = None


class Path:
    """Stores the complete state of a generated path."""

    def __init__(
        self, initial_chord: Chord | None, weights_config: dict[str, Any] | None = None
    ):
        self.initial_chord = initial_chord
        self.steps: list[PathStep] = []
        self.weights_config = weights_config if weights_config is not None else {}

        # State needed for step computation
        self._voice_map: list[int] = []  # which voice is at each position
        self._cumulative_trans: Pitch | None = None  # cumulative transposition
        self._graph_nodes: set = set()  # all graph nodes for visit tracking
        self._num_voices: int = 0  # number of voices

    def init_state(
        self, graph_nodes: set, num_voices: int, initial_chord: Chord
    ) -> None:
        """Initialize state after graph is known."""
        self._graph_nodes = graph_nodes
        self._num_voices = num_voices
        self._voice_map = list(range(num_voices))  # voice i at position i

        dims = initial_chord.dims
        self._cumulative_trans = Pitch(tuple(0 for _ in range(len(dims))), dims)

    def _get_last_visited_counts(self) -> dict:
        """Get last visited counts from the last step, or initialize fresh."""
        if self.steps:
            last_step = self.steps[-1]
            if last_step.last_visited_count_after is not None:
                return dict(last_step.last_visited_count_after)

        # Initialize fresh: all nodes start at 0 (except initial which we set to 0 explicitly)
        return {node: 0 for node in self._graph_nodes}

    def _get_sustain_counts(self) -> tuple:
        """Get sustain counts from the last step, or initialize fresh."""
        if self.steps:
            last_step = self.steps[-1]
            if last_step.sustain_count_after is not None:
                return last_step.sustain_count_after

        # Initialize fresh: all voices start at 0
        return tuple(0 for _ in range(self._num_voices))

    def step(
        self,
        edge: tuple,
        candidates: list["Candidate"],
        chosen_scores: dict[str, float] | None = None,
    ) -> PathStep:
        """Process a step: update state, compute output, return step.

        Takes edge (source_node, destination_node, edge_data), handles all voice-leading internally.
        """
        source_node = edge[0]
        destination_node = edge[1]
        edge_data = edge[2]

        trans = edge_data.get("transposition")
        movement = edge_data.get("movements", {})

        # Update cumulative transposition
        if trans is not None:
            self._cumulative_trans = self._cumulative_trans.transpose(trans)

        # Transpose the destination node
        transposed = destination_node.transpose(self._cumulative_trans)

        # Update voice map based on movement
        new_voice_map = [None] * len(self._voice_map)
        for src_idx, dest_idx in movement.items():
            new_voice_map[dest_idx] = self._voice_map[src_idx]
        self._voice_map = new_voice_map

        # Reorder pitches according to voice map
        reordered_pitches = tuple(
            transposed.pitches[self._voice_map[i]] for i in range(len(self._voice_map))
        )
        destination_chord = Chord(reordered_pitches, destination_node.dims)

        # Get previous output chord
        source_chord = self.output_chords[-1]

        # Get BEFORE state from last step (or initialize fresh)
        last_visited_before = self._get_last_visited_counts()
        sustain_before = self._get_sustain_counts()

        # Compute AFTER state
        last_visited_after = dict(last_visited_before)
        for node in last_visited_after:
            last_visited_after[node] += 1
        last_visited_after[destination_node] = 0

        sustain_after = list(sustain_before)
        for voice_idx in range(len(sustain_after)):
            curr_cents = source_chord.pitches[voice_idx].to_cents()
            next_cents = destination_chord.pitches[voice_idx].to_cents()
            if curr_cents == next_cents:
                sustain_after[voice_idx] += 1
            else:
                sustain_after[voice_idx] = 0

        # Create step with before and after state
        step = PathStep(
            source_node=source_node,
            destination_node=destination_node,
            source_chord=source_chord,
            destination_chord=destination_chord,
            transposition=trans,
            movements=movement,
            scores=chosen_scores if chosen_scores is not None else {},
            candidates=candidates,
            last_visited_count_before=last_visited_before,
            last_visited_count_after=last_visited_after,
            sustain_count_before=sustain_before,
            sustain_count_after=tuple(sustain_after),
        )

        self.steps.append(step)
        return step

    @property
    def graph_chords(self) -> list[Chord]:
        """Get list of destination graph nodes."""
        return [self.initial_chord] + [step.destination_node for step in self.steps]

    @property
    def output_chords(self) -> list[Chord]:
        """Get list of destination chords (transposed)."""
        return [self.initial_chord] + [step.destination_chord for step in self.steps]

    def __len__(self) -> int:
        """Total number of chords in path."""
        return len(self.steps) + 1

    def __iter__(self):
        """Iterate over output chords."""
        return iter(self.output_chords)

    def get_influence(self, weights: dict[str, Any]) -> dict[str, float]:
        """Compute weighted score contribution per factor for chosen candidates.

        Returns a dict mapping factor name to accumulated influence (weight * score)
        for all steps in the path.
        """
        influence = {
            "melodic": 0.0,
            "contrary_motion": 0.0,
            "dca_hamiltonian": 0.0,
            "dca_voice_movement": 0.0,
            "target_range": 0.0,
        }

        for step in self.steps:
            scores = step.scores
            w_melodic = weights.get("weight_melodic", 1)
            w_contrary = weights.get("weight_contrary_motion", 0)
            w_hamiltonian = weights.get("weight_dca_hamiltonian", 1)
            w_dca = weights.get("weight_dca_voice_movement", 1)
            w_target = weights.get("weight_target_range", 1)

            influence["melodic"] += scores.get("melodic_threshold", 0) * w_melodic
            influence["contrary_motion"] += (
                scores.get("contrary_motion", 0) * w_contrary
            )
            influence["dca_hamiltonian"] += (
                scores.get("dca_hamiltonian", 0) * w_hamiltonian
            )
            influence["dca_voice_movement"] += (
                scores.get("dca_voice_movement", 0) * w_dca
            )
            influence["target_range"] += scores.get("target_range", 0) * w_target

        return influence