import time
import os
import pickle

import numpy as np
import tensorflow as tf

from django.core.management.base import BaseCommand

from Key import Key
from plot import plot_direct

from cycling.models import id_dict_from_id_list
from machine_learning.DegradationModelBlackbox import DegradationModel
from machine_learning.LossRecordBlackbox import LossRecord

# TODO(sam): For each cell_id, needs a multigrid of (S, V, I, T) (current
#  needs to be adjusted)
# TODO(sam): Each cycle must have an index mapping to the nearest reference
#  cycle.
# TODO(sam): to evaluate a cycle, there must be the multigrid, the reference
#  cycle scaled by cycle number, the cell features, pasted together and ran
#  through a neural net.

NEIGH_MIN_CYC = 0
NEIGH_MAX_CYC = 1
NEIGH_RATE = 2
NEIGH_CELL_ID = 3
NEIGH_ABSOLUTE_CYC = 4
NEIGH_VALID_CYC = 5
NEIGH_SIGN_GRID = 6
NEIGH_VOLTAGE_GRID = 7
NEIGH_CURRENT_GRID = 8
NEIGH_TMP_GRID = 9
NEIGH_ABSOLUTE_REFERENCE = 10
NEIGH_REFERENCE = 11

NEIGH_TOTAL = 12

def ml_smoothing(options):
    """
    The main function to carry out the machine learning training and evaluation
    procedure.

    Todo(harvey): Add more description about what this does.

    Args:
        options: Dictionary defining various fitting-related arguments.

    Returns: None
    """
    if len(tf.config.experimental.list_physical_devices("GPU")) == 1:
        strategy = tf.distribute.OneDeviceStrategy(device = "/gpu:0")
    elif len(tf.config.experimental.list_physical_devices("GPU")) > 1:
        strategy = tf.distribute.MirroredStrategy()
    else:
        strategy = tf.distribute.OneDeviceStrategy("/cpu:0")

    if not os.path.exists(options[Key.PATH_PLOTS]):
        os.makedirs(options[Key.PATH_PLOTS])

    with open(
        os.path.join(options[Key.PATH_PLOTS], "fit_args_log.txt"), "w",
    ) as f:
        my_str = ""
        for k in options:
            my_str = "{} \n {}: {}".format(my_str, k, str(options[k]))
        f.write(my_str)

    dataset_path = os.path.join(
        options[Key.PATH_DATASET],
        "dataset_ver_{}.file".format(options[Key.DATA_VERSION])
    )

    dataset_names_path = os.path.join(
        options[Key.PATH_DATASET],
        "dataset_ver_{}_names.file".format(options[Key.DATA_VERSION])
    )

    if not os.path.exists(dataset_path):
        print("Path \"" + dataset_path + "\" does not exist.")
        return

    with open(dataset_path, "rb") as f:
        dataset = pickle.load(f)

    dataset_names = None
    if os.path.exists(dataset_names_path):
        with open(dataset_names_path, "rb") as f:
            dataset_names = pickle.load(f)

    cell_ids = list(dataset[Key.ALL_DATA].keys())

    if len(options[Key.CELL_IDS]) != 0:
        cell_ids = list(set(cell_ids).intersection(set(options[Key.CELL_IDS])))

    if len(cell_ids) == 0:
        print("no cell_ids")
        return

    train_and_evaluate(
        initial_processing(
            dataset, dataset_names, cell_ids, options, strategy = strategy,
        ),
        cell_ids,
        options,
    )

# TODO(sam): these huge tensors would be much easier to understand with
#  ragged tensors. Right now, I am just flattening everything.
def numpy_acc(dict, key, data):
    if key in dict.keys():
        dict[key] = np.concatenate((dict[key], data))
    else:
        dict[key] = data

    return dict

def three_level_flatten(iterables):
    for it1 in iterables:
        for it2 in it1:
            for element in it2:
                yield element

def initial_processing(
    dataset: dict, dataset_names, cell_ids, options: dict, strategy,
) -> dict:
    """ Handle the initial data processing

    Args:
        dataset: Contains the quantities given in the dataset.
        dataset_names: TODO(harvey)
        cell_ids: TODO(harvey)
        options: Parameters used to tune the machine learning fitting process.
        strategy: TODO(harvey)

    Returns:
        { Key.STRAT, Key.MODEL, Key.TENSORS, Key.TRAIN_DS, Key.CYC_M,
          Key.CYC_V, Key.OPT, Key.MY_DATA }

    """
    # TODO (harvey): Cleanup Docstring, maybe put detailed description elsewhere
    #   An appropriate place might be in the docstring for
    #   classes inside cycling.Key

    compiled_data = {}
    number_of_compiled_cycs = 0
    number_of_reference_cycs = 0

    dataset[Key.Q_MAX] = 250
    max_cap = dataset[Key.Q_MAX]

    keys = [Key.V_GRID, Key.TEMP_GRID, Key.SIGN_GRID]
    for key in keys:
        numpy_acc(compiled_data, key, np.array([dataset[key]]))

    dataset[Key.I_GRID] = dataset[Key.I_GRID] - np.log(max_cap)
    # the current grid is adjusted by the max capacity of the cell_id. It is
    # in log space, so I/q becomes log(I) - log(q)
    numpy_acc(compiled_data, Key.I_GRID, np.array([dataset[Key.I_GRID]]))

    # TODO (harvey): simplify the following using loops
    # cell ID array
    cell_id_array = np.array(cell_ids)
    # cell ID to positive electrode ID
    cell_id_to_pos_id = {}
    # cell ID to negative electrode ID
    cell_id_to_neg_id = {}
    # cell ID to electrolyte ID
    cell_id_to_lyte_id = {}
    # cell ID to dry cell ID
    cell_id_to_dry_cell_id = {}
    # dry cell ID to meta
    dry_cell_id_to_meta = {}
    # cell ID to latent
    cell_id_to_latent = {}

    # electrolyte ID to latent ID weight
    lyte_to_latent = {}
    # electrolyte ID to solvent ID weight
    lyte_to_sol_weight = {}
    # electrolyte ID to salt ID weight
    lyte_to_salt_weight = {}
    # electrolyte ID to additive ID weight
    lyte_to_addi_weight = {}

    for cell_id in cell_id_array:
        if cell_id in dataset[Key.CELL_TO_POS].keys():
            cell_id_to_pos_id[cell_id] = dataset[Key.CELL_TO_POS][cell_id]
        if cell_id in dataset[Key.CELL_TO_NEG].keys():
            cell_id_to_neg_id[cell_id] = dataset[Key.CELL_TO_NEG][cell_id]
        if cell_id in dataset[Key.CELL_TO_LYTE].keys():
            cell_id_to_lyte_id[cell_id] = dataset[Key.CELL_TO_LYTE][cell_id]
        if cell_id in dataset[Key.CELL_TO_DRY].keys():
            dry_cell_id = dataset[Key.CELL_TO_DRY][cell_id]
            cell_id_to_dry_cell_id[cell_id] = dry_cell_id

            if dry_cell_id in dataset[Key.DRY_TO_META].keys():
                dry_cell_id_to_meta[dry_cell_id]\
                    = dataset[Key.DRY_TO_META][dry_cell_id]

        if cell_id in dataset[Key.CELL_TO_LAT].keys():
            cell_id_to_latent[cell_id] = dataset[Key.CELL_TO_LAT][cell_id]

        if cell_id_to_latent[cell_id] < 0.5:
            lyte_id = cell_id_to_lyte_id[cell_id]
            if lyte_id in dataset[Key.LYTE_TO_SOL].keys():
                lyte_to_sol_weight[lyte_id] = dataset[Key.LYTE_TO_SOL][lyte_id]
            if lyte_id in dataset[Key.LYTE_TO_SALT].keys():
                lyte_to_salt_weight[lyte_id] = dataset[Key.LYTE_TO_SALT][
                    lyte_id]
            if lyte_id in dataset[Key.LYTE_TO_ADD].keys():
                lyte_to_addi_weight[lyte_id] = dataset[Key.LYTE_TO_ADD][lyte_id]
            if lyte_id in dataset[Key.LYTE_TO_LAT].keys():
                lyte_to_latent[lyte_id] = dataset[Key.LYTE_TO_LAT][lyte_id]

    mess = [
        [
            [s[0] for s in siw] for siw in lyte_to_sol_weight.values()
        ], [
            [s[0] for s in siw] for siw in lyte_to_salt_weight.values()
        ], [
            [s[0] for s in siw] for siw in lyte_to_addi_weight.values()
        ],
    ]

    mol_ids = to_sorted_array(list(three_level_flatten(mess)))
    dry_cell_ids = to_sorted_array(cell_id_to_dry_cell_id.values())
    pos_ids = to_sorted_array(cell_id_to_pos_id.values())
    neg_ids = to_sorted_array(cell_id_to_neg_id.values())
    lyte_id_list = to_sorted_array(cell_id_to_lyte_id.values())

    for cell_id_count, cell_id in enumerate(cell_ids):

        all_data = dataset[Key.ALL_DATA][cell_id]
        cyc_grp_dict = all_data[Key.CYC_GRP_DICT]

        for k_count, k in enumerate(cyc_grp_dict.keys()):

            if any([
                abs(cyc_grp_dict[k][Key.I_PREV_END_AVG]) < 1e-5,
                abs(cyc_grp_dict[k][Key.I_CC_AVG]) < 1e-5,
                abs(cyc_grp_dict[k][Key.I_END_AVG]) < 1e-5,
                abs(cyc_grp_dict[k][Key.V_PREV_END_AVG]) < 1e-1,
                abs(cyc_grp_dict[k][Key.V_END_AVG]) < 1e-1,
            ]):
                continue

            main_data = cyc_grp_dict[k][Key.MAIN]

            # normalize capacity_vector with max_cap
            normalize_keys = [
                Key.Q_CC_VEC, Key.Q_CV_VEC, Key.Q_CC_LAST, Key.Q_CV_LAST,
                Key.I_CV_VEC, Key.I_CC, Key.I_PREV_END,
            ]
            for key in normalize_keys:
                main_data[key] = 1. / max_cap * main_data[key]

            normalize_keys = [Key.I_CC_AVG, Key.I_END_AVG, Key.I_PREV_END_AVG]
            for key in normalize_keys:
                cyc_grp_dict[k][key] = 1. / max_cap * cyc_grp_dict[k][key]

            # range of cycles which exist for this cycle group
            min_cyc = min(main_data[Key.N])
            max_cyc = max(main_data[Key.N])

            """
            - now create neighborhoods, which contains the cycles,
              grouped by proximity
            - want to sample neighborhoods equally
            - neighborhoods have a central cycle and a delta on each side
            - to a first approximation, we want a delta_cyc = 300, but we have
              to vary this near the beginning of data and near the end.
            """
            number_of_centers = 10

            # the centers of neighborhoods we will try to create
            all_neigh_center_cycs = np.linspace(
                min_cyc, max_cyc, number_of_centers,
            )
            delta = (
                1.2 * (all_neigh_center_cycs[1] - all_neigh_center_cycs[0]) + 10
            )
            # check all tentative neighborhood centers and
            # commit the ones that contain good data to the dataset
            neigh_data = []

            valid_cycs = 0
            for cyc in all_neigh_center_cycs:
                # max_cyc and min_cyc are the limits of existing cycles.

                below_cyc = cyc - delta
                above_cyc = cyc + delta

                # numpy array of True and False; same length as cyc_grp_dict[k]
                # False when cycle_number falls outside out of
                # [below_cyc, above_cyc] interval
                mask = np.logical_and(
                    below_cyc <= main_data[Key.N],
                    main_data[Key.N] <= above_cyc,
                )

                # the indices for the cyc_grp_dict[k] array which correspond
                # to a True mask
                all_valid_indices = np.arange(len(mask))[mask]

                # if there are less than 1 valid cycles, skip that neighborhood
                if len(all_valid_indices) == 0:
                    continue

                """
                at this point, we know that this neighborhood
                will be added to the dataset.
                """

                min_cyc_index = all_valid_indices[0]
                max_cyc_index = all_valid_indices[-1]

                valid_cycs += 1

                """
                this commits the neighborhood to the dataset

                - record the info about the center of the neighborhood
                  (cycle number, voltage, rate of charge, rate of discharge)
                - record the relative index (within the cycle group)
                  of the min cycle, max cycle
                - record a voltage index, and a cycle group index,
                  and a cell index
                - record the absolute index into the table of cycles
                  (len(cycles_full)).
                - keep a slot empty for later

                """

                neigh_data_i = np.zeros(NEIGH_TOTAL, dtype = np.int32)

                neigh_data_i[NEIGH_MIN_CYC] = min_cyc_index
                neigh_data_i[NEIGH_MAX_CYC] = max_cyc_index
                neigh_data_i[NEIGH_RATE] = k_count
                neigh_data_i[NEIGH_CELL_ID] = cell_id_count
                neigh_data_i[NEIGH_ABSOLUTE_CYC] = number_of_compiled_cycs
                # a weight based on prevalence. Set later
                neigh_data_i[NEIGH_VALID_CYC] = 0
                neigh_data_i[NEIGH_SIGN_GRID] = 0
                neigh_data_i[NEIGH_VOLTAGE_GRID] = 0
                neigh_data_i[NEIGH_CURRENT_GRID] = 0
                neigh_data_i[NEIGH_TMP_GRID] = 0

                center_cyc = float(cyc)
                reference_cycs = all_data[Key.REF_ALL_MATS][Key.N]

                index_of_closest_reference = np.argmin(
                    abs(center_cyc - reference_cycs)
                )

                neigh_data_i[NEIGH_ABSOLUTE_REFERENCE]\
                    = number_of_reference_cycs
                neigh_data_i[NEIGH_REFERENCE] = index_of_closest_reference

                neigh_data.append(neigh_data_i)

            if valid_cycs != 0:
                neigh_data = np.array(neigh_data, dtype = np.int32)

                # the empty slot becomes the count of added neighborhoods, which
                # are used to counterbalance the bias toward longer cycle life
                neigh_data[:, NEIGH_VALID_CYC] = valid_cycs

                numpy_acc(compiled_data, Key.NEIGH_DATA, neigh_data)

            number_of_compiled_cycs += len(main_data[Key.N])
            number_of_reference_cycs += len(all_data[Key.REF_ALL_MATS][Key.N])

            dict_to_acc = {
                Key.REF_CYC: all_data[Key.REF_ALL_MATS][Key.N],
                Key.COUNT_MATRIX: all_data[Key.REF_ALL_MATS][Key.COUNT_MATRIX],
                Key.CYC: main_data[Key.N],
                Key.V_CC_VEC: main_data[Key.V_CC_VEC],
                Key.Q_CC_VEC: main_data[Key.Q_CC_VEC],
                Key.MASK_CC_VEC: main_data[Key.MASK_CC_VEC],
                Key.I_CV_VEC: main_data[Key.I_CV_VEC],
                Key.Q_CV_VEC: main_data[Key.Q_CV_VEC],
                Key.MASK_CV_VEC: main_data[Key.MASK_CV_VEC],
                Key.I_CC: main_data[Key.I_CC],
                Key.I_PREV_END: main_data[Key.I_PREV_END],
                Key.V_PREV_END: main_data[Key.V_PREV_END],
                Key.V_END: main_data[Key.V_END],
            }

            for key in dict_to_acc:
                numpy_acc(compiled_data, key, dict_to_acc[key])

    neigh_data = tf.constant(compiled_data[Key.NEIGH_DATA])

    compiled_tensors = {}
    # cycles go from 0 to 6000, but nn prefers normally distributed variables
    # so cycle numbers is normalized with mean and variance
    cycle_tensor = tf.constant(compiled_data[Key.CYC])
    cycle_m, cycle_v = tf.nn.moments(cycle_tensor, axes = [0])
    cycle_m = 0.  # we shall leave the cycle 0 at 0
    cycle_v = cycle_v.numpy()
    cycle_tensor = (cycle_tensor - cycle_m) / tf.sqrt(cycle_v)
    compiled_tensors[Key.CYC] = cycle_tensor

    labels = [
        Key.V_CC_VEC, Key.Q_CC_VEC, Key.MASK_CC_VEC, Key.Q_CV_VEC, Key.I_CV_VEC,
        Key.MASK_CV_VEC, Key.I_CC, Key.I_PREV_END, Key.V_PREV_END, Key.V_END,
        Key.COUNT_MATRIX, Key.SIGN_GRID, Key.V_GRID, Key.I_GRID, Key.TEMP_GRID,
    ]
    for label in labels:
        compiled_tensors[label] = tf.constant(compiled_data[label])

    batch_size = options[Key.BATCH]

    with strategy.scope():
        train_ds_ = tf.data.Dataset.from_tensor_slices(
            neigh_data
        ).repeat(2).shuffle(100000).batch(batch_size)

        train_ds = strategy.experimental_distribute_dataset(train_ds_)

        dry_cell_to_dry_cell_name = {}
        pos_to_pos_name = {}
        neg_to_neg_name = {}
        lyte_to_lyte_name = {}
        mol_to_mol_name = {}

        if dataset_names is not None:
            pos_to_pos_name = dataset_names[Key.NAME_POS]
            neg_to_neg_name = dataset_names[Key.NAME_NEG]
            lyte_to_lyte_name = dataset_names[Key.NAME_LYTE]
            mol_to_mol_name = dataset_names[Key.NAME_MOL]
            dry_cell_to_dry_cell_name = dataset_names[Key.NAME_DRY]

        degradation_model = DegradationModel(
            width = options[Key.WIDTH],
            depth = options[Key.DEPTH],
            cell_dict = id_dict_from_id_list(cell_id_array),
            pos_dict = id_dict_from_id_list(pos_ids),
            neg_dict = id_dict_from_id_list(neg_ids),
            lyte_dict = id_dict_from_id_list(lyte_id_list),
            mol_dict = id_dict_from_id_list(mol_ids),
            dry_cell_dict = id_dict_from_id_list(dry_cell_ids),

            cell_to_pos = cell_id_to_pos_id,
            cell_to_neg = cell_id_to_neg_id,
            cell_to_lyte = cell_id_to_lyte_id,
            cell_to_dry_cell = cell_id_to_dry_cell_id,
            dry_cell_to_meta = dry_cell_id_to_meta,

            cell_latent_flags = cell_id_to_latent,

            lyte_to_solvent = lyte_to_sol_weight,
            lyte_to_salt = lyte_to_salt_weight,
            lyte_to_additive = lyte_to_addi_weight,
            lyte_latent_flags = lyte_to_latent,

            names = (
                pos_to_pos_name,
                neg_to_neg_name,
                lyte_to_lyte_name,
                mol_to_mol_name,
                dry_cell_to_dry_cell_name,
            ),
            n_sample = options[Key.N_SAMPLE],
            options = options,
            min_latent = options[Key.MIN_LAT],
        )

        optimizer = tf.keras.optimizers.Adam(
            learning_rate = options[Key.LRN_RATE],
        )

    return {
        Key.STRAT: strategy,
        Key.MODEL: degradation_model,
        Key.TENSORS: compiled_tensors,
        Key.TRAIN_DS: train_ds,
        Key.CYC_M: cycle_m,
        Key.CYC_V: cycle_v,
        Key.OPT: optimizer,
        Key.DATASET: dataset,
    }

