"""
This script only works with guitar tabs provided by Songsterr.com.
Check you have all the modules installed.
"""


import bs4, requests, sys, math, os
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd

link = 'https://www.songsterr.com/a/wsa/mastodon-crack-the-skye-tab-s24134t0'
name = 'Some name for the title of the heatmap'
# get the tab from songster
try:
    res = requests.get(link)
except:
    print('request failed')
    sys.exit()

soup = bs4.BeautifulSoup(res.content, 'html5lib')
text = soup.prettify()



# the following represents the strings
# index 0 is string 1 = e, index 5 = E
strings = [{}, {}, {}, {}, {}, {}]
failed = []

# add all the frets. -1 = muted note
for i in range(0, 6):
    for j in range(-1, 25):
        strings[i][j] = 0

def string_and_fret_finder(text, marker):
    # find tab sentence
    p0 = text.find('<text class="_3ti_Wm _1NQZqB"', marker) + len('<text class="_3ti_Wm _1NQZqB"')
    
    # get the string
    p1 = text.find('y="', p0)+len('y="')
    p2 = text.find('"',p1)
    string = text[p1:p2]
    string = int(string)
    string = int((string)/12)

    # get the fret
    p3 = text.find('>', p2)+len('>')
    p4 = text.find('<',p3)
    rawfret = text[p3:p4]
    
    # get rid of all the empty space and lines from rawfret
    fret = []
    for i in rawfret:
        if i == 'X':
            fret.append(i)
        else:
            try:
                int(i)
                fret.append(i)
            except:
                continue
    fret = ''.join(fret)
    if fret == 'X':
        fret = -1
    else:
        fret = int(fret)
    return (string, fret, p4)


table = pd.DataFrame(columns = ("String", "Fret", "Count"))
marker = 0
while text.find('<text class="_3ti_Wm _1NQZqB"', marker) != -1:
    (string, fret, p4) = string_and_fret_finder(text, marker)
    marker = p4
    strings[string][fret] += 1

row = 1
# log version
for i in range(0,6):
    for j in strings[i]:
        x = strings[i][j]
        
        # comment out the next two lines if you don't want the log of the count
        if x != 0:
            x = math.log10(x)
        
        table.loc[row] = int(i+1), int(j), x
        row +=1
L = 485
W = 50
scale = 19
heat = table.pivot("String", "Fret", "Count")
plt.figure(figsize = (L/scale, W/scale))
ax = sns.heatmap(heat)
ax = sns.heatmap(heat, linewidths=.7).set_title(name + " - Log Count")
fig = ax.get_figure()

TotalNotes = 0
for i in range(0,6):
    for j in strings[i]:
        TotalNotes += strings[i][j]
print('Total notes: ' + str(TotalNotes))


#import os
#os.chdir('C:\\Users\\Rory\\.spyder-py3\\tab analysis')
#fig = ax.get_figure()
#fig.savefig(str(name)+".png")