#!/usr/bin/env python
"""
Harmonic space - multidimensional lattice where dimensions = prime factors.

HS_l = harmonic space with first l primes
CHS_l = collapsed harmonic space
"""

from __future__ import annotations
from typing import Iterator

import networkx as nx

from .pitch import Pitch
from .chord import Chord
from .dims import DIMS_4, DIMS_5, DIMS_7, DIMS_8


class HarmonicSpace:
    """
    Harmonic space HS_l or collapsed harmonic space CHS_l.

    A multidimensional lattice where each dimension corresponds to a prime factor.
    """

    def __init__(self, dims: tuple[int, ...] = DIMS_7, collapsed: bool = True):
        self.dims = dims
        self.collapsed = collapsed

    def __repr__(self) -> str:
        suffix = " (collapsed)" if self.collapsed else ""
        return f"HarmonicSpace({self.dims}{suffix})"

    def pitch(self, hs_array: tuple[int, ...]) -> Pitch:
        """Create a Pitch in this space."""
        return Pitch(hs_array, self.dims)

    def chord(self, pitches: tuple[Pitch, ...]) -> Chord:
        """Create a Chord in this space."""
        return Chord(pitches, self.dims)

    def root(self) -> Pitch:
        """Get the root pitch (1/1)."""
        return self.pitch(tuple(0 for _ in self.dims))

    def _branch_from(self, vertex: tuple[int, ...]) -> set[tuple[int, ...]]:
        """
        Get all vertices adjacent to the given vertex.

        For collapsed harmonic space, skip dimension 0 (the octave dimension).
        """
        branches = set()
        start_dim = 1 if self.collapsed else 0

        for i in range(start_dim, len(self.dims)):
            for delta in (-1, 1):
                branch = list(vertex)
                branch[i] += delta
                branches.add(tuple(branch))

        return branches

    def generate_connected_sets(
        self, min_size: int, max_size: int, collapsed: bool = True
    ) -> set[Chord]:
        """
        Generate all unique connected sets of a given size.

        Args:
            min_size: Minimum number of pitches in a chord
            max_size: Maximum number of pitches in a chord
            collapsed: If True, use CHS (skip dim 0 in branching).
                      If False (default), include dim 0 in branching.

        Returns:
            Set of unique Chord objects
        """
        root = tuple(0 for _ in self.dims)

        def branch_from(vertex):
            """Get adjacent vertices. Skip dim 0 for CHS."""
            branches = set()
            start_dim = 1 if collapsed else 0
            for i in range(start_dim, len(self.dims)):
                for delta in (-1, 1):
                    branch = list(vertex)
                    branch[i] += delta
                    branches.add(tuple(branch))
            return branches

        def grow(
            chord: tuple[tuple[int, ...], ...],
            connected: set[tuple[int, ...]],
            visited: set[tuple[int, ...]],
        ) -> Iterator[tuple[tuple[int, ...], ...]]:
            """Recursively grow connected sets."""
            if min_size <= len(chord) <= max_size:
                if collapsed:
                    projected = []
                    for arr in chord:
                        p = self.pitch(arr)
                        projected.append(p.project().hs_array)
                    yield tuple(projected)
                else:
                    yield chord

            if len(chord) < max_size:
                visited = set(visited)
                for b in connected:
                    if b not in visited:
                        extended = chord + (b,)
                        new_connected = connected | branch_from(b)
                        visited.add(b)
                        yield from grow(extended, new_connected, visited)

        connected = branch_from(root)
        visited = {root}

        results = set()
        for chord_arrays in grow((root,), connected, visited):
            pitches = tuple(self.pitch(arr) for arr in chord_arrays)
            sorted_pitches = tuple(sorted(pitches, key=lambda p: p.to_fraction()))
            results.add(Chord(sorted_pitches, self.dims))

        return results

    def build_voice_leading_graph(
        self,
        chords: set[Chord],
        symdiff_min: int = 2,
        symdiff_max: int = 2,
    ) -> nx.MultiDiGraph:
        """
        Build a voice leading graph from a set of chords.

        Args:
            chords: Set of Chord objects
            symdiff_min: Minimum symmetric difference between chords
            symdiff_max: Maximum symmetric difference between chords

        Returns:
            NetworkX MultiDiGraph
        """
        from itertools import combinations

        symdiff_range = (symdiff_min, symdiff_max)
        graph = nx.MultiDiGraph()

        for chord in chords:
            graph.add_node(chord)

        for c1, c2 in combinations(chords, 2):
            # Compute c1 -> c2
            edges = self._find_valid_edges(c1, c2, symdiff_range)
            for edge_data in edges:
                (
                    trans,
                    weight,
                    movements,
                    cent_diffs,
                    voice_crossing,
                    is_directly_tunable,
                    symdiff,
                ) = edge_data
                graph.add_edge(
                    c1,
                    c2,
                    transposition=trans,
                    weight=weight,
                    movements=movements,
                    cent_diffs=cent_diffs,
                    voice_crossing=voice_crossing,
                    is_directly_tunable=is_directly_tunable,
                    symdiff=symdiff,
                )

            # Compute c2 -> c1 independently
            edges = self._find_valid_edges(c2, c1, symdiff_range)
            for edge_data in edges:
                (
                    trans,
                    weight,
                    movements,
                    cent_diffs,
                    voice_crossing,
                    is_directly_tunable,
                    symdiff,
                ) = edge_data
                graph.add_edge(
                    c2,
                    c1,
                    transposition=trans,
                    weight=weight,
                    movements=movements,
                    cent_diffs=cent_diffs,
                    voice_crossing=voice_crossing,
                    is_directly_tunable=is_directly_tunable,
                    symdiff=symdiff,
                )

        return graph

    def _reverse_movements(self, movements: dict) -> dict:
        """Reverse the movement mappings (index to index)."""
        reversed_movements = {}
        for src_idx, dest_idx in movements.items():
            reversed_movements[dest_idx] = src_idx
        return reversed_movements

    def _is_directly_tunable(
        self,
        c1_pitches: tuple[Pitch, ...],
        c2_transposed_pitches: tuple[Pitch, ...],
        movements: dict,
    ) -> bool:
        """Check if all changing pitches are adjacent (directly tunable) to a staying pitch.

        1. Reorder c2_transposed using movement map to get destination pitches per voice
        2. Compare c1[i] with reordered[i]: if equal, pitch stays; if different, pitch changes
        3. For each changed position, check if new pitch is adjacent to any staying pitch
        """
        # Reorder c2_transposed using movement map
        # movements maps: source_index -> destination_index
        # After reordering, reordered[i] = c2_transposed at position movements[i]
        reordered = [None] * len(c1_pitches)
        for i in range(len(c1_pitches)):
            dest_idx = movements.get(i, i)
            reordered[i] = c2_transposed_pitches[dest_idx]

        # Find staying positions (where pitch doesn't change)
        staying_indices = [
            i for i in range(len(c1_pitches)) if c1_pitches[i] == reordered[i]
        ]

        if not staying_indices:
            return False

        # Find changing positions
        changing_indices = [
            i for i in range(len(c1_pitches)) if i not in staying_indices
        ]

        if not changing_indices:
            return True

        # For each changing position, check if new pitch is adjacent to any staying pitch
        for ch_idx in changing_indices:
            ch_pitch = reordered[ch_idx]
            is_adjacent_to_staying = False

            for st_idx in staying_indices:
                st_pitch = c1_pitches[st_idx]
                if self._is_adjacent_pitches(st_pitch, ch_pitch):
                    is_adjacent_to_staying = True
                    break

            if not is_adjacent_to_staying:
                return False

        return True

    def _find_valid_edges(
        self,
        c1: Chord,
        c2: Chord,
        symdiff_range: tuple[int, int],
    ) -> list[tuple[Pitch, float, dict, list[float], bool, bool]]:
        """Find all valid edges between two chords."""
        from itertools import combinations as iter_combinations

        edges = []

        transpositions = {
            p1.pitch_difference(p2) for p1 in c1.pitches for p2 in c2.pitches
        }

        for trans in transpositions:
            c2_transposed = c2.transpose(trans)

            symdiff = self._calc_symdiff(c1, c2_transposed)

            if not (symdiff_range[0] <= symdiff <= symdiff_range[1]):
                continue

            voice_lead_ok = self._check_voice_leading_connectivity(c1, c2_transposed)

            if not voice_lead_ok:
                continue

            movement_maps = self._build_movement_maps(c1.pitches, c2_transposed.pitches)

            for movements in movement_maps:
                cent_diffs = []
                for src_idx, dest_idx in movements.items():
                    src_pitch = c1.pitches[src_idx]
                    dst_pitch = c2_transposed.pitches[dest_idx]
                    cents = src_pitch.to_cents() - dst_pitch.to_cents()
                    cent_diffs.append(cents)

                source = list(c1.pitches)
                destination = [None] * len(source)
                for src_idx, dest_idx in movements.items():
                    destination[dest_idx] = c2_transposed.pitches[src_idx]

                voice_crossing = False
                for i in range(len(destination) - 1):
                    if destination[i] is None or destination[i + 1] is None:
                        continue
                    if destination[i].to_fraction() >= destination[i + 1].to_fraction():
                        voice_crossing = True
                        break

                is_directly_tunable = self._is_directly_tunable(
                    c1.pitches, c2_transposed.pitches, movements
                )

                edges.append(
                    (
                        trans,
                        1.0,
                        movements,
                        cent_diffs,
                        voice_crossing,
                        is_directly_tunable,
                        symdiff,
                    )
                )

        return edges

    def _build_movement_maps(
        self, c1_pitches: tuple[Pitch, ...], c2_transposed_pitches: tuple[Pitch, ...]
    ) -> list[dict]:
        """Build all valid movement maps for c1 -> c2_transposed."""
        from itertools import permutations

        c1_collapsed = [p.collapse() for p in c1_pitches]
        c2_collapsed = [p.collapse() for p in c2_transposed_pitches]

        common_indices_c1 = []
        common_indices_c2 = []
        for i, pc1 in enumerate(c1_collapsed):
            for j, pc2 in enumerate(c2_collapsed):
                if pc1 == pc2:
                    common_indices_c1.append(i)
                    common_indices_c2.append(j)
                    break

        changing_indices_c1 = [
            i for i in range(len(c1_pitches)) if i not in common_indices_c1
        ]
        changing_indices_c2 = [
            i for i in range(len(c2_transposed_pitches)) if i not in common_indices_c2
        ]

        base_map = {}
        for i in common_indices_c1:
            dest_idx = common_indices_c2[common_indices_c1.index(i)]
            base_map[i] = dest_idx

        if not changing_indices_c1:
            return [base_map]

        c1_changing = [c1_pitches[i] for i in changing_indices_c1]
        c2_changing = [c2_transposed_pitches[i] for i in changing_indices_c2]

        all_maps = []
        num_changing = len(c2_changing)

        for perm in permutations(range(num_changing)):
            new_map = dict(base_map)
            for i, c1_idx in enumerate(changing_indices_c1):
                dest_idx = changing_indices_c2[perm[i]]
                new_map[c1_idx] = dest_idx
            all_maps.append(new_map)

        return all_maps

    def _calc_symdiff(self, c1: Chord, c2: Chord) -> int:
        """Calculate symmetric difference between two chords."""
        set1 = set(c1.pitches)
        set2 = set(c2.pitches)
        return len(set1.symmetric_difference(set2))

    def _check_voice_leading_connectivity(self, c1: Chord, c2: Chord) -> bool:
        """Check that each pitch that changes is connected (adjacent in lattice) to some pitch in the previous chord."""
        c1_pitches = set(c1.pitches)
        c2_pitches = set(c2.pitches)

        common = c1_pitches & c2_pitches
        changing = c2_pitches - c1_pitches

        if not changing:
            return False

        for p2 in changing:
            is_adjacent = False
            for p1 in c1_pitches:
                if self._is_adjacent_pitches(p1, p2):
                    is_adjacent = True
                    break
            if not is_adjacent:
                return False

        return True

    def _is_adjacent_pitches(self, p1: Pitch, p2: Pitch) -> bool:
        """Check if two collapsed pitches are adjacent (differ by ±1 in one dimension)."""
        diff_count = 0
        for d in range(1, len(self.dims)):
            diff = abs(p1[d] - p2[d])
            if diff > 1:
                return False
            if diff == 1:
                diff_count += 1
        return diff_count == 1

    def _invert_transposition(self, trans: Pitch) -> Pitch:
        """Invert a transposition."""
        return Pitch(tuple(-t for t in trans.hs_array), self.dims)
