from matplotlib import pyplot as plt
from matplotlib import font_manager as plt_fnt
from matplotlib import rcParams as mplParams
from matplotlib import lines
from matplotlib.legend import Legend
import numpy as np


plt.style.use("ggplot")
prop = plt_fnt.FontProperties(family=['serif'])
mplParams["font.family"] = prop.get_name()
mplParams['hatch.linewidth'] = 0.5
mplParams['mathtext.default'] = "regular"
mplParams['lines.antialiased'] = False
mplParams['patch.antialiased'] = False


def _gen_dummy_data():
    out = {}
    for i in range(5):
        key = "key" + str(i)
        x = np.arange(0, 5, 0.5).reshape((-1, 1))
        data1 = np.random.normal(np.sin(x), 0.5)
        data2 = data1-0.2
        data3 = data1+0.2
        data = np.hstack((x, data2, data3, data1))
        out[key] = data
    return out


def _cm_to_inch(*tupl):
    inch = 2.54
    if isinstance(tupl[0], tuple):
        return tuple(k/inch for k in tupl[0])
    else:
        return tuple(k/inch for k in tupl)


def _plot_avg_velocity(velocity_arr):

    aoa_lim = 10
    v_mean = 5

    fig = plt.figure()
    fig.set_size_inches(_cm_to_inch(12, 8))
    ax = fig.add_subplot(111)
    title_string = "Radial Velocity of Airflow During CFD Simulations"
    ax.set_title(title_string, fontsize="small")

    marker_style = dict(linestyle='-', marker='', markersize=5, fillstyle='none', color='C0')
    legend1_style = dict(marker='', color='C0', linestyle='-', fillstyle='none')
    fill_style = dict(alpha=0.5, edgecolor='C0', facecolor='C0')

    vel_err = {}
    for key in velocity_arr:
        vel_err[key] = _get_err(velocity_arr[key])

    mean_handles = {}
    leg_names = []
    leg_handles = []
    i = 0

    for key in vel_err:

        legend1_style['color'] = "C" + str(i)
        leg_handle = lines.Line2D([], [], **legend1_style)
        leg_handles.append(leg_handle)
        leg_str = "Radius: %s m" % key
        leg_names.append(str(leg_str))

        marker_style["color"] = "C" + str(i)
        fill_style["edgecolor"] = "C" + str(i)
        fill_style["facecolor"] = "C" + str(i)
        mean_handles[key] = ax.plot(vel_err[key][:, 0], vel_err[key][:, -1], **marker_style)
        ax.fill_between(vel_err[key][:, 0],  vel_err[key][:, 1], vel_err[key][:, 2], **fill_style)
        i += 1

    ax2 = ax.twinx()
    ax_limits = ax.get_ylim()
    ax.set_ylim(ax_limits[0], ax_limits[1] + 0.2)
    ax_limits = ax.get_ylim()
    ax2_limits = _vel_to_aoa(v_mean, ax_limits)
    ax2.set_ylim(ax2_limits[0], ax2_limits[1])
    ax2.grid(b=None)

    ax.autoscale(False)
    ax2.autoscale(False)

    ax2_xlim = ax2.get_xlim()
    ax2.fill([ax2_xlim[0], ax2_xlim[1], ax2_xlim[1], ax2_xlim[0]], [aoa_lim, aoa_lim, ax2_limits[1], ax2_limits[1]],
             fill=False, hatch='\\', edgecolor='0')

    ax2.plot([-10, 10], [0, 0], linestyle='-.', color=[0, 0, 0], linewidth=0.5)

    leg1 = Legend(ax, leg_handles, leg_names, fontsize="small", frameon=False,
                  loc=5, ncol=1)
    ax.add_artist(leg1)

    ax.set_xlabel("Distance Above Rotors (m)", fontsize="small")
    ax.set_ylabel("Radial Velocity ($ms^{-1}$)", fontsize="small")
    ax2.set_ylabel("Equivalent Airflow Angle of Attack ($^{\circ}$)", fontsize="small")

    plt.tight_layout()
    plt.show()
    return


def _get_err(vel_arr):
    err_arr = vel_arr[:, 1:(-1-1)]
    avg_arr = vel_arr[:, -1]
    z_arr = vel_arr[:, 0]
    rows, cols = err_arr.shape
    min_arr = np.zeros([rows])
    max_arr = np.zeros([rows])
    vel_err_arr = np.zeros([rows, 4])
    for i in range(rows):
        min_arr[i] = np.min(err_arr[i, :])
        max_arr[i] = np.max(err_arr[i, :])
    vel_err_arr[:, 0] = z_arr
    vel_err_arr[:, 1] = min_arr
    vel_err_arr[:, 2] = max_arr
    vel_err_arr[:, 3] = avg_arr
    return vel_err_arr


def _vel_to_aoa(v_mean, v):
    if isinstance(v, tuple):
        return tuple(np.arctan(i/v_mean)*180/np.pi for i in v)
    else:
        return np.arctan(v/v_mean)


dat = _gen_dummy_data()
_plot_avg_velocity(dat)