"""
Relocation subroutines: DIFLOC, DIFCLUST, GET_TT_FAST, FIX_ORIGIN_TIME
(converted from growclust.f90 subroutines, vectorized with NumPy).
"""

import numpy as np

from .stats import robomean2


DEGRAD = 180.0 / np.pi
DEGKM = 111.19493


def get_tt_fast(ttab, xtab, dtab, nx, nd, ip, delo, qdep8):
    """Obtain travel time by interpolating from travel-time table (scalar version).

    Parameters
    ----------
    ttab : ndarray, shape (nx, nd, 2)
        Travel time table for P (index 0) and S (index 1).
    xtab : ndarray, shape (nx,)
        Distance table.
    dtab : ndarray, shape (nd,)
        Depth table.
    nx, nd : int
        Number of distance/depth points.
    ip : int
        Phase: 1=P, 2=S.
    delo : float
        Source-station distance (km).
    qdep8 : float
        Source depth (km).

    Returns
    -------
    tt : float
        Travel time (s).
    iflag : int
        -1 if outside table bounds, 0 for interpolation, 1 for extrapolation.
    """
    ip_idx = ip - 1
    dx = xtab[1] - xtab[0]
    dd = dtab[1] - dtab[0]

    if qdep8 < 0.0:
        qdep = 0.0
        dt01 = ttab[0, 1, ip_idx] - ttab[0, 0, ip_idx]
        velsurf = dd / dt01 if dt01 != 0 else 1.0
        tdep = abs(qdep8) / velsurf
    else:
        qdep = qdep8
        tdep = 0.0

    if qdep < dtab[0] or qdep > dtab[nd - 1] or delo < xtab[0] or delo > xtab[nx - 1]:
        return 999.0, -1

    id1 = min(int((qdep - dtab[0]) / dd), nd - 2)
    id2 = id1 + 1
    ix1 = min(int((delo - xtab[0]) / dx), nx - 2)
    ix2 = ix1 + 1

    t11 = ttab[ix1, id1, ip_idx]
    t12 = ttab[ix1, id2, ip_idx]
    t21 = ttab[ix2, id1, ip_idx]
    t22 = ttab[ix2, id2, ip_idx]

    if t11 != 0.0 and t12 != 0.0 and t21 != 0.0 and t22 != 0.0:
        xfrac = (delo - xtab[ix1]) / dx
        t1 = t11 + xfrac * (t21 - t11)
        t2 = t12 + xfrac * (t22 - t12)
        dfrac = (qdep - dtab[id1]) / dd
        return t1 + dfrac * (t2 - t1) + tdep, 0

    # Extrapolation fallback (rare)
    xmid = (xtab[:-1] + xtab[1:]) / 2.0
    xoff_all = np.abs(xmid - delo)

    valid1 = (ttab[:-1, id1, ip_idx] != 0) & (ttab[1:, id2, ip_idx] != 0)
    valid2 = (ttab[:-1, id2, ip_idx] != 0) & (ttab[1:, id2, ip_idx] != 0)

    if not np.any(valid1) or not np.any(valid2):
        return 999.0, -1

    xoff1 = np.where(valid1, xoff_all, 999.0)
    xoff2 = np.where(valid2, xoff_all, 999.0)
    ixbest1 = int(np.argmin(xoff1)) + 1
    ixbest2 = int(np.argmin(xoff2)) + 1

    xfrac1 = (delo - xtab[ixbest1 - 1]) / dx
    tt1 = ttab[ixbest1 - 1, id1, ip_idx] + xfrac1 * (ttab[ixbest1, id1, ip_idx] - ttab[ixbest1 - 1, id1, ip_idx])

    xfrac2 = (delo - xtab[ixbest2 - 1]) / dx
    tt2 = ttab[ixbest2 - 1, id2, ip_idx] + xfrac2 * (ttab[ixbest2, id2, ip_idx] - ttab[ixbest2 - 1, id2, ip_idx])

    dfrac = (qdep - dtab[id1]) / dd
    return tt1 + dfrac * (tt2 - tt1) + tdep, 1


