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

import numpy as num

from pyrockoeost.guts import Float, Int
from pyrockoeost import gf

from .base import SourceGenerator

km = 1e3
guts_prefix = 'pf.scenario'


class PseudoDynamicRuptureGenerator(SourceGenerator):
    depth_min = Float.T(
        default=0.0)
    depth_max = Float.T(
        default=5*km)
    decimation_factor = Int.T(
        default=1)

    slip_min = Float.T(
        optional=True)
    slip_max = Float.T(
        optional=True)

    strike = Float.T(
        optional=True)
    dip = Float.T(
        optional=True)
    rake = Float.T(
        optional=True)
    depth = Float.T(
        optional=True)
    nx = Int.T(
        default=5,
        optional=True)
    ny = Int.T(
        default=5,
        optional=True)
    nucleation_x = Float.T(
        optional=True)
    nucleation_y = Float.T(
        optional=True)

    width = Float.T(
        optional=True)
    length = Float.T(
        optional=True)

    gamma = Float.T(
        default=0.8)

    def draw_slip(self, rstate):
        if self.slip_min is not None and self.slip_max is not None:
            return rstate.uniform(self.slip_min, self.slip_max)
        else:
            return rstate.rand()

    def get_source(self, ievent):
        rstate = self.get_rstate(ievent)
        time = rstate.uniform(self.time_min, self.time_max)
        lat, lon = self.get_latlon(ievent)
        depth = rstate.uniform(self.depth_min, self.depth_max)
        nucleation_x = self.nucleation_x if self.nucleation_x is not None \
            else rstate.uniform(-1., 1.)
        nucleation_y = self.nucleation_y if self.nucleation_y is not None \
            else rstate.uniform(-1., 1.)

        def scale_from_slip(slip, a, b):
            return 10**((num.log10(slip) - a) / b)

        if self.slip_min is not None and self.slip_max is not None:
            slip = self.draw_slip(rstate)
            # After K. Thingbaijam et al. (2017) - Tab. 2, Normal faulting
            length = scale_from_slip(slip, a=-2.302, b=1.302)
            width = scale_from_slip(slip, a=-3.698, b=2.512)

        else:
            slip = self.draw_slip(rstate)
            length = scale_from_slip(slip, a=-2.302, b=1.302)
            width = scale_from_slip(slip, a=-3.698, b=2.512)

        length = length if not self.length else self.length
        width = width if not self.width else self.width
        depth = depth if not self.depth else self.depth

        if self.strike is None and self.dip is None and self.rake is None:
            strike, rake = rstate.uniform(-180., 180., 2)
            dip = rstate.uniform(0., 90.)
        else:
            if None in (self.strike, self.dip, self.rake):
                raise ValueError(
                    'PseudoDynamicRuptureGenerator: '
                    'strike, dip, rake must be used in combination.')

            strike = self.strike
            dip = self.dip
            rake = self.rake

        source = gf.PseudoDynamicRupture(
            time=float(time),
            lat=float(lat),
            lon=float(lon),
            anchor='top',
            depth=float(depth),
            length=float(length),
            width=float(width),
            strike=float(strike),
            dip=float(dip),
            rake=float(rake),
            slip=slip,
            nucleation_x=float(nucleation_x),
            nucleation_y=float(nucleation_y),
            nx=self.nx,
            ny=self.ny,
            decimation_factor=self.decimation_factor,
            smooth_rupture=True,
            gamma=self.gamma)

        return source

    def add_map_artists(self, automap):
        for source in self.get_sources():
            automap.gmt.psxy(
                in_rows=source.outline(cs='lonlat'),
                L='+p2p,black',
                W='1p,black',
                G='black',
                t=50,
                *automap.jxyr)
