#!/usr/bin/env python
# coding: utf-8

# In[89]:


from itertools import chain, combinations, permutations, product
from math import prod, log
from copy import deepcopy
import networkx as nx
from fractions import Fraction
import json
from operator import add

def hs_array_to_fr(hs_array):
    return prod([pow(dims[d], hs_array[d]) for d in range(len(dims))])

def hs_array_to_cents(hs_array):
    return (1200 * log(hs_array_to_fr(hs_array), 2))

def expand_pitch(hs_array):
    expanded_pitch = list(hs_array)
    frequency_ratio = hs_array_to_fr(hs_array)
    if frequency_ratio < 1:
        while frequency_ratio < 1:
            frequency_ratio *= 2
            expanded_pitch[0] += 1
    elif frequency_ratio >= 2:
        while frequency_ratio >= 2:
            frequency_ratio *= 1/2
            expanded_pitch[0] += -1
    return tuple(expanded_pitch)

def expand_chord(chord):
    return tuple(expand_pitch(p) for p in chord)

def collapse_pitch(hs_array):
    collapsed_pitch = list(hs_array)
    collapsed_pitch[0] = 0
    return tuple(collapsed_pitch)

def collapse_chord(chord):
    return tuple(collapse_pitch(p) for p in chord)

def transpose_pitch(pitch, trans):
    return tuple(map(add, pitch, trans))

def transpose_chord(chord, trans):
    return tuple(transpose_pitch(p, trans) for p in chord)

def cent_difference(hs_array1, hs_array2):
    return hs_array_to_cents(hs_array2) - hs_array_to_cents(hs_array1)

def pitch_difference(hs_array1, hs_array2):
    return transpose_pitch(hs_array1, [p * -1 for p in hs_array2])

# this is modified for different chord sizes like original version
def grow_chords(chord, root, min_chord_size, max_chord_size):
    #this could use the tranpose_pitch function
    branches = [branch for alt in [-1, 1] for d in range(1, len(root)) if (branch:=(*(r:=root)[:d], r[d] + alt, *r[(d + 1):])) not in chord]
    subsets = chain.from_iterable(combinations(branches, r) for r in range(1, max_chord_size - len(chord) + 1))
    for subset in subsets:
        extended_chord = chord + subset
        if(len(extended_chord) < max_chord_size):
            for branch in subset:
                yield from grow_chords(extended_chord, branch, min_chord_size, max_chord_size)
        if(len(extended_chord) >= min_chord_size):
            yield tuple(sorted(extended_chord, key=hs_array_to_fr))

def chords(chord, root, min_chord_size, max_chord_size):
    # this will filter out the 4x dups of paths that are loops, there might be a faster way to test this
    return set(grow_chords(chord, root, min_chord_size, max_chord_size))

