"""
Velocity model management and travel-time table generation (converted from vel_subs.f90).
"""

import numpy as np

from .params import ERAD


def flatten(z_s, vel_s, erad=ERAD):
    """Flat earth transformation.

    Parameters
    ----------
    z_s : float
        Spherical depth (km).
    vel_s : float
        Spherical velocity (km/s).
    erad : float
        Earth radius (km).

    Returns
    -------
    z_f : float
        Flat earth depth.
    vel_f : float
        Flat earth velocity.
    """
    r = erad - z_s
    z_f = -erad * np.log(r / erad)
    vel_f = vel_s * (erad / r)
    return z_f, vel_f


def unflatten(z_f, vel_f, erad=ERAD):
    """Inverse of flatten.

    Parameters
    ----------
    z_f : float
        Flat earth depth.
    vel_f : float
        Flat earth velocity.
    erad : float
        Earth radius (km).

    Returns
    -------
    z_s : float
        Spherical depth.
    vel_s : float
        Spherical velocity.
    """
    r = erad * np.exp(-z_f / erad)
    z_s = erad - r
    vel_s = vel_f * (r / erad)
    return z_s, vel_s


def layertrace(p, h, utop, ubot, imth=3):
    """Calculate travel time and range offset for ray tracing through a single layer.

    Parameters
    ----------
    p : float
        Horizontal slowness (ray parameter).
    h : float
        Layer thickness (flat earth).
    utop : float
        Slowness at top of layer.
    ubot : float
        Slowness at bottom of layer.
    imth : int
        Interpolation method:
        1 = v(z) = 1/sqrt(a - 2*b*z) (fastest)
        2 = v(z) = a - b*z (linear gradient)
        3 = v(z) = a*exp(-b*z) (preferred for flat earth)

    Returns
    -------
    dx : float
        Range offset (km).
    dt : float
        Travel time (s).
    irtr : int
        Return code:
        -1 = zero thickness layer
         0 = ray turned above layer
         1 = ray passed through layer
         2 = ray turned within layer
    """
    # Check zero thickness layer
    if h == 0.0:
        return 0.0, 0.0, -1

    u = utop

    # Check for ray turning above layer
    if u - p <= 0.0:
        return 0.0, 0.0, 0

    # Vertical slowness at top
    eta2 = (u - p) * (u + p)
    eta = np.sqrt(eta2)

    # Special function at top of layer
    if imth == 2:
        y = u + eta
        if p != 0.0:
            y = y / p
        qr = np.log(y)
    elif imth == 3:
        qr = np.arctan2(eta, p)

    # b factor
    if imth == 1:
        b = -(utop**2 - ubot**2) / (2.0 * h)
    elif imth == 2:
        vtop = 1.0 / utop
        vbot = 1.0 / ubot
        b = -(vtop - vbot) / h
    else:
        b = -np.log(ubot / utop) / h

    # Constant velocity layer
    if b == 0.0:
        b = 1.0 / h
        etau = eta
        ex = p / eta
        irtr = 1

        dx = ex / b
        dtau = etau / b
        dt = dtau + p * dx
        return float(dx), float(dt), irtr

    # Ray tracing integral at upper limit
    if imth == 1:
        etau = -eta2 * eta / 3.0
        ex = -eta * p
    elif imth == 2:
        ex = eta / u
        etau = qr - ex
        if p != 0.0:
            ex = ex / p
    else:
        etau = eta - p * qr
        ex = qr

    # Check lower limit for turning point
    u = ubot
    if u <= p:
        irtr = 2

        dx = ex / b
        dtau = etau / b
        dt = dtau + p * dx
        return float(dx), float(dt), irtr

    # No turning point: ray passes through
    irtr = 1
    eta2 = (u - p) * (u + p)
    eta = np.sqrt(eta2)

    if imth == 1:
        etau = etau + eta2 * eta / 3.0
        ex = ex + eta * p
    elif imth == 2:
        y = u + eta
        z = eta / u
        etau = etau + z
        if p != 0.0:
            y = y / p
            z = z / p
        qr = np.log(y)
        etau = etau - qr
        ex = ex - z
    else:
        qr = np.arctan2(eta, p)
        etau = etau - eta + p * qr
        ex = ex - qr

    dx = ex / b
    dtau = etau / b
    dt = dtau + p * dx
    return float(dx), float(dt), irtr


