#!/usr/bin/env python3

import shapefile
import numpy as np
import matplotlib.pyplot as plt
import itertools

def plot_distr_funct(x, label = None, ax = None):
    """Sort and plot distribution function. x should be a Numpy array."""
    
    x.sort()
    nx = np.size(x)
    
    if nx != 0:
        if ax is None: ax = plt.gca()
        ax.plot(x, (1 + np.arange(nx)) / nx, label = label)
        print("minimum value:", x[0])
        print("maximum value:", x[-1])
        
    print("number of values:", nx, "\n")

def read():
    """Read the three dbf files and return speed, radius and amplitude of
    outermost contour, radius and amplitude of maximum speed contour,
    as Numpy arrays.

    Select valid speed values and valid outermost contour (with
    sufficient area).

    """
    
    with shapefile.Reader("extremum") as extremum, \
         shapefile.Reader("outermost_contour") as outerm_cont, \
         shapefile.Reader("max_speed_contour") as max_speed_cont:
        speed = []
        rad_outer = []
        rad_speed = []
        amp_outer = []
        amp_speed = []

        for rec_extr, rec_outer, rec_max in zip(extremum.iterRecords(),
                                                outerm_cont.iterRecords(),
                                                max_speed_cont.iterRecords()):
            if rec_extr.speed != 1e4: speed.append(rec_extr.speed)

            if rec_extr.valid == 1:
                rad_outer.append(rec_outer.r_eq_area)
                amp_outer.append(rec_extr.ssh - rec_outer.ssh)
                
            if rec_max.r_eq_area != - 100:
                rad_speed.append(rec_max.r_eq_area)
                amp_speed.append(rec_extr.ssh - rec_max.ssh)

    return {"speed": np.array(speed), "rad_outer": np.array(rad_outer),
            "rad_speed": np.array(rad_speed), "amp_outer": np.array(amp_outer),
            "amp_speed": np.array(amp_speed)}

def fig_hist(xlabel, x, label):
    """Complete figure with histogram. x and label can also be sequences."""
    
    plt.figure()
    plt.hist(x, bins = "auto", histtype = "step", label = label)
    plt.xlabel(xlabel)
    plt.ylabel("number of occurrences")
    if label is not None: plt.legend()

def fig_distr_funct(xlabel, x, label = None):
    """Complete figure with distribution function. x must be a sequence of
    Numpy arrays and label None or a sequence of strings.

    """
    
    plt.figure()
    print(xlabel + ":")

    label_None = label is None
    if label_None: label = itertools.repeat(None)
        
    for xi, li in zip(x, label):
        plot_distr_funct(xi, label = li)
        
    plt.xlabel(xlabel)
    plt.ylabel("distribution function")
    if not label_None: plt.legend()
    
def plot_all(dict_list, label = None):
    """dict_list: list of dictionaries. label: list of labels, one label
    for each dictionary.

    """
    
    fig_hist("speed, in m s-1", [d["speed"] for d in dict_list], label)
    fig_distr_funct("speed, when positive, in m s-1",
                    [d["speed"][d["speed"] >= 0] for d in dict_list], label)
    fig_distr_funct("speed, when negative, in m s-1",
                    [- d["speed"][d["speed"] < 0] for d in dict_list], label)
    
    fig_hist("equivalent radius of outermost contour, in km",
             [d["rad_outer"] for d in dict_list], label)
    fig_distr_funct("equivalent radius of outermost contour, in km",
                    [d["rad_outer"] for d in dict_list], label)

    fig_hist("equivalent radius of max-speed contour, in km",
             [d["rad_speed"] for d in dict_list], label)
    fig_distr_funct("equivalent radius of max-speed contour, in km",
                    [d["rad_speed"] for d in dict_list], label)

    fig_hist("amplitude of outermost contour, in m",
             [d["amp_outer"] for d in dict_list], label)
    fig_distr_funct("amplitude of outermost contour, anticyclones, in m",
                    [d["amp_outer"][d["amp_outer"] >= 0] for d in dict_list],
                    label)
    fig_distr_funct("amplitude of outermost contour, cyclones, in m",
                    [- d["amp_outer"][d["amp_outer"] < 0] for d in dict_list],
                    label)

    fig_hist("amplitude of max-speed contour, in m",
             [d["amp_speed"] for d in dict_list], label)
    fig_distr_funct("amplitude of max-speed contour, anticyclones, in m",
                    [d["amp_speed"][d["amp_speed"] >= 0] for d in dict_list],
                    label)
    fig_distr_funct("amplitude of max-speed contour, cyclones, in m",
                    [- d["amp_speed"][d["amp_speed"] < 0] for d in dict_list],
                    label)
    
if __name__ == "__main__":
    d = read()
    plot_all([d])
    plt.show()