# this is very slow, I have an idea in mind that my be faster by simply growing the chords to max_chord_size + max_sim_diff
# technically at that point you have generated both chords and can get the second chord from the first
def edges(chords, min_symdiff, max_symdiff, max_chord_size):  

    def reverse_movements(movements):
        return {value['destination']:{'destination':key, 'cent_difference':value['cent_difference'] * -1} for key, value in movements.items()}

    def is_directly_tunable(intersection, diff):
        # this only works for now when intersection if one element - need to fix that
        return max([sum(abs(p) for p in collapse_pitch(pitch_difference(d, list(intersection)[0]))) for d in diff]) == 1

    for combination in combinations(chords, 2):
        [expanded_base, expanded_comp] = [expand_chord(chord) for chord in combination]
        edges = []
        transpositions = set(pitch_difference(pair[0], pair[1]) for pair in set(product(expanded_base, expanded_comp)))
        for trans in transpositions:
            expanded_comp_transposed = transpose_chord(expanded_comp, trans)
            intersection = set(expanded_base) & set(expanded_comp_transposed)
            symdiff_len = sum([len(chord) - len(intersection) for chord in [expanded_base, expanded_comp_transposed]])
            if (min_symdiff <= symdiff_len <= max_symdiff):
                rev_trans = tuple(t * -1 for t in trans)
                [diff1, diff2] = [list(set(chord) - intersection) for chord in [expanded_base, expanded_comp_transposed]]
                base_map = {val: {'destination':transpose_pitch(val, rev_trans), 'cent_difference': 0} for val in intersection}
                base_map_rev = reverse_movements(base_map)
                maps = []
                diff1 += [None] * (max_chord_size - len(diff1) - len(intersection))
                perms = [list(perm) + [None] * (max_chord_size - len(perm) - len(intersection)) for perm in set(permutations(diff2))]
                for p in perms:
                    appended_map = {
                        diff1[index]:
                        {
                            'destination': transpose_pitch(val, rev_trans) if val != None else None, 
                            'cent_difference': cent_difference(diff1[index], val) if None not in [diff1[index], val] else None
                        } for index, val in enumerate(p)}
                    yield (tuple(expanded_base), tuple(expanded_comp), {
                        'transposition': trans,
                        'symmetric_difference': symdiff_len, 
                        'is_directly_tunable': is_directly_tunable(intersection, diff2),
                        'movements': base_map | appended_map
                    },)
                    yield (tuple(expanded_comp), tuple(expanded_base), {
                        'transposition': rev_trans,
                        'symmetric_difference': symdiff_len, 
                        'is_directly_tunable': is_directly_tunable(intersection, diff1),
                        'movements': base_map_rev | reverse_movements(appended_map)
                    },)

def graph_from_edges(edges):
    g = nx.MultiDiGraph()
    g.add_edges_from(edges)
    return g

def generate_graph(chord_set, min_symdiff, max_symdiff, max_chord_size):
    #chord_set = chords(pitch_set, min_chord_size, max_chord_size)
    edge_set = edges(chord_set, min_symdiff, max_symdiff, max_chord_size)
    res_graph = graph_from_edges(edge_set)
    return res_graph

def display_graph(graph):
    show_graph = nx.Graph(graph)
    pos = nx.draw_spring(show_graph, node_size=5, width=0.1)
    plt.figure(1, figsize=(12,12)) 
    nx.draw(show_graph, pos, node_size=5, width=0.1)
    plt.show()
    #plt.savefig('compact_sets.png', dpi=150)

def path_to_chords(path, start_root):
    current_root = start_root
    start_chord = tuple(sorted(path[0][0], key=hs_array_to_fr))
    chords = ((start_chord, start_chord,),)
    for edge in path:
        trans = edge[2]['transposition']
        movements = edge[2]['movements']
        current_root = transpose_pitch(current_root, trans)
        current_ref_chord = chords[-1][0]
        next_ref_chord = tuple(movements[pitch]['destination'] for pitch in current_ref_chord)
        next_transposed_chord = tuple(transpose_pitch(pitch, current_root) for pitch in next_ref_chord)
        chords += ((next_ref_chord, next_transposed_chord,),)
    return tuple(chord[1] for chord in chords)

def write_chord_sequence(path):
    file = open("seq.txt", "w+")
    content = json.dumps(path)
    content = content.replace("[[[", "[\n\t[[")
    content = content.replace(", [[", ",\n\t[[")
    content = content.replace("]]]", "]]\n]")
    file.write(content)
    file.close()


# In[214]:


dims = (2, 3, 5, 7, 11)
root = (0, 0, 0, 0, 0)
chord = (root,)
chord_set = chords(chord, root, 3, 3)
graph = generate_graph(chord_set, 4, 4, 3)


# In[215]:


from random import choice, choices

