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