def get_tt_batch(ttab, xtab, dtab, nx, nd, ip_arr, delo_arr, qdep_arr):
    """Vectorized batch travel-time lookup for arrays of distances and depths.

    Parameters
    ----------
    ttab : ndarray, shape (nx, nd, 2)
    xtab : ndarray, shape (nx,)
    dtab : ndarray, shape (nd,)
    nx, nd : int
    ip_arr : ndarray of int, shape (n,)
        Phase indices (1=P, 2=S).
    delo_arr : ndarray, shape (n,)
        Source-station distances (km).
    qdep_arr : ndarray, shape (n,)
        Source depths (km).

    Returns
    -------
    tt_out : ndarray, shape (n,)
        Travel times (s). 999.0 for out-of-bounds.
    """
    n = len(delo_arr)
    dx = xtab[1] - xtab[0]
    dd = dtab[1] - dtab[0]
    ip_idx = ip_arr.astype(np.intp) - 1  # 0-based

    # Handle negative depths
    qdep = np.where(qdep_arr < 0.0, 0.0, qdep_arr)
    neg_mask = qdep_arr < 0.0
    tdep = np.zeros(n)
    if np.any(neg_mask):
        # Compute surface velocity for each phase
        for phase in (0, 1):
            pmask = neg_mask & (ip_idx == phase)
            if np.any(pmask):
                dt01 = ttab[0, 1, phase] - ttab[0, 0, phase]
                velsurf = dd / dt01 if dt01 != 0 else 1.0
                tdep[pmask] = np.abs(qdep_arr[pmask]) / velsurf

    # Range check
    valid = (qdep >= dtab[0]) & (qdep <= dtab[nd - 1]) & (delo_arr >= xtab[0]) & (delo_arr <= xtab[nx - 1])

    tt_out = np.full(n, 999.0)

    if not np.any(valid):
        return tt_out

    # Work only on valid entries
    v_delo = delo_arr[valid]
    v_qdep = qdep[valid]
    v_ip = ip_idx[valid]
    v_tdep = tdep[valid]

    # Compute interpolation indices
    id1 = np.clip(((v_qdep - dtab[0]) / dd).astype(np.intp), 0, nd - 2)
    id2 = id1 + 1
    ix1 = np.clip(((v_delo - xtab[0]) / dx).astype(np.intp), 0, nx - 2)
    ix2 = ix1 + 1

    # Gather corner values
    t11 = ttab[ix1, id1, v_ip]
    t12 = ttab[ix1, id2, v_ip]
    t21 = ttab[ix2, id1, v_ip]
    t22 = ttab[ix2, id2, v_ip]

    # Bilinear interpolation (only where all corners are nonzero)
    corners_ok = (t11 != 0) & (t12 != 0) & (t21 != 0) & (t22 != 0)

    xfrac = (v_delo - xtab[ix1]) / dx
    dfrac = (v_qdep - dtab[id1]) / dd

    t1 = t11 + xfrac * (t21 - t11)
    t2 = t12 + xfrac * (t22 - t12)
    tt_interp = t1 + dfrac * (t2 - t1) + v_tdep

    # For entries where corners are ok, use interpolated values
    result = np.where(corners_ok, tt_interp, 999.0)

    # Extrapolation fallback for entries where corners are not ok
    extrap_mask = ~corners_ok
    if np.any(extrap_mask):
        extrap_indices = np.where(extrap_mask)[0]
        for idx in extrap_indices:
            tt_val, iflag = get_tt_fast(ttab, xtab, dtab, nx, nd,
                                        int(v_ip[idx]) + 1, float(v_delo[idx]),
                                        float(v_qdep[idx]))
            result[idx] = tt_val + v_tdep[idx] if iflag >= 0 else 999.0

    tt_out[valid] = result
    return tt_out


def _compute_fit(resid, inorm):
    """Compute fit metric and origin time estimate from residuals."""
    if inorm == 1:
        residmed = float(np.median(resid))
        fit = float(np.sum(np.abs(resid - residmed)))
    elif inorm == 2:
        residmed = float(np.mean(resid))
        fit = float(np.sum((resid - residmed) ** 2))
    elif inorm == 3:
        residmed, fit = robomean2(resid, 0.1, 10)
    else:
        raise ValueError(f"Undefined inorm: {inorm}")
    return fit, residmed