def vzfillin(infile, outfile, sfact, dz, zmin_in, zmax_in):
    """Read (z, Vp, Vs) velocity model and resample at a finer depth interval.

    Parameters
    ----------
    infile : str
        Input velocity model file.
    outfile : str
        Output interpolated model file.
    sfact : float
        Vp/Vs ratio if Vs is unlisted (0.0).
    dz : float
        Interpolation depth interval.
    zmin_in : float
        Minimum depth (overwritten on output).
    zmax_in : float
        Maximum depth (used to extend model).

    Returns
    -------
    zmin : float
        Minimum depth in input model.
    zmax : float
        Maximum depth in input model.
    """
    # Read input file
    d_list, vp_list, vs_list = [], [], []
    with open(infile, "r") as f:
        for line in f:
            parts = line.split()
            if len(parts) < 3:
                continue
            depth = float(parts[0])
            vp = float(parts[1])
            vs = float(parts[2])
            if vs == 0.0:
                vs = vp / sfact
            d_list.append(depth)
            vp_list.append(vp)
            vs_list.append(vs)

    d = np.array(d_list)
    vp = np.array(vp_list)
    vs = np.array(vs_list)
    npts = len(d)

    # Extend to zmax if needed
    if d[-1] < zmax_in:
        d = np.append(d, zmax_in)
        vp = np.append(vp, vp[-1])
        vs = np.append(vs, vs[-1])
        npts = len(d)

    zmin = d[0]
    zmax = d[-1]

    # Interpolate
    nz = int(np.floor(d[-1] / dz)) + 1
    buf = []
    for iz in range(nz):
        z = iz * dz
        for i in range(npts - 1):
            z1 = d[i]
            z2 = d[i + 1]
            if z1 <= z and z2 >= z:
                if z1 == z2:
                    fact = 0.0
                else:
                    fact = (z - z1) / (z2 - z1)
                v0_p = vp[i] + fact * (vp[i + 1] - vp[i])
                v0_s = vs[i] + fact * (vs[i + 1] - vs[i])
                buf.append((z, v0_p, v0_s))

                # Add extra points at layer boundaries
                for k in range(i + 1, npts):
                    if d[k] - z < dz:
                        buf.append((d[k], vp[k], vs[k]))
                break

    # Remove duplicates
    unique_buf = [buf[0]]
    for k in range(1, len(buf)):
        if buf[k] != buf[k - 1]:
            unique_buf.append(buf[k])

    # Write output file
    with open(outfile, "w") as f:
        for z, v1, v2 in unique_buf:
            f.write(f"{z:10.3f}{v1:8.4f}{v2:8.4f}\n")

    return zmin, zmax


