from builtins import str
from pyrocko.gui.snuffling import Snuffling, Param, Switch, Choice
from pyrocko.gui.marker import PhaseMarker
from pyrocko import orthodrome
from pyrocko import gf
from pyrocko import cake
import numpy as num


class CakePhase(Snuffling):
    '''
    <html>
    <head>
    <style type="text/css">
        body { margin-left:10px };
    </style>
    </head>
    <body>
    <h1 align="center">Theoretical Phase Arrivals</h1>
    <p>
    This snuffling uses pyrocko's
    <a href="http://emolch.github.io/pyrocko/v0.3/cake_doc.html">Cake</a>
    module to calculate seismic rays for layered earth models. </p>
    <p>
    <b>Parameters:</b><br />
        <b>&middot; Global shift</b>  -  Add time onset to phases. <br />
        <b>&middot; Add Model</b>  -  Add a model to drop down menu. <br />
        <b>&middot; Add Phase</b>  -  Add a phase definition.
            (GUI reset required)<br />
    </p>
    <p>
    Instructions and information on Cake's syntax of seismic rays can be
    found in the
    <a href="http://emolch.github.io/pyrocko/
    v0.3/cake_doc.html#cmdoption-cake--phase">Cake documentation</a>.
    </p>
    </body>
    </html>
    '''

    def setup(self):
        self.set_name('Cake Phase')

        # self._phase_names = ('PmP ~S ~P ~P(moho)s ~P(sill-top)s'
        #                       ' ~P(sill-bottom)s Pdiff').split()
        self._phase_names = ('~P Pg Sg pP p P Pdiff PKP PcP PcS PKIKP pPKIKP'
                             ' SSP PPS SPP PSP SP PS ~PS ~SP Pn s S Sn PP PPP'
                             ' ScS Sdiff SS SSS PcP SKS SKIKS').split()

        for iphase, name in enumerate(self._phase_names):
            self.add_parameter(Switch(name, 'wantphase_%i' % iphase,
                                      iphase == 0))

        model_names = cake.builtin_models()
        model_names = [
                'Cake builtin: %s' % model_name for model_name in model_names]

        self._engine = gf.LocalEngine(use_config=True)
        store_ids = self._engine.get_store_ids()

        for store_id in store_ids:
            model_names.append('GF Store: %s' % store_id)

        self._models = model_names

        self.model_choice = Choice('Model', 'chosen_model',
                                   'ak135-f-continental.m', self._models)

        self.add_parameter(self.model_choice)

        self.add_parameter(Param('Global shift', 'tshift', 0., -20., 20.))
        self.add_parameter(Switch('Use station depth',
                                  'use_station_depth', False))
        self.add_trigger('Add Phase', self.add_phase_definition)
        self.add_trigger('Add Model', self.add_model_to_choice)
        self.add_trigger('Plot Model', self.plot_model)
        self.add_trigger('Plot Rays', self.plot_rays)

        self._phases = {}
        self._model = None

    def panel_visibility_changed(self, bool):
        pass

    def wanted_phases(self):
        try:
            wanted = []
            for iphase, name in enumerate(self._phase_names):
                if getattr(self, 'wantphase_%i' % iphase):
                    if name in self._phases:
                        phases = self._phases[name]
                    else:
                        if name.startswith('~'):
                            phases = [cake.PhaseDef(name[1:])]
                        else:
                            phases = cake.PhaseDef.classic(name)

                        self._phases[name] = phases
                        for pha in phases:
                            pha.name = name

                    wanted.extend(phases)
        except (cake.UnknownClassicPhase, cake.PhaseDefParseError) as e:
            self.fail(str(e))

        return wanted

    def call(self, plot_rays=False):

        self.cleanup()
        wanted = self.wanted_phases()

        if not wanted:
            return

        event, stations = self.get_active_event_and_stations()

        if not stations:
            self.fail('No station information available.')

        self.update_model()
        model = self._model[1]

        depth = event.depth
        if depth is None:
            depth = 0.0

        allrays = []
        alldists = []
        for station in stations:
            dist = orthodrome.distance_accurate50m(event, station)
            alldists.append(dist)

            if self.use_station_depth:
                rdepth = station.depth
            else:
                rdepth = 0.0

            multi_dists = []
            nmax = 1
            for i in range(0, nmax):
                multi_dists.append(dist*cake.m2d + 360.*i)
                multi_dists.append((i+1)*360. - dist*cake.m2d)

            rays = model.arrivals(
                phases=wanted,
                distances=multi_dists,
                zstart=depth,
                zstop=rdepth)

            for ray in rays:
                time = ray.t
                name = ray.given_phase().name
                incidence_angle = ray.incidence_angle()
                takeoff_angle = ray.takeoff_angle()

                time += event.time + self.tshift
                m = PhaseMarker([(station.network, station.station, '*', '*')],
                                time, time, 2,
                                phasename=name,
                                event=event,
                                incidence_angle=incidence_angle,
                                takeoff_angle=takeoff_angle)
                self.add_marker(m)

            allrays.extend(rays)

        if plot_rays:
            fig = self.figure(name='Ray Paths')
            from pyrocko import cake_plot
            cake_plot.my_rays_plot(model, None, allrays, depth, 0.0,
                                   num.array(alldists)*cake.m2d,
                                   axes=fig.gca())

            fig.canvas.draw()

    def update_model(self):
        if not self._model or self._model[0] != self.chosen_model:
            if self.chosen_model.startswith('Cake builtin: '):
                load_model = cake.load_model(
                    self.chosen_model.split(': ', 1)[1])

            elif self.chosen_model.startswith('GF Store: '):
                store_id = self.chosen_model.split(': ', 1)[1]

                load_model = self._engine.get_store(store_id)\
                    .config.earthmodel_1d
            else:
                load_model = cake.load_model(self.chosen_model)

            self._model = (self.chosen_model, load_model)

    def update_model_choices(self):
        self.set_parameter_choices('chosen_model', self._models)

    def add_model_to_choice(self):
        '''Called from trigger 'Add Model'.

        Adds another choice to the drop down 'Model' menu.
        '''

        in_model = self.input_filename('Load Model')
        if in_model not in self._models:
            self._models.append(in_model)
            self.update_model_choices()

        self.set_parameter('chosen_model', in_model)
        self.call()

    def add_phase_definition(self):
        ''' Called from trigger 'Add Phase Definition'.

        Adds another phase option.
        Requires a reset of the GUI.
        '''
        phase_def = str(self.input_dialog('Add New Phase',
                                          'Enter Phase Definition'))
        self._phase_names.append(phase_def)

        self.add_parameter(
            Switch(phase_def,
                   'wantphase_%s' % str(len(self._phase_names)-1), True))
        self.reset_gui(reloaded=True)
        self.call()

    def plot_model(self):
        self.update_model()

        from pyrocko import cake_plot

        fig = self.figure(name='Model: %s' % self._model[0])

        cake_plot.my_model_plot(self._model[1], axes=fig.gca())

        fig.canvas.draw()

    def plot_rays(self):
        self.call(plot_rays=True)


def __snufflings__():
    '''Returns a list of snufflings to be exported by this module.'''

    return [CakePhase()]