def to_sorted_array(unsorted) -> np.array:
    """ Remove duplicates from a list or a view object (of a dictionary) and
        turn it into to a sorted array.

    Args:
        unsorted: Unsorted list or view object (of a dictionary).

    Returns:
        Sorted array of `unsorted`.
    """
    return np.array(sorted(list(set(unsorted))))

def train_and_evaluate(
    init_returns: dict, cell_ids: list, options: dict,
) -> None:
    """

    Args:
        init_returns: Return value of `initial_processing`.
        cell_ids: Specified cell IDs (identifiers for different cells).
        options:
    """
    strategy = init_returns[Key.STRAT]

    epochs = 100000
    count = 0

    end = time.time()

    train_step_params = {
        Key.TENSORS: init_returns[Key.TENSORS],
        Key.OPT: init_returns[Key.OPT],
        Key.MODEL: init_returns[Key.MODEL],
    }

    @tf.function
    def dist_train_step(strategy, neigh):
        return strategy.experimental_run_v2(
            lambda neigh: train_step(neigh, train_step_params, options),
            args = (neigh,),
        )

    # TODO(harvey, confusion): what is `l`?
    l = None
    loss_record = LossRecord()
    with strategy.scope():
        for epoch in range(epochs):
            sub_count = 0
            for neigh in init_returns[Key.TRAIN_DS]:
                count += 1
                sub_count += 1

                l_ = dist_train_step(strategy, neigh)
                if l is None:
                    l = l_
                else:
                    l += l_

                if count != 0:
                    if (count % options[Key.PRINT_LOSS]) == 0:
                        tot = l / tf.cast(sub_count, dtype = tf.float32)
                        l = None
                        sub_count = 0
                        loss_record.record(count, tot.numpy())
                        loss_record.print_recent(options)

                    plot_params = {
                        "cell_ids": cell_ids,
                        "count": count,
                        Key.OPTIONS: options,
                    }

                    if (count % options[Key.VIS_FIT]) == 0:
                        start = time.time()
                        print("time to simulate: ", start - end)
                        loss_record.plot(count, options)
                        plot_direct(
                            "generic_vs_cycle", plot_params, init_returns,
                        )
                        plot_direct(
                            "generic_vs_capacity", plot_params, init_returns,
                        )

                        end = time.time()
                        print("time to plot: ", end - start)
                        print()

                if count >= options[Key.STOP]:
                    return

def train_step(neigh, params: dict, options: dict):
    """ One training step.

    Args:
        neigh: Neighbourhood.
        params: Contains all necessary parameters.
        options: Options for `ml_smoothing`.
    """
    # need to split the range
    batch_size2 = neigh.shape[0]

    degradation_model = params[Key.MODEL]
    optimizer = params[Key.OPT]
    compiled_tensors = params[Key.TENSORS]

    sign_grid_tensor = compiled_tensors[Key.SIGN_GRID]
    voltage_grid_tensor = compiled_tensors[Key.V_GRID]
    current_grid_tensor = compiled_tensors[Key.I_GRID]
    tmp_grid_tensor = compiled_tensors[Key.TEMP_GRID]

    count_matrix_tensor = compiled_tensors[Key.COUNT_MATRIX]

    cycle_tensor = compiled_tensors[Key.CYC]
    constant_current_tensor = compiled_tensors[Key.I_CC]
    end_current_prev_tensor = compiled_tensors[Key.I_PREV_END]
    end_voltage_prev_tensor = compiled_tensors[Key.V_PREV_END]
    end_voltage_tensor = compiled_tensors[Key.V_END]

    cc_voltage_tensor = compiled_tensors[Key.V_CC_VEC]
    cc_capacity_tensor = compiled_tensors[Key.Q_CC_VEC]
    cc_mask_tensor = compiled_tensors[Key.MASK_CC_VEC]
    cv_capacity_tensor = compiled_tensors[Key.Q_CV_VEC]
    cv_current_tensor = compiled_tensors[Key.I_CV_VEC]
    cv_mask_tensor = compiled_tensors[Key.MASK_CV_VEC]

    """
    if you have the minimum cycle and maximum cycle for a neighborhood,
    you can sample cycle from this neighborhood by sampling real numbers
    x from [0,1] and computing min_cyc*(1.-x) + max_cyc*x,
    but here this computation is done in index space,
    then cycle numbers and vq curves are gathered
    """

    cyc_indices_lerp = tf.random.uniform(
        [batch_size2], minval = 0., maxval = 1., dtype = tf.float32,
    )
    cyc_indices = tf.cast(
        (1. - cyc_indices_lerp) * tf.cast(
            neigh[:, NEIGH_MIN_CYC] + neigh[:, NEIGH_ABSOLUTE_CYC],
            tf.float32,
        ) + cyc_indices_lerp * tf.cast(
            neigh[:, NEIGH_MAX_CYC] + neigh[:, NEIGH_ABSOLUTE_CYC],
            tf.float32,
        ),
        tf.int32,
    )

    sign_grid = tf.gather(
        sign_grid_tensor, indices = neigh[:, NEIGH_SIGN_GRID], axis = 0,
    )
    sign_grid_dim = sign_grid.shape[1]

    voltage_grid = tf.gather(
        voltage_grid_tensor, indices = neigh[:, NEIGH_VOLTAGE_GRID], axis = 0,
    )
    voltage_grid_dim = voltage_grid.shape[1]

    current_grid = tf.gather(
        current_grid_tensor, indices = neigh[:, NEIGH_CURRENT_GRID], axis = 0,
    )
    current_grid_dim = current_grid.shape[1]

    tmp_grid = tf.gather(
        tmp_grid_tensor, indices = neigh[:, NEIGH_TMP_GRID], axis = 0,
    )
    tmp_grid_dim = tmp_grid.shape[1]

    svit_tuple = (
        tf.tile(
            tf.reshape(
                sign_grid, [batch_size2, sign_grid_dim, 1, 1, 1, 1],
            ),
            [1, 1, voltage_grid_dim, current_grid_dim, tmp_grid_dim, 1],
        ),
        tf.tile(
            tf.reshape(
                voltage_grid, [batch_size2, 1, voltage_grid_dim, 1, 1, 1],
            ),
            [1, sign_grid_dim, 1, current_grid_dim, tmp_grid_dim, 1],
        ),
        tf.tile(
            tf.reshape(
                current_grid, [batch_size2, 1, 1, current_grid_dim, 1, 1],
            ),
            [1, sign_grid_dim, voltage_grid_dim, 1, tmp_grid_dim, 1],
        ),
        tf.tile(
            tf.reshape(
                tmp_grid, [batch_size2, 1, 1, 1, tmp_grid_dim, 1],
            ),
            [1, sign_grid_dim, voltage_grid_dim, current_grid_dim, 1, 1],
        ),
    )
    svit_grid = tf.concat(svit_tuple, axis = -1)

    count_matrix = tf.reshape(
        tf.gather(
            count_matrix_tensor,
            neigh[:, NEIGH_ABSOLUTE_REFERENCE] + neigh[:, NEIGH_REFERENCE],
            axis = 0,
        ),
        [
            batch_size2, sign_grid_dim, voltage_grid_dim, current_grid_dim,
            tmp_grid_dim, 1,
        ],
    )

    cycle = tf.gather(cycle_tensor, indices = cyc_indices, axis = 0)
    constant_current = tf.gather(
        constant_current_tensor, indices = cyc_indices, axis = 0,
    )
    end_current_prev = tf.gather(
        end_current_prev_tensor, indices = cyc_indices, axis = 0,
    )
    end_voltage_prev = tf.gather(
        end_voltage_prev_tensor, indices = cyc_indices, axis = 0,
    )
    end_voltage = tf.gather(
        end_voltage_tensor, indices = cyc_indices, axis = 0,
    )

    cc_capacity = tf.gather(cc_capacity_tensor, indices = cyc_indices)
    cc_voltage = tf.gather(cc_voltage_tensor, indices = cyc_indices)
    cc_mask = tf.gather(cc_mask_tensor, indices = cyc_indices)
    cc_mask_2 = tf.tile(
        tf.reshape(
            1. / tf.cast(neigh[:, NEIGH_VALID_CYC], tf.float32),
            [batch_size2, 1],
        ),
        [1, cc_voltage.shape[1]],
    )

    cv_capacity = tf.gather(cv_capacity_tensor, indices = cyc_indices)
    cv_current = tf.gather(cv_current_tensor, indices = cyc_indices)
    cv_mask = tf.gather(cv_mask_tensor, indices = cyc_indices)
    cv_mask_2 = tf.tile(
        tf.reshape(
            1. / tf.cast(neigh[:, NEIGH_VALID_CYC], tf.float32),
            [batch_size2, 1],
        ),
        [1, cv_current.shape[1]],
    )

    cell_indices = neigh[:, NEIGH_CELL_ID]

    with tf.GradientTape() as tape:
        train_results = degradation_model(
            {
                Key.CYC: tf.expand_dims(cycle, axis = 1),
                Key.I_CC: tf.expand_dims(constant_current, axis = 1),
                Key.I_PREV_END: tf.expand_dims(end_current_prev, axis = 1),
                Key.V_PREV_END: tf.expand_dims(end_voltage_prev, axis = 1),
                Key.V_END: tf.expand_dims(end_voltage, axis = 1),
                Key.INDICES: cell_indices,
                Key.V_TENSOR: cc_voltage,
                Key.I_TENSOR: cv_current,
                Key.SVIT_GRID: svit_grid,
                Key.COUNT_MATRIX: count_matrix,
            },
            training = True,
        )

        pred_cc_capacity = train_results[Key.Pred.I_CC]
        pred_cv_capacity = train_results[Key.Pred.I_CV]

        cc_capacity_loss = get_loss(
            cc_capacity, pred_cc_capacity, cc_mask, cc_mask_2,
        )
        cv_capacity_loss = get_loss(
            cv_capacity, pred_cv_capacity, cv_mask, cv_mask_2,
        )

        main_losses = (
            options[Key.Coeff.Q_CV] * cv_capacity_loss
            + options[Key.Coeff.Q_CC] * cc_capacity_loss
        )
        loss = main_losses + tf.stop_gradient(main_losses) * (
            options[Key.Coeff.Q] * train_results[Key.Loss.Q]
            + options[Key.Coeff.CELL] * train_results[Key.Loss.CELL]
        )

    gradients = tape.gradient(loss, degradation_model.trainable_variables)

    gradients_no_nans = [
        tf.where(tf.math.is_nan(x), tf.zeros_like(x), x) for x in gradients
    ]

    gradients_norm_clipped, _ = tf.clip_by_global_norm(
        gradients_no_nans, options[Key.GLB_NORM_CLIP],
    )

    optimizer.apply_gradients(
        zip(gradients_norm_clipped, degradation_model.trainable_variables)
    )

    return tf.stack(
        [
            cc_capacity_loss,
            cv_capacity_loss,
            train_results[Key.Loss.Q],
            train_results[Key.Loss.CELL],
        ],
    )

