import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import LinearSegmentedColormap

from scipy.interpolate import griddata, SmoothBivariateSpline
from scipy.stats import norm
from scipy.optimize import curve_fit


def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name="shiftedcmap"):
    """
    Function to offset the "center" of a colormap. Useful for
    data with a negative min and positive max and you want the
    middle of the colormap's dynamic range to be at zero

    Input
    -----
      cmap : The matplotlib colormap to be altered
      start : Offset from lowest point in the colormap's range.
          Defaults to 0.0 (no lower ofset). Should be between
          0.0 and `midpoint`.
      midpoint : The new center of the colormap. Defaults to
          0.5 (no shift). Should be between 0.0 and 1.0. In
          general, this should be  1 - vmax/(vmax + abs(vmin))
          For example if your data range from -15.0 to +5.0 and
          you want the center of the colormap at 0.0, `midpoint`
          should be set to  1 - 5/(5 + 15)) or 0.75
      stop : Offset from highets point in the colormap's range.
          Defaults to 1.0 (no upper ofset). Should be between
          `midpoint` and 1.0.
    """
    cdict = {"red": [], "green": [], "blue": [], "alpha": []}

    # regular index to compute the colors
    reg_index = np.linspace(start, stop, 257)

    # shifted index to match the data
    shift_index = np.hstack(
        [
            np.linspace(0.0, midpoint, 128, endpoint=False),
            np.linspace(midpoint, 1.0, 129, endpoint=True),
        ]
    )

    for ri, si in zip(reg_index, shift_index):
        r, g, b, a = cmap(ri)

        cdict["red"].append((si, r, r))
        cdict["green"].append((si, g, g))
        cdict["blue"].append((si, b, b))
        cdict["alpha"].append((si, a, a))

    newcmap = LinearSegmentedColormap(name, cdict)
    plt.register_cmap(cmap=newcmap)

    return newcmap


def plot_color_map(series, sort_order=[True, True], color_map=cm.RdYlGn, centered=True):
    """
    2D heat map - if x-axis is time translate to days instead

    Parameters
    -----
    series: Series with multilevel index (x: first level, y: second level)
    sort_order: sorting in the x,y axis
    color_map: default Red-Yellow-Green
    centered: center yellow as 0 of the series.
    """

    x = series.index.get_level_values(0)
    y = series.index.get_level_values(1)
    if x.dtype == "<M8[ns]":
        x = (x - x[0]).days
        x.name = "Days"
    series.sort_index(ascending=sort_order, inplace=True)

    fig, ax = plt.subplots()
    midpoint = (
        1 - series.max() / (series.max() + abs(series.min()))
        if centered is True
        else 0.5
    )
    color_map = shiftedColorMap(color_map, midpoint=midpoint)

    chart = ax.imshow(
        series.values.reshape(x.unique().size, y.unique().size).T,
        extent=(x.min(), x.max(), y.min(), y.max()),
        aspect="auto",
        interpolation="bilinear",
        cmap=color_map,
    )

    ax.set_xlabel(x.name)
    ax.set_ylabel(y.name)
    ax.set_title("{} of Trade".format(series.name))
    fig.colorbar(chart, shrink=0.8)


def plot_prob_map(df, attr="pnl", path=".", color_map=cm.RdYlGn, index="IG"):

    val_date = df.index[0].date()
    df = df.reset_index()
    df["days"] = (df["date"] - val_date).dt.days
    series = df[attr]

    days_defined = np.linspace(df.days.min(), df.days.max(), 1000)
    prob_defined = np.linspace(0.001, 0.999, 1000)

    midpoint = 1 - series.max() / (series.max() + abs(series.min()))
    shifted_cmap = shiftedColorMap(color_map, midpoint=midpoint, name="shifted")

    resampled = griddata(
        (df.days, df.prob),
        series,
        (days_defined[None, :], prob_defined[:, None]),
        method="linear",
    )

    # plot
    fig, ax = plt.subplots()
    chart = ax.imshow(
        resampled.reshape(days_defined.size, prob_defined.size),
        extent=(df.days.min(), df.days.max(), 0, 1),
        aspect="auto",
        interpolation="bilinear",
        cmap=shifted_cmap,
    )

    ax.set_xlabel("Days")
    ax.set_ylabel("Probability")
    ax.set_title("{} of Trade".format(attr.title()))

    fig.colorbar(chart, shrink=0.8)


def plot_swaption_df(df, spread_shock, vol_shock, attr="pnl"):
    val_date = df.index[0].date()
    fig = plt.figure()

    ax = fig.gca(projection="3d")
    ## use smoothing spline on a finer grid
    series = df[attr]
    f = SmoothBivariateSpline(
        df.vol_shock.values, df.spread_shock.values, series.values
    )
    xx, yy = np.meshgrid(vol_shock, spread_shock)
    surf = ax.plot_surface(xx, yy, f(vol_shock, spread_shock).T, cmap=cm.viridis)
    ax.set_xlabel("Volatility shock")
    ax.set_ylabel("Spread")
    ax.set_zlabel("PnL")
    ax.set_title("{} of Trade on {}".format(attr.title(), val_date))