def difloc(qlat0, qlon0, qdep0, npick, tt_obs, ip, slat, slon,
           ttab, xtab, dtab, nx, nd, boxwid, nit, inorm):
    """Differential location for an event pair relative to its centroid.

    Uses iterative "shrinking-box" grid search with vectorized travel-time lookups.
    """
    cosqlat = np.cos(qlat0 / DEGRAD)

    dlat0 = 0.0
    dlon0 = 0.0
    ddep0 = 0.0
    dlat = 0.5 * boxwid / DEGKM
    dlon = dlat / cosqlat

    zboxwid = boxwid
    if qdep0 < 0:
        zboxwid = min(1.0, boxwid)
    elif qdep0 < boxwid / 2.0:
        zboxwid = max(1.0, 2.0 * qdep0)
    ddep = 0.5 * zboxwid

    torgdif = 0.0

    # Pre-build the 27 grid offsets: (iy, ix, iz) for -1,0,1
    grid = np.array(np.meshgrid([-1, 0, 1], [-1, 0, 1], [-1, 0, 1])).T.reshape(-1, 3)
    # grid[:, 0] = iy, grid[:, 1] = ix, grid[:, 2] = iz

    # Ensure arrays are numpy
    tt_obs = np.asarray(tt_obs, dtype=np.float64)
    ip_arr = np.asarray(ip, dtype=np.int32)
    slat_arr = np.asarray(slat, dtype=np.float64)
    slon_arr = np.asarray(slon, dtype=np.float64)

    for it in range(nit):
        # Build 27 trial location pairs
        flat1 = qlat0 + dlat0 + dlat * grid[:, 0]  # (27,)
        flat2 = qlat0 - dlat0 - dlat * grid[:, 0]
        flon1 = qlon0 + dlon0 + dlon * grid[:, 1]
        flon2 = qlon0 - dlon0 - dlon * grid[:, 1]
        fdep1 = qdep0 + ddep0 + ddep * grid[:, 2]
        fdep2 = qdep0 - ddep0 - ddep * grid[:, 2]

        # For each grid point, compute distances to all stations: shape (27, npick)
        dy1 = flat1[:, np.newaxis] - slat_arr[np.newaxis, :]  # (27, npick)
        dx1 = (flon1[:, np.newaxis] - slon_arr[np.newaxis, :]) * cosqlat
        delkm1 = np.sqrt(dx1**2 + dy1**2) * DEGKM

        dy2 = flat2[:, np.newaxis] - slat_arr[np.newaxis, :]
        dx2 = (flon2[:, np.newaxis] - slon_arr[np.newaxis, :]) * cosqlat
        delkm2 = np.sqrt(dx2**2 + dy2**2) * DEGKM

        # Flatten for batch TT lookup: (27*npick,)
        ngrid = 27
        delo1_flat = delkm1.ravel()
        delo2_flat = delkm2.ravel()
        dep1_flat = np.repeat(fdep1, npick)
        dep2_flat = np.repeat(fdep2, npick)
        ip_flat = np.tile(ip_arr, ngrid)

        tsec1_flat = get_tt_batch(ttab, xtab, dtab, nx, nd, ip_flat, delo1_flat, dep1_flat)
        tsec2_flat = get_tt_batch(ttab, xtab, dtab, nx, nd, ip_flat, delo2_flat, dep2_flat)

        # Reshape: (27, npick)
        tsec1 = tsec1_flat.reshape(ngrid, npick)
        tsec2 = tsec2_flat.reshape(ngrid, npick)

        # Residuals for all 27 grid points
        resid_all = tt_obs[np.newaxis, :] - (tsec2 - tsec1)  # (27, npick)

        # Compute fit for each grid point
        fitbest = 9.0e20
        best_idx = 0
        tbest = 0.0
        for g in range(ngrid):
            fit, residmed = _compute_fit(resid_all[g], inorm)
            if fit < fitbest:
                fitbest = fit
                best_idx = g
                tbest = residmed

        dlat0 = flat1[best_idx] - qlat0
        dlon0 = flon1[best_idx] - qlon0
        ddep0 = fdep1[best_idx] - qdep0
        torgdif = tbest

        dlat *= 0.67
        dlon *= 0.67
        ddep *= 0.67

    qlat1 = qlat0 + dlat0
    qlon1 = qlon0 + dlon0
    qdep1 = qdep0 + ddep0
    qlat2 = qlat0 - dlat0
    qlon2 = qlon0 - dlon0
    qdep2 = qdep0 - ddep0
    resol = (dlat / 0.67) * DEGKM

    # Final misfits (vectorized)
    dy1 = qlat1 - slat_arr
    dx1 = (qlon1 - slon_arr) * cosqlat
    delkm1 = np.sqrt(dx1**2 + dy1**2) * DEGKM
    tsec1_final = get_tt_batch(ttab, xtab, dtab, nx, nd, ip_arr, delkm1, np.full(npick, qdep1))

    dy2 = qlat2 - slat_arr
    dx2 = (qlon2 - slon_arr) * cosqlat
    delkm2 = np.sqrt(dx2**2 + dy2**2) * DEGKM
    tsec2_final = get_tt_batch(ttab, xtab, dtab, nx, nd, ip_arr, delkm2, np.full(npick, qdep2))

    resid = tt_obs - (tsec2_final - tsec1_final) - torgdif
    rms = float(np.sqrt(np.mean(resid**2)))
    rmed = float(np.median(np.abs(resid)))

    return qlat1, qlon1, qdep1, qlat2, qlon2, qdep2, torgdif, resid, rms, rmed, resol


