# http://pyrocko.org - GPLv3
#
# The pyrockoeost Developers, 21st Century
# ---|P------/S----------~Lg----------

import math
import numpy as num
from pyrockoeost import cake
from pyrockoeost.util import mpl_show
from . import mpl_labelspace as labelspace, mpl_init,\
    mpl_color as str_to_mpl_color, InvalidColorDef

str_to_mpl_color
InvalidColorDef

d2r = cake.d2r
r2d = cake.r2d


def globe_cross_section():
    # modified from http://stackoverflow.com/questions/2417794/
    # how-to-make-the-angles-in-a-matplotlib-polar-plot-go-clockwise-with-0-at-the-top

    from matplotlib.projections import PolarAxes, register_projection
    from matplotlib.transforms import Affine2D, Bbox, IdentityTransform

    class GlobeCrossSectionAxes(PolarAxes):
        '''
        A variant of PolarAxes where theta starts pointing north and goes
        clockwise and the radial axis is reversed.
        '''
        name = 'globe_cross_section'

        class GlobeCrossSectionTransform(PolarAxes.PolarTransform):

            def transform(self, tr):
                xy = num.zeros(tr.shape, num.float_)
                t = tr[:, 0:1]*d2r
                r = cake.earthradius - tr[:, 1:2]
                x = xy[:, 0:1]
                y = xy[:, 1:2]
                x[:] = r * num.sin(t)
                y[:] = r * num.cos(t)
                return xy

            transform_non_affine = transform

            def inverted(self):
                return GlobeCrossSectionAxes.\
                    InvertedGlobeCrossSectionTransform()

        class InvertedGlobeCrossSectionTransform(
                PolarAxes.InvertedPolarTransform):

            def transform(self, xy):
                x = xy[:, 0:1]
                y = xy[:, 1:]
                r = num.sqrt(x*x + y*y)
                theta = num.arctan2(y, x)*r2d
                return num.concatenate((theta, cake.earthradius-r), 1)

            def inverted(self):
                return GlobeCrossSectionAxes.GlobeCrossSectionTransform()

        def _set_lim_and_transforms(self):
            PolarAxes._set_lim_and_transforms(self)
            try:
                theta_position = self._theta_label1_position
            except AttributeError:
                theta_position = self.get_theta_offset()

            self.transProjection = self.GlobeCrossSectionTransform()
            self.transData = (
                self.transScale +
                self.transProjection +
                (self.transProjectionAffine + self.transAxes))
            self._xaxis_transform = (
                self.transProjection +
                self.PolarAffine(IdentityTransform(), Bbox.unit()) +
                self.transAxes)
            self._xaxis_text1_transform = (
                theta_position +
                self._xaxis_transform)
            self._yaxis_transform = (
                Affine2D().scale(num.pi * 2.0, 1.0) +
                self.transData)

            try:
                rlp = getattr(self, '_r_label1_position')
            except AttributeError:
                rlp = getattr(self, '_r_label_position')

            self._yaxis_text1_transform = (
                rlp +
                Affine2D().scale(1.0 / 360.0, 1.0) +
                self._yaxis_transform)

    register_projection(GlobeCrossSectionAxes)


tango_colors = {
    'butter1': (252, 233, 79),
    'butter2': (237, 212, 0),
    'butter3': (196, 160, 0),
    'chameleon1': (138, 226, 52),
    'chameleon2': (115, 210, 22),
    'chameleon3': (78, 154, 6),
    'orange1': (252, 175, 62),
    'orange2': (245, 121, 0),
    'orange3': (206, 92, 0),
    'skyblue1': (114, 159, 207),
    'skyblue2': (52, 101, 164),
    'skyblue3': (32, 74, 135),
    'plum1': (173, 127, 168),
    'plum2': (117, 80, 123),
    'plum3': (92, 53, 102),
    'chocolate1': (233, 185, 110),
    'chocolate2': (193, 125, 17),
    'chocolate3': (143, 89, 2),
    'scarletred1': (239, 41, 41),
    'scarletred2': (204, 0, 0),
    'scarletred3': (164, 0, 0),
    'aluminium1': (238, 238, 236),
    'aluminium2': (211, 215, 207),
    'aluminium3': (186, 189, 182),
    'aluminium4': (136, 138, 133),
    'aluminium5': (85, 87, 83),
    'aluminium6': (46, 52, 54)
}


def path2colorint(path):
    '''
    Calculate an integer representation deduced from path's given name.
    '''
    s = sum([ord(char) for char in path.phase.given_name()])
    return s


def light(color, factor=0.2):
    return tuple(1-(1-c)*factor for c in color)


