import os
import pandas as pd
import numpy as np
import matplotlib as mpl
from matplotlib import rcParams
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pgf import FigureCanvasPgf
from math import pi
mpl.backend_bases.register_backend('pdf', FigureCanvasPgf)
 
plt.style.use('tableau-colorblind10')
 
rcParams['text.usetex'] = True
rcParams['font.size'] = 18
rcParams['lines.linewidth'] = 1.0
rcParams['axes.grid'] = True
rcParams['grid.linewidth'] = 0.5
rcParams['grid.linestyle'] = 'dotted'

def set_iteration_fig_properties():
    rcParams["figure.figsize"] = (3.5,2.5)
    rcParams["text.usetex"] = True
    rcParams["figure.dpi"] = 200
    rcParams["xtick.labelsize"] = 7
    rcParams["font.size"] = 8


def plot_iterations(dFrame, testName, titleName, outputVar="GEOP", plotPolar=False):

    set_iteration_fig_properties()

    # Output directory 
    outputDir = os.curdir
    if (os.environ.get(outputVar) is not None):
        outputDir = os.environ[outputVar] 

    # Iterations (\alpha_c)
    dFrameMean = dFrame.groupby(['ALPHA_STAR']).ITERATIONS.mean()
    dFrameMedian = dFrame.groupby(['ALPHA_STAR']).ITERATIONS.median()
    dFrameMax = dFrame.groupby('ALPHA_STAR').ITERATIONS.max()
    dFrameMin = dFrame.groupby('ALPHA_STAR').ITERATIONS.min()
    
    fig1 = plt.figure()
    ax1 = plt.subplot(111)
    
    ax1.set_title(titleName + " ITERATIONS")

    ax1.plot(dFrameMax.index, dFrameMax, label=r"$\sup(N_{\alpha_c})$") 
    ax1.plot(dFrameMean.index, dFrameMean, linestyle='-', marker='*', label=r"$mean(N_{\alpha_c})$") 
    ax1.plot(dFrameMedian.index, dFrameMedian, label=r"$median(N_{\alpha_c})$") 
    ax1.plot(dFrameMin.index, dFrameMin, label=r"$\inf(N_{\alpha_c})$")
    ax1.set_ylabel(r"$N_{\alpha_c}$")
    ax1.set_xlabel(r"$\alpha_c$")
    
    # Shrink current ax1is's height by 10% on the bottom
    box1 = ax1.get_position()
    ax1.set_position([box1.x0, box1.y0 + box1.height * 0.2,
                     box1.width, box1.height * 0.8])

    # Put a legend below current ax1is
    ax1.legend(loc='upper center', bbox_to_anchor=(0.5, -0.2),
              fancybox=True, shadow=True, ncol=2)

    figPath = os.path.join(outputDir, "figures")

    if (os.path.exists(figPath)):
        fig1.savefig(os.path.join(figPath, "%s-ITERATIONS-ALPHA.pdf" % testName), 
                     bbox_inches="tight")
    

    # Plottin Iterations = Iterations(\phi)
    dFrameMaxPhi = dFrame.groupby('PHI').ITERATIONS.max()
    dFrameMeanPhi = dFrame.groupby('PHI').ITERATIONS.mean()
    dFrameMedianPhi = dFrame.groupby('PHI').ITERATIONS.median()
    dFrameMinPhi = dFrame.groupby('PHI').ITERATIONS.min()
    
    fig2 = plt.figure()
    ax2 = plt.subplot(111, projection='polar')
    
    ax2.plot(dFrameMaxPhi.index, dFrameMaxPhi, label=r"$\sup(N_\phi)$")
    ax2.plot(dFrameMeanPhi.index, dFrameMeanPhi,label=r"$mean(N_\phi)$")
    ax2.plot(dFrameMedianPhi.index, dFrameMedianPhi,label=r"$median(N_\phi)$")
    ax2.plot(dFrameMinPhi.index, dFrameMinPhi, label=r"$\inf(N_\phi)$")
    ax2.set_title(titleName + " ITERATIONS $(\\phi)$", y = 1.14)

    ax2.set_xlabel(r"$\phi$")
    ax2.xaxis.set_label_coords(-0.1, 0.6)
    ax2.set_ylabel(r"$N_\phi$", rotation='horizontal')
    ax2.yaxis.set_label_coords(1.1, 0.7)
       
    # Shrink current axis by 20%
    box2 = ax2.get_position()
    ax2.set_position([box2.x0, box2.y0, box2.width * 0.8, box2.height * 0.8])

    # Put a legend to the right of the current ax2is
    ax2.legend(loc='center left', bbox_to_anchor=(1.1, 0.25))
        
    figPathName = os.path.join(outputDir, "figures")
    if (os.path.exists(figPathName)):
        phiFileName = os.path.join(figPathName, "%s-ITERATIONS-PHI.png" % testName)
        fig2.savefig(phiFileName, bbox_inches="tight", dpi=200)

    # Iterations (\theta)
    dFrameMaxTheta = dFrame.groupby('THETA').ITERATIONS.max()
    dFrameMeanTheta = dFrame.groupby('THETA').ITERATIONS.mean()
    dFrameMedianTheta = dFrame.groupby('THETA').ITERATIONS.median()
    dFrameMinTheta = dFrame.groupby('THETA').ITERATIONS.min()

    fig3 = plt.figure()
    ax3 = plt.subplot(111, projection='polar')
    
    ax3.set_title(titleName + " ITERATIONS $(\\theta)$", y=0.9)
    ax3.plot(dFrameMaxTheta.index, dFrameMaxTheta, label=r"$\sup(N_\theta)$")
    ax3.plot(dFrameMeanTheta.index, dFrameMeanTheta, label=r"$mean(N_\theta)$")
    ax3.plot(dFrameMedianTheta.index, dFrameMedianTheta, label=r"$median(N_\theta)$")
    ax3.plot(dFrameMinTheta.index, dFrameMinTheta,label=r"$\inf(N_\theta)$")
    ax3.set_xlabel(r"$\theta$") 
    ax3.xaxis.set_label_coords(0.95, 0.75)
    ax3.set_ylabel(r"$N_\theta$", rotation='horizontal') 
    ax3.yaxis.set_label_coords(0.55,0.1) 
    # Shrink current axis by 20%
    box3 = ax3.get_position()
    ax3.set_position([box3.x0, box3.y0, box3.width * 0.8, box3.height])
    
    # Put a legend to the right of the current ax3is
    ax3.legend(loc='center left', bbox_to_anchor=(1.1, 0.5))
    
    ax3.set_thetamin(0)
    ax3.set_thetamax(180)
    ax3.set_xticks([0,pi/4,pi/2,3*pi/4,pi])

    if (os.path.exists(figPathName)):
        fig3.savefig(os.path.join(figPathName, "%s-ITERATIONS-THETA.pdf" % testName), 
                     bbox_inches="tight")
    
