#!/usr/bin/python
# _param_help.py

# flake8: noqa


'''
Helper functions for assessing the effects of different parameter settings.
'''

# geonomics imports
from geonomics.utils.viz import _check_display

# other imports
import os
import numpy as np
import matplotlib as mpl
_check_display()
import matplotlib.pyplot as plt
from scipy.stats import levy as _s_levy
from copy import deepcopy

# ------------------------------------
# CLASSES ---------------------------
# ------------------------------------


# --------------------------------------
# FUNCTIONS ---------------------------
# --------------------------------------



def plot_movement(spp, land, num_timesteps, lyr_num=None, params=None,
                  direction_distr_mu=None, direction_distr_kappa=None,
                  distance_distr_param1=None, distance_distr_param2=None,
                  distance_distr='levy',
                  move_surf=None, subset_spp=None, color='black',
                  alpha=None, land_cmap='plasma',  size=10, ticks=False,
                  color_by_individ=True, increasing_linewidth=True,
                  include_start_points=False):

    '''
    Plot example movement tracks generated by the input movement parameters

    Input a ParametersDict object with movement parameter values of interest,
    or input those values in piecemeal fashion, to generate a plot of example
    movement tracks generated by that parameterization. Useful for visual
    exploration of the movement trends in a Geonomics model.

    Parameters
    ----------
    spp : :class:`geonomics.structs.Species`
        A Species object whose movement behavior should be plotted

    land : :class:`geonomics.structs.Landscape`
        A Land object upon which the Species' movement should be plotted

    num_timesteps : int
        A number of timesteps (i.e. individual movement steps) to be plotted in
        each individual's movement track.

    lyr_num : int, optional, default: None
        The number of the Layer to be plotted (i.e. its key in the Landscape
        dict). If None, all Layers will be plotted as a stack of transparent,
        differentially colormapped rasters.

    params : {:class:`geonomics.sim.params.ParametersDict`}, optional,
    default: None
        If a ParametersDict object is passed, then the movement resulting
        from the movement parameter values it contains will be plotted (except
        for any values that are manually overridden by new values passed
        individually to any of the following four parameters). If None is
        passed instead, then all four of the following movement parameters'
        values must be passed, and their resulting movement behavior will be
        plotted.

    direction_distr_mu,
    direction_distr_kappa,
    distance_distr_param1,
    distance_distr_param2 : float, optional default: None
        These are the parameters defining the distributions from which the
        direction and distance values for individual movements are drawn. If
        passed, their values will be used to determine the movement behavior
        that is plotted (even if a ParametersDict was passed to *params*,
        because these values will override any values therein). For more
        details, see the 'Parameters' section of the online documentation. 

    distance_distr: str, optional default: 'levy'
        This indicates the distribution to use for drawing movement
        distancs. Defaults to 'levy', but can also be set to 'wald'.

    move_surf : :class:`geonomics.utils.spatial._ConductanceSurface`, optional,
    default: None
        If a _ConductanceSurface (i.e. a Species' movement or dispersal
        surface) is passed, then it will be used to determine the plotted
        movement behavior, instead of any movement or dispersal surface already
        associated with the Species object passed to *spp*.

    subset_spp: int, optional, default: None
        The number of Individuals to subset from the Species and plot (rather
        than plotting movement tracks for all Individuals in the Species, which
        can create a jumbled plot that is hard to inspect).

    color : valid mpl.plt color value, optional, default: 'black'
        Face color for the lines in the movement-track plot.  Passed to the
        *c* argument of matplotlib.pyplot.scatter.

    alpha : valid mpl.plt alpha value, optional, default: None
        Transparency value for the lines in the movement-track plot. Passed
        to the *alpha* argument of matplotlib.pyplot.scatter.

    size : scalar or array_like, optional, default: 10
        Size of the scatter points. Passed to the *s* argument of
        matplotlib.pyplot.scatter.

    land_cmap : {valid string, None}, optional, default: 'plasma'
        Colormap to use for plotting Layer rasters. If None, will default
        to the colormap automatically assigned to the Layer by Geonomics
        based on the Layer's index number in the Landscape dict. Can be
        passed any string that references a colormap in matplotlib.pyplot.cm.

    ticks : bool, optional, default: False
        If True, x- and y-axis ticks will be added to the plot.

    color_by_individ : bool, optional, default: True
        If True, each plotted Individual's movement track will have its own
        distinct color, making it easier to follow the track and distinguish it
        from others in the plot.

    increasing_linewidth : bool, optional, default: True
        If True, each Individual's movement track will get wider with each
        successive movement step, making it easier to distinguish
        directionality in the movement track.

    include_start_points : bool, optional, default: False
        If True, a scatter of points will be plotted to signify the starting
        points of individuals' movement tracks.
    '''

    #assert that either a params dict was provided or parameters were
    #stipulated
    assert (params is not None
        or (distance_distr_param1 is not None
            and distance_distr_param2 is not None
            and direction_distr_mu is not None
            and direction_distr_kappa is not None)), ("Either a parameters "
        "dictionary must be provided, or parameters must be individually "
        "defined as arguments.")

    #if a params dictionary was fed in, copy it as toy_params
    if params is not None:
        toy_params = deepcopy(params)
        #copy the movement params from toy_params
        toy_move_params = toy_params.comm.species[spp.name].movement
        #then override existing values in the movement params dictionary,
        #if they have new values that have been entered for this trial
        if direction_distr_mu != None:
            toy_move_params['direction_distr_mu'] = direction_distr_mu
        if direction_distr_kappa != None:
            toy_move_params['direction_distr_kappa'] = direction_distr_kappa
        if distance_distr_param1  != None:
            toy_move_params['movement_distance_distr_param1'] = distance_distr_param1
        if distance_distr_param2 != None:
            toy_move_params['movement_distance_distr_param2'] = distance_distr_param2
        if move_surf != None:
            toy_move_params['move_surf'] = move_surf
        toy_params.comm.species[spp.name].movement = toy_move_params

    #if no params dictionary provided, create toy_params from
    #the individual parameter values that were fed in
    elif params is None:
        toy_move_params = {}
        toy_move_params['direction_distr_mu'] = direction_distr_mu
        toy_move_params['direction_distr_kappa'] = direction_distr_kappa
        toy_move_params['movement_distance_distr_param1'] = distance_distr_param1
        toy_move_params['movement_distance_distr_param2'] = distance_distr_param2
        if move_surf != None:
            toy_move_params['move_surf'] = move_surf
        elif spp.move_surf == None:
            toy_move_params['move_surf'] = False
        else:
            toy_move_params['move_surf'] = True
            toy_params = {'comm':{'species':{'movement':toy_move_params}}}

    #now copy the species to a toy object
    toy_spp = deepcopy(spp)
    #replace its movement-parameter values
    for param in ['direction_distr_mu', 'direction_distr_kappa',
        'movement_distance_distr_param1', 'movement_distance_distr_param2']:
        setattr(toy_spp, param,
            toy_params['comm']['species'][spp.name]['movement'][param])

    #replace its movement surface, if necessary
    if move_surf is not None:
        toy_spp._move_surf = move_surf

    #subset individuals, if necessary
    if subset_spp:
        cull_individs = np.random.choice([*toy_spp],
            len(toy_spp) - subset_spp, replace = False)
        [toy_spp.pop(ind) for ind in cull_individs]
        #reset the coords before plotting
        toy_spp._set_coords_and_cells()

    #set the plotting linewidths to increase over runtime
    if increasing_linewidth:
        linewidths = np.linspace(1, 5, num=num_timesteps)
    else:
        linewidths = np.array([1] * num_timesteps)

    #Create a colors list, to plot different-colored paths for individuals
    #colors = ['black', 'red', 'orange', 'yellow', 'green', 'blue', 'white']
    #colors = [plt.cm.Greys_r(i) for i in np.linspace(0, 0.9, 9)]
    colors = ['#000000', '#2F4F4F', '#778899', '#696969', '#A9A9A9', '#C0C0C0']

    # get list of order in which to plot individuals
    ind_idx_list = [i.idx for i in toy_spp.values()]
    ind_idx_list = sorted(ind_idx_list)

    #plot the species at its starting locations
    ind_colors = [colors[i % len(colors)] for i in range(len(ind_idx_list))]
    toy_spp._plot(lyr_num=lyr_num, land=land, color=ind_colors,
                  size=25 * include_start_points, land_cmap=land_cmap,
                  ticks=ticks)

    #set the new_x and new_y objects to the current locations, before for loop
    #that will iteratively update them to the newer locations after movement
    new_x = [getattr(toy_spp[idx], 'x') for idx in ind_idx_list]
    new_y = [getattr(toy_spp[idx], 'y') for idx in ind_idx_list]

    #loop for the number of timesteps, iteratively moving individuals and
    #plotting the vectors between their previous and new positions
    for t in range(num_timesteps):
        old_x = [x for x in new_x]
        old_y = [y for y in new_y]

        toy_spp._do_movement(land)

        new_x = [getattr(toy_spp[idx], 'x') for idx in ind_idx_list]
        new_y = [getattr(toy_spp[idx], 'y') for idx in ind_idx_list]

        #plot with separate colors for each individual, if necessary
        if color_by_individ:
            [plt.plot((old_x[i], new_x[i]), (old_y[i], new_y[i]), '-',
                scalex = False, scaley = False, linewidth = linewidths[t],
                color = colors[i%len(colors)],
                alpha = 0.5) for i in range(len(old_x))];

        #or else plot without individuals colored differently
        else:
            plt.plot((old_x, new_x), (old_y, new_y), '-',
                scalex = False, scaley = False, linewidth = linewidths[t],
                color = color, alpha = alpha)