# This is for the static version
def stochastic_hamiltonian(graph):
        
    def movement_size_weights(edges):
        
        def max_cent_diff(edge):
            res = max([abs(v) for val in edge[2]['movements'].values() if (v:=val['cent_difference']) is not None])
            return res
            
        def min_cent_diff(edge):
            res = [abs(v) for val in edge[2]['movements'].values() if (v:=val['cent_difference']) is not None]
            res.remove(0)
            return min(res)
            
        for e in edges:
            yield 1000 if ((max_cent_diff(e) < 200) and (min_cent_diff(e)) > 50) else 1

    def hamiltonian_weights(edges):
        for e in edges:
            yield 10 if e[1] not in [path_edge[0] for path_edge in path] else 1 / graph.nodes[e[1]]['count']
    
    def contrary_motion_weights(edges):

        def is_contrary(edge):
            cent_diffs = [v for val in edge[2]['movements'].values() if (v:=val['cent_difference']) is not None]
            cent_diffs.sort()
            return (cent_diffs[0] < 0) and (cent_diffs[1] == 0) and (cent_diffs[2] > 0)

        for e in edges:
            yield 10 if is_contrary(e) else 1
    
    def is_directly_tunable_weights(edges):
        for e in edges:
            yield 10 if e[2]['is_directly_tunable'] else 0

    def voice_crossing_weights(edges):
        
        def has_voice_crossing(edge):
            source = list(edge[0])
            ordered_source = sorted(source, key=hs_array_to_fr)    
            source_order = [ordered_source.index(p) for p in source]
            destination = [transpose_pitch(edge[2]['movements'][p]['destination'], edge[2]['transposition']) for p in source]
            ordered_destination = sorted(destination, key=hs_array_to_fr)
            destination_order = [ordered_destination.index(p) for p in destination]
            return source_order != destination_order

        for e in edges:
            yield 10 if not has_voice_crossing(e) else 0

    def is_bass_rooted(chord):
        return max([sum(abs(p) for p in collapse_pitch(pitch_difference(chord[0], p))) for p in chord[1:]]) == 1
        
    check_graph = graph.copy()
    next_node = choice([node for node in graph.nodes() if is_bass_rooted(node)])
    check_graph.remove_node(next_node)
    for node in graph.nodes(data=True):
        node[1]['count'] = 1
    path = []
    while (nx.number_of_nodes(check_graph) > 0) and (len(path) < 5000):
        out_edges = list(graph.out_edges(next_node, data=True))
        #print([l for l in zip(movement_size_weights(out_edges), hamiltonian_weights(out_edges))])
        factors = [
            movement_size_weights(out_edges), 
            hamiltonian_weights(out_edges), 
            contrary_motion_weights(out_edges), 
            is_directly_tunable_weights(out_edges),
            voice_crossing_weights(out_edges)
        ]
        weights = [prod(a) for a in zip(*factors)]
        edge = choices(out_edges, weights=weights)[0]
        #edge = random.choice(out_edges)
        next_node = edge[1]
        node[1]['count'] += 1
        path.append(edge)
        if next_node in check_graph.nodes:
            check_graph.remove_node(next_node)
    return path
    
path = stochastic_hamiltonian(graph)
#for edge in path:
#    print(edge)
write_chord_sequence(path_to_chords(path, root))
len(path)


# In[25]:


def is_super_compact(chord):
    return max([sum(abs(p) for p in collapse_pitch(pitch_difference(c[0], c[1]))) for c in combinations(chord, 2)]) == 1
        
[node for node in graph.nodes() if is_super_compact(node)]


# In[11]:


get_ipython().run_line_magic('load_ext', 'line_profiler')


# In[134]:


chord_set = chords(chord, root, 3, 3)


# In[136]:


lprun -f edge_data edges(chord_set, 3, 3, 4)


# In[180]:


dims = (2, 3, 5, 7, 11)
root = (0, 0, 0, 0, 0)
chord = (root,)
chord_set = chords(chord, root, 3, 3)
graph = generate_graph(chord_set, 2, 2, 3)


# In[213]:


from random import choice, choices

