import matplotlib.pyplot as plt

import numpy as np


def plot_taylor_axes(axes, cax, option):
    """
    Plot axes for Taylor diagram.

    Plots the x & y axes for a Taylor diagram using the information
    provided in the AXES dictionary returned by the
    GET_TAYLOR_DIAGRAM_AXES function.

    INPUTS:
    axes   : data structure containing axes information for target diagram
    cax    : handle for plot axes
    option : data structure containing option values. (Refer to
             GET_TAYLOR_DIAGRAM_OPTIONS function for more information.)
    option['colcor']       : CORs grid and tick labels color (Default: blue)
    option['colrms']       : RMS grid and tick labels color (Default: green)
    option['colstd']       : STD grid and tick labels color (Default: black)
    option['numberpanels'] : number of panels (quadrants) to use for Taylor
                          diagram
    option['tickrms']      : RMS values to plot gridding circles from
                          observation point
    option['titlecor']     : title for CORRELATION axis
    option['titlerms']     : title for RMS axis
    option['titlestd']     : title fot STD axis

    OUTPUTS:
    ax: returns a list of handles of axis labels

    Author: Peter A. Rochford
    Acorn Science & Innovation
    prochford@acornsi.com

    Created on Dec 3, 2016

    Author: Peter A. Rochford
        Symplectic, LLC
        www.thesymplectic.com
        prochford@thesymplectic.com
    """
    ax = []
    axlabweight = 'bold'
    if option['numberpanels'] == 1:
        # Single panel

        if option['titlestd'] == 'on':
            ttt = plt.ylabel('test', fontsize=14)
            x = -0.15 * axes['rmax']
            y = 0.8 * axes['rmax']
            handle = plt.text(x, y, 'Standard Deviation', rotation=90,
                              color=option['colstd'],
                              fontweight=axlabweight,
                              fontsize=plt.get(ttt, 'fontsize'),
                              horizontalalignment='center')
            ax.append(handle)

        if option['titlecor'] == 'on':
            pos1 = 45
            DA = 15
            lab = 'Correlation Coefficient'
            c = np.fliplr([np.linspace(pos1 - DA, pos1 + DA, len(lab))])[0]
            dd = 1.1 * axes['rmax']
            for ii, ith in enumerate(c):
                handle = plt.text(dd * np.cos(ith * np.pi / 180),
                                  dd * np.sin(ith * np.pi / 180),
                                  lab[ii])
                handle.set(rotation=ith - 90, color=option['colcor'],
                           horizontalalignment='center',
                           verticalalignment='bottom',
                           fontsize=plt.get(ax[0], 'fontsize'),
                           fontweight=axlabweight)
                ax.append(handle)

        if option['titlerms'] == 'on':
            pos1 = option['tickrmsangle'] + (180 - option['tickrmsangle']) / 2
            DA = 15
            pos1 = 160
            lab = 'RMSD'
            c = np.fliplr([np.linspace(pos1 - DA, pos1 + DA, len(lab))])[0]

            # Find optimal placement of label
            itick = -1
            ratio = 1.0
            while (ratio > 0.7):
                itick += 1
                ratio = (option['axismax'] - option['tickrms']
                         [itick]) / option['axismax']
            dd = 0.7 * option['tickrms'][itick] + \
                0.3 * option['tickrms'][itick + 1]

            # Write label in a circular arc
            for ii, ith in enumerate(c):
                xtextpos = axes['dx'] + dd * np.cos(ith * np.pi / 180)
                ytextpos = dd * np.sin(ith * np.pi / 180)
                handle = plt.text(xtextpos, ytextpos, lab[ii])
                handle.set(rotation=ith - 90, color=option['colrms'],
                           horizontalalignment='center',
                           verticalalignment='top',
                           fontsize=plt.get(ax[0], 'fontsize'),
                           fontweight=axlabweight)
                ax.append(handle)

    else:
        # Double panel

        if option['titlestd'] == 'on':
            ttt = plt.ylabel('test', fontsize=14)
            x = 0
            y = -0.15 * axes['rmax']
            handle = plt.text(x, y, 'Standard Deviation', rotation=0,
                              color=option['colstd'],
                              fontweight=axlabweight,
                              fontsize=plt.get(ttt, 'fontsize'),
                              horizontalalignment='center')
            ax.append(handle)

        if option['titlecor'] == 'on':
            pos1 = 90
            DA = 25
            lab = 'Correlation Coefficient'
            c = np.fliplr([np.linspace(pos1 - DA, pos1 + DA, len(lab))])[0]
            dd = 1.1 * axes['rmax']
            for ii, ith in enumerate(c):
                handle = plt.text(dd * np.cos(ith * np.pi / 180),
                                  dd * np.sin(ith * np.pi / 180), lab[ii])
                handle.set(rotation=ith - 90, color=option['colcor'],
                           horizontalalignment='center',
                           verticalalignment='bottom',
                           fontsize=plt.get(ax[0], 'fontsize'),
                           fontweight=axlabweight)
                ax.append(handle)

        if option['titlerms'] == 'on':
            pos1 = 160
            DA = 10
            lab = 'RMSD'
            c = np.fliplr([np.linspace(pos1 - DA, pos1 + DA, len(lab))])[0]
            dd = 1.05 * option['tickrms'][0]
            for ii, ith in enumerate(c):
                xtextpos = axes['dx'] + dd * np.cos(ith * np.pi / 180)
                ytextpos = dd * np.sin(ith * np.pi / 180)
                handle = plt.text(xtextpos, ytextpos, lab[ii])
                handle.set(rotation=ith - 90, color=option['colrms'],
                           horizontalalignment='center',
                           verticalalignment='bottom',
                           fontsize=plt.get(ax[0], 'fontsize'),
                           fontweight=axlabweight)
                ax.append(handle)

    # VARIOUS ADJUSTMENTS TO THE PLOT:
    cax.set_aspect('equal')
    plt.axis('off')
    plt.gcf().patch.set_facecolor('w')

    # set axis limits
    if option['numberpanels'] == 2:
        axislim = [axes['rmax'] * x for x in [-1.15, 1.15, 0, 1.15]]
        plt.axis(axislim)
        plt.plot([-axes['rmax'], axes['rmax']], [0, 0],
                 color=axes['tc'], linewidth=2)
        plt.plot([0, 0], [0, axes['rmax']], color=axes['tc'])
    else:
        axislim = [axes['rmax'] * x for x in [0, 1.15, 0, 1.15]]
        plt.axis(axislim)
        plt.plot([0, axes['rmax']], [0, 0], color=axes['tc'],
                 linewidth=2)
        plt.plot([0, 0], [0, axes['rmax']], color=axes['tc'],
                 linewidth=2)

    return ax