def get_loss(measured, predicted, mask, mask_2):
    return tf.reduce_mean(
        (1e-10 + mask * mask_2) * tf.square(measured - predicted)
    ) / (1e-10 + tf.reduce_mean(mask * mask_2))

class Command(BaseCommand):

    def add_arguments(self, parser):

        required_args = [
            Key.PATH_DATASET,
            Key.DATA_VERSION,
            Key.PATH_PLOTS,
        ]

        float_args = {
            Key.GLB_NORM_CLIP: 10.,

            Key.LRN_RATE: 5e-4,
            Key.MIN_LAT: .05,

            Key.COEFF_FEAT_CELL_DER: .001,
            Key.COEFF_FEAT_CELL_DER2: .01,

            Key.COEFF_CELL: 1.,
            Key.COEFF_CELL_OUT: .1,
            Key.COEFF_CELL_IN: .1,
            Key.COEFF_CELL_DER: .1,
            Key.COEFF_CELL_EQ: 10.,

            Key.COEFF_LYTE: 1.,
            Key.COEFF_LYTE_OUT: .1,
            Key.COEFF_LYTE_IN: .1,
            Key.COEFF_LYTE_DER: .1,
            Key.COEFF_LYTE_EQ: 10.,

            Key.COEFF_Q_CV: 1.,
            Key.COEFF_Q_CC: 1.,

            Key.COEFF_Q: 1.,
            Key.COEFF_Q_GEQ: 1.,
            Key.COEFF_Q_LEQ: 1.,
            Key.COEFF_Q_V_MONO: .1,
            Key.COEFF_Q_DER3_V: 1.,
            Key.COEFF_Q_DER3_I: 1.,
            Key.COEFF_Q_DER3_N: 1.,
            Key.COEFF_Q_DER_I: 1.,
            Key.COEFF_Q_DER_N: 10.,
        }

        vis = 10000
        int_args = {
            Key.N_SAMPLE: 8 * 16,

            Key.DEPTH: 3,
            Key.WIDTH: 50,
            Key.BATCH: 4 * 16,

            Key.PRINT_LOSS: 500,
            Key.VIS_FIT: vis,
            Key.VIS_VQ: vis,

            Key.STOP: 1000004,
            Key.CELL_ID_SHOW: 10,
        }

        for arg in required_args:
            parser.add_argument("--" + arg, required = True)
        for arg in float_args:
            parser.add_argument(
                "--" + arg, type = float, default = float_args[arg],
            )
        for arg in int_args:
            parser.add_argument("--" + arg, type = int, default = int_args[arg])

        cell_ids = [
            57706, 57707, 57710, 57711, 57714, 57715, 64260, 64268, 83010,
            83011, 83012, 83013, 83014, 83015, 83016, 81602, 81603, 81604,
            81605, 81606, 81607, 81608, 81609, 81610, 81611, 81612, 81613,
            81614, 81615, 81616, 81617, 81618, 81619, 81620, 81621, 81622,
            81623, 81624, 81625, 81626, 81627, 81712, 81713, 82300, 82301,
            82302, 82303, 82304, 82305, 82306, 82307, 82308, 82309, 82310,
            82311, 82406, 82407, 82410, 82411, 82769, 82770, 82771, 82775,
            82776, 82777, 82779, 82992, 82993, 83083, 83092, 83101, 83106,
            83107, 83220, 83221, 83222, 83223, 83224, 83225, 83226, 83227,
            83228, 83229, 83230, 83231, 83232, 83233, 83234, 83235, 83236,
            83237, 83239, 83240, 83241, 83242, 83243, 83310, 83311, 83312,
            83317, 83318, 83593, 83594, 83595, 83596, 83741, 83742, 83743,
            83744, 83745, 83746, 83747, 83748,
        ]

        parser.add_argument(
            "--" + Key.CELL_IDS, type = int, nargs = "+", default = cell_ids,
        )

    def handle(self, *args, **options):
        ml_smoothing(options)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""EEG-driven deep learning classifications and model generation
Software tools used generate results in submitted work to (tentative citation pending review):

David O. Nahmias, Eugene F. Civillico, and Kimberly L. Kontson.
Deep Learning and Feature Based Medication Classifications from EEG in a Large Clinical Data Set
In review (2020)

If you have found this software useful please consider citing our publication.

