#windrose - Computes wind frequencies and plots windroses
#Lionel Roubeyrie - 04/2006
#Licence : this code is released under the matplotlib license

from matplotlib.patches import Rectangle
from matplotlib.axes import PolarSubplot
import matplotlib.cm as cm
from pylab import *
from numpy import *
from scipy import stats

def windfreq( val, vbins, wind, nbwbins = 18 ):
    ''' Computes frequencies of wind classes
    val : ordered list of datas to computes (ex: wind speed)
    vbins : bins of datas (ex:[0,1,2,3,4,5])
    wind : ordered list of wind directions (same order than val)
    nbwbins : number of wind bins (ex: 18 or 36 generally)
    Returns an array of shape=( nbwbins, len(vbins) ) which
    contains the number of time each vbins is reached by wind bins'''

    freq = []
    if len( val ) != len( wind):
        raise ValueError, "values and winds not same length"

    angle = 360./nbwbins
    windbins = arange( 0, 360, angle )
    windbins[0] = 360
    windbins = windbins - angle/2 # To be centred on the North
    windbins.sort()

    vbins.append(inf)
    vbins.sort()

    #We start by selecting values by wind classes, i.e.
    #all values where the wind is between windbins(n) and windbins(n+1)
    last = []
    for i in range( len( windbins ) ):
        if i == len( windbins ) - 1:
            values = select( [greater_equal( wind, windbins[i] )],[val], default=-1.e20 )
        else:
            values = select( [greater_equal( wind, windbins[i] ) * less( wind, windbins[i+1] )],[val], default=-1.e20 )
        # Now, for each winclass, we compute the number of time each class of value is reached
        res = stats.histogram2( values, vbins )
        freq.append( res )
    freq = array( freq )
    # add the last value to the first to have the freq of North winds
    freq[:,0] = freq[:,0] + freq[:,-1]
    # and remove the last col
    freq = freq[:,:-1]
    freq = array( freq )
    # To be directly in the format adapted to windplot
    return transpose( freq )


def windplot( freq, vbins, ptitle="windrose", colormap=cm.autumn ):
    '''Plots a windrose of wind frequencies.
    freq : see windfreq
    vbins : bins of datas (only used for the legend)
    ptitle : title of the plot'''

    shp = freq.shape
    nbrow= shp[0]
    nbcol = shp[1] # == Number of windbins, normally :-)
    if nbrow != len( vbins ):
        raise ValueError, "frequency table not corresponding with bins"
    dtheta = nx.pi/8
    theta = nx.pi/2 - nx.pi/nbcol #To be centred on the North
    rectpatch = []
    fig = figure(figsize=(8,8))
    ax = axes([0.1, 0.1, 0.8, 0.8], polar=True, axisbg='#FFFFFF')

    for col in range( nbcol ):
        for row in range( nbrow ):
            if row == 0: r = 0
            else: r = freq[row-1, col]
            color = colormap( 1-(float(row)/nbcol) )
            val = freq[row, col]
            rect = Rectangle( (theta, r), dtheta, val, facecolor=color)
            ax.add_patch( rect )
            rectpatch.append( rect )
        theta -= dtheta #plot counterclockwise

    # construct the legend
    vbins.append(inf)
    vbins = map( str, vbins )
    leg = []
    leg = [ "[%s : %s["%(vbins[i], vbins[i+1]) for i in range( len(vbins)-1 ) ]

    # plot
    ax.set_thetagrids( angles=arange(0,360,45), labels=['E','N-E','N','N-W','W','S-W','S','S-E'] )
    title(ptitle, fontsize=18)
    figlegend(rectpatch, leg, 'lower left' )
    ax.autoscale_view()
    show()