def dark(color, factor=0.5):
    return tuple(c*factor for c in color)


def to01(c):
    return tuple(x/255. for x in c)


colors = [to01(tango_colors[x+i]) for i in '321' for x in
          'scarletred chameleon skyblue chocolate orange plum'.split()]
shades = [light(to01(tango_colors['chocolate1']), i*0.1) for i in range(1, 9)]
shades2 = [light(to01(tango_colors['orange1']), i*0.1) for i in range(1, 9)]

shades_water = [
    light(to01(tango_colors['skyblue1']), i*0.1) for i in range(1, 9)]


def plot_xt(
        paths, zstart, zstop,
        axes=None,
        vred=None,
        distances=None,
        coloring='by_phase_definition',
        avoid_same_colors=True,
        phase_colors={}):

    if distances is not None:
        xmin, xmax = distances.min(), distances.max()

    axes = getaxes(axes)
    all_x = []
    all_t = []
    path_to_color = make_path_to_color(
        paths, coloring, avoid_same_colors, phase_colors=phase_colors)

    for ipath, path in enumerate(paths):
        if distances is not None:
            if path.xmax() < xmin or path.xmin() > xmax:
                continue

        color = path_to_color[path]
        p, x, t = path.draft_pxt(path.endgaps(zstart, zstop))
        if p.size == 0:
            continue

        all_x.append(x)
        all_t.append(t)
        if vred is not None:
            axes.plot(x, t-x/vred, linewidth=2, color=color)
            axes.plot([x[0]], [t[0]-x[0]/vred], 'o', color=color)
            axes.plot([x[-1]], [t[-1]-x[-1]/vred], 'o', color=color)
            axes.text(
                x[len(x)//2],
                t[len(x)//2]-x[len(x)//2]/vred,
                path.used_phase().used_repr(),
                color=color,
                va='center',
                ha='center',
                clip_on=True,
                bbox=dict(
                    ec=color,
                    fc=light(color),
                    pad=8,
                    lw=1),
                fontsize=10)
        else:
            axes.plot(x, t, linewidth=2, color=color)
            axes.plot([x[0]], [t[0]], 'o', color=color)
            axes.plot([x[-1]], [t[-1]], 'o', color=color)
            axes.text(
                x[len(x)//2],
                t[len(x)//2],
                path.used_phase().used_repr(),
                color=color,
                va='center',
                ha='center',
                clip_on=True,
                bbox=dict(
                    ec=color,
                    fc=light(color),
                    pad=8,
                    lw=1),
                fontsize=10)

    all_x = num.concatenate(all_x)
    all_t = num.concatenate(all_t)
    if vred is not None:
        all_t -= all_x/vred
    xxx = num.sort(all_x)
    ttt = num.sort(all_t)
    return xxx.min(), xxx[99*len(xxx)//100], ttt.min(), ttt[99*len(ttt)//100]


def labels_xt(axes=None, vred=None, as_degrees=False):
    axes = getaxes(axes)
    if as_degrees:
        axes.set_xlabel('Distance [deg]')
    else:
        axes.set_xlabel('Distance [km]')
        xscaled(d2r*cake.earthradius/cake.km, axes)

    if vred is None:
        axes.set_ylabel('Time [s]')
    else:
        if as_degrees:
            axes.set_ylabel('Time - Distance / %g deg/s [ s ]' % (
                vred))
        else:
            axes.set_ylabel('Time - Distance / %g km/s [ s ]' % (
                d2r*vred*cake.earthradius/cake.km))


def troffset(dx, dy, axes=None):
    axes = getaxes(axes)
    from matplotlib import transforms
    return axes.transData + transforms.ScaledTranslation(
        dx/72., dy/72., axes.gcf().dpi_scale_trans)


def plot_xp(
        paths,
        zstart,
        zstop,
        axes=None,
        coloring='by_phase_definition',
        avoid_same_colors=True,
        phase_colors={}):

    path_to_color = make_path_to_color(
        paths, coloring, avoid_same_colors, phase_colors=phase_colors)

    axes = getaxes(axes)
    all_x = []
    for ipath, path in enumerate(paths):
        color = path_to_color[path]
        p, x, t = path.draft_pxt(path.endgaps(zstart, zstop))
        # converting ray parameters from s/rad to s/deg
        p /= r2d
        axes.plot(x, p, linewidth=2, color=color)
        axes.plot(x[:1], p[:1], 'o', color=color)
        axes.plot(x[-1:], p[-1:], 'o', color=color)
        axes.text(
            x[len(x)//2],
            p[len(x)//2],
            path.used_phase().used_repr(),
            color=color,
            va='center',
            ha='center',
            clip_on=True,
            bbox=dict(
                ec=color,
                fc=light(color),
                pad=8,
                lw=1))

        all_x.append(x)

    xxx = num.sort(num.concatenate(all_x))
    return xxx.min(), xxx[99*len(xxx)//100]


def labels_xp(axes=None, as_degrees=False):
    axes = getaxes(axes)
    if as_degrees:
        axes.set_xlabel('Distance [deg]')
    else:
        axes.set_xlabel('Distance [km]')
        xscaled(d2r*cake.earthradius*0.001, axes)
    axes.set_ylabel('Ray Parameter [s/deg]')


def labels_model(axes=None):
    axes = getaxes(axes)
    axes.set_xlabel('S-wave and P-wave velocity [km/s]')
    xscaled(0.001, axes)
    axes.set_ylabel('Depth [km]')
    yscaled(0.001, axes)


def make_path_to_color(
        paths,
        coloring='by_phase_definition',
        avoid_same_colors=True,
        phase_colors={}):

    assert coloring in ['by_phase_definition', 'by_path']

    path_to_color = {}
    definition_to_color = phase_colors.copy()
    available_colors = set()

    for ipath, path in enumerate(paths):
        if coloring == 'by_phase_definition':
            given_name = path.phase.given_name()
            int_rep = path2colorint(path)
            color_id = int_rep % len(colors)

            if given_name not in definition_to_color:
                if avoid_same_colors:
                    if len(available_colors) == 0:
                        available_colors = set(range(0, len(colors)-1))
                    if color_id in available_colors:
                        available_colors.remove(color_id)
                    else:
                        color_id = available_colors.pop()

                    assert color_id not in available_colors

                definition_to_color[given_name] = colors[color_id]

            path_to_color[path] = definition_to_color[given_name]
        else:
            path_to_color[path] = colors[ipath % len(colors)]

    return path_to_color


def plot_rays(paths, rays, zstart, zstop,
              axes=None,
              coloring='by_phase_definition',
              legend=True,
              avoid_same_colors=True,
              aspect=None,
              phase_colors={}):

    axes = getaxes(axes)
    if aspect is not None:
        axes.set_aspect(aspect/(d2r*cake.earthradius))

    path_to_color = make_path_to_color(
        paths, coloring, avoid_same_colors, phase_colors=phase_colors)

    if rays is None:
        rays = paths

    labels = set()

    for iray, ray in enumerate(rays):
        if isinstance(ray, cake.RayPath):
            path = ray
            pmin, pmax, xmin, xmax, tmin, tmax = path.ranges(
                path.endgaps(zstart, zstop))

            if not path._is_headwave:
                p = num.linspace(pmin, pmax, 6)
                x = None

            else:
                x = num.linspace(xmin, xmin*10, 6)
                p = num.atleast_1d(pmin)

            fanz, fanx, _ = path.zxt_path_subdivided(
                p, path.endgaps(zstart, zstop), x_for_headwave=x)

        else:
            fanz, fanx, _ = ray.zxt_path_subdivided()
            path = ray.path

        color = path_to_color[path]

        if coloring == 'by_phase_definition':
            phase_label = path.phase.given_name()

        else:
            phase_label = path

        for zs, xs in zip(fanz, fanx):
            if phase_label in labels:
                phase_label = ''

            axes.plot(xs, zs, color=color, label=phase_label)
            if legend:
                labels.add(phase_label)

    if legend:
        axes.legend(loc=4, prop={'size': 11})


def sketch_model(mod, axes=None, shade=True):
    from matplotlib import transforms
    axes = getaxes(axes)
    trans = transforms.BlendedGenericTransform(axes.transAxes, axes.transData)

    for dis in mod.discontinuities():
        color = shades[-1]
        axes.axhline(dis.z, color=dark(color), lw=1.5)
        if dis.name is not None:
            axes.text(
                0.90, dis.z, dis.name,
                transform=trans,
                va='center',
                ha='right',
                color=dark(color),
                bbox=dict(ec=dark(color), fc=light(color, 0.3), pad=8, lw=1))

    for ilay, lay in enumerate(mod.layers()):
        if lay.mtop.vs == 0.0 and lay.mbot.vs == 0.0:
            tab = shades_water
        else:
            if isinstance(lay, cake.GradientLayer):
                tab = shades
            else:
                tab = shades2

        color = tab[ilay % len(tab)]
        if shade:
            axes.axhspan(
                lay.ztop, lay.zbot, fc=color, ec=dark(color), label='abc')

        if lay.name is not None:
            axes.text(
                0.95, (lay.ztop + lay.zbot)*0.5,
                lay.name,
                transform=trans,
                va='center',
                ha='right',
                color=dark(color),
                bbox=dict(ec=dark(color), fc=light(color, 0.3), pad=8, lw=1))


def plot_source(zstart, axes=None):
    axes = getaxes(axes)
    axes.plot([0], [zstart], 'o', color='black')


def plot_receivers(zstop, distances, axes=None):
    axes = getaxes(axes)
    axes.plot(
        distances, cake.filled(zstop, len(distances)), '^', color='black')


def getaxes(axes=None):
    from matplotlib import pyplot as plt
    if axes is None:
        return plt.gca()
    else:
        return axes


def mk_sc_classes():
    from matplotlib.ticker import FormatStrFormatter, AutoLocator

    class Scaled(FormatStrFormatter):
        def __init__(self, factor):
            FormatStrFormatter.__init__(self, '%g')
            self._factor = factor

        def __call__(self, v, i=0):
            return FormatStrFormatter.__call__(self, v*self._factor, i)

    class ScaledLocator(AutoLocator):
        def __init__(self, factor):
            AutoLocator.__init__(self)
            self._factor = factor

        def bin_boundaries(self, vmin, vmax):
            return [x/self._factor for x in AutoLocator.bin_boundaries(
                self, vmin*self._factor, vmax*self._factor)]

    return Scaled, ScaledLocator


def xscaled(factor, axes):
    Scaled, ScaledLocator = mk_sc_classes()
    xaxis = axes.get_xaxis()
    xaxis.set_major_formatter(Scaled(factor))
    xaxis.set_major_locator(ScaledLocator(factor))


def yscaled(factor, axes):
    Scaled, ScaledLocator = mk_sc_classes()
    yaxis = axes.get_yaxis()
    yaxis.set_major_formatter(Scaled(factor))
    yaxis.set_major_locator(ScaledLocator(factor))


def labels_rays(axes=None, as_degrees=False):
    axes = getaxes(axes)
    if as_degrees:
        axes.set_xlabel('Distance [deg]')
    else:
        axes.set_xlabel('Distance [km]')
        xscaled(d2r*cake.earthradius/cake.km, axes)
    axes.set_ylabel('Depth [km]')
    yscaled(1./cake.km, axes)


def plot_surface_efficiency(mat):
    from matplotlib import pyplot as plt
    data = []
    for angle in num.linspace(0., 90., 910.):
        pp = math.sin(angle*d2r)/mat.vp
        ps = math.sin(angle*d2r)/mat.vs

        escp = cake.psv_surface(mat, pp, energy=True)
        escs = cake.psv_surface(mat, ps, energy=True)
        data.append((
            angle,
            escp[cake.psv_surface_ind(cake.P, cake.P)],
            escp[cake.psv_surface_ind(cake.P, cake.S)],
            escs[cake.psv_surface_ind(cake.S, cake.S)],
            escs[cake.psv_surface_ind(cake.S, cake.P)]))

    a, pp, ps, ss, sp = num.array(data).T

    plt.plot(a, pp, label='PP')
    plt.plot(a, ps, label='PS')
    plt.plot(a, ss, label='SS')
    plt.plot(a, sp, label='SP')
    plt.xlabel('Incident Angle')
    plt.ylabel('Energy Normalized Coefficient', position=(-2., 0.5))
    plt.legend()
    mpl_show(plt)


def my_xp_plot(
        paths, zstart, zstop,
        distances=None,
        as_degrees=False,
        axes=None,
        show=True,
        phase_colors={}):

    if axes is None:
        from matplotlib import pyplot as plt
        mpl_init()
        axes = plt.gca()
    else:
        plt = None

    labelspace(axes)
    xmin, xmax = plot_xp(
        paths, zstart, zstop, axes=axes, phase_colors=phase_colors)

    if distances is not None:
        xmin, xmax = distances.min(), distances.max()

    axes.set_xlim(xmin, xmax)
    labels_xp(as_degrees=as_degrees, axes=axes)

    if plt:
        if show is True:
            mpl_show(plt)


def my_xt_plot(
        paths, zstart, zstop,
        distances=None,
        as_degrees=False,
        vred=None,
        axes=None,
        show=True,
        phase_colors={}):

    if axes is None:
        from matplotlib import pyplot as plt
        mpl_init()
        axes = plt.gca()
    else:
        plt = None

    labelspace(axes)
    xmin, xmax, ymin, ymax = plot_xt(
        paths,
        zstart,
        zstop,
        vred=vred,
        distances=distances,
        axes=axes,
        phase_colors=phase_colors)

    if distances is not None:
        xmin, xmax = distances.min(), distances.max()

    axes.set_xlim(xmin, xmax)
    axes.set_ylim(ymin, ymax)
    labels_xt(as_degrees=as_degrees, vred=vred, axes=axes)
    if plt:
        if show is True:
            mpl_show(plt)


def my_rays_plot_gcs(
        mod, paths, rays, zstart, zstop,
        distances=None,
        show=True,
        phase_colors={}):

    from matplotlib import pyplot as plt
    mpl_init()

    globe_cross_section()
    axes = plt.subplot(1, 1, 1, projection='globe_cross_section')
    plot_rays(paths, rays, zstart, zstop, axes=axes, phase_colors=phase_colors)
    plot_source(zstart, axes=axes)
    if distances is not None:
        plot_receivers(zstop, distances, axes=axes)

    axes.set_ylim(0., cake.earthradius)
    axes.get_yaxis().set_visible(False)

    if plt:
        if show is True:
            mpl_show(plt)


def my_rays_plot(
        mod, paths, rays, zstart, zstop,
        distances=None,
        as_degrees=False,
        axes=None,
        show=True,
        aspect=None,
        shade_model=True,
        phase_colors={}):

    if axes is None:
        from matplotlib import pyplot as plt
        mpl_init()
        axes = plt.gca()
    else:
        plt = None

    if paths is None:
        paths = list(set([x.path for x in rays]))

    labelspace(axes)
    plot_rays(
        paths, rays, zstart, zstop,
        axes=axes, aspect=aspect, phase_colors=phase_colors)

    xmin, xmax = axes.get_xlim()
    ymin, ymax = axes.get_ylim()
    sketch_model(mod, axes=axes, shade=shade_model)

    plot_source(zstart, axes=axes)
    if distances is not None:
        plot_receivers(zstop, distances, axes=axes)
    labels_rays(as_degrees=as_degrees, axes=axes)
    mx = (xmax-xmin)*0.05
    my = (ymax-ymin)*0.05
    axes.set_xlim(xmin-mx, xmax+mx)
    axes.set_ylim(ymax+my, ymin-my)

    if plt:
        if show is True:
            mpl_show(plt)


def my_combi_plot(
        mod, paths, rays, zstart, zstop,
        distances=None,
        as_degrees=False,
        show=True,
        vred=None,
        phase_colors={}):

    from matplotlib import pyplot as plt
    mpl_init()
    ax1 = plt.subplot(211)
    labelspace(plt.gca())

    xmin, xmax, ymin, ymax = plot_xt(
        paths, zstart, zstop,
        vred=vred,
        distances=distances,
        phase_colors=phase_colors)

    if distances is None:
        plt.xlim(xmin, xmax)

    labels_xt(vred=vred, as_degrees=as_degrees)
    plt.setp(ax1.get_xticklabels(), visible=False)
    plt.xlabel('')

    ax2 = plt.subplot(212, sharex=ax1)
    labelspace(plt.gca())
    plot_rays(paths, rays, zstart, zstop, phase_colors=phase_colors)
    xmin, xmax = plt.xlim()
    ymin, ymax = plt.ylim()
    sketch_model(mod)

    plot_source(zstart)
    if distances is not None:
        plot_receivers(zstop, distances)
    labels_rays(as_degrees=as_degrees)
    mx = (xmax-xmin)*0.05
    my = (ymax-ymin)*0.05
    ax2.set_xlim(xmin-mx, xmax+mx)
    ax2.set_ylim(ymax+my, ymin-my)

    if show is True:
        mpl_show(plt)


def my_model_plot(mod, axes=None, show=True):

    if axes is None:
        from matplotlib import pyplot as plt
        mpl_init()
        axes = plt.gca()
    else:
        plt = None

    labelspace(axes)
    labels_model(axes=axes)
    sketch_model(mod, axes=axes)
    z = mod.profile('z')
    vp = mod.profile('vp')
    vs = mod.profile('vs')
    axes.plot(vp, z, color=colors[0], lw=2.)
    axes.plot(vs, z, color=colors[2], lw=2.)
    ymin, ymax = axes.get_ylim()
    xmin, xmax = axes.get_xlim()
    xmin = 0.
    my = (ymax-ymin)*0.05
    mx = (xmax-xmin)*0.2
    axes.set_ylim(ymax+my, ymin-my)
    axes.set_xlim(xmin, xmax+mx)
    if plt:
        if show is True:
            mpl_show(plt)