def deptable(vmodel, iw, plongcut, dep1, dep2, dep3, del1, del2, del3, ideprad, ttfile=None):
    """Create tables of seismic travel times as a function of distance and source depth.

    Parameters
    ----------
    vmodel : str
        Filename of interpolated velocity model.
    iw : int
        Phase index: 1 = P, 2 = S.
    plongcut : float
        Minimum ray parameter.
    dep1, dep2, dep3 : float
        Min depth, max depth, depth spacing.
    del1, del2, del3 : float
        Min distance, max distance, distance spacing.
    ideprad : int
        Velocity model format: 1=depth, 2=radius.
    ttfile : str, optional
        Output travel-time table file. If None, no file written.

    Returns
    -------
    outTT : ndarray, shape (nx, nz)
        Travel time table for this phase.
    outX : ndarray, shape (nx,)
        Distance array.
    outZ : ndarray, shape (nz,)
        Depth array.
    """
    pi = np.pi
    erad = ERAD
    ecircum = 2.0 * pi * erad
    kmdeg = ecircum / 360.0

    nray = 10000
    itype = 1  # del in km, t in sec

    # Read velocity model
    z_s, alpha_s, beta_s = [], [], []
    with open(vmodel, "r") as f:
        for line in f:
            parts = line.split()
            if len(parts) < 3:
                continue
            zz = float(parts[0])
            aa = float(parts[1])
            bb = float(parts[2])
            if ideprad == 2:
                zz = erad - zz
            if zz == erad:
                break
            z_s.append(zz)
            alpha_s.append(aa)
            beta_s.append(bb)

    # Set up dummy interface at bottom
    z_s.append(z_s[-1])
    alpha_s.append(alpha_s[-1])
    beta_s.append(beta_s[-1])

    npts = len(z_s)
    z_s = np.array(z_s)
    alpha_s = np.array(alpha_s)
    beta_s = np.array(beta_s)

    # Flat earth transform
    z = np.zeros(npts)
    alpha = np.zeros(npts)
    beta = np.zeros(npts)
    for i in range(npts):
        z[i], alpha[i] = flatten(z_s[i], alpha_s[i], erad)
        _, beta[i] = flatten(z_s[i], beta_s[i], erad)

    # Compute slowness
    slow = np.zeros((npts, 2))
    for i in range(npts):
        slow[i, 0] = 1.0 / alpha[i]
        if beta[i] != 0.0:
            slow[i, 1] = 1.0 / beta[i]
        else:
            slow[i, 1] = 1.0 / alpha[i]

    # iw is 1 or 2, but python is 0-indexed
    iw_idx = iw - 1

    # Depth table
    ndep = int(np.floor((dep2 + dep3 / 10.0 - dep1) / dep3)) + 1
    qdeptab = dep1 + np.arange(ndep) * dep3

    # Ray parameters
    pmin = 0.0
    pmax = slow[0, iw_idx]
    pstep = (pmax - pmin) / float(nray)

    # Arrays for ray tracing
    del2W = np.full(nray + 1, -999.0)
    tt2W = np.full(nray + 1, -999.0)
    ptab = np.zeros(nray + 1)
    qdepxcor = np.full((nray + 1, ndep), -999.0)
    qdeptcor = np.full((nray + 1, ndep), -999.0)

    # Initialize surface source depth
    for idep in range(ndep):
        if qdeptab[idep] == 0.0:
            qdepxcor[:, idep] = 0.0
            qdeptcor[:, idep] = 0.0

    # Ray tracing loop
    imth = 3  # preferred for flat-earth
    dep0_mask = qdeptab == 0.0  # precompute mask for zero-depth entries

    for ip in range(nray):
        p = pmin + ip * pstep
        ptab[ip] = p
        x = 0.0
        t = 0.0

        # Initialize for this ray (vectorized)
        qdepxcor[ip, dep0_mask] = 0.0
        qdeptcor[ip, dep0_mask] = 0.0
        qdepxcor[ip, ~dep0_mask] = -999.0
        qdeptcor[ip, ~dep0_mask] = -999.0

        exceeded_zmax = False
        for i in range(npts - 1):
            if z_s[i] >= 9999.0:
                del2W[ip] = -999.0
                tt2W[ip] = -999.0
                exceeded_zmax = True
                break

            h = z[i + 1] - z[i]
            if h == 0.0:
                continue

            dx_ray, dt_ray, irtr = layertrace(p, h, slow[i, iw_idx], slow[i + 1, iw_idx], imth)
            x += dx_ray
            t += dt_ray

            if irtr == 0 or irtr == 2:
                break

            # Save x, t for matching source depths (vectorized)
            match_mask = np.abs(z_s[i + 1] - qdeptab) < dep3 / 1.99
            qdepxcor[ip, match_mask] = x
            qdeptcor[ip, match_mask] = t

        if not exceeded_zmax:
            del2W[ip] = 2 * x
            tt2W[ip] = 2 * t

    # Special point for p = pmax (vertical ray at surface)
    ip = nray
    ptab[ip] = slow[0, iw_idx]
    del2W[ip] = 0.0
    tt2W[ip] = 0.0
    for idep in range(ndep):
        if qdeptab[idep] == 0.0:
            qdepxcor[ip, idep] = 0.0
            qdeptcor[ip, idep] = 0.0
        else:
            qdepxcor[ip, idep] = -999.0
            qdeptcor[ip, idep] = -999.0

    total_nray = nray + 1

    # Compute T(X,Z) for first arriving rays
    ndel = int(np.floor((del2 + del3 / 10.0 - del1) / del3)) + 1
    sdeltab = del1 + np.arange(ndel) * del3
    tt = np.full((ndel, ndep), 99999.0)

    for idep in range(ndep):
        xsave = []
        tsave = []
        psave = []

        if qdeptab[idep] == 0.0:
            i2 = total_nray - 1  # skip upgoing loop, start downgoing from last ray
        else:
            # Upgoing rays
            xold = -999.0
            i2 = 0
            for i in range(total_nray):
                x2 = qdepxcor[i, idep]
                if x2 == -999.0:
                    i2 = i - 1
                    break
                if x2 <= xold:
                    i2 = i - 1
                    break
                t2 = qdeptcor[i, idep]
                xsave.append(x2)
                tsave.append(t2)
                psave.append(-ptab[i])
                xold = x2
            else:
                i2 = total_nray - 1

        # Downgoing rays
        for i in range(i2, -1, -1):
            if qdepxcor[i, idep] == -999.0:
                continue
            if del2W[i] == -999.0:
                continue
            x2 = del2W[i] - qdepxcor[i, idep]
            t2 = tt2W[i] - qdeptcor[i, idep]
            xsave.append(x2)
            tsave.append(t2)
            psave.append(ptab[i])

        ncount = len(xsave)
        if ncount < 2:
            continue

        xsave = np.array(xsave)
        tsave = np.array(tsave)
        psave = np.array(psave)

        for idel in range(ndel):
            deldel = sdeltab[idel]
            delo = deldel
            if itype == 2:
                delo = deldel * kmdeg

            for i in range(1, ncount):
                x1 = xsave[i - 1]
                x2 = xsave[i]

                if x1 > delo or x2 < delo:
                    continue

                # Skip downgoing refracted rays
                if psave[i] > 0.0 and psave[i] < plongcut:
                    continue

                frac = (delo - x1) / (x2 - x1) if x2 != x1 else 0.0
                tbest = tsave[i - 1] + frac * (tsave[i] - tsave[i - 1])

                if tbest < tt[idel, idep]:
                    tt[idel, idep] = tbest

            if tt[idel, idep] == 99999.0:
                tt[idel, idep] = 0.0
            if itype == 2:
                tt[idel, idep] = tt[idel, idep] / 60.0

    # Fix edge case
    if ndel > 0 and ndep > 0:
        if sdeltab[0] == 0.0 and qdeptab[0] == 0.0:
            tt[0, 0] = 0.0

    # Build output arrays
    outX = sdeltab[:ndel].copy()
    outZ = qdeptab[:ndep].copy()
    outTT = tt[:ndel, :ndep].copy()

    # Write output file if requested
    if ttfile is not None:
        with open(ttfile, "w") as f:
            header1 = f"From deptable, file= {vmodel[:20]:20s} iw ={iw:2d} pmin={pmin:8.5f} pmax={pmax:8.5f} nray={total_nray:6d}\n"
            f.write(header1)
            f.write(f"{ndel:5d}{ndep:5d}\n")
            # Depth header row
            f.write(" " * 8 + "".join(f"{qdeptab[j]:9.1f}" for j in range(ndep)) + "\n")
            # Data rows
            for i in range(ndel):
                f.write(f"{sdeltab[i]:9.4f}" + "".join(f"{tt[i, j]:9.4f}" for j in range(ndep)) + "\n")

    return outTT, outX, outZ