def plot_iteration_boxplots(cubicDf, newtonDf, polyhedronName, annotateMedian=True):
    """Plots a box plot that compares the consecutive cubic and newton cubic method in terms 
    of the iteration distribution for a polyhedron.""" 

    rcParams["figure.figsize"] = (2.5,2)
    rcParams["text.usetex"] = True
    rcParams["figure.dpi"] = 300
    rcParams["xtick.labelsize"] = 7
    rcParams["font.size"] = 8
    
    fig, ax1 = plt.subplots()

    props = dict(widths=0.7,patch_artist=True, medianprops=dict(ls='--', color="lightgray"))
    box1 = ax1.boxplot(newtonDf['ITERATIONS'], positions=[0], **props)
    for patch in box1['boxes']:
        patch.set(facecolor='C0')
    ax1.set_ylabel("NCS ITERATIONS")


    if (annotateMedian):
        for line in box1['medians']:
            # get position data for median line
            x0, y0 = line.get_xydata()[0] # start of the median line
            x1, y1 = line.get_xydata()[1] # end of the median line
            # overlay median value
            ax1.annotate(r'$\mathbf{%.1f}$' % y0, (0.5*(x0 + x1),y0))

    ax2 = ax1.twinx()
    box2 = ax2.boxplot(cubicDf['ITERATIONS'], positions=[1], **props)
    for patch in box2['boxes']:
        patch.set(facecolor='C1')
    ax2.set_ylabel('CCS ITERATIONS')

    if ("NON" in polyhedronName.upper()):
        ax1.set_yscale("log")
        ax2.set_yscale("log")
        
    if (annotateMedian):
        for line in box2['medians']:
            # get position data for median line
            x0, y0 = line.get_xydata()[0] # start of the median line
            x1, y1 = line.get_xydata()[1] # end of the median line
            # overlay median value
            ax2.annotate(r'$\mathbf{%.1f}$' % y0, (0.5*(x0+x1),y0))

    ax1.set_xticklabels(['NCS', 'CCS'])
    ax1.set_title('%s iterations' % polyhedronName)

    polyName = polyhedronName.upper().split(' ')

    figName = ""
    for pName in polyName: 
        figName = figName + pName + '_' 

    figName = figName + "ITERATIONS_BOX_PLOT.png"

    plt.savefig(figName, dpi=200)

    if "GEOP" in os.environ:
        figPath = os.path.join(os.environ["GEOP"], "figures")
        if (os.path.exists(figPath)):
            plt.savefig(os.path.join(figPath, figName),
                        bbox_inches='tight', dpi=400)