Public domain license
"""

""" Disclaimer:
This software and documentation (the "Software") were developed at the Food and Drug Administration (FDA) by employees
of the Federal Government in the course of their official duties. Pursuant to Title 17, Section 105 of the United States Code,
this work is not subject to copyright protection and is in the public domain. Permission is hereby granted, free of charge,
to any person obtaining a copy of the Software, to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense, or sell copies of the Software or derivatives,
and to permit persons to whom the Software is furnished to do so. FDA assumes no responsibility whatsoever for use by other
parties of the Software, its source code, documentation or compiled executables, and makes no guarantees, expressed or implied,
about its quality, reliability, or any other characteristic. Further, use of this code in no way implies endorsement by the FDA
or confers any advantage in regulatory decisions. Although this software can be redistributed and/or modified freely, we ask that
any derivative works bear some notice that they are derived from it, and any modified versions bear some notice that they have been modified.
"""

__author__ = 'David Nahmias'
__copyright__ = 'No copyright - US Government, 2020 , DeepEEG classification'
__credits__ = ['David Nahmias']
__license__ = 'Public domain'
__version__ = '0.0.1'
__maintainer__ = 'David Nahmias'
__email__ = 'david.nahmias@fda.hhs.gov'
__status__ = 'alpha'

import logging
import time
from copy import copy
import sys

from collections import Counter
import random
import numpy as np
from numpy.random import RandomState
import resampy
from torch import optim
import torch.nn.functional as F
import torch as th
from torch.nn.functional import elu
from torch import nn

from braindecode.datautil.signal_target import SignalAndTarget
from braindecode.torch_ext.util import np_to_var
from braindecode.torch_ext.util import set_random_seeds
from braindecode.torch_ext.modules import Expression
from braindecode.experiments.experiment import Experiment
from braindecode.datautil.iterators import CropsFromTrialsIterator
from braindecode.experiments.monitors import (RuntimeMonitor, LossMonitor,
                                              MisclassMonitor)
from braindecode.experiments.stopcriteria import MaxEpochs
from braindecode.models.shallow_fbcsp import ShallowFBCSPNet
from braindecode.models.deep4 import Deep4Net
from braindecode.models.util import to_dense_prediction_model
from braindecode.datautil.iterators import get_balanced_batches
from braindecode.torch_ext.constraints import MaxNormDefaultConstraint
from braindecode.torch_ext.util import var_to_np
from braindecode.torch_ext.functions import identity

from dataset import DiagnosisSet
from monitors import compute_preds_per_trial, CroppedDiagnosisMonitor

log = logging.getLogger(__name__)
log.setLevel('DEBUG')

import pdb
from loadNEDC import loadNEDCdata,loadSubNormData,addDataNoise

if len(sys.argv)>1:
    if ',' in sys.argv[1]:
        CLASSY = sys.argv[1].split(',')
        CLASSY[2] = int(CLASSY[2])
        if len(CLASSY) == 4:
            CLASSY[3] = int(CLASSY[3])
    else:
        CLASSY = str(sys.argv[1])

def splitDataRandom(allData,allLabels,setNum=0,shuffle=0):
    numberEqSamples = min(Counter(allLabels).values())
    trainSamplesNum = int(np.ceil(numberEqSamples*0.9))
    testSamplesNum = numberEqSamples-trainSamplesNum

    labels0 = allLabels[allLabels == 0]
    labels1 = allLabels[allLabels == 1]
    data0 = allData[allLabels == 0]
    data1 = allData[allLabels == 1]

    fullRange = list(range(numberEqSamples))
    random.shuffle(fullRange)

    testIndecies = fullRange[trainSamplesNum:]
    trainIndecies = fullRange[:trainSamplesNum]

    allDataTrain = np.concatenate((data0[trainIndecies],data1[trainIndecies]),axis=0)
    allLabelsTrain = np.concatenate((labels0[trainIndecies],labels1[trainIndecies]),axis=0)

    allDataTest = np.concatenate((data0[testIndecies],data1[testIndecies]),axis=0)
    allLabelsTest = np.concatenate((labels0[testIndecies],labels1[testIndecies]),axis=0)

    if shuffle == 1:
        random.shuffle(allLabelsTrain)
        np.save('sessionsData/trainLabels%s-shuffle-sessions'%CLASSY,allLabelsTrain)
        np.save('sessionsData/dataRangeOrder%s-shuffle-sessions'%CLASSY,fullRange)
    else:
        np.save('sessionsData/dataRangeOrder%s-sessions'%CLASSY,fullRange)

    return allDataTrain,allLabelsTrain,allDataTest,allLabelsTest

def create_set(X, y, inds):
    """
    X list and y nparray
    :return:
    """
    new_X = []
    for i in inds:
        new_X.append(X[i])
    new_y = y[inds]
    return SignalAndTarget(new_X, new_y)

class TrainValidTestSplitter(object):
    def __init__(self, n_folds, i_test_fold, shuffle):
        self.n_folds = n_folds
        self.i_test_fold = i_test_fold
        self.rng = RandomState(39483948)
        self.shuffle = shuffle

    def split(self, X, y,):
        if len(X) < self.n_folds:
            raise ValueError("Less Trials: {:d} than folds: {:d}".format(
                len(X), self.n_folds
            ))
        folds = get_balanced_batches(len(X), self.rng, self.shuffle,
                                     n_batches=self.n_folds)
        test_inds = folds[self.i_test_fold]
        valid_inds = folds[self.i_test_fold - 1]
        all_inds = list(range(len(X)))
        train_inds = np.setdiff1d(all_inds, np.union1d(test_inds, valid_inds))
        assert np.intersect1d(train_inds, valid_inds).size == 0
        assert np.intersect1d(train_inds, test_inds).size == 0
        assert np.intersect1d(valid_inds, test_inds).size == 0
        assert np.array_equal(np.sort(
            np.union1d(train_inds, np.union1d(valid_inds, test_inds))),
            all_inds)

        train_set = create_set(X, y, train_inds)
        valid_set = create_set(X, y, valid_inds)
        test_set = create_set(X, y, test_inds)

        return train_set, valid_set, test_set

class TrainValidSplitter(object):
    def __init__(self, n_folds, i_valid_fold, shuffle):
        self.n_folds = n_folds
        self.i_valid_fold = i_valid_fold
        self.rng = RandomState(39483948)
        self.shuffle = shuffle

    def split(self, X, y):
        if len(X) < self.n_folds:
            raise ValueError("Less Trials: {:d} than folds: {:d}".format(
                len(X), self.n_folds
            ))
        folds = get_balanced_batches(len(X), self.rng, self.shuffle,
                                     n_batches=self.n_folds)
        valid_inds = folds[self.i_valid_fold]
        all_inds = list(range(len(X)))
        train_inds = np.setdiff1d(all_inds, valid_inds)
        assert np.intersect1d(train_inds, valid_inds).size == 0
        assert np.array_equal(np.sort(np.union1d(train_inds, valid_inds)),
            all_inds)

        train_set = create_set(X, y, train_inds)
        valid_set = create_set(X, y, valid_inds)
        return train_set, valid_set

def run_exp(data_folders,
            n_recordings,
            sensor_types,
            n_chans,
            max_recording_mins,
            sec_to_cut, duration_recording_mins,
            test_recording_mins,
            max_abs_val,
            sampling_freq,
            divisor,
            test_on_eval,
            n_folds, i_test_fold,
            shuffle,
            model_name,
            n_start_chans, n_chan_factor,
            input_time_length, final_conv_length,
            model_constraint,
            init_lr,
            batch_size, max_epochs,cuda,):

    import torch.backends.cudnn as cudnn
    cudnn.benchmark = True
    preproc_functions = []
    preproc_functions.append(
        lambda data, fs: (data[:, int(sec_to_cut * fs):-int(
            sec_to_cut * fs)], fs))
    preproc_functions.append(
        lambda data, fs: (data[:, :int(duration_recording_mins * 60 * fs)], fs))
    if max_abs_val is not None:
        preproc_functions.append(lambda data, fs:
                                 (np.clip(data, -max_abs_val, max_abs_val), fs))

    preproc_functions.append(lambda data, fs: (resampy.resample(data, fs,
                                                                sampling_freq,
                                                                axis=1,
                                                                filter='kaiser_fast'),
                                               sampling_freq))

    if divisor is not None:
        preproc_functions.append(lambda data, fs: (data / divisor, fs))

    dataset = DiagnosisSet(n_recordings=n_recordings,
                           max_recording_mins=max_recording_mins,
                           preproc_functions=preproc_functions,
                           data_folders=data_folders,
                           train_or_eval='train',
                           sensor_types=sensor_types)
    if test_on_eval:
        if test_recording_mins is None:
            test_recording_mins = duration_recording_mins
        test_preproc_functions = copy(preproc_functions)
        test_preproc_functions[1] = lambda data, fs: (
            data[:, :int(test_recording_mins * 60 * fs)], fs)
        test_dataset = DiagnosisSet(n_recordings=n_recordings,
                                max_recording_mins=None,
                                preproc_functions=test_preproc_functions,
                                data_folders=data_folders,
                                train_or_eval='eval',
                                sensor_types=sensor_types)
    #X,y = dataset.load()
    #test_X, test_y = test_dataset.load()

    #X,y = loadNEDCdata(mode='train')
    print(CLASSY)
    #X,y,test_X,test_y = loadNEDCdata(mode='all',classy=CLASSY)

    data = np.load('sessionsData/data%s-sessions.npy'%CLASSY[:3])
    labels = np.load('sessionsData/labels%s-sessions.npy'%CLASSY[:3])

    #pdb.set_trace()

    if len(CLASSY) > 3:
        random.seed(11081992 + int(CLASSY[3]))
    else:
        random.seed(11081992)

    X,y,test_X,test_y = splitDataRandom(data,labels,shuffle=0)

    #X = np.load('trainData%s.npy'%CLASSY)
    #y = np.load('trainLabels%s.npy'%CLASSY)
    #test_X = np.load('testData%s.npy'%CLASSY)
    #test_y = np.load('testLabels%s.npy'%CLASSY)

    #return False

    #pdb.set_trace()

    #X,y,test_X,test_y = loadSubNormData(mode='all')

    #X = addDataNoise(X,band=[1,4])
    #test_X = addDataNoise(test_X,band=[1,4])
    #pdb.set_trace()

    max_shape = np.max([list(x.shape) for x in X],
                       axis=0)
    assert max_shape[1] == int(duration_recording_mins *
                               sampling_freq * 60)
    if test_on_eval:
        #test_X, test_y = test_dataset.load()
        #test_X, test_y = loadNEDCdata(mode='eval')
        max_shape = np.max([list(x.shape) for x in test_X],
                           axis=0)
        assert max_shape[1] == int(test_recording_mins *
                                   sampling_freq * 60)
    if not test_on_eval:
        splitter = TrainValidTestSplitter(n_folds, i_test_fold,
                                          shuffle=shuffle)
        train_set, valid_set, test_set = splitter.split(X, y)
    else:
        splitter = TrainValidSplitter(n_folds, i_valid_fold=i_test_fold,
                                          shuffle=shuffle)
        train_set, valid_set = splitter.split(X, y)
        test_set = SignalAndTarget(test_X, test_y)
        del test_X, test_y
    del X,y # shouldn't be necessary, but just to make sure

    #pdb.set_trace()

    set_random_seeds(seed=20170629, cuda=cuda)
    n_classes = 2
    if model_name == 'shallow':
        model = ShallowFBCSPNet(in_chans=n_chans, n_classes=n_classes,
                                n_filters_time=n_start_chans,
                                n_filters_spat=n_start_chans,
                                input_time_length=input_time_length,
                                final_conv_length=final_conv_length).create_network()
    elif model_name == 'deep':
        model = Deep4Net(n_chans, n_classes,
                         n_filters_time=n_start_chans,
                         n_filters_spat=n_start_chans,
                         input_time_length=input_time_length,
                         n_filters_2 = int(n_start_chans * n_chan_factor),
                         n_filters_3 = int(n_start_chans * (n_chan_factor ** 2.0)),
                         n_filters_4 = int(n_start_chans * (n_chan_factor ** 3.0)),
                         final_conv_length=final_conv_length,
                        stride_before_pool=True).create_network()
    elif (model_name == 'deep_smac'):
        if model_name == 'deep_smac':
            do_batch_norm = False
        else:
            assert model_name == 'deep_smac_bnorm'
            do_batch_norm = True
        double_time_convs = False
        drop_prob = 0.244445
        filter_length_2 = 12
        filter_length_3 = 14
        filter_length_4 = 12
        filter_time_length = 21
        final_conv_length = 1
        first_nonlin = elu
        first_pool_mode = 'mean'
        first_pool_nonlin = identity
        later_nonlin = elu
        later_pool_mode = 'mean'
        later_pool_nonlin = identity
        n_filters_factor = 1.679066
        n_filters_start = 32
        pool_time_length = 1
        pool_time_stride = 2
        split_first_layer = True
        n_chan_factor = n_filters_factor
        n_start_chans = n_filters_start
        model = Deep4Net(n_chans, n_classes,
                 n_filters_time=n_start_chans,
                 n_filters_spat=n_start_chans,
                 input_time_length=input_time_length,
                 n_filters_2=int(n_start_chans * n_chan_factor),
                 n_filters_3=int(n_start_chans * (n_chan_factor ** 2.0)),
                 n_filters_4=int(n_start_chans * (n_chan_factor ** 3.0)),
                 final_conv_length=final_conv_length,
                 batch_norm=do_batch_norm,
                 double_time_convs=double_time_convs,
                 drop_prob=drop_prob,
                 filter_length_2=filter_length_2,
                 filter_length_3=filter_length_3,
                 filter_length_4=filter_length_4,
                 filter_time_length=filter_time_length,
                 first_nonlin=first_nonlin,
                 first_pool_mode=first_pool_mode,
                 first_pool_nonlin=first_pool_nonlin,
                 later_nonlin=later_nonlin,
                 later_pool_mode=later_pool_mode,
                 later_pool_nonlin=later_pool_nonlin,
                 pool_time_length=pool_time_length,
                 pool_time_stride=pool_time_stride,
                 split_first_layer=split_first_layer,
                 stride_before_pool=True).create_network()
    elif model_name == 'shallow_smac':
        conv_nonlin = identity
        do_batch_norm = True
        drop_prob = 0.328794
        filter_time_length = 56
        final_conv_length = 22
        n_filters_spat = 73
        n_filters_time = 24
        pool_mode = 'max'
        pool_nonlin = identity
        pool_time_length = 84
        pool_time_stride = 3
        split_first_layer = True
        model = ShallowFBCSPNet(in_chans=n_chans, n_classes=n_classes,
                                n_filters_time=n_filters_time,
                                n_filters_spat=n_filters_spat,
                                input_time_length=input_time_length,
                                final_conv_length=final_conv_length,
                                conv_nonlin=conv_nonlin,
                                batch_norm=do_batch_norm,
                                drop_prob=drop_prob,
                                filter_time_length=filter_time_length,
                                pool_mode=pool_mode,
                                pool_nonlin=pool_nonlin,
                                pool_time_length=pool_time_length,
                                pool_time_stride=pool_time_stride,
                                split_first_layer=split_first_layer,
                                ).create_network()
    elif model_name == 'linear':
        model = nn.Sequential()
        model.add_module("conv_classifier",
                         nn.Conv2d(n_chans, n_classes, (600,1)))
        model.add_module('softmax', nn.LogSoftmax(dim=1))
        model.add_module('squeeze', Expression(lambda x: x.squeeze(3)))
    else:
        assert False, "unknown model name {:s}".format(model_name)
    to_dense_prediction_model(model)
    log.info("Model:\n{:s}".format(str(model)))
    if cuda:
        model.cuda()
    # determine output size
    test_input = np_to_var(
        np.ones((2, n_chans, input_time_length, 1), dtype=np.float32))
    if cuda:
        test_input = test_input.cuda()
    log.info("In shape: {:s}".format(str(test_input.cpu().data.numpy().shape)))

    out = model(test_input)
    log.info("Out shape: {:s}".format(str(out.cpu().data.numpy().shape)))
    n_preds_per_input = out.cpu().data.numpy().shape[2]
    log.info("{:d} predictions per input/trial".format(n_preds_per_input))
    iterator = CropsFromTrialsIterator(batch_size=batch_size,
                                       input_time_length=input_time_length,
                                       n_preds_per_input=n_preds_per_input)
    optimizer = optim.Adam(model.parameters(), lr=init_lr)

    loss_function = lambda preds, targets: F.nll_loss(
        th.mean(preds, dim=2, keepdim=False), targets)

    if model_constraint is not None:
        assert model_constraint == 'defaultnorm'
        model_constraint = MaxNormDefaultConstraint()
    monitors = [LossMonitor(), MisclassMonitor(col_suffix='sample_misclass'),
                CroppedDiagnosisMonitor(input_time_length, n_preds_per_input),
                RuntimeMonitor(),]
    stop_criterion = MaxEpochs(max_epochs)
    batch_modifier = None
    run_after_early_stop = True
    exp = Experiment(model, train_set, valid_set, test_set, iterator,
                     loss_function, optimizer, model_constraint,
                     monitors, stop_criterion,
                     remember_best_column='valid_misclass',
                     run_after_early_stop=run_after_early_stop,
                     batch_modifier=batch_modifier,
                     cuda=cuda)
    exp.run()
    return exp

if __name__ == "__main__":
    import config
    print('Classifying: {}'.format(CLASSY))
    #pdb.set_trace()
    start_time = time.time()
    logging.basicConfig(format='%(asctime)s %(levelname)s : %(message)s',
                     level=logging.DEBUG, stream=sys.stdout)
    exp = run_exp(
        config.data_folders,
        config.n_recordings,
        config.sensor_types,
        config.n_chans,
        config.max_recording_mins,
        config.sec_to_cut, config.duration_recording_mins,
        config.test_recording_mins,
        config.max_abs_val,
        config.sampling_freq,
        config.divisor,
        config.test_on_eval,
        config.n_folds, config.i_test_fold,
        config.shuffle,
        config.model_name,
        config.n_start_chans, config.n_chan_factor,
        config.input_time_length, config.final_conv_length,
        config.model_constraint,
        config.init_lr,
        config.batch_size, config.max_epochs,config.cuda,)
    end_time = time.time()
    run_time = end_time - start_time

    log.info("Experiment runtime: {:.2f} sec".format(run_time))

    #pdb.set_trace()
    if len(CLASSY) == 3:
        CLASSY[2] = str(CLASSY[2])
        CLASSY = '-'.join(CLASSY)
    elif len(CLASSY) == 4:
        CLASSY[2] = str(CLASSY[2])
        CLASSY[3] = str(CLASSY[3])
        CLASSY = '-'.join(CLASSY)

    th.save(exp.model.state_dict(),'sessionsData/'+CLASSY+'Model-sessions.pt')
    #th.save(exp.model.state_dict(),CLASSY+'Model-shuffle.pt')

    #savedModel = th.load('subNormModel.pt')

    # In case you want to recompute predictions for further analysis:
    exp.model.eval()
    for setname in ('train', 'valid', 'test'):
        log.info("Compute predictions for {:s}...".format(
            setname))
        dataset = exp.datasets[setname]
        if config.cuda:
            preds_per_batch = [var_to_np(exp.model(np_to_var(b[0]).cuda()))
                      for b in exp.iterator.get_batches(dataset, shuffle=False)]
        else:
            preds_per_batch = [var_to_np(exp.model(np_to_var(b[0])))
                      for b in exp.iterator.get_batches(dataset, shuffle=False)]
        preds_per_trial = compute_preds_per_trial(
            preds_per_batch, dataset,
            input_time_length=exp.iterator.input_time_length,
            n_stride=exp.iterator.n_preds_per_input)
        mean_preds_per_trial = [np.mean(preds, axis=(0, 2)) for preds in
                                    preds_per_trial]
        mean_preds_per_trial = np.array(mean_preds_per_trial)


import os
import datetime as dt
from collections import OrderedDict
import shutil
from copy import copy

import matplotlib
import matplotlib.dates
import pylab as plt
import numpy as np
from mpl_toolkits.basemap import Basemap
from scipy import stats
from matplotlib_venn import venn2, venn2_circles

from ..processing.c20data import C20Data
from plotting import lon_convert
from ..results import StormtracksResultsManager
from ..analysis import StormtracksAnalysis, ClassificationAnalysis
from .. import analysis
from ..ibtracsdata import IbtracsData
import plotting
from ..load_settings import settings
from ..analysis.classification import SCATTER_ATTRS
from .. import classification
from .. import processing.matching

def load_20th_century():
    cla_analysis = ClassificationAnalysis()
    cal_cd = cla_analysis.load_cal_classification_data()
    val_cd = cla_analysis.load_val_classification_data()

    sgdc = classification.SGDClassifier()
    sgdc_best = sgdc.load('sgdc_best')
    sgdc.train(cal_cd, **sgdc_best)

    sgdc.predict(cal_cd)
    cal_res = copy(sgdc.res)
    cal_sens, cal_ppv = sgdc.sensitivity, sgdc.ppv

    sgdc.predict(val_cd)
    val_sens, val_ppv = sgdc.sensitivity, sgdc.ppv
    val_res = copy(sgdc.res)

    sens = (cal_sens + val_sens) / 2
    ppv = (cal_ppv + val_ppv) / 2
    adjustment_ratio_old = ppv / sens
    print('OLD sens: {0}, ppv: {1}, ar: {2}'.format(sens, ppv, adjustment_ratio_old))

    tp = cal_res['tp'] + val_res['tp']
    fn = cal_res['fn'] + val_res['fn']
    fp = cal_res['fp'] + val_res['fp']
    sens = 1. * tp / (tp + fn)
    ppv = 1. * tp / (tp + fp)
    adjustment_ratio = ppv / sens
    print('NEW sens: {0}, ppv: {1}, ar: {2}'.format(sens, ppv, adjustment_ratio))

    ib_hurrs, ib_pdis, cla_hurrs, cla_pdis = cla_analysis.run_yearly_analysis(sgdc)
    return ib_hurrs, ib_pdis, cla_hurrs, cla_pdis, adjustment_ratio

def load_thresh_metric_vs_vort_lo():
    cla_analysis = ClassificationAnalysis()
    cal_cd = cla_analysis.load_cal_classification_data()
    cc = classification.CutoffClassifier()
    cc_best = cc.load('cc_best')
    print(cc_best)

    metrics = []
    for vort_lo in np.arange(0.00001, 0.0005, 0.00001):
        cc_best['vort_lo'] = vort_lo
        print(vort_lo)
        cc.train(cal_cd, **cc_best)
        cc.predict(cal_cd)
        metrics.append((vort_lo, cc.sensitivity, cc.ppv, cc.fpr))

    return np.array(metrics)

def load_classifier_metrics():
    cla_analysis = ClassificationAnalysis()

    cal_cd, val_cd, clas = cla_analysis.get_trained_classifiers()

    metrics = OrderedDict()
    for (cla, settings, fmt, name) in clas:
        metrics[name] = {'fmt': fmt}

        cla.predict(cal_cd)
        metrics[name]['cal'] = (cla.sensitivity, cla.ppv, cla.fpr)

        cla.predict(val_cd)
        metrics[name]['val'] = (cla.sensitivity, cla.ppv, cla.fpr)

    return metrics

def load_ibtracs_info():
    yearly_hurr_distribution, hurr_per_year = analysis.analyse_ibtracs_data(False)
    return yearly_hurr_distribution, hurr_per_year

def plot_20th_century(ib_hurrs, cla_hurrs, adjustment_ratio):
    ar = adjustment_ratio
    fig = plt.figure()
    ax = plt.subplot(211)
    plt.plot(range(1890, 2010), ib_hurrs, 'r-')
    plt.plot(range(1890, 2010), cla_hurrs.mean(axis=1) * ar, 'b--')
    plt.xlim((1890, 2010))
    plt.ylim((0, 350))
    plt.fill_between(range(1890, 2010),
                     cla_hurrs.min(axis=1) * ar,
                     cla_hurrs.max(axis=1) * ar,
                     color=(0, 0, 1, 0.2))
    plt.plot((1914, 1914), (0, 350), 'k--')
    plt.plot((1944, 1944), (0, 350), 'k--')
    plt.plot((1966, 1966), (0, 350), 'k--')

    plt.annotate('Panama\nCanal', xy=(1914, 290), xytext=(1915, 290), fontsize=10)
    plt.annotate('Aircraft\nRecon.', xy=(1944, 290), xytext=(1945, 290), fontsize=10)
    plt.annotate('Satellite', xy=(1966, 290), xytext=(1967, 290), fontsize=10)
    plt.setp(ax.get_xticklabels(), visible=False)

    plt.ylabel('Hurricane-timesteps')
    ax.yaxis.set_label_coords(-0.12, 0.5)

    ax = plt.subplot(212)
    plt.plot(range(1890, 2010), np.array(ib_hurrs) - cla_hurrs.mean(axis=1) * ar, 'b-')
    plt.plot((1890, 2010), (0, 0), 'k-')

    plt.xlim((1890, 2010))
    ylim = (-150, 200)
    plt.ylim(ylim)

    plt.plot((1914, 1914), ylim, 'k--')
    plt.plot((1944, 1944), ylim, 'k--')
    plt.plot((1966, 1966), ylim, 'k--')

    yoffset = 60
    plt.annotate('Panama\nCanal', xy=(1914, ylim[1] - yoffset),
                 xytext=(1915, ylim[1] - yoffset), fontsize=10)
    plt.annotate('Aircraft\nRecon.', xy=(1944, ylim[1] - yoffset),
                 xytext=(1945, ylim[1] - yoffset), fontsize=10)
    plt.annotate('Satellite', xy=(1966, ylim[1] - yoffset),
                 xytext=(1967, ylim[1] - yoffset), fontsize=10)

    plt.ylabel('$\Delta$ Hurricane-timesteps')
    ax.yaxis.set_label_coords(-0.12, 0.5)
    fig.set_size_inches(6.3, 6)
    _save_figure('20th_century_hurricane_timesteps.png')

def plot_poster_20th_century(ib_hurrs, cla_hurrs, adjustment_ratio):
    ar = adjustment_ratio
    fig = plt.figure()
    ax = plt.subplot(211)
    plt.plot(range(1890, 2010), ib_hurrs, 'r-')
    plt.plot(range(1890, 2010), cla_hurrs.mean(axis=1) * ar, 'b--')
    plt.xlim((1890, 2010))
    plt.ylim((0, 350))
    plt.fill_between(range(1890, 2010),
                     cla_hurrs.min(axis=1) * ar,
                     cla_hurrs.max(axis=1) * ar,
                     color=(0, 0, 1, 0.2))
    plt.plot((1914, 1914), (0, 350), 'k--')
    plt.plot((1944, 1944), (0, 350), 'k--')
    plt.plot((1966, 1966), (0, 350), 'k--')

    plt.annotate('Panama\nCanal', xy=(1914, 290), xytext=(1915, 290), fontsize=10)
    plt.annotate('Aircraft\nRecon.', xy=(1944, 290), xytext=(1945, 290), fontsize=10)
    plt.annotate('Satellite', xy=(1966, 290), xytext=(1967, 290), fontsize=10)
    plt.setp(ax.get_xticklabels(), visible=False)

    plt.ylabel('Hurricane-timesteps', fontsize=16)
    ax.yaxis.set_label_coords(-0.12, 0.5)

    ax = plt.subplot(212)
    plt.plot(range(1890, 2010), np.array(ib_hurrs) - cla_hurrs.mean(axis=1) * ar, 'b-')
    plt.plot((1890, 2010), (0, 0), 'k-')

    plt.xlim((1890, 2010))
    ylim = (-150, 200)
    plt.ylim(ylim)

    plt.plot((1914, 1914), ylim, 'k--')
    plt.plot((1944, 1944), ylim, 'k--')
    plt.plot((1966, 1966), ylim, 'k--')

    yoffset = 60
    plt.annotate('Panama\nCanal', xy=(1914, ylim[1] - yoffset),
                 xytext=(1915, ylim[1] - yoffset), fontsize=10)
    plt.annotate('Aircraft\nRecon.', xy=(1944, ylim[1] - yoffset),
                 xytext=(1945, ylim[1] - yoffset), fontsize=10)
    plt.annotate('Satellite', xy=(1966, ylim[1] - yoffset),
                 xytext=(1967, ylim[1] - yoffset), fontsize=10)

    plt.ylabel('$\Delta$ Hurricane-timesteps', fontsize=16)
    ax.yaxis.set_label_coords(-0.12, 0.5)
    fig.set_size_inches(6.3, 6)
    _save_figure('20th_century_hurricane_timesteps_poster.png')

def plot_poster_20th_century_trends(ib_hurrs, cla_hurrs, adjustment_ratio):
    ar = adjustment_ratio
    fig = plt.figure()
    ax = plt.subplot(211)
    plt.plot(range(1890, 2010), ib_hurrs, 'r-')
    plt.xlim((1890, 2010))
    plt.ylim((0, 350))
    plt.plot((1914, 1914), (0, 350), 'k--')
    plt.plot((1944, 1944), (0, 350), 'k--')
    plt.plot((1966, 1966), (0, 350), 'k--')

    plt.annotate('Panama\nCanal', xy=(1914, 290), xytext=(1915, 290), fontsize=10)
    plt.annotate('Aircraft\nRecon.', xy=(1944, 290), xytext=(1945, 290), fontsize=10)
    plt.annotate('Satellite', xy=(1966, 290), xytext=(1967, 290), fontsize=10)
    plt.setp(ax.get_xticklabels(), visible=False)

    plt.ylabel('Hurricane-timesteps', fontsize=16)
    ax.yaxis.set_label_coords(-0.12, 0.5)

    ax = plt.subplot(212)
    plt.plot(range(1890, 2010), cla_hurrs.mean(axis=1) * ar, 'b-')
    plt.xlim((1890, 2010))
    ylim = (0, 350)
    plt.ylim(ylim)

    plt.plot((1914, 1914), ylim, 'k--')
    plt.plot((1944, 1944), ylim, 'k--')
    plt.plot((1966, 1966), ylim, 'k--')

    yoffset = 60
    plt.annotate('Panama\nCanal', xy=(1914, ylim[1] - yoffset),
                 xytext=(1915, ylim[1] - yoffset), fontsize=10)
    plt.annotate('Aircraft\nRecon.', xy=(1944, ylim[1] - yoffset),
                 xytext=(1945, ylim[1] - yoffset), fontsize=10)
    plt.annotate('Satellite', xy=(1966, ylim[1] - yoffset),
                 xytext=(1967, ylim[1] - yoffset), fontsize=10)

    ar = adjustment_ratio

    for i, name, sl in ((1, 'pre-Panama Canal', slice(0, 24)),
                        (2, 'Panama Canal', slice(24, 54)),
                        (3, 'Aircraft', slice(54, 76)),
                        (4, 'Satellite', slice(76, 120))):
        corr_ib = stats.linregress(range(sl.stop - sl.start), ib_hurrs[sl] * ar)
        corr_c2 = stats.linregress(range(sl.stop - sl.start), cla_hurrs.mean(axis=1)[sl] * ar)
        plt.subplot(211)
        if corr_ib[3] < 0.05:
            fmt = 'r-'
        else:
            fmt = 'r--'

        plt.plot((sl.start + 1890, sl.stop + 1890),
                 (corr_ib[1], (sl.stop - sl.start) * corr_ib[0] + corr_ib[1]), fmt)
        plt.subplot(212)
        if corr_c2[3] < 0.05:
            fmt = 'b-'
        else:
            fmt = 'b--'
        plt.plot((sl.start + 1890, sl.stop + 1890),
                 (corr_c2[1], (sl.stop - sl.start) * corr_c2[0] + corr_c2[1]), fmt)

        latex = False
        if not latex:
            print(name)
            print('IB: gradient: {0}, inter: {1}, r2: {2}, p: {3}'.format(corr_ib[0],
                                                                          corr_ib[1],
                                                                          corr_ib[2] ** 2,
                                                                          corr_ib[3]))
            print('C2: gradient: {0}, inter: {1}, r2: {2}, p: {3}'.format(corr_c2[0],
                                                                          corr_c2[1],
                                                                          corr_c2[2] ** 2,
                                                                          corr_c2[3]))
        else:
            print('{0} & {1:.2f} & {2:.3f} & {3:.2f} & {4:.3f} \\\\'.format(name,
                                                                            corr_ib[0],
                                                                            corr_ib[3],
                                                                            corr_c2[0],
                                                                            corr_c2[3]))
    plt.ylabel('Hurricane-timesteps', fontsize=16)
    ax.yaxis.set_label_coords(-0.12, 0.5)

    fig.set_size_inches(6.3, 6)
    _save_figure('20th_century_trends.png')

def print_20th_century_trends(ib_hurrs, cla_hurrs, adjustment_ratio, latex=True):
    ar = adjustment_ratio
    for i, name, sl in ((1, 'pre-Panama Canal', slice(0, 24)),
                        (2, 'Panama Canal', slice(24, 54)),
                        (3, 'Aircraft', slice(54, 76)),
                        (4, 'Satellite', slice(76, 120))):
        corr_ib = stats.linregress(range(sl.stop - sl.start), ib_hurrs[sl] * ar)
        corr_c2 = stats.linregress(range(sl.stop - sl.start), cla_hurrs.mean(axis=1)[sl] * ar)
        if not latex:
            print(name)
            print('IB: gradient: {0}, inter: {1}, r2: {2}, p: {3}'.format(corr_ib[0],
                                                                          corr_ib[1],
                                                                          corr_ib[2] ** 2,
                                                                          corr_ib[3]))
            print('C2: gradient: {0}, inter: {1}, r2: {2}, p: {3}'.format(corr_c2[0],
                                                                          corr_c2[1],
                                                                          corr_c2[2] ** 2,
                                                                          corr_c2[3]))
        else:
            print('{0} & {1:.2f} & {2:.3f} & {3:.2f} & {4:.3f} \\\\'.format(name,
                                                                            corr_ib[0],
                                                                            corr_ib[3],
                                                                            corr_c2[0],
                                                                            corr_c2[3]))

def plot_20th_century_corr(ib_hurrs, cla_hurrs, adjustment_ratio):
    ar = adjustment_ratio
    fig = plt.figure()
    plt.subplot(121)
    plt.plot(ib_hurrs, cla_hurrs.mean(axis=1) * ar, 'k+')
    corr = stats.linregress(ib_hurrs, cla_hurrs.mean(axis=1) * ar)
    print(corr)
    label = 'grad.: {0:.2f}\nintercept: {1:.2f}\nr$^2$: {2:.2f}'.format(corr[0], corr[1],
                                                                        corr[2] ** 2)
    plt.plot((0, 300), (corr[0] * 0 + corr[1], corr[0] * 300 + corr[1]), 'k--', label=label)
    plt.xlim((0, 300))
    plt.ylim((0, 300))
    plt.legend(bbox_to_anchor=(0.75, 1.1), numpoints=1, prop={'size': 10})

    plt.xlabel('IBTrACS Hurricane-timesteps')
    plt.ylabel('Estimated Hurricane-timesteps')
    # fig.set_size_inches(4, 4)

    ax = plt.subplot(122)
    for label, sl, c, fmt in (('pre-1944', slice(0, 54), 'r', '+'),
                              ('post-1944', slice(54, 120), 'b', 'x')):
        plt.plot(ib_hurrs[sl], cla_hurrs.mean(axis=1)[sl] * ar, '{0}{1}'.format(c, fmt),
                 label=label)
        corr = stats.linregress(ib_hurrs[sl], cla_hurrs.mean(axis=1)[sl] * ar)
        print(corr)
        # label = 'grad.: {0:.2f}\nintercept: {1:.2f}\nr$^2$: {2:.2f}'.format(corr[0],
        # corr[1], corr[2] ** 2)
        plt.plot((0, 300), (corr[0] * 0 + corr[1], corr[0] * 300 + corr[1]), '{0}--'.format(c))

    plt.setp(ax.get_yticklabels(), visible=False)
    plt.xlim((0, 300))
    plt.ylim((0, 300))
    plt.legend(bbox_to_anchor=(0.60, 1.1), numpoints=1, prop={'size': 10})

    plt.xlabel('IBTrACS Hurricane-timesteps')
    # plt.ylabel('Estimated Hurricane timesteps')
    fig.set_size_inches(6.3, 3)
    _save_figure('20th_century_corr.png')

    fig = plt.figure()

    for i, label, sl, c, fmt, bbox in ((1, 'pre-Panama', slice(0, 24), 'r', '+', (0.97, 0.35)),
                                       (2, 'Panama', slice(24, 54), 'b', 'x', (0.97, 0.35)),
                                       (3, 'aircraft', slice(54, 76), 'g', '^', (0.65, 1.1)),
                                       (4, 'satellite', slice(76, 120), 'c', 'o', (0.71, 1.1))):
        ax = plt.subplot(2, 2, i)
        plt.plot(ib_hurrs[sl], cla_hurrs.mean(axis=1)[sl] * ar, '{0}{1}'.format(c, fmt),
                 label=label)
        corr = stats.linregress(ib_hurrs[sl], cla_hurrs.mean(axis=1)[sl] * ar)
        print(label)
        print(corr[2] ** 2)
        label = 'grad.: {0:.2f}\ninter.: {1:.2f}'.format(corr[0], corr[1])
        plt.plot((0, 300), (corr[0] * 0 + corr[1], corr[0] * 300 + corr[1]), '{0}--'.format(c),
                 label=label)

        if i in (1, 2):
            plt.setp(ax.get_xticklabels(), visible=False)
        else:
            plt.xlabel('IBTrACS Hurricane-timesteps')

        if i in (1, 3):
            plt.ylabel('Estimated Hurricane-timesteps')
        else:
            plt.setp(ax.get_yticklabels(), visible=False)

        plt.xlim((0, 300))
        plt.ylim((0, 300))

        plt.legend(bbox_to_anchor=bbox, numpoints=1, prop={'size': 10})

    # plt.xlabel('IBTrACS Hurricane timesteps')
    # plt.ylabel('Estimated Hurricane timesteps')
    fig.set_size_inches(6.3, 6.3)
    _save_figure('20th_century_corr_split.png')

def plot_galveston():
    c20data = C20Data(1900, fields=['prmsl', 'u', 'v'])
    c20data.set_date(dt.datetime(1900, 9, 7, 18, 0))
    loc = {'llcrnrlat': 15, 'urcrnrlat': 35, 'llcrnrlon': -100, 'urcrnrlon': -70}

    fig = plt.figure()
    plt.subplot(121)
    m = _raster_on_earth(c20data.lons, c20data.lats, c20data.vort * 10000, loc=loc, colorbar=None)
    m.colorbar(location='bottom', pad='7%', ticks=(-1, 0, 1, 2))
    plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)', labelpad=30)
    plt.subplot(122)
    m = _raster_on_earth(c20data.lons, c20data.lats, c20data.prmsl / 100, loc=loc, colorbar=None)
    m.colorbar(location='bottom', pad='7%', ticks=(970, 1000, 1030))
    plt.xlabel('Pressure (hPa)', labelpad=30)

    fig.set_size_inches(6.3, 3)
    _save_figure('galveston_1900-9-7_18-00_em0.png')

def plot_venn():
    fig = plt.figure()
    plt.clf()
    fig.set_size_inches(6, 5)

    v = venn2(subsets=(5, 6, 7), set_labels=('', ''))
    v.get_label_by_id('10').set_text('False\nPositives')
    v.get_label_by_id('01').set_text('False\nNegatives')
    v.get_label_by_id('11').set_text('True\nPositives')
    c = venn2_circles(subsets=(5, 6, 7), linestyle='dashed')
    c[0].set_lw(1.0)
    c[0].set_ls('dotted')

    plt.annotate('True Negatives', xy=np.array([0, 0]),
                 xytext=(0, -140),
                 ha='center', textcoords='offset points')
    _save_figure('tf_np_venn.png')

def plot_thresh_metric_vs_vort_lo(metrics):
    fig = plt.figure()
    ax = plt.subplot(211)
    plt.plot(metrics[:40, 0] * 10000, metrics[:40, 1] + metrics[:40, 2], 'k-', label='Sum')
    plt.legend(loc='best', prop={'size': 10})
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.ylim((1, 1.4))
    ax.set_yticks((1, 1.1, 1.2, 1.3, 1.4))

    plt.subplot(212)
    plt.plot(metrics[:40, 0] * 10000, metrics[:40, 1], 'b--', label='sensitivity')
    plt.plot(metrics[:40, 0] * 10000, metrics[:40, 2], 'g--', label='PPV')
    plt.legend(loc='best', prop={'size': 10})
    plt.xlabel('Vorticity threshold ($10^{-4}$ s$^{-1}$)')

    fig.set_size_inches(6.3, 4)

    _save_figure('threshold_sens_ppv_vort.png')

def plot_classifer_metrics(metrics):
    plt.figure()

    for name in metrics:
        metric = metrics[name]
        fmt = metric['fmt']
        plt.subplot(2, 2, 1)
        plt.plot(metric['cal'][1], metric['cal'][0], fmt, label=name)

        plt.subplot(2, 2, 2)
        plt.plot(metric['cal'][2], metric['cal'][0], fmt, label=name)

        plt.subplot(2, 2, 3)
        plt.plot(metric['val'][1], metric['cal'][0], fmt, label=name)

        plt.subplot(2, 2, 4)
        plt.plot(metric['val'][2], metric['cal'][0], fmt, label=name)

    plt.subplot(2, 2, 1)
    # plt.xlabel('PPV')
    plt.ylabel('Calibration\nsensitivity')
    plt.xlim((0, 1))
    plt.ylim((0, 1))

    plt.subplot(2, 2, 2)
    # plt.xlabel('FPR')
    # plt.ylabel('sensitivity')
    plt.xlim((0, 0.1))
    plt.ylim((0, 1))
    plt.legend(bbox_to_anchor=(1.1, 1.1), numpoints=1)

    plt.subplot(2, 2, 3)
    plt.xlabel('PPV')
    plt.ylabel('Validation\nsensitivity')
    plt.xlim((0, 1))
    plt.ylim((0, 1))

    plt.subplot(2, 2, 4)
    plt.xlabel('FPR')
    # plt.ylabel('sensitivity')
    plt.xlim((0, 0.1))
    plt.ylim((0, 1))

    _save_figure('ppv_and_fpr_vs_sens.png')
    return metrics

def plot_wld(stormtracks_analysis=None, years=None):
    if not stormtracks_analysis:
        stormtracks_analysis = StormtracksAnalysis(2000)

    if not years:
        years = range(2000, 2010)

    wlds = {}
    wlds['draws'] = []

    for year in years:
        print(year)
        stormtracks_analysis.set_year(year)
        k0, w0, k1, w1, d = stormtracks_analysis.run_wld_analysis(active_configs={'scale': 3})

        if k0 not in wlds:
            wlds[k0] = []
        wlds[k0].append(w0)

        if k1 not in wlds:
            wlds[k1] = []
        wlds[k1].append(w1)

        wlds['draws'].append(d)

    print(wlds)

    plt.figure()
    plt.title('wld')
    for k in wlds:
        plt.plot(years, wlds[k], label=k)
    plt.legend(loc='best')

def plot_tracking_stats(stormtracks_analysis=None, years=None, sort_col='cumoveroverlap'):
    if not stormtracks_analysis:
        stormtracks_analysis = StormtracksAnalysis(2000)

    k995 = []
    for scale in (1, 2, 3):
        config = {'pressure_level': 995, 'scale': scale, 'tracker': 'nearest_neighbour'}
        k995.append(stormtracks_analysis.good_matches_key(config))

    keys = ('pl995', 'pl850', 'scale3')

    config_keys = {}
    for key in keys:
        config_keys[key] = []

    for scale in (1, 2, 3):
        config = {'pressure_level': 995, 'scale': scale, 'tracker': 'nearest_neighbour'}
        config_keys['pl995'].append(stormtracks_analysis.good_matches_key(config))

        config = {'pressure_level': 850, 'scale': scale, 'tracker': 'nearest_neighbour'}
        config_keys['pl850'].append(stormtracks_analysis.good_matches_key(config))

    for pl in (995, 850):
        config = {'pressure_level': pl, 'scale': 3, 'tracker': 'nearest_neighbour'}
        config_keys['scale3'].append(stormtracks_analysis.good_matches_key(config))

    all_wins = {}
    for key in keys:
        all_wins[key] = []

    if not years:
        years = range(2000, 2010)

    for year in years:
        print(year)
        stormtracks_analysis.set_year(year)

        res = {}
        res['pl995'] =\
            stormtracks_analysis.run_position_analysis(sort_on=sort_col,
                                                       active_configs={'pressure_level': 995})
        res['pl850'] =\
            stormtracks_analysis.run_position_analysis(sort_on=sort_col,
                                                       active_configs={'pressure_level': 850})
        res['scale3'] =\
            stormtracks_analysis.run_position_analysis(sort_on=sort_col,
                                                       active_configs={'scale': 3})

        for key in keys:
            print(key)
            wins = []
            for config_key in config_keys[key]:
                print('  {0}'.format(config_key))
                try:
                    wins.append(res[key][config_key][0])
                except KeyError:
                    wins.append(0)

            all_wins[key].append(wins)

    for key in keys:
        all_wins[key] = np.array(all_wins[key])

    _plot_all_wins(years, all_wins)
    return years, all_wins

def _plot_all_wins(years, all_wins):
    fmt = matplotlib.ticker.ScalarFormatter(useOffset=False)
    fig = plt.figure()
    fmt.set_scientific(False)
    ax = plt.subplot(311)
    plt.title('Near Surface Pressure Level (NSPL)')
    plt.plot(years, all_wins['pl995'][:, 0], 'r-', label='Scale 1')  # scale 1
    plt.plot(years, all_wins['pl995'][:, 1], 'g-', label='Scale 2')  # scale 2
    plt.plot(years, all_wins['pl995'][:, 2], 'b-', label='Scale 3')  # scale 3
    plt.ylim(0, 60)
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.legend(bbox_to_anchor=(1.1, 1.14), numpoints=1, prop={'size': 10})
    ax.xaxis.set_major_formatter(fmt)

    ax = plt.subplot(312)
    plt.title('850 hPa Pressure Level')
    plt.plot(years, all_wins['pl850'][:, 0], 'r--', label='Scale 1')  # scale 1
    plt.plot(years, all_wins['pl850'][:, 1], 'g--', label='Scale 2')  # scale 2
    plt.plot(years, all_wins['pl850'][:, 2], 'b--', label='Scale 3')  # scale 3
    plt.ylim(0, 60)
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.legend(bbox_to_anchor=(1.1, 1.14), numpoints=1, prop={'size': 10})
    ax.xaxis.set_major_formatter(fmt)

    ax = plt.subplot(313)
    plt.title('Scale 3')
    plt.plot(years, all_wins['scale3'][:, 0], 'b-', label='NSPL')  # pl 995
    plt.plot(years, all_wins['scale3'][:, 1], 'b--', label='850 hPa')  # pl 850
    plt.ylim(0, 60)
    plt.legend(bbox_to_anchor=(1.1, 1.14), numpoints=1, prop={'size': 10})
    ax.xaxis.set_major_formatter(fmt)

    fig.set_size_inches(6.3, 6)
    _save_figure('tracking_wins_losses.png')

def plot_poster_katrina():
    loc = {'llcrnrlat': 15, 'urcrnrlat': 35, 'llcrnrlon': -100, 'urcrnrlon': -70}
    fig = plt.figure(1)
    plt.clf()
    c20data = C20Data(2005, fields=['u', 'v'])
    c20data.set_date(dt.datetime(2005, 8, 27, 18))

    m = _raster_on_earth(c20data.lons, c20data.lats, c20data.vort * 10000, loc=loc, colorbar=False)
    fig.set_size_inches(6.4, 3)
    _save_figure('katrina.png')

def plot_poster_2005_best_tracks():
    loc = {'llcrnrlat': 10, 'urcrnrlat': 60, 'llcrnrlon': -100, 'urcrnrlon': -20}
    fig = plt.figure(2)
    plt.clf()
    c20data = C20Data(2005, fields=['u', 'v'])
    ibdata = IbtracsData()
    bts = ibdata.load_ibtracks_year(2005)
    m = _raster_on_earth(c20data.lons, c20data.lats, None, loc=loc, colorbar=False, labels=False)
    fig.set_size_inches(6.4, 3)

    for bt in bts:
        plotting.plot_track(bt)
    _save_figure('2005_best_tracks.png')

def plot_data_processing_figures():
    _plot_katrina()
    _plot_katrina_maxs_mins()

def _plot_katrina():
    c20data = C20Data(2005, fields=['u', 'v'])
    c20data.set_date(dt.datetime(2005, 8, 27, 18))

    fig = plt.figure(1)
    plt.clf()

    loc = {'llcrnrlat': 15, 'urcrnrlat': 35, 'llcrnrlon': -100, 'urcrnrlon': -70}
    ax = plt.subplot(131)
    # plt.title('Wind')

    m = _vec_plot_on_earth(c20data.lons, c20data.lats, -c20data.u, c20data.v, loc=loc)
    m.colorbar(location='bottom', pad='7%', ticks=(0, 8, 16, 24))
    plt.xlabel('Wind speed (ms$^{-1}$)')
    ax.xaxis.set_label_coords(0.5, -0.33)

    ax = plt.subplot(132)
    # plt.title('Vorticity')
    m = _raster_on_earth(c20data.lons, c20data.lats, c20data.vort * 10000, loc=loc, colorbar=False)
    m.colorbar(location='bottom', pad='7%', ticks=(-1, 0, 1, 2))
    plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')
    ax.xaxis.set_label_coords(0.5, -0.33)

    c20data = C20Data(2005, fields=['u', 'v'], upscaling=True, scale_factor=3)
    c20data.set_date(dt.datetime(2005, 8, 27, 18))
    ax = plt.subplot(133)
    # plt.title('Downscaled\nVorticity')
    m = _raster_on_earth(c20data.up_lons, c20data.up_lats, c20data.up_vort * 10000,
                         loc=loc, colorbar=False)
    m.colorbar(location='bottom', pad='7%', ticks=(-1, 0, 1, 2))
    plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')
    ax.xaxis.set_label_coords(0.5, -0.33)

    fig.set_size_inches(7, 4)
    _save_figure('katrina_data_proc.png')

def _plot_katrina_maxs_mins():
    c20data = C20Data(2005, fields=['prmsl', 'u', 'v'])
    c20data.set_date(dt.datetime(2005, 8, 27, 18))
    loc = {'llcrnrlat': 0, 'urcrnrlat': 45, 'llcrnrlon': -120, 'urcrnrlon': -60}

    fig = plt.figure(2)
    plt.clf()

    plt.subplot(221)
    _raster_on_earth(c20data.lons, c20data.lats, c20data.vort, loc=loc, colorbar=False)
    plt.ylabel('Vorticity')

    plt.subplot(222)
    _raster_on_earth(c20data.lons, c20data.lats, None, loc=loc)
    points = c20data.vmaxs
    for p_val, p_loc in points:
        _plot_point_on_earth(p_loc[0] + 1, p_loc[1] + 1, 'ro')
    points = c20data.vmins
    for p_val, p_loc in points:
        _plot_point_on_earth(p_loc[0] + 1, p_loc[1] + 1, 'kx')

    plt.subplot(223)
    _raster_on_earth(c20data.lons, c20data.lats, c20data.prmsl,
                     vmin=99000, vmax=103000, loc=loc, colorbar=False)
    plt.ylabel('Pressure')

    plt.subplot(224)
    _raster_on_earth(c20data.lons, c20data.lats, None, loc=loc)
    points = c20data.pmaxs
    for p_val, p_loc in points:
        _plot_point_on_earth(p_loc[0] + 1, p_loc[1] + 1, 'ro')
    points = c20data.pmins
    for p_val, p_loc in points:
        _plot_point_on_earth(p_loc[0] + 1, p_loc[1] + 1, 'kx')

    _save_figure('katrina_max_mins.png')

def plot_matching_figures():
    _plot_six_configs_figure()
    _plot_individual_katrina_figure()

def plot_yearly_hurr_dist(yearly_hurr_distribution):
    start_doy = dt.datetime(2001, 6, 1).timetuple().tm_yday
    end_doy = dt.datetime(2001, 12, 1).timetuple().tm_yday

    fig = plt.figure()
    # plt.title('Hurricane Distribution over the Year')
    plt.plot(yearly_hurr_distribution.keys(), yearly_hurr_distribution.values())
    plt.plot((start_doy, start_doy), (0, 250), 'k--')
    plt.plot((end_doy, end_doy), (0, 250), 'k--')
    plt.xlabel('Day of Year')
    plt.ylabel('Hurricane-timesteps')
    fig.set_size_inches(6.3, 3)
    _save_figure('yearly_hurr_dist.png')

def plot_hurr_per_year(hurr_per_year):
    fig = plt.figure()
    # plt.title('Hurricanes per Year')
    plt.plot(hurr_per_year.keys(), hurr_per_year.values())
    plt.xlim((1890, 2010))
    plt.ylim((0, 300))

    plt.plot((1914, 1914), (0, 300), 'k--')
    plt.plot((1944, 1944), (0, 300), 'k--')
    plt.plot((1966, 1966), (0, 300), 'k--')

    plt.annotate('Panama\nCanal', xy=(1914, 250), xytext=(1915, 250), fontsize=10)
    plt.annotate('Aircraft\nRecon.', xy=(1944, 250), xytext=(1945, 250), fontsize=10)
    plt.annotate('Satellite', xy=(1966, 250), xytext=(1967, 250), fontsize=10)

    plt.xlabel('Year')
    plt.ylabel('Hurricane-timesteps')
    fig.set_size_inches(6.3, 3)
    _save_figure('hurr_per_year.png')

def plot_cdp_with_hurr_info(cla_analysis=None):
    if not cla_analysis:
        cla_analysis = ClassificationAnalysis()

    val_cd = cla_analysis.load_val_classification_data()
    _plot_2005_cdp(val_cd)
    return val_cd

def plot_cal_cd(cal_cd):
    m = cal_cd.are_hurr_actual

    i1 = SCATTER_ATTRS['vort']['index']
    i2 = SCATTER_ATTRS['pmin']['index']

    fig = plt.figure()

    ax = plt.subplot(131)
    plt.plot(cal_cd.data[:, i1][m] * 10000, cal_cd.data[:, i2][m] / 100.,
             'ro', zorder=0, label='hurricane')
    plt.plot(cal_cd.data[:, i1][~m] * 10000, cal_cd.data[:, i2][~m] / 100.,
             'bx', zorder=1, label='not hurricane')
    plt.xlim((0, 4.5))
    plt.ylim((920, 1040))
    ax.set_xticks((0, 1.5, 3, 4.5))

    plt.ylabel('Pressure (hPa)')
    plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')

    ax = plt.subplot(132)
    plt.plot(cal_cd.data[:, i1][m] * 10000, cal_cd.data[:, i2][m] / 100.,
             'ro', zorder=0, label='hurricane')
    plt.xlim((0, 4.5))
    plt.ylim((920, 1040))
    ax.set_xticks((0, 1.5, 3, 4.5))

    plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')
    plt.setp(ax.get_yticklabels(), visible=False)

    ax = plt.subplot(133)
    plt.plot(cal_cd.data[:, i1][~m] * 10000, cal_cd.data[:, i2][~m] / 100.,
             'bx', zorder=1, label='not hurricane')
    plt.plot(-10, -10, 'ro', zorder=1, label='hurricane')  # dummy point.
    plt.xlim((0, 4.5))
    plt.ylim((920, 1040))
    ax.set_xticks((0, 1.5, 3, 4.5))

    plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')
    plt.legend(bbox_to_anchor=(1.3, 1.16), numpoints=1, prop={'size': 10})
    plt.setp(ax.get_yticklabels(), visible=False)

    fig.set_size_inches(6.3, 2.3)

    _save_figure('cal_cdp_with_hurrs.png')

def _plot_2005_cdp(val_cd):
    m = val_cd.are_hurr_actual
    m2005 = val_cd.data[:, -2] == 2005  # -2 is year col.
    data = val_cd.data[m2005]
    m = m[m2005]

    i1 = SCATTER_ATTRS['vort']['index']
    i2 = SCATTER_ATTRS['pmin']['index']

    fig = plt.figure()

    labelx = -0.35
    ax = plt.subplot(231)
    plt.plot(data[:, i1][m] * 10000, data[:, i2][m] / 100., 'ro', zorder=0, label='hurricane')
    plt.plot(data[:, i1][~m] * 10000, data[:, i2][~m] / 100., 'bx', zorder=1, label='not hurricane')
    plt.xlim((0, 3))
    plt.ylim((940, 1040))
    ax.set_xticks((0, 1, 2, 3))
    ax.yaxis.set_label_coords(labelx, 0.5)

    plt.ylabel('Pressure (hPa)')
    # plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')

    ax = plt.subplot(232)
    plt.plot(data[:, i1][m] * 10000, data[:, i2][m] / 100., 'ro', zorder=0, label='hurricane')
    plt.xlim((0, 3))
    plt.ylim((940, 1040))
    ax.set_xticks((0, 1, 2, 3))

    # plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')
    plt.setp(ax.get_yticklabels(), visible=False)

    ax = plt.subplot(233)
    plt.plot(data[:, i1][~m] * 10000, data[:, i2][~m] / 100., 'bx', zorder=1, label='not hurricane')
    plt.plot(-10, -10, 'ro', zorder=1, label='hurricane')  # dummy point.
    plt.xlim((0, 3))
    plt.ylim((940, 1040))
    ax.set_xticks((0, 1, 2, 3))

    # plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')
    plt.legend(bbox_to_anchor=(1.3, 1.16), numpoints=1, prop={'size': 10})
    plt.setp(ax.get_yticklabels(), visible=False)

    i2 = SCATTER_ATTRS['t850']['index']

    ax = plt.subplot(234)
    plt.plot(data[:, i1][m] * 10000, data[:, i2][m], 'ro', zorder=0, label='hurricane')
    plt.plot(data[:, i1][~m] * 10000, data[:, i2][~m], 'bx', zorder=1, label='not hurricane')
    plt.xlim((0, 3))
    plt.ylim((250, 310))
    ax.set_xticks((0, 1, 2, 3))
    ax.yaxis.set_label_coords(labelx, 0.5)

    plt.ylabel('Temp. at 850 hPa (K)')
    # plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')

    ax = plt.subplot(235)
    plt.plot(data[:, i1][m] * 10000, data[:, i2][m], 'ro', zorder=0, label='hurricane')
    plt.xlim((0, 3))
    plt.ylim((250, 310))
    ax.set_xticks((0, 1, 2, 3))

    plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')
    plt.setp(ax.get_yticklabels(), visible=False)

    ax = plt.subplot(236)
    plt.plot(data[:, i1][~m] * 10000, data[:, i2][~m], 'bx', zorder=1, label='not hurricane')
    plt.plot(-10, -10, 'ro', zorder=1, label='hurricane')  # dummy point.
    plt.xlim((0, 3))
    plt.ylim((250, 310))
    ax.set_xticks((0, 1, 2, 3))

    # plt.xlabel('Vorticity ($10^{-4}$ s$^{-1}$)')
    # plt.legend(bbox_to_anchor=(1.3, 1.16), numpoints=1, prop={'size': 10})
    plt.setp(ax.get_yticklabels(), visible=False)

    fig.set_size_inches(6.3, 4.6)

    _save_figure('cdp_2005_with_hurrs.png')

def _plot_individual_katrina_figure(em=7):
    c20data = C20Data(2005, fields=['prmsl', 'u', 'v'])
    srm = StormtracksResultsManager('pyro_tracking_analysis')
    ta = StormtracksAnalysis(2005)
    ibdata = IbtracsData()
    w, k = ibdata.load_wilma_katrina()
    bt = k
    loc = {'llcrnrlat': 10, 'urcrnrlat': 45, 'llcrnrlon': -100, 'urcrnrlon': -65}

    config = {'pressure_level': 850, 'scale': 1, 'tracker': 'nearest_neighbour'}

    fig = plt.figure()
    plt.clf()

    print(config)
    all_matches = []

    key = ta.good_matches_key(config)
    good_matches = srm.get_result(2005, em, key)

    for good_match in good_matches:
        if good_match.best_track.name == bt.name:
            break

    _raster_on_earth(c20data.lons, c20data.lats, None, loc=loc)
    plotting.plot_match_with_date(good_match, None)
    _save_figure('katrina_individual_match_em7')

def _plot_six_configs_figure():
    c20data = C20Data(2005, fields=['prmsl', 'u', 'v'])
    srm = StormtracksResultsManager('pyro_tracking_analysis')
    ta = StormtracksAnalysis(2005)
    ibdata = IbtracsData()
    w, k = ibdata.load_wilma_katrina()
    bt = k
    loc = {'llcrnrlat': 10, 'urcrnrlat': 45, 'llcrnrlon': -100, 'urcrnrlon': -65}
    fig3 = plt.figure(3)
    plt.clf()

    for j, config in enumerate(ta.analysis_config_options):
        plt.subplot(3, 2, j + 1)
        print(config)
        if config['pressure_level'] == 995:
            title = 'NSLP, Scale: {scale}'.format(**config)
        else:
            title = '850 hPa, Scale: {scale}'.format(**config)
        plt.title(title, fontsize=10)
        all_matches = []
        _raster_on_earth(c20data.lons, c20data.lats, None, loc=loc, labels=False)
        plotting.plot_track(bt, zorder=2)

        for i in range(56):
            key = ta.good_matches_key(config)
            good_matches = srm.get_result(2005, i, key)
            matches = []
            for good_match in good_matches:
                if good_match.best_track.name == bt.name:
                    matches.append(good_match)

            if matches:
                all_matches.append(matches)
                for match in matches:
                    vt = match.vort_track
                    mask = (vt.dates >= bt.dates[0]) & (vt.dates <= bt.dates[-1])
                    plotting.plot_path_on_earth(vt.lons[mask], vt.lats[mask], 'b--')

                    # plotting.plot_track(vt, 'b--')
                    # return vt, bt
            else:
                print('Could not find wilma in {0}-{1}'.format(i, key))

    fig3.set_size_inches(5.4, 7.8)

    _save_figure('katrina_six_tracking_configs')

def plot_katrina_correlation(ca=None, em=7):
    fig = plt.figure()
    ibdata = IbtracsData()
    w, k = ibdata.load_wilma_katrina()
    if not ca:
        ca = ClassificationAnalysis()
    cs, ms, ums = ca.run_individual_cla_analysis(2005, em)

    pressures = []
    winds = []
    dates = []

    for cm in ms:
        if cm.best_track.name == k.name:
            for date, bt_pres, bt_wind in zip(cm.best_track.dates,
                                              cm.best_track.pressures,
                                              cm.best_track.winds):
                if date in cm.cyclone.pmins and cm.cyclone.pmins[date]:
                    dates.append(date)
                    pressures.append((bt_pres, cm.cyclone.pmins[date] / 100.))
                    winds.append((bt_wind, cm.cyclone.max_windspeeds[date]))

    pressures, winds = np.array(pressures), np.array(winds)

    labelx = -0.1
    ax = plt.subplot(211)
    plt.plot_date(dates, pressures[:, 0], 'b-', label='best track')
    plt.plot_date(dates, pressures[:, 1], 'b--', label='derived track')
    plt.ylabel('pressure (hPa)')
    plt.legend(bbox_to_anchor=(1.07, 1.16), numpoints=1, prop={'size': 10})
    # ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%b %d'))
    plt.setp(ax.get_xticklabels(), visible=False)
    ax.yaxis.set_label_coords(labelx, 0.5)

    ax = plt.subplot(212)
    plt.plot_date(dates, winds[:, 0], 'r-', label='best track')
    plt.plot_date(dates, winds[:, 1], 'r--', label='derived track')
    plt.ylabel('max. wind speed (ms$^{-1}$)')
    plt.legend(bbox_to_anchor=(1.07, 1.07), numpoints=1, prop={'size': 10})
    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%b %d'))
    ax.yaxis.set_label_coords(-0.09, 0.5)

    fig.set_size_inches(6.3, 5)
    _save_figure('katrina_best_derived_comparison.png')

def plot_2005_pressure_wind_corr(ca=None):
    if not ca:
        ca = ClassificationAnalysis()

    # c20data = C20Data(2005, fields=['prmsl', 'u', 'v'])
    key = 'cyclones'

    pressures = []
    winds = []

    for em in range(56):
        print(em)
        cs, ms, ums = ca.run_individual_cla_analysis(2005, em)

        for cm in ms:
            for date, bt_pres, bt_wind in zip(cm.best_track.dates,
                                              cm.best_track.pressures,
                                              cm.best_track.winds):
                if date in cm.cyclone.pmins and cm.cyclone.pmins[date]:
                    pressures.append((bt_pres, cm.cyclone.pmins[date] / 100.))
                    winds.append((bt_wind, cm.cyclone.max_windspeeds[date]))
    pressures, winds = np.array(pressures), np.array(winds)
    _plot_pres_wind(pressures, winds)
    return pressures, winds

def _plot_pres_wind(pressures, winds):
    fig = plt.figure()
    ax = plt.subplot(121)
    plt.plot(pressures[:, 0], pressures[:, 1], 'b+')
    rp = stats.linregress(pressures)
    label = 'grad.: {0:.2f}\nintercept: {1:.1f}\n r$^2$: {2:.2f}'.format(rp[0], rp[1], rp[2] ** 2)
    plt.plot((880, 1040), (880 * rp[0] + rp[1], 1040 * rp[0] + rp[1]), 'r--', label=label)
    plt.ylabel('derived track pressure (hPa)')
    plt.xlabel('best track\npressure (hPa)')
    plt.legend(bbox_to_anchor=(0.9, 1.23), numpoints=1, prop={'size': 10})
    ax.set_xticks((880, 920, 960, 1000, 1040))

    ax = plt.subplot(122)
    plt.plot(winds[:, 0], winds[:, 1], 'b+')
    rw = stats.linregress(winds)
    label = 'grad.: {0:.2f}\nintercept: {1:.1f}\n r$^2$: {2:.2f}'.format(rw[0], rw[1], rw[2] ** 2)
    plt.plot((0, 160), (0 * rw[0] + rw[1], 160 * rw[0] + rw[1]), 'r--', label=label)
    plt.ylabel('derived track max. wind speed (ms$^{-1}$)')
    plt.xlabel('best track\nmax. wind speed (ms$^{-1}$)')
    plt.legend(bbox_to_anchor=(0.9, 1.23), numpoints=1, prop={'size': 10})
    ax.set_xticks((0, 40, 80, 120, 160))
    ax.yaxis.tick_right()
    ax.yaxis.set_label_position("right")

    fig.set_size_inches(6.3, 3)
    _save_figure('press_max_ws_corr_2005.png')

def _plot_point_on_earth(lon, lat, plot_fmt=None):
    if plot_fmt:
        plt.plot(lon_convert(lon), lat, plot_fmt)
    else:
        plt.plot(lon_convert(lon), lat)

def _raster_on_earth(lons, lats, data, vmin=None, vmax=None, loc=None, colorbar=True, labels=True):
    if not loc:
        m = Basemap(projection='cyl', resolution='c',
                    llcrnrlat=-90, urcrnrlat=90, llcrnrlon=-180, urcrnrlon=180)
    else:
        m = Basemap(projection='cyl', resolution='c', **loc)

    if data is not None:
        plot_lons, plot_data = _extend_data(lons, lats, data)
        lons, lats = np.meshgrid(plot_lons, lats)
        x, y = m(lons, lats)
        if vmin:
            m.pcolormesh(x, y, plot_data, vmin=vmin, vmax=vmax)
        else:
            m.pcolormesh(x, y, plot_data)

    m.drawcoastlines()

    if labels:
        p_labels = [0, 1, 0, 0]
        m.drawparallels(np.arange(-90., 90.1, 45.), labels=p_labels, fontsize=10)
        m.drawmeridians(np.arange(-180., 180., 60.), labels=[0, 0, 0, 1], fontsize=10)

    if colorbar and data is not None:
        m.colorbar(location='right', pad='7%')
    return m

def _vec_plot_on_earth(lons, lats, x_data, y_data, vmin=-4, vmax=12, loc=None, colorbar=False):
    plot_lons, plot_x_data = _extend_data(lons, lats, x_data)
    plot_lons, plot_y_data = _extend_data(lons, lats, y_data)

    lons, lats = np.meshgrid(plot_lons, lats)

    if not loc:
        m = Basemap(projection='cyl', resolution='c',
                    llcrnrlat=-90, urcrnrlat=90, llcrnrlon=-180, urcrnrlon=180)
    else:
        m = Basemap(projection='cyl', resolution='c', **loc)
    x, y = m(lons, lats)

    mag = np.sqrt(plot_x_data**2 + plot_y_data**2)
    vmin, vmax = mag.min(), mag.max()
    m.contourf(x, y, mag)
    # m.pcolormesh(x, y, mag, vmin=vmin, vmax=vmax)
    # m.quiver(x, y, plot_x_data, plot_y_data)
    skip = 1
    m.quiver(x[::skip, ::skip], y[::skip, ::skip],
             plot_x_data[::skip, ::skip], plot_y_data[::skip, ::skip], scale=500)

    m.drawcoastlines()
    m.drawparallels(np.arange(-90., 90., 45.), labels=[1, 0, 0, 0], fontsize=10)
    m.drawmeridians(np.arange(-180., 180., 60.), labels=[0, 0, 0, 1], fontsize=10)

    if colorbar:
        m.colorbar(location='right', pad='7%')
    return m

def _extend_data(lons, lats, data):
    if False:
        # TODO: probably doesn't work!
        # Adds extra data at the end.
        plot_offset = 2
        plot_lons = np.zeros((lons.shape[0] + plot_offset,))
        plot_lons[:-plot_offset] = lons
        plot_lons[-plot_offset:] = lons[-plot_offset:] + 3.75 * plot_offset

        plot_data = np.zeros((data.shape[0], data.shape[1] + plot_offset))
        plot_data[:, :-plot_offset] = data
        plot_data[:, -plot_offset:] = data[:, :plot_offset]
    else:
        # Adds extra data before the start.
        delta = lons[1] - lons[0]
        plot_offset = 180
        plot_lons = np.ma.zeros((lons.shape[0] + plot_offset,))
        plot_lons[plot_offset:] = lons
        plot_lons[:plot_offset] = lons[-plot_offset:] - delta * (lons.shape[0])

        plot_data = np.ma.zeros((data.shape[0], data.shape[1] + plot_offset))
        plot_data[:, plot_offset:] = data
        plot_data[:, :plot_offset] = data[:, -plot_offset:]

    return plot_lons, plot_data

def _save_figure(name):
    if not os.path.exists(settings.FIGURE_OUTPUT_DIR):
        os.makedirs(settings.FIGURE_OUTPUT_DIR)
    plt.savefig(os.path.join(settings.FIGURE_OUTPUT_DIR, name), bbox_inches='tight')

def main():
    import sys
    import inspect

    print('Loading all data from: {}'.format(settings.DATA_DIR))
    ib_hurrs, ib_pdis, cla_hurrs, cla_pdis, adjustment_ratio = load_20th_century()
    thresh_metrics = load_thresh_metric_vs_vort_lo()
    classifier_metrics = load_classifier_metrics()
    yearly_hurr_distribution, hurr_per_year = load_ibtracs_info()
    print('Data loaded')

    functions = inspect.getmembers(sys.modules[__name__], inspect.isfunction)
    for name, fn in functions:
        if name[:5] == 'plot_':
            try:
                print(name)
                if name in ('plot_20th_century',
                            'plot_poster_20th_century',
                            'plot_poster_20th_century_trends',
                            'plot_20th_century_corr'):
                    fn(ib_hurrs, cla_hurrs, adjustment_ratio)
                elif name == 'plot_thresh_metric_vs_vort_lo':
                    fn(thresh_metrics)
                elif name == 'plot_classifer_metrics':
                    fn(classifier_metrics)
                elif name == 'plot_yearly_hurr_dist':
                    fn(yearly_hurr_distribution)
                elif name == 'plot_hurr_per_year':
                    fn(hurr_per_year)
                elif name == 'plot_cal_cd':
                    cla_analysis = ClassificationAnalysis()
                    cal_cd = cla_analysis.load_cal_classification_data()
                    plot_cal_cd(cal_cd)
                else:
                    fn()
            except:
                print('PROBLEM PLOTTING {0}, SKIPPING'.format(name))

if __name__ == '__main__':
    main()