# functions to visualize the different distributions that can be
# parameterized from the parameters file
def plot_distr_movement_distance(spp=None, distance_distr_param1=None,
                                 distance_distr_param2=None,
                                 distance_dist='levy'):
    if spp is not None:
        distance_distr_param1 = spp.movement_distance_distr_param1
        distance_distr_param2 = spp.movement_distance_distr_param2
    else:
        assert (distance_distr_param1 is not None and
                distance_distr_param2 is not None), ('If a Species object '
                                                     'is not provided then '
                                                     'the parameter values '
                                                     'must be provided.')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    if distance_distr == 'levy':
        fig.suptitle(('movement distance: '
                  '~Levy($\loc$=%.4E, $\scale$=%.4E)') % (distance_distr_param1,
                                                        distance_distr_param2))
        vals = _s_levy.rvs(loc=distance_distr_param1,
                              scale=distance_distr_param2, size=10000)
    elif distance_distr == 'wald':
        fig.suptitle(('movement distance: '
                  '~Wald($\mu$=%.4E, $\scale$=%.4E)') % (distance_distr_param1,
                                                         distance_distr_param2))
        vals = np.random.wald(mean=distance_distr_param1,
                              scale=distance_distr_param2, size=10000)
    ax.hist(vals, bins = 25)
    ax.set_xlim((min(vals), max(vals)))
    plt.show()


