
try:
    import psyco
    psyco.profile()
except:
    pass

# TODO: break parts -- or just unroll everything

import abc, fraction

import sys, os, math, optparse

def notes(text):
    return [ item[0] for item in abc.read(text) ]

def note_range(text, extra=0):
    low, high = notes(text)
    return range(low-extra,high+1+extra)

# Recorder ranges
#TODO: make this configurable
#treble = note_range("C d'", 0)
#alto = note_range("F, g", 0)
#tenor = note_range("C, d", 0)
#bass = note_range("F,, F", 0)

# Voice ranges
treble = note_range("C g'", 0)
alto = note_range("G, c", 0)
tenor = note_range("C, g", 0)
bass = note_range("A,, C", 0)


def possible_chords(top):
    for item1 in alto:
        if top != None and item1 >= top: continue
#        if item1 == final: continue
        for item2 in tenor:
            if item2 >= item1: continue
#            if item2 == final: continue
            for item3 in bass:
                if item3 >= item2: continue
#                if item3 == final: continue
                yield (top, item1, item2, item3)


minifactor = 32
minitone = 2**(1.0/(12*minifactor))

r = minifactor / 7.0 #6.0
#pure = N.ceil( minitone**(N.arange(10000)-r) ) <= N.floor( minitone**(N.arange(10000)+r) )
pure = [ math.ceil(minitone**(i-r)) <= math.floor(minitone**(i+r)) for i in xrange(10000) ]
del r

pure_indicies = [ i for i in xrange(len(pure)) if pure[i] ]

def root(chord):
    chord = [ int(item*minifactor) for item in chord ]
    base = reduce(min, chord)
    for index in pure_indicies:
        possible_root = base-index
        for tone in chord:
            if not pure[tone-possible_root]: break
        else:
            return float(possible_root)/minifactor
    
def chord_cost(chord):
    chord = [ item for item in chord if item != None ]
    neg_chord = [ -item for item in chord ]
    return -(root(neg_chord)+root(chord)) \
           #+ (-reduce(min,chord)+reduce(max,chord)) *0.5 #/1.5

def chord_cost_etc_a(chord, final):
    chord = [ item for item in chord if item != None ]
    
    #sum = 0.0
    #for i in xrange(len(chord)-1):
    #    sum += ((chord[i]-chord[i+1])/6.0) ** 2
    sum = (-reduce(min,chord)+reduce(max,chord)) * 0.5 #0.7

    chord = chord + [final]
    neg_chord = [ -item for item in chord ]
    r = root(chord)
    nr = root(neg_chord)
#    return -(root(neg_chord)+root(chord)) + sum\
#           , r,nr #/1.5
    return -(nr+r) + sum, r,nr

def chord_cost_etc(chord, final):
    c1, r1, nr1 = chord_cost_etc_a(chord, final)
    c2, r2, nr2 = chord_cost_etc_a(chord, final+7)
#    return (c1*c2)**0.5, (r1, nr1, r2, nr2)
    return (c1+c2)*0.5, (r1, nr1, r2, nr2)

#    r = min(r1,r2)
#    nr = max(nr1, nr2)
#    return (-reduce(min,chord)+reduce(max,chord)) * 0.5 - r - nr, (r, nr, r, nr)


if 0:
 pure = [ ]
 for i in xrange(1,41):
    r = minifactor*4.0/i #6.0
    pure.append([ math.ceil(minitone**(i-r)) <= math.floor(minitone**(i+r)) for i in xrange(10000) ])

 def root(chord, pure):
    chord = [ item*minifactor for item in chord ]
    base = reduce(min, chord)
    for index in pure_indicies:
        possible_root = base-index
        for tone in chord:
            if not pure[tone-possible_root]: break
        else:
            return float(possible_root)/minifactor

 def chord_cost(chord):
    chord = [ item for item in chord if item != None ]
    neg_chord = [ -item for item in chord ]
    
    total = 0.0
    a = -reduce(min,chord)+reduce(max,chord)
    for item in pure:
        total += -(root(neg_chord, item)+root(chord, item)) #+ a
    #total /= len(pure)
    return total/16.0 + a
    #return -(root(neg_chord)+root(chord)) \
    #       + (-reduce(min,chord)+reduce(max,chord)) *0.5 #/1.5