def plot_total_data(totalData):
    """Plot the total CPU time and total iterations, comparing Newton Cubic Spline and Successive Cubic Spline."""

    rcParams["figure.figsize"] = 4.5,3
    rcParams["text.usetex"] = True
    rcParams["figure.dpi"] = 200
    rcParams["xtick.labelsize"] = 7
    rcParams["font.size"] = 8

    from matplotlib.ticker import FormatStrFormatter
    
    newtonData = totalData[totalData["TEST"].str.contains("NEWTON_CUBIC_SPLINE")]
    newtonData = newtonData.reset_index(drop=True)

    cubicData = totalData[totalData["TEST"].str.contains("CONSECUTIVE_CUBIC_SPLINE")]
    cubicData = cubicData.reset_index(drop=True)
    
    # Get the colors from the color style. 
    colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
    
    cpuFig, cpuAx = plt.subplots()
    cpuAx.yaxis.set_major_formatter(FormatStrFormatter('$%.1e$'))
    iterFig, iterAx = plt.subplots()
    iterAx.yaxis.set_major_formatter(FormatStrFormatter('$%1d$'))

    cpuOffset = 0.1 * min(newtonData["AVERAGE_CPU_TIME_MICROSECONDS"].mean(), 
                          cubicData["AVERAGE_CPU_TIME_MICROSECONDS"].mean());

    iterOffset = 0.1 * min(newtonData["AVERAGE_ITERATIONS"].mean(), 
                           cubicData["AVERAGE_ITERATIONS"].mean());
    
    for i in range(len(newtonData["TEST"])):
        
        # Plot the total CPU time
        newtonCPU = newtonData["AVERAGE_CPU_TIME_MICROSECONDS"][i]
        cubicCPU = cubicData["AVERAGE_CPU_TIME_MICROSECONDS"][i]
        
        cpuAx.bar(2*i, newtonCPU, color=colors[0], label="NEWTON CUBIC SPLINE")
        cpuAx.bar(2*i+1, cubicCPU, color=colors[1], label="CONSECUTIVE CUBIC SPLINE")

        cpuAx.annotate("speedup", (2*i+0.6,cubicCPU + 3*cpuOffset), fontsize=5, color='k', 
                       weight='bold')
        cpuAx.annotate("$\mathbf{%.2f}$" % (newtonCPU / cubicCPU), 
                       (2*i+0.6,cubicCPU + cpuOffset), fontsize=8, color='k')
        
        # Plot the total number of iterations
        newtonIter = newtonData["AVERAGE_ITERATIONS"][i]
        cubicIter = cubicData["AVERAGE_ITERATIONS"][i]
        
        iterAx.bar(2*i, newtonIter, color=colors[0], label="NEWTON CUBIC SPLINE")
        iterAx.bar(2*i+1, cubicIter, color=colors[1], label="CONSECUTIVE CUBIC SPLINE")
        iterAx.annotate("speedup", (2*i+0.6,cubicIter + 3*iterOffset), fontsize=5, color='k', 
                        weight='bold')
        iterAx.annotate("$\mathbf{%.2f}$" % (newtonIter / cubicIter), 
                       (2*i+0.6,cubicIter + iterOffset), fontsize=8, color='k')
    
    
    cpuAx.set_xlabel("TET = tetrahedron, DOD = dodecahedron, \nENDO = endo dodecahedron, NPDO = non-planar dodecahedron",
                 wrap=True, ha='left', x=-0.05)
    cpuAx.set_xticks(np.arange(2*len(newtonData["TEST"]))) 
    cpuAx.set_xticklabels(["TET", "TET", "CUBE", "CUBE", "DOD", "DOD", "ENDO", "ENDO", "NPDO", "NPDO"])
    cpuAx.set_ylabel("Average CPU time in microseconds")
    cpuAx.legend(('NEWTON CUBIC SPLINE', 'CONSECUTIVE CUBIC SPLINE'))
    cpuFig.savefig("POSITIONING_CPU_TIME_MICROSECONDS_TOTAL.pdf", bbox_inches='tight')


    figPathName = "" 
    if "GEOP" in os.environ:
        figPathName = os.path.join(os.environ["GEOP"], "figures") 

    if (os.path.exists(figPathName)):
        cpuFig.savefig(os.path.join(figPathName,"POSITIONING_AVERAGE_CPU_TIME_MICROSECONDS.pdf"),
                       bbox_inches='tight')
        
        
    iterAx.set_xlabel("TET = tetrahedron, DOD = dodecahedron, \nENDO = endo dodecahedron, NPDO = non-planar dodecahedron",
                 wrap=True, ha='left', x=-0.05)
    iterAx.set_xticks(np.arange(2*len(newtonData["TEST"]))) 
    iterAx.set_xticklabels(["TET", "TET", "CUBE", "CUBE", "DOD", "DOD", "ENDO", "ENDO", "NPDO", "NPDO"])
    iterAx.set_ylabel("Average number of iterations")
    iterAx.legend(('NEWTON CUBIC SPLINE', 'CONSECUTIVE CUBIC SPLINE'))
    iterFig.savefig("POSITIONING_ITERATIONS_TOTAL.pdf", bbox_inches='tight')
    
    if (os.path.exists(figPathName)):
        iterFig.savefig(os.path.join(figPathName,"POSITIONING_AVERAGE_ITERATIONS.pdf"),
                        bbox_inches='tight')

    testNameDict = {"TETRAHEDRON_NEWTON_CUBIC_SPLINE" : "TET NCS",
                "CUBE_NEWTON_CUBIC_SPLINE" : "CUBE NCS",
                "DODECAHEDRON_NEWTON_CUBIC_SPLINE" : "DOD NCS",
                "ENDO_DODECAHEDRON_NEWTON_CUBIC_SPLINE" : "ENDO NCS",
                "NON_PLANAR_DODECAHEDRON_NEWTON_CUBIC_SPLINE" : "NPDO NCS",
                "TETRAHEDRON_CONSECUTIVE_CUBIC_SPLINE" : "TET CCS",
                "CUBE_CONSECUTIVE_CUBIC_SPLINE" : "CUBE CCS",
                "DODECAHEDRON_CONSECUTIVE_CUBIC_SPLINE" : "DOD CCS",
                "ENDO_DODECAHEDRON_CONSECUTIVE_CUBIC_SPLINE" : "ENDO CCS",
                "NON_PLANAR_DODECAHEDRON_CONSECUTIVE_CUBIC_SPLINE" : "NPDO CCS"
               };

    barLabels = ["ROOT", "POLY", "GEOM"]

    barFig = plt.figure(figsize=(5,6))
    for i,testName in enumerate(totalData["TEST"]):
        ax = plt.subplot(5,2, i+1)
        ax.set_title(testNameDict[testName])
        testData = totalData[totalData["TEST"] == testName] 
        barData = testData[["ROOT_FINDING_MICROSECONDS_AVERAGE",
                           "POLYNOMIAL_INTERPOLATION_MICROSECONDS_AVERAGE",
                           "GEOMETRICAL_INTERSECTION_MICROSECONDS_AVERAGE"]] 
        x = np.arange(len(barData.columns)) 
        ax.semilogy()
        color = "C1"
        if ("NEWTON" in testName):
            color = "C0"
        ax.bar(x, barData.to_numpy().reshape(-1,1).flatten() * 1e-06, color=color)
        ax.set_xticks(x)
        ax.set_xticklabels(barLabels)
        ax.set_ylabel("CPU time in seconds")
    barFig.tight_layout()

    barFig.savefig("POSITIONING_CPU_TIME_DISTRIBUTION.png", bbox_inches='tight')
    if (os.path.exists(figPathName)):
        barFig.savefig(os.path.join(figPathName,"POSITIONINING_CPU_TIME_DISTRIBUTION.png"),
                       bbox_inches='tight')