def plot_distr_movement_direction(spp=None, direction_distr_mu=None,
                            direction_distr_kappa=None):
    if spp is not None:
        direction_distr_mu = spp.direction_distr_mu
        direction_distr_sigma = spp.direction_distr_kappa
    else:
        assert (direction_distr_mu is not None and
                direction_distr_kappa is not None), ('If a Species object '
                                                     'is not provided then '
                                                     'the parameter values '
                                                     'must be provided.')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    fig.suptitle(('movement direction: '
                 '~von Mises($\mu$ = %.4E, '
                  '$\kappa$ = %.4E)') % (direction_distr_mu,
                                         direction_distr_kappa))
    vals = np.random.vonmises(direction_distr_mu, direction_distr_kappa, 10000)
    ax.hist(vals, bins = 25)
    ax.set_xlim((min(vals), max(vals)))
    plt.show()


def plot_distr_dispersal_distance(spp=None, dispersal_distr_param1=None,
                                  dispersal_distr_param2=None,
                                  dispersal_distance_distr='levy'):
    if spp is not None:
        dispersal_distr_param1 = spp.dispersal_distance_distr_param1
        dispersal_distr_param2 = spp.dispersal_distance_distr_param2
    else:
        assert (dispersal_distr_param1 is not None and
                dispersal_distr_param2 is not None), ('If a Species object '
                                                     'is not provided then '
                                                     'the parameter values '
                                                     'must be provided.')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    if dispersal_distance_distr == 'levy':
        fig.suptitle(('dispersal distance: '
                      '~Levy($\loc$=%.4E, $\scale$=%.4E)') % (dispersal_distr_param1,
                                                           dispersal_distr_param2))
        vals = _s_levy.rvs(dispersal_distr_param1, dispersal_distr_param2, 10000)
    elif dispersal_distance_distr == 'wald':
        fig.suptitle(('dispersal distance: '
                      '~Wald($\mean$=%.4E, $\scale$=%.4E)') % (dispersal_distr_param1,
                                                           dispersal_distr_param2))
        vals = np.random.wald(dispersal_distr_param1, dispersal_distr_param2, 10000)
    ax.hist(vals, bins = 25)
    ax.set_xlim((min(vals), max(vals)))
    plt.show()


