import sys
import csv
import os
import glob
import pandas as pd
import numpy as np
from math import pi
from datetime import datetime as dt
from bokeh.io import output_file,show
from bokeh.models import DatetimeTickFormatter,ColumnDataSource,CircleCross
from bokeh.models import HoverTool
import bokeh.plotting as plot
from bokeh.transform import jitter
from bokeh.layouts import column,gridplot

def extract_and_seperate(path,outIdiomatic,outNonidiomatic):

    # 'arg[] path' is the path that contains Siamese's CSV output
    # 'arg[] outIdiomatic' is the outputted idiomatic.csv
    # 'arg[] outNonidiomatic' is the outputted nonidiomatic
    
    pattern_tokens=[]
    siameseCSVsFiles = []  # List of the CSVs generated by Siamese

    # idiomName:marker Dictionary
    marker_dict = {
        'dict-comprehension':'circle',
        'enumerate':'triangle',
        'file-reading-statement':'square',
        'list-comprehension':'diamond',
        'if-statement':'hex',
        'string-formatting':'asterisk',
        'code-formatting':'circle_cross',
        'set':'cross',
        'tuple':'inverted_triangle',
        'variable-swapping':'square_cross'
    }

    # Get all the CSV file names in the specified path
    # r=root, d=directories, f = files
    for r, d, f in os.walk(path):
        for file in f:
            if '.csv' in file:
                siameseCSVsFiles.append(os.path.join(r, file))
    
    # Creating the header row for outIdiomatic and outNonidiomatic
    with open(outIdiomatic, 'w', newline='') as idiomaticCSV:
        print("style|idiomName|fileName|color|marker|CommitNO|CommitID|methodName|startLine|endLine",file=idiomaticCSV)
    with open(outNonidiomatic, 'w', newline='') as nonidiomaticCSV:
        print("style|idiomName|fileName|color|marker|CommitNO|CommitID|methodName|startLine|endLine|recommend",file=nonidiomaticCSV)
    
    for f in siameseCSVsFiles:

        # Extracting repoName, commitNo, and commitID from Siamese's CSV file names
        # i.e.: horovod_447_1gtfvx.csv
        #  > file_name_splitted[0]: 'horovod'
        #  > file_name_splitted[1]: '447'
        #  > file_name_spllited[2]: '1gtfvx'
        file_name_splitted = os.path.basename(f).replace(".csv","").split("_")
        (repoName,commitNO,commitID) = file_name_splitted

        # Reading rows from a CSV
        with open(f, "r") as w:
            # data = list(csv.reader(w))
            data = w.read()
            data_splitted = data.split("@")    # Used the '@' symbol to split between each record in the CSV
            data_splitted.pop()                # An extra element is tailed at the end because of the last '@' in the last row    
         
        # Loop through each row in a CSV object 'f'
        for row in data_splitted:
            
            processed_row = []
            idiom_name_tokens = []
            
            row_splitted = row.split(',/')      # Split between the query columm and the pattern column
            path_query = row_splitted[0]        # The query part of the row
            path_pattern = row_splitted[1]      # The pattern part of the row

            # Extract the pattern type (pi/npi) and style from path_pattern
            path_pattern_tokens = path_pattern.split('.py')[0].split('/')         
            # Extract the type (first element in path_pattern_tokens)

            # ['pi','if','statement'] 
            pattern_tokens = path_pattern_tokens[len(path_pattern_tokens)-1].split('-') 
            pi_or_npi = pattern_tokens[0]
            
            # Extract the style (the remaining elements in path_pattern_tokens)
            for i in range(1,len(pattern_tokens)):
                idiom_name_tokens.append(pattern_tokens[i])
            idiom_name = '-'.join(idiom_name_tokens)
        
            # Extract the query's fileName, methodName, startLine and endLine
            path_query_tokens = path_query.split('/')
            query_file_name = path_query_tokens[len(path_query_tokens)-1].split('.py')[0]+'.py'
            query_method_name = path_query_tokens[len(path_query_tokens)-1].split('.py_')[1].split('#')[0]
            query_method_startLine = path_query_tokens[len(path_query_tokens)-1].split('.py_')[1].split('#')[1]
            query_method_endLine = path_query_tokens[len(path_query_tokens)-1].split('.py_')[1].split('#')[2]

            
            # "CommitNO,CommitID,fileName,methodName,startLine,endLine,style,idiomName"
            processed_row.append(pi_or_npi)
            processed_row.append(idiom_name)
            processed_row.append(query_file_name)
            
            # Label the color (red/green) to be used for the markers
            # 'green' for pi, 'red' for npi
            if processed_row[0] == 'pi':
                processed_row.append('green')
            else:
                processed_row.append('red')
                
            # Add the marker based on the idiomName (processed_row[1]) to processed_row[]
            processed_row.append(marker_dict.get(processed_row[1]))
            processed_row.append(commitNO)
            processed_row.append(commitID)            
            processed_row.append(query_method_name)
            processed_row.append(query_method_startLine)
            processed_row.append(query_method_endLine)

            # Final check if the row is of a PI or NPI
            # if it's a PI, append the processed_row to outIdiomatic.csv
            # if it's a NPI, add recommendPattern first then append to outNonidiomatic.csv
            if processed_row[0] == 'pi':
                with open(outIdiomatic,'a',newline='') as outIdiomaticCSV:
                    processed_row_string = '|'.join(processed_row)
                    print(processed_row_string,file=outIdiomaticCSV)
            else:
                # Extract the recommendation pattern
                recommend_pattern = path_pattern.split('#')
                # Encode tabes and newline as alphabetical '\t' and '\n' and remove the tailing '@' inherited from the Siamese's output
                recommend_pattern = recommend_pattern[len(recommend_pattern)-1].replace('\n','\\n').replace('\t','\\t').replace('@','')
                # Append recommend_pattern to process_row
                processed_row.append(recommend_pattern)

                # Append the completed processed_row to outNonidiomatic
                with open(outNonidiomatic,'a',newline='') as outNonidiomaticCSV:
                    processed_row_string = '|'.join(processed_row)
                    print(processed_row_string,file=outNonidiomaticCSV)
            