def difclust(qlat0, qlon0, qdep0, npick, tt_obs, ip, slat, slon,
             qlat1_off, qlon1_off, qdep1_off, qorg1_off,
             qlat2_off, qlon2_off, qdep2_off, qorg2_off,
             ttab, xtab, dtab, nx, nd, boxwid, nit, inorm):
    """Relative relocation of two clusters (vectorized).

    Analogous to DIFLOC but for clusters with individual event offsets.
    """
    cosqlat = np.cos(qlat0 / DEGRAD)

    dlat0 = 0.0
    dlon0 = 0.0
    ddep0 = 0.0
    dlat = 0.5 * boxwid / DEGKM
    dlon = dlat / cosqlat

    zboxwid = boxwid
    if qdep0 < 0:
        zboxwid = min(1.0, boxwid)
    elif qdep0 < boxwid / 2.0:
        zboxwid = max(1.0, 2.0 * qdep0)
    ddep = 0.5 * zboxwid

    torgdif = 0.0

    # 27 grid offsets
    grid = np.array(np.meshgrid([-1, 0, 1], [-1, 0, 1], [-1, 0, 1])).T.reshape(-1, 3)
    ngrid = 27

    tt_obs = np.asarray(tt_obs, dtype=np.float64)
    ip_arr = np.asarray(ip, dtype=np.int32)
    slat_arr = np.asarray(slat, dtype=np.float64)
    slon_arr = np.asarray(slon, dtype=np.float64)
    qlat1_off = np.asarray(qlat1_off, dtype=np.float64)
    qlon1_off = np.asarray(qlon1_off, dtype=np.float64)
    qdep1_off = np.asarray(qdep1_off, dtype=np.float32)
    qorg1_off = np.asarray(qorg1_off, dtype=np.float32)
    qlat2_off = np.asarray(qlat2_off, dtype=np.float64)
    qlon2_off = np.asarray(qlon2_off, dtype=np.float64)
    qdep2_off = np.asarray(qdep2_off, dtype=np.float32)
    qorg2_off = np.asarray(qorg2_off, dtype=np.float32)

    for it in range(nit):
        flat1 = qlat0 + dlat0 + dlat * grid[:, 0]
        flat2 = qlat0 - dlat0 - dlat * grid[:, 0]
        flon1 = qlon0 + dlon0 + dlon * grid[:, 1]
        flon2 = qlon0 - dlon0 - dlon * grid[:, 1]
        fdep1 = qdep0 + ddep0 + ddep * grid[:, 2]
        fdep2 = qdep0 - ddep0 - ddep * grid[:, 2]

        # Event 1: station distances (27, npick)
        dy1 = flat1[:, np.newaxis] + qlat1_off[np.newaxis, :] - slat_arr[np.newaxis, :]
        dx1 = (flon1[:, np.newaxis] + qlon1_off[np.newaxis, :] - slon_arr[np.newaxis, :]) * cosqlat
        delkm1 = np.sqrt(dx1**2 + dy1**2) * DEGKM

        # Event 2: station distances
        dy2 = flat2[:, np.newaxis] + qlat2_off[np.newaxis, :] - slat_arr[np.newaxis, :]
        dx2 = (flon2[:, np.newaxis] + qlon2_off[np.newaxis, :] - slon_arr[np.newaxis, :]) * cosqlat
        delkm2 = np.sqrt(dx2**2 + dy2**2) * DEGKM

        # Depths including offsets (27, npick)
        dep1_all = fdep1[:, np.newaxis] + qdep1_off[np.newaxis, :]
        dep2_all = fdep2[:, np.newaxis] + qdep2_off[np.newaxis, :]

        # Flatten for batch lookup
        ip_flat = np.tile(ip_arr, ngrid)
        tsec1_flat = get_tt_batch(ttab, xtab, dtab, nx, nd, ip_flat, delkm1.ravel(), dep1_all.ravel())
        tsec2_flat = get_tt_batch(ttab, xtab, dtab, nx, nd, ip_flat, delkm2.ravel(), dep2_all.ravel())

        tsec1 = tsec1_flat.reshape(ngrid, npick)
        tsec2 = tsec2_flat.reshape(ngrid, npick)

        # Differential time with origin time offsets
        tdif_pred = tsec2 + qorg2_off[np.newaxis, :] - (tsec1 + qorg1_off[np.newaxis, :])
        resid_all = tt_obs[np.newaxis, :] - tdif_pred

        fitbest = 9.0e20
        best_idx = 0
        tbest = 0.0
        for g in range(ngrid):
            fit, residmed = _compute_fit(resid_all[g], inorm)
            if fit < fitbest:
                fitbest = fit
                best_idx = g
                tbest = residmed

        dlat0 = flat1[best_idx] - qlat0
        dlon0 = flon1[best_idx] - qlon0
        ddep0 = fdep1[best_idx] - qdep0
        torgdif = tbest

        dlat *= 0.67
        dlon *= 0.67
        ddep *= 0.67

    clat1 = qlat0 + dlat0
    clon1 = qlon0 + dlon0
    cdep1 = qdep0 + ddep0
    clat2 = qlat0 - dlat0
    clon2 = qlon0 - dlon0
    cdep2 = qdep0 - ddep0
    resol = (dlat / 0.67) * DEGKM

    # Cluster distance
    dy = clat2 - clat1
    dx = (clon2 - clon1) * cosqlat
    delkm = np.sqrt(dx**2 + dy**2) * DEGKM
    cdist = np.sqrt(delkm**2 + (cdep2 - cdep1) ** 2)

    # Final misfits (vectorized)
    dy1 = clat1 + qlat1_off - slat_arr
    dx1 = (clon1 + qlon1_off - slon_arr) * cosqlat
    delkm1 = np.sqrt(dx1**2 + dy1**2) * DEGKM
    dep1_final = np.full(npick, cdep1) + qdep1_off
    tsec1_final = get_tt_batch(ttab, xtab, dtab, nx, nd, ip_arr, delkm1, dep1_final)

    dy2 = clat2 + qlat2_off - slat_arr
    dx2 = (clon2 + qlon2_off - slon_arr) * cosqlat
    delkm2 = np.sqrt(dx2**2 + dy2**2) * DEGKM
    dep2_final = np.full(npick, cdep2) + qdep2_off
    tsec2_final = get_tt_batch(ttab, xtab, dtab, nx, nd, ip_arr, delkm2, dep2_final)

    tdif_pred = tsec2_final + qorg2_off - (tsec1_final + qorg1_off)
    resid = tt_obs - tdif_pred - torgdif
    rms = float(np.sqrt(np.mean(resid**2)))
    rmed = float(np.median(np.abs(resid)))

    return clat1, clon1, cdep1, clat2, clon2, cdep2, cdist, torgdif, resid, rms, rmed, resol