# This is for the rising version / yitgadal
def stochastic_hamiltonian(graph):
        
    def movement_size_weights(edges):
        
        def max_cent_diff(edge):
            res = max([v for val in edge[2]['movements'].values() if (v:=val['cent_difference']) is not None])
            return res
            
        def min_cent_diff(edge):
            res = [v for val in edge[2]['movements'].values() if (v:=val['cent_difference']) is not None]
            res.remove(0)
            return min(res)
            
        for e in edges:
            yield 1000 if ((max_cent_diff(e) < 175) and (min_cent_diff(e)) >= 0) else 1

    def hamiltonian_weights(edges):
        for e in edges:
            yield 10 if e[1] not in [path_edge[0] for path_edge in path] else 1 / graph.nodes[e[1]]['count']
    
    def contrary_motion_weights(edges):

        def is_contrary(edge):
            cent_diffs = [v for val in edge[2]['movements'].values() if (v:=val['cent_difference']) is not None]
            cent_diffs.sort()
            return (cent_diffs[0] < 0) and (cent_diffs[1] == 0) and (cent_diffs[2] > 0)

        for e in edges:
            yield 2 if is_contrary(e) else 1
    
    def is_directly_tunable_weights(edges):
        for e in edges:
            yield 10 if e[2]['is_directly_tunable'] else 0

    def transposition_weight(edges):
        for e in edges:
            yield 1000 if 0 <= hs_array_to_cents(e[2]['transposition']) < 100 else 0

    def is_sustained_voice(edges, voice):
        
        def is_sustained(edge):
            source = list(edge[0])
            ordered_source = sorted(source, key=hs_array_to_fr)    
            destination = [transpose_pitch(edge[2]['movements'][p]['destination'], edge[2]['transposition']) for p in source]
            ordered_destination = sorted(destination, key=hs_array_to_fr)
            return ordered_source[voice] == ordered_destination[voice]

        for e in edges:
            yield 10 if is_sustained(e) else 0

    def voice_crossing_weights(edges):
        
        def has_voice_crossing(edge):
            source = list(edge[0])
            ordered_source = sorted(source, key=hs_array_to_fr)    
            source_order = [ordered_source.index(p) for p in source]
            destination = [transpose_pitch(edge[2]['movements'][p]['destination'], edge[2]['transposition']) for p in source]
            ordered_destination = sorted(destination, key=hs_array_to_fr)
            destination_order = [ordered_destination.index(p) for p in destination]
            return source_order != destination_order

        for e in edges:
            yield 10 if not has_voice_crossing(e) else 0

    def is_bass_rooted(chord):
        return max([sum(abs(p) for p in collapse_pitch(pitch_difference(chord[0], p))) for p in chord[1:]]) == 1
        
    check_graph = graph.copy()
    #next_node = choice([node for node in graph.nodes() if is_bass_rooted(node)])
    next_node = choice(list(graph.nodes()))
    print(next_node)
    check_graph.remove_node(next_node)
    for node in graph.nodes(data=True):
        node[1]['count'] = 1
    path = []
    while (nx.number_of_nodes(check_graph) > 0) and (len(path) < 500):
        out_edges = list(graph.out_edges(next_node, data=True))
        #print([l for l in zip(movement_size_weights(out_edges), hamiltonian_weights(out_edges))])
        factors = [
            movement_size_weights(out_edges), 
            hamiltonian_weights(out_edges), 
            #contrary_motion_weights(out_edges), 
            #is_directly_tunable_weights(out_edges),
            voice_crossing_weights(out_edges),
            #transposition_weight(out_edges)
            #is_sustained_voice(out_edges, 0)
        ]
        weights = [prod(a) for a in zip(*factors)]
        #print(weights)
        edge = choices(out_edges, weights=weights)[0]
        #print(edge)
        #edge = random.choice(out_edges)
        next_node = edge[1]
        node[1]['count'] += 1
        path.append(edge)
        if next_node in check_graph.nodes:
            check_graph.remove_node(next_node)
    print(len(check_graph.nodes()))
    return path
    
path = stochastic_hamiltonian(graph)
#for edge in path:
#    print(edge)
write_chord_sequence(path_to_chords(path, root))
len(path)


# In[212]:


path = stochastic_hamiltonian(graph)
#for edge in path:
#    print(edge)
write_chord_sequence(path_to_chords(path, root))
len(path)