#tonic_power = 1.0 #0.9

cost_cache = { }
def get_costs(top, final):
    if (top, final) not in cost_cache:
        #result = [ (chord_cost((top,tonic)+item)*tonic_power+(1.0-tonic_power)*chord_cost((top,)+item), item) for item in possible_chords(top) ]
#        result = [ (chord_cost_etc((top,tonic)+item), item) for item in possible_chords(top) ]
        result = [ (chord_cost_etc((top,)+item, final), item) for item in possible_chords(top) ]
        result.sort()
        
        cost_cache[(top,final)] = result #[:500]
    
    return cost_cache[(top,final)]

def transition_cost(chord1, chord2, bar_position, beat, allotment, roots1,roots2):
    #return 0.0
#    chord1 = chord1 + (final,)
#    chord2 = chord2 + (final,)

#    nchord1 = [ -i for i in chord1 ]
#    nchord2 = [ -i for i in chord2 ]
    
#    a = root(chord1) #- reduce(max,chord1)
#    b = root(nchord1) #- reduce(max,chord2)
#    aa = root(chord2) #- reduce(max,nchord1)
#    bb = root(nchord2) #- reduce(max,nchord2)
    #total = ((a-aa)/8.0)**2 + ((b-bb)/8.0)**2
    #aaa = a-aa #Orthogonal to actual consonance (good!) ... er ... sideways
    #bbb = b-bb
    
    #def dfunc(d): return math.log((d/2.0)**2+1.0)*0.5
    def dfunc(d): return (d/6.0)**2
    
#    aaa = a-b
#    bbb = aa-bb
#    total = dfunc(aaa-bbb)
#    aaa = a+b
#    bbb = aa+bb
#    if aaa < bbb: total += dfunc(aaa-bbb)
    #total = (dfunc(a-aa)+dfunc(b-bb)) #/2.0
#    total = dfunc(a+b-aa-bb)
    
    total = 0.0
#    if a > aa: total += dfunc(a-aa)*2.0
#    if b > bb: total += dfunc(b-bb)*2.0
#    if a+b > aa+bb: total += 12 #(aa+bb)-(a+b)
    
    # Encourages cadences, apparently
#    if a > aa: total += (a-aa)*0.75 * 2.0
#    if b > bb: total += (b-bb)*0.75 * 2.0

#    if roots1[0]+roots1[1] < roots1[2]+roots1[3]:
#        r = [0,1]
#    else:
#        r = [2,3]
    if roots1[0] < roots1[2]: r = [0]
    else: r = [2]
    if roots1[1] < roots1[3]: r += [1]
    else: r += [3]
    
    #r = [2,3]
    for i in r: #xrange(len(roots1)):
        if roots1[i] > roots2[i]: total += (roots1[i]-roots2[i]) * 0.75
            
    def dd(chord):
        nchord = [ -item for item in chord ] 
        return -(root(chord)+root(nchord))
#    total += dd(a,aa)
#    total += dd(b,bb)

#    if chord1 != chord2:
#        total += (dd(chord1 + chord2) - dd(chord2))*0.5
#        total += (dd((a,-b,aa,-bb)) - dd((a,-b)))*0.5
        
#    total += (2*dd(chord1 + chord2)-dd(chord1)-dd(chord2))*0.5
#    total += dd((a,-b,aa,-bb))
#    for i in xrange(1, len(chord1)):
#        if chord1[i] != chord2[i]:
#            total += 2
            
#    return total
    
