#!/usr/bin/env python
"""
PathFinder - finds paths through voice leading graphs.
"""

from __future__ import annotations
import networkx as nx
from random import choices, seed
from typing import Iterator


class PathFinder:
    """Finds paths through voice leading graphs."""

    def __init__(self, graph: nx.MultiDiGraph):
        self.graph = graph

    def find_stochastic_path(
        self,
        start_chord: "Chord | None" = None,
        max_length: int = 100,
        weights_config: dict | None = None,
    ) -> list["Chord"]:
        """Find a stochastic path through the graph."""
        if weights_config is None:
            weights_config = self._default_weights_config()

        chord = self._initialize_chords(start_chord)
        if not chord or chord[0] is None or len(self.graph.nodes()) == 0:
            return []

        original_chord = chord[0]
        graph_node = original_chord
        output_chord = original_chord

        path = [output_chord]
        last_graph_nodes = (graph_node,)

        graph_path = [graph_node]

        from .pitch import Pitch

        dims = output_chord.dims
        cumulative_trans = Pitch(tuple(0 for _ in range(len(dims))), dims)

        num_voices = len(output_chord.pitches)
        voice_map = list(range(num_voices))

        voice_stay_count = [0] * num_voices

        for _ in range(max_length):
            out_edges = list(self.graph.out_edges(graph_node, data=True))

            if not out_edges:
                break

            weights = self._calculate_edge_weights(
                out_edges,
                path,
                last_graph_nodes,
                weights_config,
                tuple(voice_stay_count),
                graph_path,
            )

            edge = choices(out_edges, weights=weights)[0]
            next_graph_node = edge[1]
            trans = edge[2].get("transposition")
            movement = edge[2].get("movements", {})

            for src_idx, dest_idx in movement.items():
                if src_idx == dest_idx:
                    voice_stay_count[src_idx] += 1
                else:
                    voice_stay_count[src_idx] = 0

            new_voice_map = [None] * num_voices
            for src_idx, dest_idx in movement.items():
                new_voice_map[dest_idx] = voice_map[src_idx]
            voice_map = new_voice_map

            if trans is not None:
                cumulative_trans = cumulative_trans.transpose(trans)

            transposed = next_graph_node.transpose(cumulative_trans)

            reordered_pitches = tuple(
                transposed.pitches[voice_map[i]] for i in range(num_voices)
            )
            from .chord import Chord

            output_chord = Chord(reordered_pitches, dims)

            graph_node = next_graph_node
            graph_path.append(graph_node)

            path.append(output_chord)
            last_graph_nodes = last_graph_nodes + (graph_node,)
            if len(last_graph_nodes) > 2:
                last_graph_nodes = last_graph_nodes[-2:]

        return path

    def _initialize_chords(self, start_chord: "Chord | None") -> tuple:
        """Initialize chord sequence."""
        if start_chord is not None:
            return (start_chord,)

        nodes = list(self.graph.nodes())
        if nodes:
            import random

            random.shuffle(nodes)

            weights_config = self._default_weights_config()
            weights_config["voice_crossing_allowed"] = False

            for chord in nodes[:50]:
                out_edges = list(self.graph.out_edges(chord, data=True))
                if len(out_edges) == 0:
                    continue

                weights = self._calculate_edge_weights(
                    out_edges, [chord], (chord,), weights_config, None
                )
                nonzero = sum(1 for w in weights if w > 0)

                if nonzero > 0:
                    return (chord,)

            return (nodes[0],)

        return (None,)

    def _default_weights_config(self) -> dict:
        """Default weights configuration."""
        return {
            "contrary_motion": True,
            "direct_tuning": True,
            "voice_crossing_allowed": False,
            "melodic_threshold_min": 0,
            "melodic_threshold_max": 500,
            "hamiltonian": True,
            "dca": 2.0,
        }

    def _calculate_edge_weights(
        self,
        out_edges: list,
        path: list["Chord"],
        last_chords: tuple["Chord", ...],
        config: dict,
        voice_stay_count: tuple[int, ...] | None = None,
        graph_path: list["Chord"] | None = None,
    ) -> list[float]:
        """Calculate weights for edges based on configuration."""
        weights = []

        dca_multiplier = config.get("dca", 0)
        if dca_multiplier is None:
            dca_multiplier = 0

        melodic_min = config.get("melodic_threshold_min", 0)
        melodic_max = config.get("melodic_threshold_max", float("inf"))

        for edge in out_edges:
            w = 1.0
            edge_data = edge[2]

            cent_diffs = edge_data.get("cent_diffs", [])
            voice_crossing = edge_data.get("voice_crossing", False)
            is_directly_tunable = edge_data.get("is_directly_tunable", False)

            if melodic_min is not None or melodic_max is not None:
                all_within_range = True
                for cents in cent_diffs:
                    if melodic_min is not None and cents < melodic_min:
                        all_within_range = False
                        break
                    if melodic_max is not None and cents > melodic_max:
                        all_within_range = False
                        break

                if all_within_range:
                    w *= 10
                else:
                    w = 0.0

                if w == 0.0:
                    weights.append(w)
                    continue

            if config.get("contrary_motion", False):
                if len(cent_diffs) >= 3:
                    sorted_diffs = sorted(cent_diffs)
                    if sorted_diffs[0] < 0 and sorted_diffs[-1] > 0:
                        w *= 100

            if config.get("direct_tuning", False):
                if is_directly_tunable:
                    w *= 10

            if not config.get("voice_crossing_allowed", False):
                if edge_data.get("voice_crossing", False):
                    w = 0.0

            if config.get("hamiltonian", False):
                destination = edge[1]
                if graph_path and destination in graph_path:
                    w *= 0.1
                else:
                    w *= 10

            if dca_multiplier > 0 and voice_stay_count is not None and len(path) > 0:
                source_chord = path[-1]
                movements = edge_data.get("movements", {})

                move_boost = 1.0
                for voice_idx in range(len(voice_stay_count)):
                    if voice_idx in movements:
                        dest_idx = movements[voice_idx]
                        if dest_idx != voice_idx:
                            stay_count = voice_stay_count[voice_idx]
                            move_boost *= dca_multiplier**stay_count

                w *= move_boost

            weights.append(w)

        return weights

    def is_hamiltonian(self, path: list["Chord"]) -> bool:
        """Check if a path is Hamiltonian (visits all nodes exactly once)."""
        return len(path) == len(self.graph.nodes()) and len(set(path)) == len(path)