def plot_graph(idiomaticCSV,nonidiomaticCSV,siameseCSVsloc):

    outputHTMLpath = siameseCSVsloc+"/teddyPlus_result.html"
    plot.output_file(outputHTMLpath)

    idiomin = pd.read_csv(idiomaticCSV,sep='|')
    nonidiomin = pd.read_csv(nonidiomaticCSV,sep='|')

    sourceipy = ColumnDataSource(data=dict(
        styleipy = idiomin['style'].astype('str'),
        inameipy = idiomin['idiomName'].astype('str'),
        fnameipy =  idiomin['fileName'].astype('str'),
        coloripy = idiomin['color'].astype('str'),
        markeripy = idiomin['marker'].astype('str'),
        commitNOipy = idiomin['CommitNO'],
        commitIDipy = idiomin['CommitID'].astype('str'),
        methodNameipy = idiomin['methodName'].astype('str'),
        startlineipy = idiomin['startLine'],
        endlineipy = idiomin['endLine']
    ))

    recommend_whitespace =[]
    for i in nonidiomin['recommend']:
        j = i.replace("\\n","\n")
        k = j.replace("\\t","\t")
        recommend_whitespace.append(k)

    sourcenipy = ColumnDataSource(data=dict(
        stylenipy = nonidiomin['style'].astype('str'),
        inamenipy = nonidiomin['idiomName'].astype('str'),
        fnamenipy =  nonidiomin['fileName'].astype('str'),
        colornipy = nonidiomin['color'].astype('str'),
        markernipy = nonidiomin['marker'].astype('str'),
        commitNOnipy = nonidiomin['CommitNO'],
        commitIDnipy = nonidiomin['CommitID'].astype('str'),
        methodNamenipy = nonidiomin['methodName'].astype('str'),
        startlinenipy = nonidiomin['startLine'],
        endlinenipy = nonidiomin['endLine'],
        # recommendnipy = nonidiomin['recommend'].astype('str')
        recommendnipy = pd.Series(recommend_whitespace).astype('str')
    ))

    filename_pi_npi = []

    for f in idiomin['fileName']:
        filename_pi_npi.append(f)

    for f in nonidiomin['fileName']:
        filename_pi_npi.append(f)

    p = plot.figure(
        y_range = pd.Series(filename_pi_npi).unique().astype('str'),
        #y_range = nonidiomin['fileName'].unique().astype('str'),
        plot_width = 1500,
        title="Analysis of Good (Idiomatic) and Bad (Nonidiomatic) Python Code Patterns",
        tools=['wheel_zoom','pan','reset']
    )

    ip_plot = p.scatter(
                x='commitNOipy',
                y=jitter('fnameipy',width=0.1, range=p.y_range),
                marker='markeripy',
                size=10,
                source=sourceipy,
                color='coloripy',
                alpha=0.2,
                legend_label='Idiomatic (IP)'
                )
    ip_hover = HoverTool(
        tooltips=[
            ("style","@styleipy"),
            ("idiom type", "@inameipy"),
            ("file name","@fnameipy"),
            ("method name","@methodNameipy"),
            ("commit ID","@commitIDipy"),
            ("start line","@startlineipy"),
            ("end line","@endlineipy")
        ],
        renderers=[ip_plot]
    )
    p.add_tools(ip_hover)
    
    nip_plot = p.scatter(
                    x='commitNOnipy',
                    y=jitter('fnamenipy', width=0.1, range=p.y_range),
                    marker = 'markernipy',
                    size = 10,
                    source=sourcenipy,
                    color='colornipy',
                    alpha=0.2,
                    legend_label="Nonidiomatic (NIP)"
                )  
    nip_hover = HoverTool(
        tooltips=[
            ("style","@stylenipy"),
            ("idiom type","@inamenipy"),
            ("file name","@fnamenipy"),
            ("method name","@methodNamenipy"),
            ("commit ID","@commitIDnipy"),
            ("start line","@startlinenipy"),
            ("end line","@endlinenipy"),
            ("recommend","@recommendnipy")          
        ],
        renderers=[nip_plot]
    )
    p.add_tools(nip_hover)
        
    p.yaxis.axis_label = "File Name"
    p.xaxis.axis_label = "Commit Number"

    p.legend.location = "top_right"
    p.legend.click_policy="hide"
    # plot.show(p) >> Uncomment to have the graph automatically launched in a web browser

siamese_input = str(sys.argv[1])
out_idiomatic_csv = str(sys.argv[2])
out_nonidiomatic_csv = str(sys.argv[3])
print("[TeddyPlus] Starting visualization using Bokeh ...")
print("[TeddyPlus] ***Configuration***")
print("[TeddyPlus] Siamese's CSV files location: ", str(sys.argv[1]))
print("[TeddyPlus] Idiomatic CSV: ", str(sys.argv[2]))
print("[TeddyPlus] Nonidiomatic CSV: ", str(sys.argv[3]))
extract_and_seperate(siamese_input,out_idiomatic_csv,out_nonidiomatic_csv)
plot_graph(out_idiomatic_csv,out_nonidiomatic_csv,siamese_input)
print("[TeddyPlus] Your plot is ready. See "+siamese_input+"/teddyPlus_result.html")
