import numpy
from matplotlib import pyplot
from matplotlib import gridspec
from matplotlib.collections import LineCollection

def plotLines(file_name):

    # -- Misc ----------------------------------------------------------------
    axis_font_size = 50.
    line_thickness = 3.

    # -- Data ----------------------------------------------------------------
    angles = numpy.linspace(-180., 180, num=201)
    amp_values = [1., 2., 3.] 

    line_values = []     
    for amp in reversed(amp_values):
        sine = amp * (numpy.sin(angles * (numpy.pi/180.)))**-1
        line_values.append(zip(angles, sine))

    amp_values = numpy.array(amp_values)
    lines = LineCollection(
        line_values, 
        array=amp_values, 
        cmap=pyplot.cm.viridis,
        linewidth=line_thickness
        )

    # -- Init Plot -----------------------------------------------------------
    pyplot.clf()  # Clear plot area
    fig = pyplot.figure()
    ax_grid = gridspec.GridSpec(
        1, 2, wspace=0.05, hspace=0.00,
        width_ratios=[20, 1]
        )

    # -- Plot ----------------------------------------------------------------
    plot = pyplot.subplot(ax_grid[0])

    plot.add_collection(lines)

    # Axis limits
    plot.set_xlim([numpy.min(angles), numpy.max(angles)])
    plot.set_ylim([-numpy.max(amp_values), numpy.max(amp_values)])

    # Axis labels
    plot.set_xlabel(
        "$\\mathrm{Angle} \\, \\left[^{\\circ}\\right]$", 
        fontsize=axis_font_size
        )
    plot.set_ylabel(
        "$\\mathrm{Amp} \\,"
        "\\times \\, \\left[\\mathrm{sine} \\,"
        "\\left(\\mathrm{Angle}\\right)\\right]^{-1}$",
        fontsize=axis_font_size,
        )

    # -- Colour Bar ----------------------------------------------------------
    cbar_axis = pyplot.subplot(ax_grid[1])

    cbar = pyplot.colorbar(lines, ticks=amp_values, cax=cbar_axis)

    # CBar Label
    cbar.ax.set_ylabel(
        "$\\mathrm{Amp}$",
        fontsize=axis_font_size,
        )

    # -- Save ----------------------------------------------------------------
    pyplot.savefig(
        "{:s}.eps".format(file_name),
        bbox_inches="tight",
        dpi=600
        )
    pyplot.close()


if __name__ == "__main__":
    plotLines("Without")
    import matplotlib as mpl
    mpl.rcParams['text.usetex'] = True  # Process all text with TeX
    plotLines("With")