import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, RegularPolygon
from matplotlib.path import Path
from matplotlib.projections.polar import PolarAxes
from matplotlib.projections import register_projection
from matplotlib.spines import Spine
from matplotlib.transforms import Affine2D
import matplotlib.colors as mcolors
import os

path_prefix = os.environ['UDG_DATA_PATH']

results_0 = [[1242.91, 849.18, 1456.2, 840.74, 1481.22, 1214.12],
            [985.76, 832.38, 1381.18, 962.47, 1534.01, 1182.93],
            [1480.1, 1050.54, 1249.99, 685.75, 1235.17, 1027.98]]
results_1 = [[1108.76, 1214.7, 870.44, 50.02, 409.31, 536.4],
            [1076.62, 1410.98, 484.67, 33, 393.84, 264.95],
            [1124.03, 854.35, 793.03, 27.03, 494.58, 329.29]]
results_2 = [[1108.76, 701.89, 85.22, 50.81, 371.68, 786.1],
            [1076.62, 390.11, 391.13, 11.7, 403.29, 780.64],
            [1124.03, 674.14, 98.84, 14.44, 223.85, 861.65]]

results = [np.array(results_0), np.array(results_1), np.array(results_2)]

res_max = max([np.max(x) for x in results])
axis_max = res_max
res_mean = [(np.mean(x, axis=0)/axis_max) for x in results]
res_std = [(np.std(x, axis=0)/axis_max) for x in results]


def radar_factory(num_vars, frame='circle'):
    """
    Create a radar chart with `num_vars` axes.

    This function creates a RadarAxes projection and registers it.

    Parameters
    ----------
    num_vars : int
        Number of variables for radar chart.
    frame : {'circle', 'polygon'}
        Shape of frame surrounding axes.

    """
    # calculate evenly-spaced axis angles
    theta = np.linspace(0, 2*np.pi, num_vars, endpoint=False)

    class RadarTransform(PolarAxes.PolarTransform):

        def transform_path_non_affine(self, path):
            # Paths with non-unit interpolation steps correspond to gridlines,
            # in which case we force interpolation (to defeat PolarTransform's
            # autoconversion to circular arcs).
            if path._interpolation_steps > 1:
                path = path.interpolated(num_vars)
            return Path(self.transform(path.vertices), path.codes)

    class RadarAxes(PolarAxes):

        name = 'radar'
        PolarTransform = RadarTransform

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            # rotate plot such that the first axis is at the top
            self.set_theta_zero_location('N')

        def fill(self, *args, closed=True, **kwargs):
            """Override fill so that line is closed by default"""
            return super().fill(closed=closed, *args, **kwargs)
        
        def fill_between(self, *args, closed=True, **kwargs):
            """Override fill so that line is closed by default"""
            return super().fill_between(closed=closed, *args, **kwargs)

        def plot(self, *args, **kwargs):
            """Override plot so that line is closed by default"""
            lines = super().plot(*args, **kwargs)
            for line in lines:
                self._close_line(line)

        def _close_line(self, line):
            x, y = line.get_data()
            # FIXME: markers at x[0], y[0] get doubled-up
            if x[0] != x[-1]:
                x = np.append(x, x[0])
                y = np.append(y, y[0])
                line.set_data(x, y)

        def set_varlabels(self, labels, fontsize):
            self.set_thetagrids(np.degrees(theta), labels, fontsize=fontsize)

        def _gen_axes_patch(self):
            # The Axes patch must be centered at (0.5, 0.5) and of radius 0.5
            # in axes coordinates.
            if frame == 'circle':
                return Circle((0.5, 0.5), 0.5)
            elif frame == 'polygon':
                return RegularPolygon((0.5, 0.5), num_vars,
                                      radius=.5, edgecolor="k")
            else:
                raise ValueError("Unknown value for 'frame': %s" % frame)

        def _gen_axes_spines(self):
            if frame == 'circle':
                return super()._gen_axes_spines()
            elif frame == 'polygon':
                # spine_type must be 'left'/'right'/'top'/'bottom'/'circle'.
                spine = Spine(axes=self,
                              spine_type='circle',
                              path=Path.unit_regular_polygon(num_vars))
                # unit_regular_polygon gives a polygon of radius 1 centered at
                # (0, 0) but we want a polygon of radius 0.5 centered at (0.5,
                # 0.5) in axes coordinates.
                spine.set_transform(Affine2D().scale(.5).translate(.5, .5)
                                    + self.transAxes)
                return {'polar': spine}
            else:
                raise ValueError("Unknown value for 'frame': %s" % frame)

    register_projection(RadarAxes)
    return theta

N = 6
theta = radar_factory(N, frame='polygon')

spoke_labels = [u'%i\u00b0'%np.round(np.degrees(x)) for x in theta]

fig, ax = plt.subplots(figsize=(9, 9), subplot_kw=dict(projection='radar'))
ref_colors = mcolors.TABLEAU_COLORS
colors = [ref_colors['tab:blue'], ref_colors['tab:red'], ref_colors['tab:orange']]

rtick_locs = np.arange(0.0, 1.0, 0.2)
rtick_labels = ['' for r in rtick_locs]

ax.set_rgrids(rtick_locs, rtick_labels, fontsize='large')
ax.set_ylim([0.0, 1.0])

for d, std, color in zip(res_mean, res_std, colors):
    ax.plot(theta, d.tolist(), color=color)

    ext_theta = theta.tolist()
    ext_theta = ext_theta + [ext_theta[0]]
    ext_d1 = (d+std).tolist()
    ext_d1 = ext_d1 + [ext_d1[0]]
    ext_d2 = (d-std).tolist()
    ext_d2 = ext_d2 + [ext_d2[0]]
    ax.fill_between(ext_theta, ext_d1, ext_d2, facecolor=color, alpha=0.25, label='_nolegend_')
ax.tick_params(pad=15)
ax.set_varlabels(spoke_labels, fontsize=17)

# add legend relative to top-left plot
labels = ('unsupervised', 'supervised w/ div', 'supervised')
legend = ax.legend(labels, loc=[0.65,0.95],
                            labelspacing=0.1, fontsize=17)

fig_path = path_prefix + 'mopo-local/fig/radar.png'
plt.savefig(fig_path)
print("Figure plotting complete!")