def fix_origin_time(qyr, qmon, qday, qhr, qmin, qsec):
    """Fix origin time when seconds field is negative or >= 60."""
    if qsec < 0.0:
        qsec += 60.0
        qmin -= 1
    elif qsec >= 60.0:
        qsec -= 60.0
        qmin += 1
    else:
        return qyr, qmon, qday, qhr, qmin, qsec

    if qmin < 0:
        qmin += 60
        qhr -= 1
    elif qmin >= 60:
        qmin -= 60
        qhr += 1
    else:
        return qyr, qmon, qday, qhr, qmin, qsec

    if qhr < 0:
        qhr += 24
        qday -= 1
    elif qhr >= 24:
        qhr -= 24
        qday += 1
    else:
        return qyr, qmon, qday, qhr, qmin, qsec

    if 0 < qday < 29:
        return qyr, qmon, qday, qhr, qmin, qsec

    if qday < 1:
        qmon -= 1
        if qmon < 1:
            qmon = 12
            qyr -= 1

    if qyr % 400 == 0:
        ileap = 1
    elif qyr % 4 == 0 and qyr % 100 != 0:
        ileap = 1
    else:
        ileap = 0

    maxday = 31
    if qmon == 2:
        maxday = 28 + ileap
    elif qmon in (4, 6, 9, 11):
        maxday = 30

    if qday < 1:
        qday = maxday
    if qday > maxday:
        qday = 1
        qmon += 1
        if qmon > 12:
            qmon = 1
            qyr += 1

    return qyr, qmon, qday, qhr, qmin, qsec
