r"""
Velocity model representation and layer geometry for 1-D layered media.

This module provides a :class:`LayerStack` data structure that encapsulates
the sequence of layers traversed by a ray between a source and receiver,
and a helper :func:`build_layer_stack` to extract this structure from a
velocity model (DataFrame or :class:`ModelArrays`).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Union

import numpy as np
import pandas as pd


@dataclass
class LayerStack:
    r"""Layers traversed by a ray between source and receiver depths.

    Parameters
    ----------
    h : numpy.ndarray
        Layer thicknesses (m), shape ``(N,)``.  Ordered from the
        shallowest traversed depth to the deepest.
    vp : numpy.ndarray
        P-wave velocities (m/s), shape ``(N,)``.
    vs : numpy.ndarray
        S-wave velocities (m/s), shape ``(N,)``.
    rho : numpy.ndarray or None
        Densities (kg/m³), shape ``(N,)``, or *None*.
    qp : numpy.ndarray or None
        P-wave quality factors, shape ``(N,)``, or *None*.
    qs : numpy.ndarray or None
        S-wave quality factors, shape ``(N,)``, or *None*.
    """

    h: np.ndarray
    vp: np.ndarray
    vs: np.ndarray
    rho: np.ndarray | None = None
    qp: np.ndarray | None = None
    qs: np.ndarray | None = None

    @property
    def n_layers(self) -> int:
        """Number of layers in the stack."""
        return len(self.h)

    def v(self, vel_type: str = "Vp") -> np.ndarray:
        """Return the velocity array for the requested wave type.

        Parameters
        ----------
        vel_type : str
            ``'Vp'`` or ``'Vs'``.
        """
        if vel_type.lower() in ("vp", "p"):
            return self.vp
        elif vel_type.lower() in ("vs", "s"):
            return self.vs
        raise ValueError(f"vel_type must be 'Vp' or 'Vs', got '{vel_type}'")

    def q_factor(self, vel_type: str = "Vp") -> np.ndarray | None:
        """Return the Q-factor array for the requested wave type."""
        if vel_type.lower() in ("vp", "p"):
            return self.qp
        return self.qs


@dataclass
class ModelArrays:
    """Pre-extracted numpy arrays from a velocity DataFrame.

    Use :meth:`from_dataframe` to construct once and pass to
    :func:`build_layer_stack` to avoid repeated DataFrame
    column access when tracing many rays.
    """

    depths: np.ndarray
    vp: np.ndarray
    vs: np.ndarray
    rho: np.ndarray | None = None
    qp: np.ndarray | None = None
    qs: np.ndarray | None = None

    @classmethod
    def from_dataframe(cls, vel_df: pd.DataFrame) -> "ModelArrays":
        """Extract arrays once from *vel_df*."""
        return cls(
            depths=vel_df["Depth"].values.astype(np.float64),
            vp=vel_df["Vp"].values.astype(np.float64),
            vs=vel_df["Vs"].values.astype(np.float64),
            rho=vel_df["Rho"].values.astype(np.float64) if "Rho" in vel_df.columns else None,
            qp=vel_df["Qp"].values.astype(np.float64) if "Qp" in vel_df.columns else None,
            qs=vel_df["Qs"].values.astype(np.float64) if "Qs" in vel_df.columns else None,
        )


def _layer_index(depth: float, boundaries: np.ndarray) -> int:
    """Return the index of the layer containing *depth*.

    Layer *k* spans from ``boundaries[k]`` (inclusive) to
    ``boundaries[k+1]`` (exclusive).  The last layer extends to
    infinity.

    Parameters
    ----------
    depth : float
        Query depth (positive downward).
    boundaries : numpy.ndarray
        Sorted layer-top depths from the velocity model.

    Returns
    -------
    int
        Layer index (0-based).
    """
    idx = int(np.searchsorted(boundaries, depth, side="right")) - 1
    return max(idx, 0)


def _extract_arrays(
    vel_model: Union[pd.DataFrame, ModelArrays],
) -> tuple[np.ndarray, np.ndarray, np.ndarray,
           np.ndarray | None, np.ndarray | None, np.ndarray | None]:
    """Return ``(depths, vp, vs, rho, qp, qs)`` from either input type."""
    if isinstance(vel_model, ModelArrays):
        return (vel_model.depths, vel_model.vp, vel_model.vs,
                vel_model.rho, vel_model.qp, vel_model.qs)
    # pd.DataFrame path
    depths = vel_model["Depth"].values.astype(np.float64)
    vp = vel_model["Vp"].values.astype(np.float64)
    vs = vel_model["Vs"].values.astype(np.float64)
    rho = vel_model["Rho"].values.astype(np.float64) if "Rho" in vel_model.columns else None
    qp = vel_model["Qp"].values.astype(np.float64) if "Qp" in vel_model.columns else None
    qs = vel_model["Qs"].values.astype(np.float64) if "Qs" in vel_model.columns else None
    return depths, vp, vs, rho, qp, qs


def build_layer_stack(
    vel_model: Union[pd.DataFrame, ModelArrays],
    z_src: float,
    z_rcv: float,
) -> LayerStack:
    r"""Extract the layer stack between source and receiver depths.

    The returned :class:`LayerStack` contains the layers traversed by a
    ray connecting *z_src* and *z_rcv*, with partial thicknesses at the
    source and receiver layers.  Layers are always ordered from the
    shallowest point to the deepest, regardless of which endpoint is the
    source.

    Parameters
    ----------
    vel_model : pandas.DataFrame or ModelArrays
        Velocity model.  A DataFrame must have columns ``Depth``,
        ``Vp``, ``Vs`` (optional: ``Rho``, ``Qp``, ``Qs``).
        A :class:`ModelArrays` instance avoids repeated column
        extraction when tracing many rays.
    z_src : float
        Source depth (positive downward, m).
    z_rcv : float
        Receiver depth (positive downward, m).

    Returns
    -------
    LayerStack
    """
    depths, vp_all, vs_all, rho_all, qp_all, qs_all = _extract_arrays(vel_model)

    has_rho = rho_all is not None
    has_qp = qp_all is not None
    has_qs = qs_all is not None

    z_top = min(z_src, z_rcv)
    z_bot = max(z_src, z_rcv)

    i_top = _layer_index(z_top, depths)
    i_bot = _layer_index(z_bot, depths)

    # ── Single-layer case ──
    if i_top == i_bot:
        h_arr = np.array([z_bot - z_top])
        slc = slice(i_top, i_top + 1)
        return LayerStack(
            h=h_arr,
            vp=vp_all[slc].copy(),
            vs=vs_all[slc].copy(),
            rho=rho_all[slc].copy() if has_rho else None,
            qp=qp_all[slc].copy() if has_qp else None,
            qs=qs_all[slc].copy() if has_qs else None,
        )

    # ── Multi-layer case ──
    n = i_bot - i_top + 1
    h_arr = np.empty(n)

    # First (shallowest) layer — partial
    if i_top + 1 < len(depths):
        h_arr[0] = depths[i_top + 1] - z_top
    else:
        h_arr[0] = z_bot - z_top  # only layer

    # Middle layers — full thickness
    for k in range(1, n - 1):
        idx = i_top + k
        if idx + 1 < len(depths):
            h_arr[k] = depths[idx + 1] - depths[idx]
        else:
            h_arr[k] = z_bot - depths[idx]

    # Last (deepest) layer — partial
    h_arr[n - 1] = z_bot - depths[i_bot]

    slc = slice(i_top, i_bot + 1)
    return LayerStack(
        h=h_arr,
        vp=vp_all[slc].copy(),
        vs=vs_all[slc].copy(),
        rho=rho_all[slc].copy() if has_rho else None,
        qp=qp_all[slc].copy() if has_qp else None,
        qs=qs_all[slc].copy() if has_qs else None,
    )


# Backward-compatible alias
build_layer_stack_fast = build_layer_stack


def discretize_gradient_layer(
    z_top: float, 
    z_bot: float, 
    v_func: Callable[[float], float], 
    dz: float = 2.0,
) -> pd.DataFrame:
    """
    Discretize a gradient layer into thin constant-velocity layers.

    Parameters
    ----------
    z_top : float
        Top depth of the gradient layer.
    z_bot : float
        Bottom depth of the gradient layer.
    v_func : Callable[[float], float]
        Function that returns the velocity at a given depth.
    dz : float, optional
        Thickness of each discretized layer, by default 2.0.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns ``Depth``, ``Vp``, ``Vs``, ``Rho``.
    """
    depths = np.arange(z_top, z_bot, dz)
    z_mid = depths + dz / 2.0
    v_vals = v_func(z_mid)
    
    df = pd.DataFrame({
        "Depth": depths,
        "Vp": v_vals,
        "Vs": v_vals / 1.732,
        "Rho": 2500.0
    })
    return df