"""Plots tornado reports, storm tracks, and linkages."""

import copy
import os.path
import argparse
import numpy
import pandas
import matplotlib
matplotlib.use('agg')
import matplotlib.colors
import matplotlib.pyplot as pyplot
from gewittergefahr.gg_io import tornado_io
from gewittergefahr.gg_utils import linkage
from gewittergefahr.gg_utils import time_conversion
from gewittergefahr.gg_utils import storm_tracking_utils as tracking_utils
from gewittergefahr.gg_utils import file_system_utils
from gewittergefahr.plotting import plotting_utils
from gewittergefahr.plotting import storm_plotting

# TODO(thunderhoser): Maybe put some of this code in linkage_plotting.py.
#  Either way, methods need unit tests.

# TODO(thunderhoser): Turn this into one plot for each tornado, rather than one
# with all.

LOG_MESSAGE_TIME_FORMAT = '%Y-%m-%d-%H%M%S'
SEPARATOR_STRING = '\n\n' + '*' * 50 + '\n\n'

SENTINEL_VALUE = -9999
LATLNG_TOLERANCE_DEG = 0.001

COLOUR_MAP_NAME = 'YlOrRd'

LINKAGE_FONT_SIZE = 12
LINKAGE_FONT_COLOUR = numpy.full(3, 0.)

TORNADO_FONT_SIZE = 18
TORNADO_START_MARKER_TYPE = 'o'
TORNADO_START_MARKER_SIZE = 8
TORNADO_START_MARKER_EDGE_WIDTH = 1
TORNADO_END_MARKER_TYPE = 'x'
TORNADO_END_MARKER_SIZE = 12
TORNADO_END_MARKER_EDGE_WIDTH = 3

TEXT_BOUNDING_BOX_DICT = {
    'alpha': 0.5, 'edgecolor': 'k', 'linewidth': 1
}

NUM_PARALLELS = 8
NUM_MERIDIANS = 6
LATLNG_BUFFER_DEG = 0.5
BORDER_COLOUR = numpy.full(3, 0.)
FIGURE_RESOLUTION_DPI = 300

SHORT_TORNADO_ID_COLUMN = 'short_tornado_id_string'

LINKAGE_DIR_ARG_NAME = 'input_linkage_dir_name'
GENESIS_ONLY_ARG_NAME = 'genesis_only'
MAX_DISTANCE_ARG_NAME = 'max_link_distance_metres'
FIRST_DATE_ARG_NAME = 'first_spc_date_string'
LAST_DATE_ARG_NAME = 'last_spc_date_string'
NUM_COLOURS_ARG_NAME = 'num_colours'
MIN_LATITUDE_ARG_NAME = 'min_plot_latitude_deg'
MAX_LATITUDE_ARG_NAME = 'max_plot_latitude_deg'
MIN_LONGITUDE_ARG_NAME = 'min_plot_longitude_deg'
MAX_LONGITUDE_ARG_NAME = 'max_plot_longitude_deg'
OUTPUT_FILE_ARG_NAME = 'output_file_name'

LINKAGE_DIR_HELP_STRING = (
    'Name of top-level directory with linkage files.  Files therein will be '
    'found by `linkage.find_linkage_file` and read by '
    '`linkage.read_linkage_file`.')

GENESIS_ONLY_HELP_STRING = (
    'Boolean flag.  If 1, will plot linkages only to tornadogenesis events.  If'
    ' 0, will plot linkages to tornado occurrences.')

MAX_DISTANCE_HELP_STRING = (
    'Max linkage distance.  Will not show linkages with greater distance.  To '
    'plot all linkages, leave this alone.')

SPC_DATE_HELP_STRING = (
    'SPC date (format "yyyymmdd").  Linkages will be plotted for the period '
    '`{0:s}`...`{1:s}`.'
).format(FIRST_DATE_ARG_NAME, LAST_DATE_ARG_NAME)

NUM_COLOURS_HELP_STRING = (
    'Number of colours in colour scheme (tornado reports and storm-track '
    'segments are coloured by time).')

LATITUDE_HELP_STRING = (
    'Latitude (deg N, in range -90...90).  Plotting area will be '
    '`{0:s}`...`{1:s}`.  To let plotting area be determined by storm tracks, '
    'leave this alone.'
).format(MIN_LATITUDE_ARG_NAME, MAX_LATITUDE_ARG_NAME, SENTINEL_VALUE)

LONGITUDE_HELP_STRING = (
    'Longitude (deg E, in range 0...360).  Plotting area will be '
    '`{0:s}`...`{1:s}`.  To let plotting area be determined by storm tracks, '
    'leave this alone.'
).format(MIN_LONGITUDE_ARG_NAME, MAX_LONGITUDE_ARG_NAME, SENTINEL_VALUE)

OUTPUT_FILE_HELP_STRING = 'Path to output file (figure will be saved here).'

INPUT_ARG_PARSER = argparse.ArgumentParser()
INPUT_ARG_PARSER.add_argument(
    '--' + LINKAGE_DIR_ARG_NAME, type=str, required=True,
    help=LINKAGE_DIR_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + GENESIS_ONLY_ARG_NAME, type=int, required=False, default=1,
    help=GENESIS_ONLY_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + MAX_DISTANCE_ARG_NAME, type=float, required=False, default=0.,
    help=MAX_DISTANCE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + FIRST_DATE_ARG_NAME, type=str, required=True,
    help=SPC_DATE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + LAST_DATE_ARG_NAME, type=str, required=True,
    help=SPC_DATE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + NUM_COLOURS_ARG_NAME, type=int, required=False, default=12,
    help=NUM_COLOURS_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + MIN_LATITUDE_ARG_NAME, type=float, required=False,
    default=SENTINEL_VALUE, help=LATITUDE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + MAX_LATITUDE_ARG_NAME, type=float, required=False,
    default=SENTINEL_VALUE, help=LATITUDE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + MIN_LONGITUDE_ARG_NAME, type=float, required=False,
    default=SENTINEL_VALUE, help=LONGITUDE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + MAX_LONGITUDE_ARG_NAME, type=float, required=False,
    default=SENTINEL_VALUE, help=LONGITUDE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + OUTPUT_FILE_ARG_NAME, type=str, required=True,
    help=OUTPUT_FILE_HELP_STRING)


def _long_to_short_tornado_ids(long_id_strings):
    """Converts long tornado IDs to short IDs.

    N = number of IDs

    :param long_id_strings: length-N list of long IDs.
    :return: short_id_strings: length-N list of short IDs.
    """

    unique_long_id_strings, orig_to_unique_indices = numpy.unique(
        numpy.array(long_id_strings), return_inverse=True
    )

    short_id_strings = [''] * len(long_id_strings)

    for i in range(len(unique_long_id_strings)):
        these_indices = numpy.where(orig_to_unique_indices == i)[0]

        for j in these_indices:
            short_id_strings[j] = '{0:d}'.format(i)

    return short_id_strings


def _truncate_colour_map(
        orig_colour_map_object, num_colours, min_normalized_value=0.25,
        max_normalized_value=1.):
    """Truncates colour map.

    :param orig_colour_map_object: Original colour map (instance of
        `matplotlib.pyplot.cm`).
    :param num_colours: Number of colours in new map.
    :param min_normalized_value: Minimum normalized value for new map (in range
        0...1).
    :param max_normalized_value: Max normalized value for new map (in range
        0...1).
    :return: colour_map_object: New colour map (instance of
        `matplotlib.pyplot.cm`).
    """