"""
Module containing various thermostat. Berendsen only for now.
"""
import numpy as np
from numba import njit

class Thermostat:
    """
    Thermostat object.

    Attributes
    ----------
    kB : float
        Boltzmann constant in correct units.

    no_species : int
        Total number of species.

    species_np : numpy.ndarray
        Number of particles of each species.

    species_masses : numpy.ndarray
        Mass of each species.

    relaxation_rate: float
        Berendsen parameter tau.

    relaxation_timestep: int
        Timestep at which thermostat is turned on.

    type: str
        Thermostat type

    berendsen_tau: float
        Berendsen parameter.

    """

    def __init__(self):
        self.temperatures = None
        self.temperatures_eV = None
        self.type = None
        self.relaxation_rate = None
        self.relaxation_timestep = None
        self.kB = None
        self.species_num = None
        self.species_masses = None
        self.berendsen_tau = None
        self.eV_temp_flag = False
        self.K_temp_flag = False

    def __repr__(self):
        sortedDict = dict(sorted(self.__dict__.items(), key=lambda x: x[0].lower()))
        disp = 'Thermostat( \n'
        for key, value in sortedDict.items():
            disp += "\t{} : {}\n".format(key, value)
        disp += ')'
        return disp

    def from_dict(self, input_dict: dict):
        """
        Update attributes from input dictionary.

        Parameters
        ----------
        input_dict: dict
            Dictionary to be copied.

        """
        self.__dict__.update(input_dict)

        # Make sure list are turned into numpy arrays
        if self.temperatures_eV:
            if not isinstance(self.temperatures_eV, np.ndarray):
                self.temperatures_eV = np.array([self.temperatures_eV])
            self.eV_temp_flag = True

        if self.temperatures:
            if not isinstance(self.temperatures, np.ndarray):
                self.temperatures = np.array([self.temperatures])
            self.K_temp_flag = True

    def pretty_print(self):
        """Print Thermostat information in a user-friendly way."""
        print('Type: {}'.format(self.type))
        print('First thermostating timestep, i.e. relaxation_timestep = {}'.format(self.relaxation_timestep))
        print("Berendsen parameter tau: {:.3f} [timesteps]".format(self.berendsen_tau))
        print("Berendsen relaxation rate: {:.3f} [1/timesteps] ".format(self.relaxation_rate))
        if not self.eV_temp_flag and not self.K_temp_flag:
            # If you forgot to give thermostating temperatures
            print("\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! WARNING !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
            print("Equilibration temperatures not defined. I will use the species's temperatures")
            print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n")
        print("Thermostating temperatures: ")
        for i, (t, t_ev) in enumerate(zip(self.temperatures, self.temperatures_eV)):
            print("Species ID {}: T_eq = {:.6e} [K] = {:.6e} [eV]".format(i, t, t_ev))

    def setup(self, params):
        """
        Assign attributes from simulation's parameters.

        Parameters
        ----------
        params: sarkas.core.Parameters
            Simulation's parameters

        """

        # Check whether you input temperatures in eV or K

        if self.eV_temp_flag:
            self.temperatures = params.eV2K * np.copy(self.temperatures_eV)
        elif self.K_temp_flag:
            self.temperatures_eV = np.copy(self.temperatures) / params.eV2K
        else:
            self.temperatures = np.copy(params.species_temperatures)
            self.temperatures_eV = np.copy(self.temperatures) / params.eV2K

        if self.berendsen_tau:
            self.relaxation_rate = 1.0 / self.berendsen_tau
        else:
            self.berendsen_tau = 1.0 / self.relaxation_rate

        if not self.temperatures.all():
            self.temperatures = np.copy(params.species_temperatures)

        self.kB = params.kB
        self.species_num = np.copy(params.species_num)
        self.species_masses = np.copy(params.species_masses)

        assert self.type.lower() == "berendsen", "Only Berendsen thermostat is supported."

    def update(self, ptcls, it):
        """
        Update particles' velocities according to the chosen thermostat

        Parameters
        ----------
        ptcls : sarkas.core.Particles
            Particles' data.

        it : int
            Current timestep.

        """
        K, T = ptcls.kinetic_temperature()
        berendsen(ptcls.vel, self.temperatures, T, self.species_num, self.relaxation_timestep,
                  self.relaxation_rate, it)

@njit
def calc_kin_temp(vel, nums, masses, kB):
    """
    Calculates the kinetic energy and temperature.

    Parameters
    ----------
    kB: float
        Boltzmann constant in chosen units.

    masses: numpy.ndarray
        Mass of each species.

    nums: numpy.ndarray
        Number of particles of each species.

    vel: numpy.ndarray
        Particles' velocities.

    Returns
    -------
    K : numpy.ndarray
        Kinetic energy of each species.

    T : numpy.ndarray
        Temperature of each species.
    """

    num_species = len(nums)

    K = np.zeros(num_species)
    T = np.zeros(num_species)
    const = 2.0 / (kB * nums * vel.shape[1])
    kinetic_energies = 0.5 * masses * (vel ** 2).transpose()

    species_start = 0
    species_end = 0
    for i, num in enumerate(nums):
        species_end += num
        K[i] = np.sum(kinetic_energies[:, species_start:species_end])
        T[i] = const[i] * K[i]
        species_start += num

    return K, T

@njit
def berendsen(vel, T_desired, T, species_np, therm_timestep, tau, it):
    """
    Update particle velocity based on Berendsen thermostat [Berendsen1984]_.

    Parameters
    ----------
    T : numpy.ndarray
        Instantaneous temperature of each species.

    vel : numpy.ndarray
        Particles' velocities to rescale.

    T_desired : numpy.ndarray
        Target temperature of each species.

    tau : float
        Scale factor.

    therm_timestep : int
        Timestep at which to turn on the thermostat.

    species_np : numpy.ndarray
        Number of each species.

    it : int
        Current timestep.

    References
    ----------
    .. [Berendsen1984] `H.J.C. Berendsen et al., J Chem Phys 81 3684 (1984) <https://doi.org/10.1063/1.448118>`_

    """

    # if it < therm_timestep:
    #     fact = np.sqrt(T_desired / T)
    # else:
    #     fact = np.sqrt(1.0 + (T_desired / T - 1.0) * tau)  # eq.(11)

    # branchless programming
    fact = 1.0 * (it < therm_timestep) + np.sqrt(1.0 + (T_desired / T - 1.0) * tau) * (it >= therm_timestep)
    species_start = 0
    species_end = 0

    for i, num in enumerate(species_np):
        species_end += num
        vel[species_start:species_end, :] *= fact[i]
        species_start += num


# -*- coding: utf-8; -*-
#
# (c) 2016 microquake development team
#
# This file is part of the microquake library
#
# microquake is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# microquake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with microquake.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
from microquake.core import Stream, Trace
from obspy.signal.trigger import recursive_sta_lta, pk_baer, \
    classic_sta_lta, coincidence_trigger
from obspy.realtime.signal import kurtosis
from scipy.signal import detrend
from scipy.ndimage.filters import gaussian_filter1d
from microquake.core import event
from microquake.core.util.decorator import deprecated
import microquake.core.nlloc as nll
#from microquake.core.util.decorator import logger
from loguru import logger

from microquake.core.event import make_pick
from microquake.core.util.tools import copy_picks_to_dict

def measure_polarity(st, catalog, site, average_time_window=1e-3,
                     hp_filter_freq=100):
    """
    Measure the P- and S-wave polarities
    :param st: seismograms
    :type st: either obspy.core.Stream or microquake.core.Stream
    :param catalog: catalog object
    :type catalog: microquake.core.Catalog
    :param site: sensor information
    :type site: microquake.core.station.Site
    :param average_time_window: time window over which the polarity is measured (s)
    :type average_time_window: float
    :rparam: returns a copy of the input catalog
    :rtype: microquake.core.Stream
    """
    cat_out = catalog.copy()

    for iev, ev in enumerate(cat_out.events):
        for ipick, pick in enumerate(ev.picks):
            sta_code = pick.waveform_id.station_code
            ch_code = pick.waveform_id.channel_code
            trs = st.select(station=sta_code, channel=ch_code)
            trs = trs.detrend('linear').detrend('demean')
            sta = site.select(station=sta_code, channel=ch_code)
            if not sta.networks[0].stations:
                continue
            if (len(sta.networks) > 1) or (len(sta.networks[0].stations) > 1):
                logger.warning("MeasurePolarity: multiple station selected "
                               "from site. Only the first will be used")
            sta = sta.networks[0].stations[0]
            motion_type = sta.motion_type
            if motion_type == 'acceleration':
                trs = trs.integrate().detrend('linear').detrend(
                    'demean').integrate()
            elif motion_type == 'velocity':
                trs.integrate()
            elif motion_type == 'displacement':
                pass
            else:
                logger.warning("MeasurePolarity: motion_type not set for "
                               "sensor %s... Displacement will be assumed"
                               %stname)

            trs.filter('highpass', freq=hp_filter_freq)
            if len(trs) > 1:
                logger.warning("number of trace for station %s and channel %s"
                 "is greater than 1. Only the first trace will be used" % (
                                   sta_code, ch_code))
            tr = trs.traces[0]

            sp_s = int((pick.time - tr.stats.starttime) * tr.stats.sampling_rate)
            sp_e = int((sp_s + average_time_window * tr.stats.sampling_rate)) + 1
            pol = np.sign(np.mean(tr.data[sp_s:sp_e]) - tr.data[sp_s])

            if pol > 0:
                cat_out.events[iev].picks[ipick].polarity = "positive"
            elif pol < 0:
                cat_out.events[iev].picks[ipick].polarity = "negative"
            else:
                cat_out.events[iev].picks[ipick].polarity = "undecidable"

    return cat_out

def _measure_pick_polarity_tr(tr, pick_time, signal_type='Acceleration', nsamp_avg=10):
    """Mesure pick polarity
    The pick polarity is measured on the displacement trace looking at the difference
    between the amplitude at the pick time and the sign of the average amplitude
    for the 10 samples following the pick time

    :param tr: signal trace seismogram
    :type tr: obspy Trace
    :param pick_time: pick time
    :type pick_time: obspy UTCDateTime
    :param signal_type: Type of signal
    :type signal_type: str [accepted values are 'Acceleration', 'Velocity', and 'Displacement']
    :param nsamp_avg: number of sample on which to calculate the average difference
    :type nsamp_avg: int
    :returns: signa of signal polarity (1 or -1)
    :rtype: int
    """

    if signal_type not in ['Acceleration', 'Velocity', 'Displacement']:
        # raise error and exit funciton
        pass

    starttime = tr.stats.starttime
    sampling_rate = tr.stats.sampling_rate
    sample_pick = int(pick_time - starttime * sampling_rate)

    polarity = np.sign(np.mean(tr.data[sample_pick + 1:sample_pick +
                                                       nsamp_avg + 2] -
                               tr.data[sample_pick]))

    return polarity

def pick_uncertainty(tr, pick_time, snr_window=10):
    """Estimate the pick uncertainty, dt, based on the following equation
    dt = 1 / (fm * log_2(1 + SNR ** 2)
    where fm is the middle frequency and SNR, the signal to noise ratio in dB.

    :param tr: signal trace seismogram
    :type tr: obspy Trace
    :param pick_time: pick time
    :type pick_time: obspy UTCDateTime
    :param snr_window: length of window in ms centered on pick_time to
    calculate SNR.
        Noise and signal energy are calculated over the first and second half of
        the pick window, respectively.
    :type snr_window: float
    :returns: pick uncertainty in second
    :rtype: float
    """

    fm = 100  # central frequency in Hz
    # fm = centralfrequency(tr, pick_time, window)
    # tr is obspy tr, pick_time same, window window from pick over wich
    # signal frequencies are calculated
    st = Stream(traces=[tr])
    snr = calculate_snr(st, pick_time, snr_window)
    return 1 / (fm * np.log(1 + snr ** 2) / np.log(2))

# noinspection PyProtectedMember
def STALTA_picker_refraction(st, nsta=1e-3, nlta=4e-2, fc=50, nc=2,
    noise_mag=1e-25, SNR_window=5e-3):
    """The STA/LTA picker provides a first order good estimate of the arrival
    times for both the P- and S- wave. The STA/LTA picker uses
    :py:class:`obspy.signal.trigger.recursive_sta_lta`.

    :param st: seismogram containing a seismic event
    :type st: :py:class:`obspy.core.stream.Stream`
    :param nsta: Length of the short term average (STA) window in seconds.
    The value
        must be smaller than the minimum separation between the P- and S-
        wave arrival (explanation to be improved).
    :param nlta: Length of the long term average (LTA) window in seconds. The value
        must be long enough to capture well the noise but not too long (
        description to be improved)
    :param fc: Corner frequency of the Gaussian kernel used to smooth the
    LTA/STA
        function in Hz. This parameter must be chosen to limit the number of
        peaks in the STA/LTA function.
    :param nc: Number of desired onset. *** PARAMETER NOT USED ***
    :param noise_mag: Amplitude of the noise that is add to the original signal
    :param SNR_window: Length of window in seconds in which the SNR is
    calculated before and after the pick
    :returns:  :py:class:`obspy.core.event.Catalog` -- a new catalog
    containing a single event with a list of picks
    """

    # nphase = 1
    st.detrend('demean')
    stations = np.unique([tr.stats.station for tr in st])

    # for station in stations:
    opicks = []
    cfs = []
    snrs = []
    for station in stations:

        trs = st.select(station=station)
        # if len(trs) == 3:
        #   trs = _RotateP_S(trs)

        # else:
        #   continue

        sta = int(nsta * trs.traces[0].stats['sampling_rate'])
        lta = int(nlta * trs.traces[0].stats['sampling_rate'])

        # Noise added to avoid zeroed trace.
        rd = np.random.randn(len(trs.traces[0].data)) * noise_mag
        trstmp = [Trace(
            data=np.hstack(
                (tr.data[::-1].reshape(len(tr)) + rd[::-1].reshape(len(rd)),
                tr.data.reshape(len(tr)) + rd.reshape(len(rd)))),
            header=tr.stats) for tr in trs]
        sttmp = Stream(traces=trstmp)

        # tr.data = tr.data + np.random.randn(len(tr.data)) * 1e-10

        StaLta = np.sum([classic_sta_lta(tr, sta, lta) for tr in sttmp], axis=0)
        StaLta = StaLta[len(StaLta) / 2::]

        sfreq = trs.traces[0].stats['sampling_rate']
        sigma = sfreq / (2 * np.pi * fc)
        StaLtaf = gaussian_filter1d(StaLta, sigma=sigma, mode='reflect')

        picks = _Pick_STALTA_refraction(trs, StaLtaf)
        snrs = [calculate_snr(trs, p, SNR_window) for p in picks]

        cfs.append(StaLtaf)
        snrs.append(snrs)

        for k, p in enumerate(picks):
            if k == 0:
                opicks.append(make_pick(p, 'P', trs.traces[0], snrs[k]))
            elif k == 1:
                opicks.append(make_pick(p, 'S', trs.traces[0], snrs[k]))
            else:
                opicks.append(make_pick(p, '?', trs.traces[0], snrs[k]))

    catalog = event.Catalog()
    catalog.events.append(event.Event(picks=opicks))

    return catalog, cfs, snrs

def STALTA_picker(st, nphase=2, nsta=1e-3, nlta=4e-2, fc=50, nc=2,
    noise_mag=1e-25, SNR_window=5e-3):
    """The STA/LTA picker provides a first order good estimate of the arrival
    times for both the P- and S- wave. The STA/LTA picker uses
    :py:class:`obspy.signal.trigger.recursive_sta_lta `.

    :param st: seismogram containing a seismic event
    :type st: :py:class:`obspy.core.stream.Stream`
    :param nphase: Number of phases to pick. The number of phase should
    represent the number of phase present in the signal. This parameter will
    be set-up to two for picking the P- and S-wave onset times.
    :param nsta: Length of the short term average (STA) window in seconds. The value
        must be smaller than the minimum separation between the P- and S-
        wave arrival (explanation to be improved).
    :param nlta: Length of the long term average (LTA) window in seconds. The value
        must be long enough to capture well the noise but not too long (
        description to be improved)
    :param fc: Corner frequency of the Gaussian kernel used to smooth the LTA/STA
        function in Hz. This parameter must be chosen to limit the number of
        peaks in the STA/LTA function (default 50 Hz).
    :param nc: Number of desired onset. *** PARAMETER NOT USED ***
    :param noise_mag: Amplitude of the noise to be added to the original signal
    :param SNR_window: Length of window in seconds in which the SNR is
    calculated before and after the pick
    :returns:  :py:class:`obspy.core.event.Catalog` -- a new catalog
    containing a single event with a list of picks
    """

    st.detrend('demean')
    stations = np.unique([tr.stats.station for tr in st])

    # for station in stations:
    opicks = []
    cfs = []
    snrs = []
    for station in stations:

        trs = st.select(station=station)
        # if len(trs) == 3:
        #   trs = _RotateP_S(trs)

        # else:
        #   continue

        sta = int(nsta * trs.traces[0].stats['sampling_rate'])
        lta = int(nlta * trs.traces[0].stats['sampling_rate'])

        # Noise added to avoid zeroed trace.
        rd = np.random.randn(len(trs.traces[0].data)) * noise_mag
        trstmp = [Trace(
            data=np.hstack(
                (tr.data[::-1].reshape(len(tr)) + rd[::-1].reshape(len(rd)),
                tr.data.reshape(len(tr)) + rd.reshape(len(rd)))),
            header=tr.stats) for tr in trs]
        sttmp = Stream(traces=trstmp)

        # tr.data = tr.data + np.random.randn(len(tr.data)) * 1e-10

        StaLta = np.sum([recursive_sta_lta(tr, sta, lta) for tr in sttmp],
                        axis=0)
        StaLta = StaLta[len(StaLta) / 2::]

        sfreq = trs.traces[0].stats['sampling_rate']
        sigma = sfreq / (2 * np.pi * fc)
        StaLtaf = gaussian_filter1d(StaLta, sigma=sigma, mode='reflect')

        picks = _Pick_STALTA(trs, StaLtaf, nphase)
        SNRs = [calculate_snr(trs, p, SNR_window) for p in picks]

        # print picks[0], (picks - trs[0].stats.starttime) * trs[0].stats.sampling_rate, SNRs

        cfs.append(StaLtaf)
        snrs.append(SNRs)

        for k, p in enumerate(picks):
            if k == 0:
                opicks.append(make_pick(p, 'P', trs.traces[0], SNRs[k]))
            elif k == 1:
                opicks.append(make_pick(p, 'S', trs.traces[0], SNRs[k]))
            else:
                opicks.append(make_pick(p, '?', trs.traces[0], SNRs[k]))

    catalog = event.Catalog()
    catalog.events.append(event.Event(picks=opicks))

    return catalog, cfs, snrs

def kurtosis_picker(st, picks, freqmin=100, freqmax=1000, pick_freqs=None,
    kurtosis_window=None, CF3_tol=10e-3, SNR_window=5e-3):
    """Kurtosis picker adapted for microseismic event processing
    from ``Baillard et al. 2014``.

    :param st: seismogram containing a seismic event
    :type st: :py:class:`obspy.core.stream.Stream`
    :param cat: catalog containing the event with previous picks
    :type cat: obspy.core.event.Catalog
    :param freqmin: Low end of frequency band used to prefilter the seismograms.
        This value depends on the type of sensor.
        The default value is optimized for 2:3 kHz accelerometers.
    :param freqmax: High end of frequency band used to prefilter the seismograms.
        This value depends on the type of sensor.
        The default value is optimized for 2:3 kHz accelerometers.
    :param pick_freqs: The smoothing frequencies for the smoothing
        window applied to the CF3 function (this is used instead of Ns).
    :param kurtosis_window: Windows lengths used to calculate CF3 up to P-S delay
    :param CF3_tol: Maximum time to move a pick
    :param SNR_window: Length of window in seconds in which the SNR is calculated before and after the pick
    :returns:  :py:class:`obspy.core.event.Catalog` -- a new catalog containing a single event with a list of picks
    """

    '''
    if not cat.events:
        return None

    evt = cat.events[0]

    if not evt['picks']:
        return None

    prevPicks = evt['picks']
    '''
    prevPicks = copy_picks_to_dict(picks)

    # # clip the signal to the existing picks
    # if prevPicks:
    #   p_early = UTCDateTime(2100, 1, 1)
    #   p_late = UTCDateTime(1970, 1, 1)

    #   for picks in prevPicks:
    #       if picks['time'] < p_early:
    #           p_early = picks['time']
    #       if picks['time'] > p_late:
    #           p_late = picks['time']

    #   p_early -= 25e-3
    #   p_late += 25e-3
    #   st.trim(starttime=p_early, endtime=p_late)

    if pick_freqs is None:
        pick_freqs = np.linspace(50, 1000, 20)
    if kurtosis_window is None:
        kurtosis_window = np.array([1, 2, 3]) * 1e-3

    st.detrend('demean')
    stations = np.unique([tr.stats.station for tr in st])

    opicks = []
    cfs = []
    snrs = []
    for station in stations:
        # find the existing P and S picks for the current station
        trs = st.select(station=station)
        #print('kurtosis: sta:%s ntr=%d' % (station, len(trs)))
        #if len(trs) < 3:
            #continue
        CF3 = np.zeros(len(trs.traces[0].data))

        for ws in kurtosis_window:
            for tr in trs:
                cf1, cf2, cf3 = _CalculateCF1_3(tr, WS=ws,
                                         BW=[freqmin, freqmax])
                CF3 += cf3

        cfs.append(CF3)

        # TODO:
        # Picker should probably work with no previous picks.
        # If the current station has not been picked before, a new pick will not be generated.
        # If the old pick has only either P or S, then only a new corresponding P or S will be generated, not both.
        if not prevPicks:
            continue

        for phase in prevPicks[station]:
            oldPick = prevPicks[station][phase]

        #for oldPick in prevPicks:
            #if oldPick['waveform_id'].station_code != station:
                #continue

            pick = _Pick_CF3(trs, CF3, oldPick['time'], pick_freqs, CF3_tol)
            SNR = calculate_snr(trs, pick, SNR_window)
            snrs.append(SNR)
            opicks.append(make_pick(pick, oldPick['phase_hint'], trs.traces[0], SNR))

            print('kurtosis: sta:%s pha:%s old_pick:%s new:%s' % (station,phase,oldPick.time,opicks[-1].time))

    #catalog = event.Catalog()
    #catalog.events.append(event.Event(picks=opicks))

    #return catalog, cfs, snrs
    return opicks

def snr_picker(st, picks, snr_dt=None, snr_window=(1e-3, 20e-3), filter=None):
    """
    Function to improve the picks based on the SNR.
    :param st: seismogram containing a seismic event
    :type st: :py:class:`obspy.core.stream.Stream`
    :param picks: list of microquake.core.event.Pick object
    picks
    :type picks: microquake.core.event.Catalog
    :param snr_dt: Window in which the picks will be improved.
    :param snr_window: Length of window in seconds in which the SNR is calculated
    before and after the pick
    :type snr_window: (tuple)
    :returns:  Tuple comprising 1) a :py:class:`microquake.core.event.Catalog`
    a new catalog containing a single event with a list of picks and 2) the SNR
    """

    function_name = 'snr_picker'

    filter_p = False
    filter_s = False
    if filter == 'S':
        filter_p = True
    elif filter == 'P':
        filter_s = True

    previous_picks = copy_picks_to_dict(picks)

    if snr_dt is None:
        snr_dt = np.linspace(-5e-3, 5e-3, 20)

    st.detrend('demean')
    stations = np.unique([tr.stats.station for tr in st])

    opicks = []
    snrs = []

    pre_window_length = snr_window[0]
    post_window_length = snr_window[1]

    for station in stations:

        tr = st.select(station=station).composite()[0]

        if station not in previous_picks:
            logger.warning('SNR_detect: station:[%s] has no previous picks'
                           % station)
            continue

        for phase in previous_picks[station]:

            if filter_p and phase == 'P':
                continue
            elif filter_s and phase == 'S':
                continue

            earliest_time = latest_time = None
            if phase == 'S' and 'P' in previous_picks[station]:
                delta = 1 / 2 * (previous_picks[station]['S'].time -
                                 previous_picks[station]['P'].time)
                earliest_time = previous_picks[station]['S'].time - delta

            elif phase == 'S' and 'P' not in previous_picks[station]:
                earliest_time = previous_picks[station]['S'].time - \
                                pre_window_length
            elif phase == 'P' and 'S' in previous_picks[station]:
                delta = 1 / 2 * (previous_picks[station]['S'].time -
                                 previous_picks[station]['P'].time)
                latest_time = previous_picks[station]['P'].time + delta
            elif phase == 'P' and 'S' not in previous_picks[station]:
                latest_time = previous_picks[station]['P'].time + \
                              post_window_length

            old_pick = previous_picks[station][phase]

            if earliest_time is None or earliest_time < tr.stats.starttime + \
                    .03:
                earliest_time = tr.stats.starttime + .03

            if latest_time is None or latest_time > tr.stats.endtime - .06:
                latest_time = tr.stats.endtime - .06

            tau = []
            for dt in snr_dt:
                taut = old_pick.time + dt

                if earliest_time <= taut <= latest_time:
                    tau.append(taut)

            if not tau:
                continue
            if tau[0] <= tr.stats.starttime:
                logger.error("Too early! tau[0]=%s <= tr.st=%s"
                             % (tau[0], tr.stats.starttime))
                logger.error("earliest_time:%s latest_time:%s"
                             % (earliest_time, latest_time))

                continue
            if tau[-1] >= tr.stats.endtime:
                logger.error("Too late! tau[-1]=%s >= tr.et=%s"
                             % (tau[-1], tr.stats.endtime))
                continue

            tau = np.array(tau)
            indices = (tau - tr.stats.starttime) * tr.stats.sampling_rate
            tmp = np.array([(taut, index, calculate_snr(tr, taut,
                                                 pre_wl=pre_window_length,
                                                 post_wl=post_window_length))
                            for taut, index in zip(tau, indices)])

        # MTH: this is a hack to try to force the solution close to the oldPick
            alpha = 0
            """
            alpha = 10.
            for i, foo in enumerate(tmp):
                time = foo[0]
                snr = foo[1]
                dt = np.abs(old_pick.time - foo[0])
                #dt = np.abs(oldPick['time'] - foo[0])
                scale = np.exp(-alpha * dt)
                tmp[i,1] *= np.exp(-alpha * dt)
                #print(time, dt, snr, scale, snr*scale, tmp[i,1])
            """

            index = np.argmax(tmp[:, 2])
            pick_time = tmp[index, 0]

            import matplotlib.pyplot as plt

            snr = calculate_snr(tr, pick_time, pre_wl=pre_window_length,
                                post_wl=post_window_length)

            # plt.plot(tmp[:, 1], tmp[:, 2] / np.max(tmp[:, 2]))
            # plt.plot(tr.data / np.max(tr.data))
            # plt.axvline(tmp[index, 1], color='r', ls='--')
            # plt.xlim([tmp[0, 1], tmp[-1, 1]])
            # plt.show()

            # from ipdb import set_trace; set_trace()
            logger.debug("%s: sta:%s [%s] time_diff:%0.3f SNR:%.2f" %
                         (function_name, station, phase, old_pick.time -
                          pick_time, snr))

            method_string = 'snr_picker preWl=%.3g postWl=%.3g alpha=%.1f' % \
                            (pre_window_length, post_window_length, alpha)
            opicks.append(make_pick(pick_time, phase=old_pick.phase_hint,
                                    wave_data=tr, snr=snr,
                                    method_string=method_string,
                                    resource_id=old_pick.resource_id))
            snrs.append(snr)

    return snrs, opicks

def calculate_snr(trace, pick, pre_wl=1e-3, post_wl=10e-3):
    """
    input :
    trs - Obspy stream
    Pick Time - in Obspy UTCDateTime
    preWl - Length of pre-window in seconds
    postWl - Length of post-window in seconds

    output:
    SNR - Signal to noise ratio
    """

    tr = trace

    sr = tr.stats.sampling_rate
    st = tr.stats.starttime
    et = tr.stats.endtime
    ps = int((pick - st) * sr)
    n_pre = int(pre_wl * sr)
    n_post = int(post_wl * sr)

    if pick + post_wl > et:
        energy_s = np.var(tr.data[ps:])
    else:
        energy_s = np.var(tr.data[ps:ps + n_post])

    if pick - pre_wl < st:
        energy_n = np.var(tr.data[:ps])
    else:
        energy_n = np.var(tr.data[ps - n_pre:ps])

    if (energy_n == 0) | (energy_s == 0):
        return 0

    snr = 10 * np.log10(energy_s / energy_n)

    return snr

def calculate_energy(stream, pick, Wl=5e-3):
    sr = stream.traces[0].stats['sampling_rate']
    st = stream.traces[0].stats['starttime']
    Ps = int((pick - st) * sr)
    Nb = int(Wl * sr)

    EnergyS = np.sum([np.var(tr.data[Ps-Nb/4:Ps+3*Nb/2]) for tr in stream])

    return EnergyS

def _CalculateCF1_3(tr1, BW=None, WS=1e-3):
    if not BW:
        BW = [100, 5000]
    tr = tr1.copy()
    tr.taper(max_percentage=0.5, type='cosine')
    tr.filter(type='bandpass', freqmin=BW[0], freqmax=BW[1])

    cf1 = kurtosis(tr, win=WS)
    cf1 /= np.max(np.abs(cf1))

    cf2 = np.zeros(cf1.shape)
    dcf1 = np.diff(cf1)
    dcf1[dcf1 < 0] = 0
    cf2[0] = cf1[0]

    for k in range(1, len(cf1)):
        cf2[k] = cf2[k - 1] + dcf1[k - 1]

    try:
        cf3 = detrend(cf2, type='linear')
    except:
        cf3 = np.zeros(cf2.shape)

    return cf1, cf2, cf3

def _Pick_STALTA(st, stalta, nphase):

    # Finding the two largest maximum of the smoothed STALTA
    # function. The two largest maximum are used as starting pick
    # values.

    sr = st.traces[0].stats['sampling_rate']
    starttime = st.traces[0].stats['starttime']
    endtime = st.traces[0].stats['endtime']

    buf = (endtime - starttime) * 0.0001  # picks cannot be in the first and last 1% of the stream

    mx = np.r_[True, stalta[1:] > stalta[:-1]] & np.r_[stalta[:-1] > stalta[1:], True]

    i1 = np.nonzero(mx)[0]

    i2 = np.argsort(stalta[i1])[::-1]
    # i2 = i1

    # EnergySs = np.array([calculate_energy(st, starttime + i1[k] / sr, Wl=5e-3) for k in i2])

    # Eratio = EnergySs / np.max(EnergySs)
    # ie = np.nonzero(Eratio > 0.05)[0]
    # i2 = i2[ie]

    picks = []
    for k in range(0, nphase):
        try:
            picks.append(i1[i2[k]])
        except:
            pass
            # logger.warning("_Pick_STALTA: station=%s phase=%d FAILED!" % (st.traces[0].stats.station, k))
            # picks.append(int((endtime - starttime)*sr))

    # if len(picks) < nphase:
    #   logger.warning("_Pick_STALTA: Not all phases were picked")

    picks = np.sort([starttime + p / sr for p in picks])
    picks = np.array([p for p in picks if ((starttime + buf) < p < (endtime - buf))])

    return picks

def _Pick_STALTA_refraction(st, stalta):

    # Finding the two largest maximum of the smoothed STALTA
    # function. The two largest maximum are used as starting pick
    # values.

    sr = st.traces[0].stats['sampling_rate']
    starttime = st.traces[0].stats['starttime']
    # endtime = st.traces[0].stats['endtime']

    mx = np.r_[True, stalta[1:] > stalta[:-1]] & np.r_[stalta[:-1] > stalta[1:], True]

    i1 = np.nonzero(mx)[0]

    i2 = 0  # np.argsort(stalta[i1])[::-1]
    # i2 = i1

    # EnergySs = np.array([calculate_energy(st,starttime+i1[k]/sr,Wl = 5e-3) for k in i2])

    # Eratio = EnergySs/np.max(EnergySs)
    # ie = np.nonzero(Eratio > 0.05)[0]
    # i2 = i2[ie]

    picks = [i1[i2]]

    picks = np.sort([starttime + p / sr for p in picks])

    return picks

# This should work with only one pick at a time regardless whether it is a P or S pick.
def _Pick_CF3(st, cf3, iniPick, f=np.linspace(50, 1000, 20), tol=10e-3):
    """
    input :
    st - stream object
    cf3 - characteristic function number 3
    f - list of frequency (used instead of the Ns parameter)
    tol - maximum time to move a pick
    """

    sr = st.traces[0].stats['sampling_rate']
    ST = st.traces[0].stats['starttime']
    # first = True
    CF3_2 = np.hstack((cf3, cf3[::-1])) - cf3[0]
    CF3_2TR = Trace(data=CF3_2, header=st.traces[0].stats)  # ??

    Pick = int((iniPick - ST) * sr)  # initial pick in sample
    Pick_stalta = Pick

    for freq in f:
        # sigma = sfreq/(2*np.pi*freq)
        CF3_2TR_F = CF3_2TR.copy()
        CF3_2TR_F.filter('lowpass', freq=freq)
        # CF3_f = gaussian_filter1d(CF3,sigma=sigma,mode='reflect')
        CF3_f = CF3_2TR_F.data
        #CF3_f = CF3_f[:len(CF3_f) / 2]
        CF3_f = CF3_f[:int(len(CF3_f) / 2)]
        s = np.r_[True, CF3_f[1:] < CF3_f[:-1]] & np.r_[CF3_f[:-1] < CF3_f[1:], True]
        indices = np.nonzero(s)[0]
        CF4 = np.zeros(CF3_f.shape)
        for i in indices[0:-1]:
            CF4[i] = CF3_f[i] - CF3_f[i + 1]

        indices = np.nonzero(CF4)[0]

        try:
            pick_tmp = indices[np.argmin(np.abs(Pick - indices[indices <= Pick_stalta]))]
        except:
            pick_tmp = Pick

        if np.abs(pick_tmp - Pick_stalta) <= (tol * sr):
                Pick = pick_tmp

    # Return pick time in samples
    pick = ST + Pick / sr
    return pick

@deprecated
def triggersByGroup(st, trigger_type="recstalta", group="station", thr_on=3, thr_off=2,
    thr_coincidence_sum=1, sta=0.01, lta=1):

    # Accentuate peaks - bug below, skip for now
    # for i,tr in enumerate(st):
    #   st2.traces[0].data = tr.data**2*np.sign(tr.data)

    gp = np.array([])

    if group == "all":
        # one trigger for all stations
        trig = coincidence_trigger(trigger_type, thr_on, thr_off, st,
                                   thr_coincidence_sum, sta=sta, lta=lta, details=True)

    elif group == "station":
        # one trigger per station
        trig = []

        SensorList = []
        for tr in st:
            SensorList.append(tr.stats.station)
        SensorList = np.unique(SensorList)

        for k, S in enumerate(SensorList):
            st2 = st.select(station=S)
            try:
                trigtmp = coincidence_trigger(trigger_type, thr_on, thr_off,
                st2, thr_coincidence_sum, sta=sta, lta=lta, details=True)
            except:
                trigtmp = []

            # plot - debug
            # for tr in st2:
            #   cft = classic_sta_lta(tr.data, int(sta * tr.stats.sampling_rate), int(lta * tr.stats.sampling_rate))
            #   plotTrigger(tr, cft, thr_on, thr_off)
            # rpdb.set_trace()

            if k == 0:
                trig = trigtmp
            else:
                trig = trig + trigtmp

            gp = np.hstack((gp, np.ones(len(trigtmp)) * k))
    else:

        # individual recursive_sta_lta for each trace
        trig = []
        pass

    return trig, gp

@deprecated
def associateTriggers(trig, gp, tolerance=25e-3):

    tme = np.sort([trg['time'] for trg in trig])

    trigger = []

    # at least 2 sensors need to be in the same tolerance window
    indices2 = np.argsort([trg['time'] for trg in trig])
    gp = gp[indices2]
    trig = np.array(trig)
    trig = trig[indices2]

    k = 0
    while k < len(tme):
        indices = np.nonzero((tme - tme[k] > 0) & (tme - tme[k] < tolerance))[0]
        if len(indices) > 0:
            k = indices[-1]
            gp2 = gp[indices]
            if len(np.unique(gp2)) > 1:
                trigs = trig[indices]
                best_trig = np.argmax([t['cft_peak_wmean'] for t in trigs])
                best_trig = trigs[best_trig]

                data = {'mean_time': best_trig['time'],
                        'data': trigs}

                trigger.append(data)
        else:
                k += 1

    return trigger

@deprecated
def picksFromTriggers(st, trg, method="by_triggers", tolerance=20e-3, clip_stream=True, filter_stream=True):

    picks = []

    if method == "all":
        SensorList = []
        for tr in st:
            SensorList.append(tr.stats.station)

        SensorList = np.unique(SensorList)
        for S in SensorList:
            st2 = st.select(station=S)

            picks2 = compute_picks(st2, trg, tolerance, clip_stream, filter_stream)
            if picks2 is not None:
                picks = np.hstack((picks, picks2))

    elif method == "by_triggers":
        for tg in trg:
            SensorList = np.unique(tg['stations'])
            for S in SensorList:
                st2 = st.select(station=S)

                picks2 = compute_picks(st2, tg, tolerance, clip_stream, filter_stream)
                if picks2 is not None:
                    picks = np.hstack((picks, picks2))

    return picks

@deprecated
def compute_picks(st2, trg, tolerance=20e-3, clip_stream=True, filter_stream=True):

    picks = []

    tg_time = trg['time']
    if clip_stream:
        st = st2.trim(tg_time - tolerance, tg_time + tolerance)
        if filter_stream:
            st = st.filter_stream(lf=100, hf=1000, copy=False)
    else:
        if filter_stream:
            st = st2.filter_stream(lf=100, hf=1000, copy=True)
        else:
            st = st2

    starttime = st.traces[0].stats.starttime
    # starttime2 = st2.traces[0].stats.starttime
    sr = st.traces[0].stats.sampling_rate
    station_type = st.traces[0].stats.station_type

    try:
        if len(st) < 3:
            data = st.traces[0].data ** 2 * np.sign(st.traces[0].data)
        else:
            data = (st.traces[0].data ** 2 +
                    st.traces[1].data ** 2 +
                    st.traces[2].data ** 2) * \
                    np.sign(st.traces[0].data)
        i = np.argmax(np.abs(data))
        data = data / data[i]

    except:
        return

    if 'A' in station_type:
        # Parameters originally estimated for Northparkes TBM project (which had a sampling rate of around 5000) -> Correction x2
        pickBaer = pk_baer(reltrc=data, samp_int=1, tdownmax=2, tupevent=40, thr1=10, thr2=20, preset_len=10, p_dur=10)
        # pickBaer = pk_baer(reltrc=data, samp_int=1, tdownmax=16, tupevent=60, thr1=10, thr2=20, preset_len=20, p_dur=20)
    else:
        # Parameters originally estimated for Northparkes TBM project (which had a sampling rate of around 2500) -> Correction x4
        pickBaer = pk_baer(reltrc=data, samp_int=1, tdownmax=4, tupevent=10, thr1=10, thr2=20, preset_len=10, p_dur=10)
        # pickBaer = pk_baer(reltrc=data, samp_int=1, tdownmax=32, tupevent=40, thr1=10, thr2=20, preset_len=10, p_dur=10)

    # plot - debug
    # if st.traces[0].stats.station == '14_A1':
    # trg_time = (trg['time'] - starttime)*sr
    # diff_samples = (starttime - starttime2)*sr
    # num_traces = len(st)
    # plt.close()
    # # for k, (tr, tr2) in enumerate(zip(st, st3.select(station=st.traces[0].stats.station))):
    # for k, (tr, tr2) in enumerate(zip(st, st2)):
    #   trace.plot_traces(tr, k, num_traces*2, pickSample=[trg_time, pickBaer[0]])
    #   trace.plot_traces(tr2, k+num_traces, num_traces*2, pickSample=[trg_time+diff_samples, pickBaer[0]+diff_samples])
    # plt.show()

    if pickBaer[0] < 10:  # pick can't occur in the first N samples
        return

    Noise = np.var(data[0:pickBaer[0]])
    Signal = np.var(data[pickBaer[0]:pickBaer[0] + 50])
    # SNR = 10*np.log((Signal-Noise)/(Noise+1e-10))
    SNR = 10 * np.log10(Signal / (Noise + 1e-10))
    # print SNR, pickBaer

    if SNR < 3:
        return

    t = starttime + pickBaer[0] / sr

    for tr in st:
        this_pick = event.Pick()
        this_pick.time = t
        this_pick.phase_hint = 'P'
        this_pick.waveform_id = event.WaveformStreamID(
            network_code=tr.stats.network,
            station_code=tr.stats.station,
            location_code=tr.stats.location,
            channel_code=tr.stats.channel)
        this_pick.evaluation_mode = 'automatic'
        # this_pick.creation_info = creation_info
        if 'E' in pickBaer[1]:
            this_pick.onset = 'emergent'
        if 'I' in pickBaer[1]:
            this_pick.onset = 'impulsive'
        this_pick.evaluation_status = 'preliminary'

        if len(st) < 3:
            if 'U' in pickBaer[1]:
                this_pick.polarity = 'positive'
            elif 'D' in pickBaer[1]:
                this_pick.polarity = 'negative'
        else:
            d_0 = tr.data[pickBaer[0]]
            d_1 = tr.data[pickBaer[0] + 3]
            if d_1 > d_0:
                this_pick.polarity = 'positive'
            else:
                this_pick.polarity = 'negative'

        this_pick.SNR = SNR

        picks.append(this_pick)

    return picks

def automatic_picking(st_in, site, params):
    """
    Automatic picking algorithm
    :param st_in: time series data
    :type st_in: microquake.core.Stream
    :param site: an object containing information on the station
    :type site: microquake.core.station.Site
    :param params: dictionnary containing the control parameters for the
    automatic picker
    :type params: AttribDict
    :return: returns an event
    :rtype: obspy.core.event.Event
    """

    st = st_in.copy()

    if not ('picker' in params.keys()):
        logger.error('params is not well formed... exiting')
        return

    #NLL_BASE = params.nll.NLL_BASE
    nll_opts_auto = nll.init_nlloc_from_params(params)
    #nll_opts_auto = nll.NLL(params.project_code, suffix=nll_suffix, base_folder=NLL_BASE)

    pcks = []

    try:
        Polarity_kwargs = params.picker.polarity
    except:
        logger.warning('Probably no parameter for polarity measurements in '
                       'the control file')
        Polarity_kwargs = None

    for stname in st.unique_stations():
        for net in site.select(station=stname).networks:
            if net:
                curr_sta = net.stations[0]
                stype = curr_sta.sensor_type

                try:
                    STALTA_picker_kwargs = params.picker.STALTA_picker[stype]
                except:
                    logger.warning('Probably no parameter for STALTA_picker '
                                   'in the control file for %s...' %stype)
                    STALTA_picker_kwargs = None

                try:
                    SNR_picker_kwargs = params.picker.SNR_picker[stype]
                except:
                    logger.warning('Probably no parameter for SNR_picker in '
                                   'the control file for %s...' %stype)
                    SNR_picker_kwargs = None

        st_station = st.select(station=stname)
        cat_STALTA1, cfs, snrs = STALTA_picker(st_station,
                                               **STALTA_picker_kwargs)

        try:
            cat_SNR1, snrs = SNR_picker(st_station, cat_STALTA1,
                                        **SNR_picker_kwargs)
        except:
            return

        snr_th = params.picker[stype].SNR_threshold

        for (pk, snr) in zip(cat_SNR1.events[0].picks, snrs):
            if snr > snr_th:
                pcks.append(pk)

    if (len(pcks) < params.picker.min_picks):
        logger.info('too few picks ... exiting')
        return

    cat_auto1 = event.Catalog(cat_SNR1)
    cat_auto1.events[0].picks = pcks

    cat_filt1 = event.Catalog(cat_auto1)
    # cat_filt2.events[0].picks = picks

    # we remove bad picks only from uniaxials
    while True:
        # creating an Origin and associating pick to arrivals
        origin = event.Origin()
        cat_filt1[0].preferred_origin_id = origin.resource_id.id
        for k, pick in enumerate(cat_filt1[0].picks):
            arrival = event.Arrival()
            arrival.pick_id = pick.resource_id.id
            arrival.phase = pick.phase_hint
            origin.arrivals.append(arrival)

        cat_filt1[0].origins.append(origin)
        cat_filt1[0].preferred_origin_id = origin.resource_id.id
        cat_auto = nll_opts_auto.run_event(cat_filt1[0], status='preliminary')
        if not cat_auto:
            logger.error('unable to locate the event')
            return
        res = np.abs([arr.time_residual for arr in cat_auto.events[0].origins[0].arrivals])

        if (np.max(res) < params.picker.rms_residual_threshold) or (len(res) \
            < params.picker.min_picks):
            break
        del cat_filt1.events[0].picks[np.argmax(res)]

    if (len(res) < params.picker.min_picks):
        return

    res = np.abs([arr.time_residual for arr in cat_auto.events[0].origins[0].arrivals])

    g = lambda x: x ** 2
    rms_residual = np.sqrt(np.mean(g(res)))

    uncertainty = cat_auto.events[0].origins[0].origin_uncertainty.confidence_ellipsoid.semi_major_axis_length
    if (rms_residual > params.picker.rms_residual_threshold) \
        or (uncertainty > params.picker.uncertainty_threshold):
        cat_auto.events[0].origins[0].evaluation_status = 'rejected'
    else:
        cat_auto.events[0].origins[0].evaluation_status = 'preliminary'

    return cat_auto[0]

def eventCategorization_polarity(catalog, site):
    """
    determine the event category by looking at the polarity of the first motion.
    It is assumed that blast will generate mostly positive first motion whereas seismic event
    will generate mixte first motion

    :param st: seismograms
    :type st: obspy.core.Stream or microquake.core.Stream
    :param catalog: events catalog
    :type catalog: obspy.core.event.Catalog or microquake.core.event.Catalog
    :param site: information on network
    :type site: microquake.core.data.station.Site
    """

    catalog = event.Catalog(cat=catalog)
    for evi, evt in enumerate(catalog.events):
        picks = evt.picks
        polarity = []
        for origin in evt.origins:
            evloc = np.array([origin.x, origin.y, origin.z])
            for pick in picks:
                if not pick.polarity:
                    continue
                if pick.polarity.lower() == "positive":
                    pick_polarity = 1
                elif pick.polarity.lower() == 'negative':
                    pick_polarity = -1
                else:
                    continue

                sta_code = pick.waveform_id.station_code
                station = site.stations(station=sta_code)[0]
                stloc = station.loc
                ev_st_vect = stloc - evloc
                for channel in station:
                    if not np.any(channel.orientation):
                        continue

                    polarity.append(np.sign(np.dot(channel.orientation, ev_st_vect)))

        polarity = np.array(polarity)
        if len(polarity[polarity == 1]) >= 0.85 * len(polarity):
            catalog.events[evi].event_type = "mining explosion"
        else:
            catalog.events[evi].event_type = "induced or triggered event"

        catalog.events[evi].event_type_certainty = "suspected"

    return catalog


"""
Model steps for lubricated contacts
"""
import typing
import warnings
from collections.abc import Sequence
from numbers import Number

import numpy as np
import slippy

from slippy.core import _NonDimensionalReynoldSolverABC
from ._model_utils import get_gap_from_model
from ._step_utils import make_interpolation_func, solve_normal_loading
from slippy.core.influence_matrix_utils import plan_convolve
from .steps import _ModelStep
from slippy.core.materials import _IMMaterial

__all__ = ['IterSemiSystem']

class IterSemiSystem(_ModelStep):
    """
    Lubrication solutions by iteration of semi systems

    Parameters
    ----------
    step_name: str
        An identifying name for the step used for errors and outputs
    reynolds_solver: _NonDimensionalReynoldSolverABC
        A reynolds solver object which will be used to solve for pressures
    rolling_speed: float or Sequence of float, optional (None)
        The mean speed of the surfaces (u1+u2)/2 parameters can be:

        * A constant (float) value, indicating a constant rolling speed.
        * A two element sequence of floats, indicating the the start and finish rolling speed, if this is used, the
          movement_interpolation_mode will be used to generate intermediate values
        * A 2 by n array of n rolling speed and n time values normalised to a 0-1 scale. array[0] should be
          position values and array[1] should be time values, time values must be between 0 and 1. The
          movement_interpolation_mode will be used to generate intermediate values
    number_of_steps: int
        The number of sub steps the problem will be split into, together with the time_period this controls the duration
        of the time steps used
    no_time: bool, optional (False)
        Set to true if there is no time dependence and the time steps can be solved in any order (no permanent changes
        between steps such as plastic deformation or heat generation), if True the model will be solved more efficiently
    time_period: float, optional (1.0)
        The total time period of this model step, used for solving sub-models and writing outputs
    off_set_x, off_set_y: float or Sequence of float, optional (0.0)
        The off set between the surfaces in the x and y directions, this can be a relative off set or an absolute off
        set, controlled by the relative_loading parameter:

        * A constant (float) value, indicating a constant offset between the surfaces (no relative movement of profiles)
        * A two element sequence of floats, indicating the the start and finish offsets, if this is used, the
          movement_interpolation_mode will be used to generate intermediate values
        * A 2 by n array of n absolute position values and n time values normalised to a 0-1 scale. array[0] should be
          position values and array[1] should be time values, time values must be between 0 and 1. The
          movement_interpolation_mode will be used to generate intermediate values
    interference, normal_load: float or Sequence of float, optional (None)
        The interference and normal load between the surfaces, only one of these can be set (the other will be solved
        for) setting neither keeps the interference as it is at the start of this model step. As above for the off sets,
        either of these parameters can be:

        * A constant (float) value, indicating a constant load/ interference between the surfaces.
        * A two element sequence of floats, indicating the the start and finish load. interference, if this is used, the
          movement_interpolation_mode will be used to generate intermediate values
        * A 2 by n array of n absolute position values and n time values normalised to a 0-1 scale. array[0] should be
          position values and array[1] should be time values, time values must be between 0 and 1. The
          movement_interpolation_mode will be used to generate intermediate values
    relative_loading: bool, optional (False)
        If True the load or displacement and off set will be applied relative to the value at the start of the step,
        otherwise the absolute value will be used. eg, if the previous step ended with a load of 10N and this step ramps
        from 0 to 10N setting relative_loading to True will ramp the total load form 10 to 20N over this step.
    movement_interpolation_mode: str or int, optional ('linear')
        Any valid input to scipy.interpolate.interp1d as the 'kind' parameter, using 'nearest', 'previous' or 'next'
        will cause a warning. This parameter controls how the offset and loading is interpolated over the step
    profile_interpolation_mode: {'nearest', 'linear'}, optional ('nearest')
        Used to generate the grid points for the second surface at the location of the grid points for the first
        surface, nearest ensures compatibility with sub models which change the profile, if the grid spacings of the
        surfaces match
    periodic_geometry: bool, optional (False)
        If True the surface profile will warp when applying the off set between the surfaces
    periodic_axes: tuple, optional ((False, False))
        For each True value the corresponding axis will be solved by circular convolution, meaning the result is
        periodic in that direction. NOTE: this only controls the deformation result from the materials, ensure that the
        reynolds equation solver used supports periodic solutions and is set to produce a periodic solution to the
        pressure equation. Additionally, ensure that the axes of periodicity match.
    periodic_im_repeats: tuple, optional (1,1)
        The number of times the influence matrix should be wrapped along periodic dimensions, only used if at least one
        of periodic axes is True. This is necessary to ensure truly periodic behaviour, no physical limit exists
    max_it_pressure: int, optional (100)
        The maximum number of iterations in the fluid pressure calculation loop
    rtol_pressure: float, optional (1e-7)
        The relative tolerance for the fluid pressure calculation loop
    max_it_interference: int, optional (100)
        The maximum number of iterations in the loop that finds the interference between the surfaces
    rtol_interference: float, optional (1e-7)
        The relative tolerance on the total load (integral of the pressure solution minus the applied load)
    initial_guess: {callable, 'previous', list}, optional ('previous')
        The initial guess for the interference, and/ or pressure profile between the surfaces, any callable will be
        called with the contact model and the undeformed nd_gap as positional arguments, it must return the interference
        and the pressure profile. 'previous' will use the result(s) from the previous step, if results are not found the
        interference and pressure profile will be set to 0. Can also be a 2 element list, the first element being the
        interference and the second being the pressure profile as an array, if this is the wrong shape zeros wil be
        used.
    no_update_warning: bool, optional (True)
        Change to False to suppress warning given when no movement or loading changes are specified

    Notes
    -----
    The solver iterates through a 'pressure loop' until the solution has converged to a set of pressure values, then the
    loading is checked, if the total pressure is too low the surfaces are brought closer together. This is continued
    until the total load has converged to the set value. This outer loop is referred to as the interference loop.

    Examples
    --------
    In this example we will model smooth surface EHL with a non newtonian fluid:

    >>> import slippy
    >>> slippy.CUDA = False  # Note: we are using a reynolds solver which does not currently support the CUDA back end
    >>> import slippy.surface as s
    >>> import slippy.contact as c
    >>>
    >>> radius = 0.01905       # The radius of the ball
    >>> load = 800             # The load on the ball in N
    >>> rolling_speed = 4      # The rolling speed in m/s (The mean speed of the surfaces)
    >>> youngs_modulus = 200e9 # The youngs modulus of the surfaces
    >>> p_ratio = 0.3          # The poission's ratio of the surfaces
    >>> grid_size = 65         # The number of points in the descretisation grid
    >>> eta_0 = 0.096          # Coefficient in the roelands pressure-viscosity equation
    >>> roelands_p_0 = 1/5.1e-9# Coefficient in the roelands pressure-viscosity equation
    >>> roelands_z = 0.68      # Coefficient in the roelands pressure-viscosity equation
    >>>
    >>> # Solving the hertzian contact to get the domain size and the initial guess
    >>> hertz_result = c.hertz_full([radius, radius], [float('inf'), float('inf')],
    >>>                             [youngs_modulus, youngs_modulus],
    >>>                             [p_ratio, p_ratio], load)
    >>> hertz_pressure = hertz_result['max_pressure']
    >>> hertz_a = hertz_result['contact_radii'][0]
    >>> hertz_deflection = hertz_result['total_deflection']
    >>> hertz_pressure_function = hertz_result['pressure_f']
    >>>
    >>> # make the surfaces
    >>> ball = s.RoundSurface((radius,)*3, shape = (grid_size, grid_size),
    >>>                       extent=(hertz_a*4,hertz_a*4), generate = True)
    >>> flat = s.FlatSurface()
    >>>
    >>> # assigning materials
    >>> steel = c.Elastic('steel', {'E' : youngs_modulus, 'v' : p_ratio})
    >>> ball.material = steel
    >>> flat.material = steel
    >>>
    >>> # make the non newtonian fluid
    >>> oil = c.Lubricant('oil') # Making a lubricant object to contain our sub models
    >>> oil.add_sub_model('nd_viscosity', c.lubricant_models.nd_roelands(eta_0, roelands_p_0,
    >>>                                                                  hertz_pressure, roelands_z))
    >>> oil.add_sub_model('nd_density', c.lubricant_models.nd_dowson_higginson(hertz_pressure))
    >>>
    >>> # make the contact model
    >>> my_model = c.ContactModel('lubrication_test', ball, flat, oil)
    >>>
    >>> # make a reynolds solver
    >>> reynolds = c.UnifiedReynoldsSolver(time_step = 0,
    >>>                                    grid_spacing = ball.grid_spacing,
    >>>                                    hertzian_pressure = hertz_pressure,
    >>>                                    radius_in_rolling_direction=radius,
    >>>                                    hertzian_half_width=hertz_a,
    >>>                                    dimentional_viscosity=eta_0,
    >>>                                    dimentional_density=872)
    >>>
    >>> # Find the hertzian pressure distribution as an initial guess
    >>> X, Y = ball.get_points_from_extent()
    >>> X, Y = X + ball._total_shift[0], Y + ball._total_shift[1]
    >>> hertzian_pressure_dist = hertz_pressure_function(X, Y)
    >>>
    >>> # Making the step
    >>> step = c.IterSemiSystem('main', reynolds, rolling_speed, 1, no_time=True, normal_load=load,
    >>>                         initial_guess=[hertz_deflection, hertzian_pressure_dist],
    >>>                         relaxation_factor=0.05, max_it_interference=3000)
    >>>
    >>> # Adding the step to the contact model
    >>> my_model.add_step(step)
    >>>
    >>> # solve the model:
    >>> state = my_model.solve()

    """
    "The minimum number of iterations in the reynolds solving loop"
    _dh = 0
    "The base change in height"
    _interferences = list()
    _load_errors = list()

    _reynolds: typing.Optional[_NonDimensionalReynoldSolverABC] = None
    initial_guess: typing.Optional[typing.Union[typing.Callable, list, str]]

    def __init__(self, step_name: str, reynolds_solver: _NonDimensionalReynoldSolverABC,
                 rolling_speed: typing.Union[float, typing.Sequence[float]],
                 number_of_steps: int = 1,
                 no_time: bool = False, time_period: float = 1.0,
                 off_set_x: typing.Union[float, typing.Sequence[float]] = 0.0,
                 off_set_y: typing.Union[float, typing.Sequence[float]] = 0.0,
                 interference: typing.Union[float, typing.Sequence[float]] = None,
                 normal_load: typing.Union[float, typing.Sequence[float]] = None,
                 relative_loading: bool = False,
                 movement_interpolation_mode: str = 'linear',
                 profile_interpolation_mode: str = 'nearest',
                 periodic_geometry: bool = False, periodic_axes: tuple = (False, False),
                 periodic_im_repeats: tuple = (1, 1),
                 max_it_pressure: int = 5000, rtol_pressure: float = 2e-6,
                 max_it_interference: int = 5000,
                 rtol_interference: float = 1e-4,
                 relaxation_factor: float = 0.1,
                 initial_guess: typing.Union[typing.Callable, str, typing.Sequence] = 'previous',
                 no_update_warning: bool = True):

        self._adjust_height_every_step = True
        self._initial_guess = initial_guess
        self._no_time = no_time
        self.total_time = time_period
        self._relative_loading = relative_loading
        self.profile_interpolation_mode = profile_interpolation_mode
        self._periodic_profile = periodic_geometry
        self._periodic_axes = periodic_axes
        self._periodic_im_repeats = periodic_im_repeats
        self._max_it_pressure = max_it_pressure
        self._max_it_interference = max_it_interference
        self._rtol_pressure = rtol_pressure
        self._rtol_interference = rtol_interference
        self._nd_max_pressure = None

        self.reynolds = reynolds_solver

        if relaxation_factor <= 0 or relaxation_factor > 1:
            raise ValueError("Relaxation factor must be greater than 0 and less than or equal to 1")
        self._relaxation_factor = relaxation_factor

        self.time_step = time_period / number_of_steps
        self.number_of_steps = number_of_steps

        self.update = set()

        if not isinstance(off_set_x, Number) or not isinstance(off_set_y, Number):
            if no_time:
                raise ValueError("Can not have no time dependence and sliding contact")
            off_set_x = [off_set_x] * 2 if isinstance(off_set_x, Number) else off_set_x
            off_set_y = [off_set_y] * 2 if isinstance(off_set_y, Number) else off_set_y
            off_set_x_func = make_interpolation_func(off_set_x, movement_interpolation_mode, 'relative_off_set_x')
            off_set_y_func = make_interpolation_func(off_set_y, movement_interpolation_mode, 'relative_off_set_y')
            self._off_set_upd = lambda time: np.array([off_set_x_func(time), off_set_y_func(time)])
            self.update.add('off_set')
            self.off_set = None
        else:
            self.off_set = np.array([off_set_x, off_set_y])

        if normal_load is not None and interference is not None:
            raise ValueError("Both normal_load and interference are set, only one of these can be set")
        if normal_load is None and interference is None:
            if relative_loading:
                interference = 0
            else:
                raise ValueError("Cannot have no set load or interference and not relative loading, set either the"
                                 "normal load, normal interference or change relative_loading to True")

        if isinstance(rolling_speed, Number):
            self.rolling_speed = rolling_speed
        else:
            self.rolling_speed = None
            self._rolling_speed_upd = make_interpolation_func(rolling_speed, movement_interpolation_mode,
                                                              'rolling_speed')
            self.update.add('rolling_speed')

        if normal_load is not None:
            if isinstance(normal_load, Number):
                self.normal_load = normal_load
            else:
                self.normal_load = None
                self._normal_load_upd = make_interpolation_func(normal_load, movement_interpolation_mode,
                                                                'normal_load')
                self.update.add('normal_load')
            self.load_controlled = True
        else:
            self.normal_load = None

        if interference is not None:
            if isinstance(interference, Number):
                self.interference = interference
            else:
                self.interference = None
                self._interference_upd = make_interpolation_func(interference, movement_interpolation_mode,
                                                                 'interference')
                self.update.add('interference')
            self.load_controlled = False

        if not self.update and no_update_warning:
            warnings.warn("Nothing set to update")

        self._provides = None

        base_provides = {'just_touching_gap', 'surface_1_points', 'surface_2_points', 'off_set', 'time_step', 'time',
                         'interference', 'total_normal_load', 'pressure', 'nd_pressure', 'loads',
                         'surface_1_displacement', 'surface_2_displacement', 'total_displacement', 'converged',
                         'gap', 'nd_gap', 'rolling_speed'}

        provides = base_provides.union(reynolds_solver.provides).union(reynolds_solver.requires)

        super().__init__(step_name, time_period, provides)

    @property
    def provides(self):
        results_set = self._provides
        if self.model is None:
            return results_set
        if self.model.lubricant_model is None:
            return results_set
        return results_set.union(set(self.model.lubricant_model.sub_models.keys()))

    @provides.setter
    def provides(self, value):
        if self._provides is None:
            self._provides = value
        else:
            raise ValueError("The provides property can only be set during instantiation")

    @property
    def reynolds(self):
        return self._reynolds

    @reynolds.setter
    def reynolds(self, value):
        if isinstance(value, _NonDimensionalReynoldSolverABC):
            self._reynolds = value
        else:
            raise ValueError("Cannot set a non reynolds solver object as the reynolds solver, to use custom solvers"
                             f"first subclass _NonDimensionalReynoldSolverABC from slippy.core, received "
                             f"type was {type(value)}")

    @reynolds.deleter
    def reynolds(self):
        self._reynolds = None

    def data_check(self, current_state):
        if self.reynolds is None:
            _data_check_error_or_warn(f"Reynolds solver not set for step {self.name}")
        if self.model.lubricant_model is None:
            _data_check_error_or_warn(f"Step {self.name} requires a lubricant to be defined in the contact model")
        if (self.reynolds.requires - set(self.model.lubricant_model.sub_models)) - {'nd_gap', 'nd_pressure'}:
            _data_check_error_or_warn(f"Reynolds solve in step {self.name} has requirements which are not provided by"
                                      f" the lubricant sub models:\n"
                                      f"Requires: {self.reynolds.requires}\n"
                                      f"Sub models provide: {set(self.model.lubricant_model.sub_models)}")

    def update_movement(self, relative_time, original):
        for name in self.update:
            if self._relative_loading:
                self.__setattr__(name, original[name] + self.__getattribute__(f'_{name}_upd')(relative_time))
            else:
                self.__setattr__(name, self.__getattribute__(f'_{name}_upd')(relative_time))

    def solve(self, previous_state: dict, output_file):
        cuda = slippy.CUDA
        slippy.CUDA = False
        start_time = previous_state['time']
        gs = self.model.surface_1.grid_spacing

        im_mats = (isinstance(self.model.surface_1.material, _IMMaterial) and
                   isinstance(self.model.surface_2.material, _IMMaterial))
        surf_1_material = self.model.surface_1.material
        surf_2_material = self.model.surface_2.material

        for s in self.sub_models:
            s.no_time = self._no_time

        relative_time = np.linspace(0, 1, self.number_of_steps+1)[1:]
        just_touching_gap = None

        original = dict()

        if self._relative_loading:
            original['normal_load'] = previous_state['total_normal_load'] if 'total_normal_load' in previous_state \
                else 0
            original['interference'] = previous_state['interference'] if 'interference' in previous_state else 0
            original['off_set'] = np.array(previous_state['off_set']) if 'off_set' in previous_state else \
                np.array([0, 0])

        previous_gap_shape = None  # shape of just touching gap array

        for i in range(self.number_of_steps):
            self.update_movement(relative_time[i], original)
            self.reynolds.rolling_speed = self.rolling_speed
            # find overlapping nodes
            if 'off_set' in self.update or just_touching_gap is None or not self._no_time:
                just_touching_gap, surface_1_points, surface_2_points \
                    = get_gap_from_model(self.model, interference=0, off_set=self.off_set,
                                         mode=self.profile_interpolation_mode, periodic=self._periodic_profile)

            time_step_current_state = dict(just_touching_gap=just_touching_gap, surface_1_points=surface_1_points,
                                           surface_2_points=surface_2_points, off_set=self.off_set,
                                           time_step=self.time_step, time=start_time+(i+1)*self.time_step)

            # make a new loads function if we need it
            if (previous_gap_shape is None or previous_gap_shape != just_touching_gap.shape) and im_mats:
                span = tuple([s*(2-pa) for s, pa in zip(just_touching_gap.shape, self._periodic_axes)])
                max_pressure = self.reynolds.dimensionalise_pressure(min([surf_1_material.max_load,
                                                                          surf_2_material.max_load]), True)
                self._nd_max_pressure = max_pressure
                im1 = surf_1_material.influence_matrix(components=['zz'], grid_spacing=[gs] * 2,
                                                       span=span, periodic_strides=self._periodic_im_repeats)['zz']
                im2 = surf_2_material.influence_matrix(components=['zz'], grid_spacing=[gs] * 2, span=span,
                                                       periodic_strides=self._periodic_im_repeats)['zz']
                total_im = im1 + im2
                loads_func = plan_convolve(just_touching_gap, total_im, circular=self._periodic_axes)
                previous_gap_shape = just_touching_gap.shape

            elif not im_mats:
                def loads_func(loads):
                    return solve_normal_loading(loads_z=loads, model=self.model,
                                                deflections='z', current_state=time_step_current_state)[0]['z']
            # sort out initial guess
            if i >= 0:
                initial_guess = 'previous'
            else:
                initial_guess = self.initial_guess

            if initial_guess is None:
                initial_guess = [self.reynolds.dimensionalise_gap(0.01),
                                 self.reynolds.dimensionalise_pressure(0.05)]
            if isinstance(initial_guess, str) and initial_guess.lower() == 'previous':
                pressure = np.zeros_like(just_touching_gap) if 'pressure' not in previous_state else \
                    previous_state['pressure']
                interference = 0.0 if 'interference' not in previous_state else previous_state['interference']
            elif isinstance(initial_guess, Sequence):
                interference = initial_guess[0]
                if isinstance(initial_guess[1], Number):
                    pressure = initial_guess[1] * np.ones_like(just_touching_gap)
                else:
                    try:
                        pressure = np.asarray(initial_guess[1], dtype=np.float)
                        assert (pressure.shape == just_touching_gap.shape)
                    except ValueError:
                        raise ValueError('Initial guess for pressure could not be converted to a numeric array')
                    except AssertionError:
                        # noinspection PyUnboundLocalVariable
                        raise ValueError("Initial guess for pressure produced an array of the wrong size:"
                                         f"expected {just_touching_gap.shape}, got: {pressure.shape}")
            elif hasattr(initial_guess, '__call__'):
                interference, pressure = initial_guess(self.model, just_touching_gap)
            else:
                raise ValueError('Unsupported type for initial guess')

            results_last_it = {'nd_pressure': self.reynolds.dimensionalise_pressure(pressure, True),
                               'just_touching_gap': just_touching_gap,
                               'interference': interference,
                               'pressure': pressure}

            # we have the interference, and the pressure initial guesses, find the initial displacement before solving

            if not (results_last_it['nd_pressure'] == 0).all():
                # noinspection PyUnboundLocalVariable
                results_last_it['total_displacement_z'] = loads_func(previous_state['pressure'])

            else:
                results_last_it['total_displacement_z'] = np.zeros_like(just_touching_gap)
            results_last_it = self.model.lubricant_model.solve_sub_models(results_last_it)
            # main loops
            it_num = 0
            # Find the gap and non denationalise it
            gap = just_touching_gap + results_last_it['total_displacement_z'] - results_last_it['interference']
            results_last_it['nd_interference'] = self.reynolds.dimensionalise_gap(results_last_it['interference'], True)
            results_last_it['gap'] = gap

            while True:
                nd_gap = self.reynolds.dimensionalise_gap(results_last_it['gap'], True)
                results_last_it['nd_gap'] = nd_gap
                # if flag:
                #     return locals()
                # else:
                #     flag = True
                # solve reynolds equation
                results_this_it = self.reynolds.solve(results_last_it, self._nd_max_pressure)

                # add just touching gap, needed for sub models
                results_this_it['just_touching_gap'] = just_touching_gap

                # check for pressure convergence
                change_in_pressures = results_this_it['nd_pressure'] - results_last_it['nd_pressure']
                total_nd_pressure = np.sum(results_last_it['nd_pressure'])  # use previous state here ... more stable
                if total_nd_pressure > 0:
                    pressure_relative_error = np.sum(np.abs(change_in_pressures)) / total_nd_pressure
                else:
                    pressure_relative_error = 1
                pressure_converged = pressure_relative_error < self._rtol_pressure

                # apply the relaxation factor to the pressure result
                if self._nd_max_pressure is not None:
                    results_this_it['nd_pressure'] = np.clip(results_last_it['nd_pressure'] +
                                                             self._relaxation_factor * change_in_pressures, None,
                                                             self._nd_max_pressure)
                else:
                    results_this_it['nd_pressure'] = (results_last_it['nd_pressure'] +
                                                      self._relaxation_factor * change_in_pressures)

                # solve contact geometry
                results_this_it['pressure'] = self.reynolds.dimensionalise_pressure(results_this_it['nd_pressure'])
                results_this_it['total_displacement_z'] = loads_func(results_this_it['pressure'])

                # find gap
                gap = just_touching_gap + results_this_it['total_displacement_z'] - results_last_it['interference']
                results_this_it['gap'] = gap

                # solve lubricant sub models
                results_this_it = self.model.lubricant_model.solve_sub_models(results_this_it)

                # check for load convergence
                total_load = np.sum(results_this_it['pressure']) * gs ** 2

                if self.load_controlled:
                    load_relative_error = (total_load / self.normal_load) - 1
                    load_converged = abs(load_relative_error) < self._rtol_pressure
                else:
                    load_converged = True
                    load_relative_error = 0.0

                results_this_it['nd_gap'] = self.reynolds.dimensionalise_gap(results_this_it['gap'], True)
                results_this_it['interference'] = results_last_it['interference']

                # escape the loop if it converged
                if pressure_converged and load_converged:
                    converged = True
                    print(f"Step {self.name} converged successfully after {it_num} iterations.")
                    print(f"Converged load is {total_load}, last change in pressure was {pressure_relative_error}\n")
                    break

                # escape the loop it if failed
                if it_num > self._max_it_pressure:  # this logic has changed used to just check the error
                    converged = False
                    print(f"Step {self.name} failed to converge after {it_num} iterations.\n")
                    print("Consider increasing the maximum number of iterations or reducing the relaxation factor")
                    print(f"Converged load is {total_load}, last change in pressure was {pressure_relative_error}")
                    break

                # adjust height for load balance

                if self._adjust_height_every_step and self.load_controlled:
                    # adjust height based on load balance
                    new_nd_interference = self.update_interference(it_num, pressure_relative_error,
                                                                   load_relative_error,
                                                                   results_last_it['nd_interference'],
                                                                   np.mean(results_this_it['nd_gap']),
                                                                   np.min(results_this_it['nd_gap']))
                    interference_updated = True

                elif pressure_converged:
                    # adjust height based on load balance
                    new_nd_interference = self.update_interference(it_num, pressure_relative_error,
                                                                   load_relative_error,
                                                                   results_last_it['nd_interference'],
                                                                   np.mean(results_this_it['nd_gap']),
                                                                   np.min(results_this_it['nd_gap']))
                    interference_updated = True
                else:
                    new_nd_interference = 0
                    interference_updated = False

                if interference_updated:
                    results_this_it['nd_interference'] = new_nd_interference
                    results_this_it['interference'] = self.reynolds.dimensionalise_gap(new_nd_interference)
                    gap = (just_touching_gap + results_this_it['total_displacement_z'] - results_this_it['interference'])
                    results_this_it['gap'] = gap
                    # print summary of iteration to log file
                    old_int = results_last_it['nd_interference']
                    print(f'{it_num}\ter_load: {load_relative_error:.4g}\t'
                          f'er_press: {pressure_relative_error:.4g}\t'
                          f'old_int: {old_int:.6g}\t'
                          f'new_int: {new_nd_interference:.6g}')
                else:
                    results_this_it['nd_interference'] = results_last_it['nd_interference']
                    results_this_it['interference'] = results_last_it['interference']

                it_num += 1
                results_last_it = results_this_it

            # clean up after it has converged
            current_state = {**time_step_current_state, **results_this_it}
            pressure = current_state['pressure']
            if im_mats:
                current_state['surface_1_displacement_z'] = plan_convolve(pressure, im1,
                                                                          circular=self._periodic_axes)(pressure)
                current_state['surface_2_displacement_z'] = plan_convolve(pressure, im2,
                                                                          circular=self._periodic_axes)(pressure)
                current_state['total_displacement_z'] = current_state['total_displacement_z']

            else:
                all_disp = solve_normal_loading(pressure, self.model, current_state, 'z')
                current_state['total_displacement_z'] = all_disp[0]
                current_state['surface_1_displacement_z'] = all_disp[1]
                current_state['surface_2_displacement_z'] = all_disp[2]

            # del current_state['total_displacement_z']
            current_state['total_normal_load'] = total_load
            current_state['loads_z'] = current_state['pressure']
            current_state['converged'] = converged
            current_state['rolling_speed'] = self.rolling_speed
            current_state = self.solve_sub_models(current_state)
            self.save_outputs(current_state, output_file)

            previous_state = current_state

        slippy.CUDA = cuda

        return current_state

    def __repr__(self):
        return "Lubrication step"

    def update_interference(self, it_num, pressure_error_rel, load_error_rel, current_interference, mean_gap, min_gap):
        """This method updates the interference between the 2 surfaces during solution

        Parameters
        ----------
        it_num: int
            The current iteration number of the solution
        pressure_error_rel
            The current relative pressure error on the solution
        load_error_rel: float
            The current relative load error
        current_interference: float
            The current interference between the two bodies (the maximum overlap between the undeformed profiles)
        mean_gap
            The average nd_gap between the surfaces
        min_gap
            The minimum nd_gap between the surfaces

        Returns
        -------
        new_interference: float
            The non dimensional interference between the surfaces

        Notes
        -----
        This method is quite basic at the moment and can definitely be improved, if executed on every loop, the height
        is just updated by a fixed proportion of the minimum nd_gap size

        If updated only when the solver converges, the Regula-Falsi method is used
        """
        if self._adjust_height_every_step:
            new_interference = current_interference - 0.1 * load_error_rel
            return new_interference

        # else:  # height only adjusted when the load has converged, in this case use the Regula-Falsi method
        self._interferences.append(current_interference)
        self._load_errors.append(load_error_rel)

        if len(self._interferences) == 1:  # if this is the first guess give a value that will probably bound it
            new_interference = current_interference + abs(mean_gap) * -np.sign(load_error_rel)
            return new_interference

        if len(self._interferences) > 2 and abs(sum(np.sign(self._load_errors))) == 1:
            # if this is true we must have bound a root

            if self._load_errors[0] * self._load_errors[2] < 0:
                del (self._interferences[1])
                del (self._load_errors[1])
            else:
                del (self._interferences[0])
                del (self._load_errors[0])

        else:  # we have not bound a root, continue with the secant method but del the first item so we progress
            del (self._interferences[0])
            del (self._load_errors[0])

        new_interference = (self._interferences[0] - self._load_errors[0] * (self._interferences[1] -
                                                                             self._interferences[0]) /
                            (self._load_errors[1] - self._load_errors[0]))

        print(f'Adjusting interference, new interference is {new_interference}')

        return new_interference

def _data_check_error_or_warn(msg: str):
    if slippy.ERROR_IN_DATA_CHECK:
        raise ValueError(msg)
    else:
        warnings.warn(msg)


import numpy as np
from astropy import units as u
from astropy import table
from scipy import interpolate as scinterp

class SNANAHostLib():
    """Class for parsing a SNANA HOSTLIB file.
    The file may contain a weight map, and must contain host galaxy data,
    following the standard SNANA HOSTLIB format.
    """
    def __init__(self, filename):
        """Read in a SNANA HOSTLIB file"""
        # find the 'VARNAMES' line, use it to define the start of the hostlib
        # section (as opposed to the wgtmap section)
        nwgtmapstart = -1
        ngaldataheader = -1
        ngaldatastart = -1
        iline = 0
        with open(filename, 'r') as read_obj:
            for line in read_obj:
                if len(line.strip().lstrip('#'))==0:
                    continue
                if line.strip().startswith('NVAR_WGTMAP:'):
                    wgtmaphdrline = line.split()
                    varnames_wgtmap = wgtmaphdrline[3:]
                if line.strip().startswith('WGT:') and nwgtmapstart<0:
                    nwgtmapstart = iline
                if line.strip().startswith('GAL:') and ngaldatastart<0:
                    ngaldatastart = iline
                if line.strip().startswith('VARNAMES:'):
                    ngaldataheader = iline
                iline += 1
        if ngaldataheader < 0:
            raise RuntimeError(r"{filename} is not an SNANA HOSTLIB file")

        if nwgtmapstart >= 0:
            self.wgtmaptable = table.Table.read(
                filename, format='ascii.basic',
                names=['label']+varnames_wgtmap+['wgt','snmagshift'],
                data_start=nwgtmapstart-1,
                data_end=ngaldataheader-2,
                comment='#'
                )
        else:
            self.wgtmaptable = None

        galdatatable = table.Table.read(filename, format='ascii.basic',
                                  header_start=ngaldataheader-1,
                                  data_start=ngaldatastart-1, comment='#'
                                  )
        galdatatable.remove_columns(['VARNAMES:'])
        self.galdatatable = galdatatable
        return

class CatalogBasedRedshiftSim():
    """Class for projecting redshift completeness from an input
    galaxy catalog.
    """

    def __init__(self):
        self.postsurvey = False
        self.galaxies = None

    def read_galaxy_catalog(self, filename):
        """Read in a catalog of galaxy properties

        Parameters
        ----------
        filename : str
          full path to the file containing galaxy properties (e.g. Mass, SFR,
          magnitudes, etc.).  May be a SNANA HOSTLIB file, or any formtat that
          can be auto-parsed by astropy.table.Table.read()
        """
        # TODO: check if it is a hostlib without try/except
        try :
            self.galaxies = table.Table.read(filename)
        except:
            try:
                hostlib = SNANAHostLib(filename)
                self.galaxies = hostlib.galdatatable
            except:
                raise RuntimeError(
                    f"Can't read in {filename}. "
                    "It may not be a valid hostlib or astropy-readable table.")
        return

    def assign_snhost_prob(self, snr_model='AH18S',
                           logmasscolname='logmass',
                           logsfrcolname='logsfr',
                           verbose=True):
        """Add a column to the 'galaxies' catalog that gives the relative
        probability for each galaxy hosting a SN in any given observer-frame
        year.  This is computed based on the predicted SN rate (number of SN
        explosions per observer-frame year) of each galaxy, adopting the
        specified SN rate model.

        Parameters
        ----------
        snr_model : str
           'A+B' : SNR = A*M + B*SFR   (Scannapieco & Bildsten 2005)
           'AH18S' : the smooth logarithmic sSFR model (Andersen & Hjorth 2018)
           'AH18PW' : the piecewise sSFR model (Andersen & Hjorth 2018)

        logmasscolname : str
           name of column in the galaxies Table containing the log10(Mass)

        logsfrcolname : str
           name of column in the galaxies Table containing the
           log10(StarFormationRate)

        verbose : bool
            Set to True to print messages.
        """
        if self.galaxies is None:
            print("No 'galaxies' catalog loaded. Use 'read_galaxy_catalog()'")

        if snr_model.lower()=='a+b':
            # Note: adopting the A and B values from Andersen & Hjorth 2018
            # but dividing by 1e-4 (so the SNR below actually counts the number
            # of SN explodiing per 10000 yrs)
            A = 4.66 * 1e-10
            B = 4.88
            snr = A * 10 ** self.galaxies[logmasscolname] + B * 10 ** self.galaxies[logsfrcolname]
            # divide by the total snr to get relative probabilities
            snr /= np.nanmax(snr)
            snrcolname = 'snr_A+B'
            snrcol = table.Column(data=snr, name='snr_A+B')
        elif snr_model.lower() == 'ah18s':
            logssfr = self.galaxies[logsfrcolname] - self.galaxies[logmasscolname]
            ssnr = ssnr_ah18_smooth(logssfr)
            snr = ssnr * 10 ** self.galaxies[logmasscolname]
            snr /= np.nanmax(snr)
            snrcolname = 'snr_AH18_smooth'
            snrcol = table.Column(data=snr, name=snrcolname)
        elif snr_model.lower() == 'ah18pw':
            logssfr = self.galaxies[logsfrcolname] - self.galaxies[logmasscolname]
            ssnr = ssnr_ah18_piecewise(logssfr)
            snr = ssnr * 10 ** self.galaxies[logmasscolname]
            snr /= np.nanmax(snr)
            snrcolname = 'snr_AH18_piecewise'
        else:
            raise RuntimeError(r"{snr_model} is not a know SN rate model.")

        snrcol = table.Column(data=snr, name=snrcolname)
        if snrcolname in self.galaxies.colnames:
            self.galaxies[snrcolname] = snr
        else:
            self.galaxies.add_column(snrcol)
        if verbose:
            print(f"Added/updated relative SN rate column using {snr_model} model")
        return

    def pick_host_galaxies(self, nsn, snrcolname='snr_AH18_piecewise',
                           replace=False, verbose=True):
        """Do a random draw to assign 'nsn' supernovae to galaxies in the
        galaxies catalog, based on the (pre-defined) relative SN rates.

        TODO: (Alternatively, read in a SNANA output file (.dump file maybe?)
        that has already run a survey simulation and picked host galaxies.)

        Parameters
        ----------
        replace
        nsn : int
          number of SN to assign to host galaxies

        snrcolname : str
           name of the column in the galaxies catalog that gives the relative
           SN rate (or 'weight') for each galaxy.  This may be created by the
           assign_snhost_prob() method.

        replace : bool
           Whether to sample with replacement.  If True, a galaxy may host
           more than one SN. If False, then assign no more than one SN to
           each galaxy (requires nsn<len(galaxies))
        """
        if ~replace and nsn > len(self.galaxies):
            raise RuntimeError(
                r'Picking hosts without replacement, but Nsn > len(galaxies)')

        # Pick SN host galaxies
        galindices = np.arange(len(self.galaxies))
        psnhost = self.galaxies[snrcolname]/np.sum(self.galaxies[snrcolname])
        snindices = np.random.choice(
            galindices, nsn, replace=replace, p=psnhost)

        # Add a boolean 'host' column to the galaxies catalog
        ishost = np.zeros(len(self.galaxies), dtype=bool)
        ishost[snindices] = True
        hostcol = table.Column(name='host', data=ishost)
        if 'host' in self.galaxies.colnames:
            self.galaxies['host'] = hostcol
        else:
            self.galaxies.add_column(hostcol, index=1)

        # TODO: Alternate approach:  read in a SNANA output file (.dump file
        #  maybe?) that has already run a survey simulation and picked hosts.

        if verbose:
            print(f"Assigned {nsn} SNe to hosts using {snrcolname} probabilities.")
        return

    def apply_specz_completeness_map(self, filename,
                                     defining_columns_galcat,
                                     defining_columns_speczmap,
                                     efficiency_columns_speczmap,
                                     fill_value = np.nan
                                     ):
        """Read in a 'map' for spectroscopic redshift completeness, which
        maps from one or more galaxy properties (mag, SFR, z...) onto a
        probability of getting a spec-z.

        Preferred format of the input file is a .ecsv file, but anything
        that astropy.table can read is OK in principle.

        Then apply the specz completeness map to the catalog
        of host galaxy properties (already read in) to define exactly which
        of the galaxies gets a redshift.

        If the method 'pick_host_galaxies' has already been run
        (so the flag postsurvey == True), then only galaxies defined as SN
        hosts are assigned a redshift.

        Parameters
        ----------
        filename : str
           path to astropy-readable file

        defining_columns_galcat : listlike
           list of strings specifying the column names in the galaxy catalog
           (self.galaxies) for parameters that are used to define the specz
           efficiency (e.g. if this is a SFR-based specz map then this may
           be ['logSFR'])

        defining_columns_speczmap : listlike, same length as above
           list of strings specifying the corresponding column names in the
           specz map file (given by 'filename').  Must be the same length as
           defining_columns_galcat, giving corresponding column names in the
           same order.

        efficiency_columns_speczmap : listlike, same length as above
           list of column names giving the specz
           efficiency (or completeness fraction) for each row in the specz
           map file.
        """
        if (len(defining_columns_galcat)!=len(defining_columns_speczmap) or
            len(defining_columns_galcat)!=len(efficiency_columns_speczmap)):
            raise RuntimeError(
                'You must specify the same number of columns from the '
                'galaxy catalog and the specz efficiency catalog.')

        # TODO : make a masked array to remove NaNs ? ?
        speczmap = table.Table.read(filename)

        # TODO : build a separate interpolating function for each of
        #  the possible input parameters ?
        interpolatordict = {}
        for i in range(len(defining_columns_galcat)):
            colname_galcat = defining_columns_galcat[i]
            xobs = self.galaxies[colname_galcat]
            colname_param = defining_columns_speczmap[i]
            x = speczmap[colname_param]
            colname_efficiency = efficiency_columns_speczmap[i]
            y = speczmap[colname_efficiency]
            interpolator = scinterp.interp1d(
                x, y, bounds_error=False, fill_value=fill_value)
            interpolatordict[colname_galcat] = interpolator

        return(interpolatordict)

    def make_photoz_accuracy_map(self):
        """For every galaxy in the catalog of galaxy properties, apply a
        photo-z function that defines the 'measured' photo-z value and
        uncertainty (photoz pdf).  Includes catastrophic outliers.
        """
        pass

    def report_redshift_completeness(self):
        """Produce a report of the overall redshift completeness, accuracy
        and precision, based on multiple spectroscopic 'filters' and the
        random assignment of photo-z values.
        """
        pass

def ssnr_ah18_smooth(logssfr):
    """ Returns the Type Ia specific SN rate per Tyr
    (number of SN Ia exploding per 10^12 yr per solar mass)
    for a galaxy, using the model of Andersen & Hjorth 2018, which is based
    on the specific star formation rate, given as log10(SSFR).
    """
    a = (1.5)*1e-13 # (1.12)*1e-13
    b = 0.5 # 0.73
    k = 0.4 # 0.49
    ssfr0 = 1.7e-10# 1.665e-10
    # logssfr0 = -9.778585762157661    # log10(ssfr0)
    ssfr = np.power(10.,logssfr)
    ssnr = (a + (a/k) * np.log10(ssfr/ssfr0 + b)) * 1e12
    #ssnr = np.max(ssnr, 0.7)
    return(ssnr)

def ssnr_ah18_piecewise(logssfr):
    """ Returns the Type Ia specific SN rate per Tyr
    (number of SN Ia exploding per 10^12 yr per solar mass)
    for a galaxy, using the piecwise linear model
    of Andersen & Hjorth 2018, which is based
    on the specific star formation rate, given as log10(SSFR).
    """
    # Note that the alpha scaling parameter
    # has been multiplied by 1e12 to get units of Tyr-1
    alpha = (1.12)* 1e5
    beta = 0.586
    ssfr2 = 1.01e-11
    ssfr1 = 1.04e-9

    S1 = np.power(ssfr1, beta)
    S2 = np.power(ssfr2, beta)

    if not np.iterable(logssfr):
        logssfr = np.array([logssfr])

    ssfr = np.power(10.,logssfr)

    ilow = np.where(ssfr<=ssfr2)[0]
    imid = np.where((ssfr>ssfr2) & (ssfr<ssfr1))[0]
    ihi = np.where(ssfr>=ssfr1)[0]

    ssnrmid = alpha * np.power(ssfr[imid], beta)

    ssnr = alpha * np.where(ssfr<=ssfr2, S2,
                            np.where(ssfr>=ssfr1, S1,
                                     np.power(ssfr, beta)))
    if len(ssnr)==1:
        ssnr = ssnr[0]
    return(ssnr)