#!/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 Any

from .pitch import Pitch
from .chord import Chord


@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)
    weight: float = 0.0  # computed later by _compute_weights
    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, step: PathStep) -> PathStep:
        """Add a completed step to the path.

        Takes a PathStep (computed as a hypothetical step), updates internal state,
        and adds it to the path.
        """
        # Update cumulative transposition
        if step.transposition is not None:
            self._cumulative_trans = self._cumulative_trans.transpose(
                step.transposition
            )

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

        # 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[step.destination_node] = 0

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

        # Update step with computed state
        step.last_visited_count_before = last_visited_before
        step.last_visited_count_after = last_visited_after
        step.sustain_count_before = sustain_before
        step.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