#    if a > aa: 
#        total += 12 #((a-aa)**2.0)*0.1
#    if b > bb: 
#        total += 12 #((b-bb)**2.0)*0.1
    #return total
    
    #total += (abs(a-aa)+abs(b-bb))*0.5
    
    #total = ((aaa-bbb)/6.0)**2
    #total = (abs(a-b)+abs(aa-bb))/3.0
    #return total

    #total = 0.0
#    for i in xrange(1, len(chord1)):
#        if chord1[i] != chord2[i]:
#            best = dfunc(chord1[i]-chord2[i])
#            for j in xrange(0,len(chord1)):
                #this = dfunc(chord1[i]-chord2[j]) #tree rooted at final
                #if this < best: best = this
#                if chord1[j] != None:
#                    this = dfunc(chord1[j]-chord2[i]) #tree rooted at start
            #        if this < best: best = this
#            total += best

#            best = dfunc(chord1[i]-chord2[i])
#            for j in xrange(0,len(chord1)):
#                this = dfunc(chord1[j]-chord2[i]) #tree rooted at start
#                if this < best: best = this
#            total += best

            
#            total += dfunc(chord1[i]-chord2[i])
            #total += ((chord1[i]-chord2[i])/8.0) ** 2
#            total += ((chord1[i]-chord2[i])/4.0) ** 2
#            total += ((chord1[i]-chord2[i])/2.0) ** 2
            
            #if bar_position.den > 2:
            #    total += 64
            #elif bar_position.den == 2:
            #    total += 2
            #elif bar_position.num not in [0, 2]:
            #    total += 4

#    for i in xrange(len(chord1)):
#        for j in xrange(i+1,len(chord1)):
#            a = chord2[i]-chord1[i]
#            b = chord2[j]-chord1[j]
#            if a and a == b:
#                total += 12

    ups = 0
    downs = 0
    for i in xrange(len(chord1)):
        if chord1[i] != chord2[i]:
            if chord2[i] != None and chord1[i] != None and chord2[i] > chord1[i]:
                ups += 1
            else:
                downs += 1

    total += 2*abs(ups-downs)
    
#    if not ups or not downs:
#        total += 12
#    sum1 = 0.0
#    sum2 = 0.0
#    for i in xrange(len(chord1)):
#        if chord1[i] != None and chord2[i] != None:
#           sum1 += chord1[i]
#           sum2 += chord1[i]
#    total += dfunc(sum1-sum2)

    return total


    if chord1[0] != chord2[0]:
        m = 1
    else:
        m = 0
            
    n = 0
    for i in xrange(len(chord1)):
        if chord1[i] != chord2[i]:
            n += 1            
    
    if m > allotment:
        allotment = m
        
    if n > allotment:
        total += (n-allotment)*24

    return total

def transition_ok(chord1, chord2, allotment):        
#    for i in xrange(len(chord1)):
#      for j in xrange(i+1,len(chord1)):
#          if chord1[i] != None and chord2[i] != None and chord1[i] == chord1[j]:
#              return False
        
    if chord1[0] != chord2[0]:
        m = 1
    else:
        m = 0
            
    n = 0
    ups = 0
    downs = 0
    for i in xrange(len(chord1)):
        if chord1[i] != chord2[i]:
            n += 1
            if chord2[i] != None and chord1[i] != None and chord2[i] > chord1[i]:
                ups += 1
            else:
                downs += 1
                
    if m > allotment:
        return n <= m
        
    #if ups and downs:
    #    n -= 1
            
    return n <= allotment
#    return (ups <= allotment and downs <= allotment) and n <= allotment+1
    
#    if n > allotment:
#        total += 64
#            
#    return total

