import numpy as np
from scipy.integrate import cumulative_trapezoid

from . import interp

class _Echo:
    z        = None
    t_min    = None
    tau_min  = None
    zeta_min = None

    t    = None
    tau  = None
    zeta = None

    t_interp    = None
    tau_interp  = None
    zeta_interp = None

    def __init__(self, el, z, t, zeta1d, alpha1d, idx_pos_prev, idx_pos_next):
        dt = t[1] - t[0]
        interp_stencil = 4

        self.z = z

        # interpolate ζ between two neighboring positive values
        # to find the true minimum
        t_interp    = np.linspace(t[idx_pos_prev], t[idx_pos_next], 1024)
        zeta_interp = interp.interp1d(t[0], dt, zeta1d, t_interp, interp_stencil)

        idx_min         = np.argmin(zeta_interp)
        self.t_min      = t_interp[idx_min]
        self.zeta_min   = zeta_interp[idx_min]

        # integrate the proper time along the z=const wordline corresponding to
        # the minimum of ζ
        tau  = cumulative_trapezoid(alpha1d, t, initial = 0.0)
        self.tau_min = interp.interp1d(t[0], dt, tau, np.array([self.t_min]), interp_stencil)

        self.t    = t
        self.tau  = tau
        self.zeta = zeta1d

        self.t_interp       = np.linspace(t[0] + 16 * dt, t[-1] - 16 * dt, 4096)
        self.tau_interp     = interp.interp1d(t[0], dt, tau, self.t_interp, interp_stencil)
        self.zeta_interp    = interp.interp1d(t[0], dt, zeta1d, self.t_interp, interp_stencil)

    def __repr__(self):
        return 'Echo(t=%g,z=%g,minζ=%g)' % (self.t_min, self.z, self.zeta_min)

class EchoLoc:

    t    = None
    z    = None
    zeta = None

    echoes = None

    def __init__(self, t, z, zeta, alpha):
        self.t    = t
        self.z    = z
        self.zeta = zeta

        remainder = zeta
        echoes    = []
        while remainder.shape[0] > 0:
            # find the global minimum in the remaining part of the data
            idx_min = np.unravel_index(np.argmin(remainder), remainder.shape)

            # out of negative minima, finish
            if remainder[idx_min] >= 0.0:
                break

            # find the index where zeta along this z=const. line becomes positive again
            idx_pos = np.argwhere(remainder[:idx_min[0], idx_min[1]] >= 0.0)
            if idx_pos.shape[0] == 0:
                break

            # find the index where zeta becomes positive _after_ the minimum
            # if there is no such index, skip this as a false positive
            idx_pos_next = np.argwhere(remainder[idx_min[0]:, idx_min[1]] >= 0.0)
            if idx_pos_next.shape[0] > 0:
                echoes.append(_Echo(self, z[idx_min[1]], t[:remainder.shape[0]],
                                    remainder[:, idx_min[1]], alpha[:remainder.shape[0], idx_min[1]],
                                    idx_pos[-1, 0], idx_min[0] + idx_pos_next[0, 0]))

            remainder = remainder[:idx_pos[-1, 0]]

        self.echoes = echoes

    def __repr__(self):
        return '%d echoes: t=[%s]' % (len(self.echoes), ','.join((str(e.t_min) for e in self.echoes)))