def plot_n_births(spp=None, n_births_distr_lambda=None):
    if spp is not None:
        n_births_distr_lambda = spp.n_births_distr_lambda
    else:
        assert n_births_distr_lambda is not None ('If a Species object '
                                                  'is not provided then '
                                                  'the parameter values '
                                                  'must be provided.')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    fig.suptitle(('number of births per mating event: '
                  '~Poisson($\lamba$=%.4E)') % n_births_distr_lambda)
    vals = np.random.poisson(n_births_distr_lambda, 10000)
    ax.hist(vals, bins = 25)
    ax.set_xlim((min(vals), max(vals)))
    plt.show()


def plot_distr_delet_effect_sizes(spp=None, delet_alpha_distr_shape=None,
                                  delet_alpha_distr_scale=None):
    if spp is not None:
        delet_alpha_distr_shape = spp.gen_arch.delet_alpha_distr_shape
        delet_alpha_distr_scale = spp.gen_arch.delet_alpha_distr_scale
    else:
        assert (delet_alpha_distr_shape is not None
                and delet_alpha_distr_scale is not None), ('If a Species '
                                                           'object is not '
                                                           'provided then '
                                                           'the parameter '
                                                           'values must be '
                                                           'provided.')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    fig.suptitle(('effect sizes of deleterious mutations: '
                  '~Gamma($shape$=%.4E, '
                  '$scale$=%.4E)') % (delet_alpha_distr_shape,
                                      delet_alpha_distr_scale))
    vals = np.random.gamma(delet_alpha_distr_shape,
                           delet_alpha_distr_scale, 10000)
    ax.hist(vals, bins = 25)
    ax.set_xlim((min(vals), max(vals)))
    plt.show()


def plot_distr_trait_effect_sizes(spp=None, r_distr_alpha=None,
                            r_distr_beta=None):
    if spp is not None:
        r_distr_alpha = spp.gen_arch.recombinations._r_distr_alpha
        r_distr_beta = spp.gen_arch.recombinations._r_distr_beta
    else:
        assert (r_distr_alpha is not None
                and r_distr_beta is not None), ('If a Species '
                                                'object is not '
                                                'provided then '
                                                'the parameter '
                                                'values must be '
                                                'provided.')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    fig.suptitle(('recombination rates: '
                  '~Beta($\alpha$=%.4E, $\beta$=%.4E)') % (r_distr_alpha,
                                                           r_distr_beta))
    vals = np.random.beta(r_distr_alpha, r_distr_beta, 10000)
    ax.hist(vals, bins = 25)
    ax.set_xlim((min(vals), max(vals)))
    plt.show()


def plot_distr_recomb_rates(spp=None, trt_num=None, alpha_distr_mu=None,
                            alpha_distr_sigma=None):
    if spp is not None:
        assert trt_num is not None, ('Trait number must be provided if a '
                                     'Species object is provided.')
        alpha_distr_mu = spp.gen_arch.traits[trt_num]._alpha_distr_mu
        alpha_distr_sigma = spp.gen_arch.traits[trt_num]._alpha_distr_sigma
    else:
        assert (alpha_distr_mu is not None
                and alpha_distr_sigma is not None), ('If a Species '
                                                     'object is not '
                                                     'provided then '
                                                     'the parameter '
                                                     'values must be '
                                                     'provided.')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    fig.suptitle(('effect sizes of loci, trait %i: '
                  '~Normal($\mu$=%.4E, $\sigma$=%.4E)') % (alpha_distr_mu,
                                                           alpha_distr_sigma))
    vals = np.random.normal(alpha_distr_mu, alpha_distr_sigma, 10000)
    ax.hist(vals, bins = 25)
    ax.set_xlim((min(vals), max(vals)))
    plt.show()