#def beam_search(music):
#    beam_size = 1000 #50
#        
#    beam = [(0.0, None, [[],[],[],[]])] #[(cost, last_chord, [musics])]
#    
#    for i in xrange(len(music)):
#        if i and music[i-1].part != music[i].part:
#            # Collapse previous section to best
#            beam = [(beam[0][0], None, beam[0][2])]
#    
#        if not isinstance(music[i], abc.Note): # or music[i].value == None:
#            for item in beam:
#                for j in xrange(4):
#                    item[2][j].append(music[i])
#            continue
#    
#        new_beam = [ ]
#        
#        n = (i*20+10) // len(music)
#        sys.stderr.write("\r   [%s|%s] " % (" "*n," "*(19-n)))
#        sys.stderr.flush()
#        
#        if music[i].bar_position.multiple_of(music[i].emphasis):
#            allotment = 4
#        elif music[i].bar_position.multiple_of(music[i].beat):
#            allotment = 2
#        elif music[i].bar_position.multiple_of(music[i].shatter):
#            allotment = 1
#        else:
#            allotment = 0
#            
#        #if music[i].bar_position.multiple_of(music[i].beat): ...use shatter unit
#        #    allotment = 4
#        #elif music[i].bar_position.den > 2:
#        #    allotment = 0
#        #elif music[i].beat.den == 1 and music[i].bar_position.den != 1:
#        #    allotment = 1
#        #else:
#        #    allotment = 2
#        
#        for (cost, r,nr), chord in get_costs(music[i].value, music[i].final):
#            cost = cost * float(music[i].length)
#            
#            if len(new_beam) == beam_size and \
#               beam[0][0]+cost > new_beam[-1][0]:
#                break
#            
#            best_value = None
#            for item in beam:
#                this_value = item[0] + cost
#                
#                if best_value != None and this_value > best_value: break
#            
#                if item[1]:
#                    if not transition_ok(item[1][0], chord, allotment):
#                        continue
#                
#                    this_value += music[i].horizontality * \
#                                  transition_cost(item[1][0], chord, music[i].bar_position, music[i].beat,
#                                  allotment, item[1][1], (r,nr))
#                    
#                if best_value == None or this_value < best_value:
#                    best_value = this_value
#                    best_item = item
#            
#            if best_value != None:
#                position = len(new_beam)
#                while position > 0 and new_beam[position-1][0] > best_value:
#                    position -= 1
#                if position < beam_size:
#                    best = (best_value, (chord, (r, nr)),
#                        [ best_item[2][j] + [music[i].but_with(value=chord[j])] 
#                          for j in xrange(4) ])
#              
#                    new_beam = (new_beam[:position] + [best] + new_beam[position:])[:beam_size]

#        beam = new_beam
#        
#    sys.stderr.write("\r" + " "*30 + "\r")
#    sys.stderr.flush()
        
#    return beam[0][2]

def beam_search(music):
    beam_size = 500
        
    beam = [(0.0, None, [[],[],[],[]])] #[(cost, last_chord, [musics])]
    
    for i in xrange(len(music)-1,-1,-1):
        if i+1 < len(music) and music[i+1].part != music[i].part:
            # Collapse previous section to best
            beam = [(beam[0][0], None, beam[0][2])]
    
        if not isinstance(music[i], abc.Note): # or music[i].value == None:
            for item in beam:
                for j in xrange(4):
                    item[2][j].insert(0, music[i])
            continue
            
        note = music[i].value
        if note == None:
            # Rests are like a continuation of the previous note
            j = i-1
            while note == None and j > 0:
                if isinstance(music[j], abc.Note):
                    note = music[j].value    
                j -= 1
    
        new_beam = [ ]
        
        n = (i*20+10) // len(music)
        sys.stderr.write("\r   [%s|%s] " % (" "*n," "*(19-n)))
        sys.stderr.flush()
        
        bpl = music[i].bar_position + music[i].length
        if bpl.multiple_of(music[i].emphasis):
            allotment = 4
        elif bpl.multiple_of(music[i].beat):
            allotment = 2
        elif bpl.multiple_of(music[i].shatter):
            allotment = 1
        else:
            allotment = 0
            
        #if music[i].bar_position.multiple_of(music[i].beat): ...use shatter unit
        #    allotment = 4
        #elif music[i].bar_position.den > 2:
        #    allotment = 0
        #elif music[i].beat.den == 1 and music[i].bar_position.den != 1:
        #    allotment = 1
        #else:
        #    allotment = 2
        
        for (cost, roots), chord in get_costs(note, music[i].final):
            cost = cost * float(music[i].length)
            
            if len(new_beam) == beam_size and \
               beam[0][0]+cost > new_beam[-1][0]:
                break
            
            best_value = None
            for item in beam:
                this_value = item[0] + cost
                
                if best_value != None and this_value > best_value: break
            
                if item[1]:
                    if not transition_ok(chord, item[1][0], allotment):
                        continue
                
                    this_value += music[i].horizontality * \
                                  transition_cost(chord, item[1][0], music[item[1][2]].bar_position, music[item[1][2]].beat,
                                  allotment, roots, item[1][1])
                    
                if best_value == None or this_value < best_value:
                    best_value = this_value
                    best_item = item
            
            if best_value != None:
                position = len(new_beam)
                while position > 0 and new_beam[position-1][0] > best_value:
                    position -= 1
                if position < beam_size:
                    best = (best_value, (chord, roots, i),
                        [ [music[i].but_with(value=chord[j])] + best_item[2][j] 
                          for j in xrange(4) ])
              
                    new_beam = (new_beam[:position] + [best] + new_beam[position:])[:beam_size]

        beam = new_beam
        
    sys.stderr.write("\r" + " "*30 + "\r")
    sys.stderr.flush()
        
    return beam[0][2]


parser = optparse.OptionParser("usage: python %prog [options] ABCFILE")
parser.add_option("-e", None, type="string", dest="indicies",
                  help="Harmonize only listed index numbers, eg -e 2,3,4")
                  
options, args = parser.parse_args()

if len(args) < 1:
    parser.print_help()
    sys.exit(1)
if len(args) > 1:
    parser.error("please give one file only")

filename = args[0]

print >> sys.stderr, "Writing to out.abc"
output_file = open("out.abc", "wt")

songs = abc.read_songs(open(filename, "rt"))

for song in songs:
    if options.indicies:
        selected = False
        for item in options.indicies.split(","):
            if song.index == item:
                selected = True
        if not selected: continue

    print song.index + ". " + song.title
    for line in song.header[:-1]:
        output_file.write(line + "\n")
        
    output_file.write("V: 1 clef=treble\nV: 2 clef=treble merge\nV: 3 clef=bass\nV: 4 clef=bass merge\n")
    
    output_file.write(song.header[-1] + "\n")
        
    music = abc.shatter(song.music)    
    result = beam_search(music)
    
    result = [ abc.join_up(item) for item in result ]

#    if not song.parts:    
#        output_file.write(song.music_to_abc(song.music).strip())
#        for i in xrange(1,4):
#            temp = [ item for item in result[i] 
#                     if not isinstance(item, abc.Annotation) ]
#            output_file.write("[V:%d] " % (i+1) + 
#                              song.music_to_abc(temp).strip())
#    else:
    for part in [None] + song.parts:
        if part != None:
            output_file.write(" [P:%s] " % part)
        temp = [ item for item in song.music if item.part == part ]
        output_file.write(" [V:1] " + song.music_to_abc(temp))
        for i in xrange(1,4):
            temp = [ item for item in result[i] 
                     if item.part == part ]
#                     if item.part == part and not isinstance(item, abc.Annotation) ]
            output_file.write(" [V:%d] " % (i+1) + 
                              song.music_to_abc(temp))
    
    output_file.write("\n")
    output_file.flush()
    
output_file.close()

#abcm2ps out.abc && abc2midi out.abc && sleep 0.25 && timidity -EI0 out1.mid
