"""
profile.py
----------
Defines ``SoilProfile`` — an ordered, contiguous stack of soil layers that
computes vertical total stress, pore water pressure, and effective stress at
arbitrary depths.

Layers are defined by ``soil`` and ``thickness`` only.  Absolute depth
coordinates (``top``, ``bottom``) are assigned automatically by
``SoilProfile._assign_depths()`` based on stacking order, so callers never
need to track or supply elevation values manually.

Modified and improved by Claude Code.
Optimized with compute_point_at_index to reduce redundant lookups.
"""

from __future__ import annotations

import bisect
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import numpy as np

from .layer import Layer
from .point import Point


class SoilProfile:
    """An ordered, contiguous stack of soil layers for vertical stress computation.

    Layers are supplied with only ``soil`` and ``thickness``; the profile
    assigns each layer's ``top`` and ``bottom`` automatically based on
    stacking order, starting at depth 0.

    :param layers: Ordered from shallowest to deepest.
    :type layers: sequence of Layer
    :param water_table_depth: Depth to the phreatic surface below the ground
        surface. Use a **negative** value for a water table (or standing
        water) *above* the ground surface. Defaults to ``total_depth`` (fully
        drained).
    :type water_table_depth: float, optional
    :param surcharge: External load applied uniformly at the ground surface.
        Default 0.
    :type surcharge: float
    :param gamma_w: Unit weight of water. Default = 62.4.
    :type gamma_w: float
    :raises TypeError: If any element of *layers* is not a ``Layer`` instance.
    :raises ValueError: If *layers* is empty.
    """

    def __init__(
        self,
        layers: Optional[Sequence[Layer]] = None,
        water_table_depth: Optional[float] = None,
        surcharge: float = 0.0,
        gamma_w: float = 62.4,
    ) -> None:
        if layers is None:
            layers = []
        if not layers:
            raise ValueError("SoilProfile requires at least one Layer.")
        if not all(isinstance(layer, Layer) for layer in layers):
            raise TypeError("All elements of 'layers' must be Layer instances.")

        self._layer_tops: List[float] = []
        self._cum_stress: List[float] = []

        self._layers: List[Layer] = list(layers)
        self._assign_depths()

        self.gamma_w: float = gamma_w
        self.surcharge: float = surcharge
        self.water_table_depth: float = (
            water_table_depth if water_table_depth is not None else self.total_depth
        )
        # ensure WT is a layer boundary for correct stress calculations
        self.split_at_water_table()
        if not self._cum_stress:
            self._update_cumulative_stress()

    def _assign_depths(self) -> None:
        """Set each layer's ``top``/``bottom`` from its position in the stack.

        Called after any mutation (construction, add, insert, split) so the
        profile always has consistent absolute depth coordinates.
        """
        current_top = 0.0
        self._layer_tops = []
        for layer in self._layers:
            self._layer_tops.append(current_top)
            layer._set_depths(current_top)
            current_top = layer.bottom

        if hasattr(self, 'surcharge') and hasattr(self, 'water_table_depth'):
            self._update_cumulative_stress()

    def _update_cumulative_stress(self) -> None:
        """Precomputes total vertical stress at the top of each layer."""
        self._cum_stress = []
        sigma_v = self.total_surcharge
        wt = self.water_table_depth

        for layer in self._layers:
            self._cum_stress.append(sigma_v)
            if layer.bottom <= wt:
                sigma_v += layer.soil.unit_weight * layer.thickness
            elif layer.top >= wt:
                sigma_v += layer.soil.unit_weight_sat * layer.thickness
            else:
                # Fallback if layer crosses WT (should not happen if split correctly)
                sigma_v += layer.soil.unit_weight * (wt - layer.top)
                sigma_v += layer.soil.unit_weight_sat * (layer.bottom - wt)

    @property
    def layers(self) -> List[Layer]:
        """All layers in profile order (read-only view).

        :rtype: list of Layer
        """
        return list(self._layers)

    @property
    def num_layers(self) -> int:
        """Number of layers in the profile.

        :rtype: int
        """
        return len(self._layers)

    @property
    def total_depth(self) -> float:
        """Depth to the base of the deepest layer.

        :rtype: float
        """
        return self._layers[-1].bottom

    @property
    def standing_water_head(self) -> float:
        """Height of standing water above ground surface, or 0 if WT is below.

        :rtype: float
        """
        return max(0.0, -self.water_table_depth)

    @property
    def standing_water_pressure(self) -> float:
        """Additional total stress from standing water above the surface.

        :rtype: float
        """
        return self.standing_water_head * self.gamma_w

    @property
    def total_surcharge(self) -> float:
        """Combined surcharge including standing-water pressure.

        :rtype: float
        """
        return self.surcharge + self.standing_water_pressure

    def layer_at(self, depth: float) -> tuple[Layer, int]:
        """Return ``(layer, index)`` for the layer containing *depth*.

        Uses half-open intervals ``(top, bottom]``; a depth exactly on a
        layer interface belongs to the *upper* layer (the one whose bottom
        equals that depth).  The ground surface (depth 0) is a special case
        that belongs to the first layer.

        :param depth: Absolute depth from ground surface.
        :type depth: float
        :returns: A tuple of ``(layer, zero-based index)``.
        :rtype: tuple[Layer, int]
        :raises ValueError: If *depth* is outside ``[0, total_depth]``.
        """
        self._assert_in_profile(depth)
        if depth == 0.0:
            return self._layers[0], 0
            
        # Binary search over layer tops
        idx = bisect.bisect_right(self._layer_tops, depth) - 1
        idx = max(0, min(idx, len(self._layers) - 1))
        
        # Exact interface floating point safety fallback
        if not self._layers[idx].contains(depth): # pragma: no cover
            for i, layer in enumerate(self._layers):
                if layer.contains(depth):
                    return layer, i
            raise AssertionError(f"Depth {depth!r} not contained by any layer.")
            
        return self._layers[idx], idx

    def _assert_in_profile(self, depth: float) -> None:
        if depth < 0 or depth > self.total_depth:
            raise ValueError(
                f"Depth {depth!r} is outside the profile [0, {self.total_depth}]."
            )

    def add_layer(self, layer: Layer) -> "SoilProfile":
        """Append a ``Layer`` to the bottom of the profile.

        Depth coordinates are reassigned automatically after appending.

        :param layer: The layer to append.
        :type layer: Layer
        :returns: ``self``, for chaining.
        :rtype: SoilProfile
        :raises TypeError: If *layer* is not a ``Layer`` instance.
        """
        if not isinstance(layer, Layer):
            raise TypeError(f"Expected a Layer instance, got {type(layer).__name__!r}.")
        self._layers.append(layer)
        self._assign_depths()
        self.split_at_water_table()
        return self

    def insert_layer(self, index: int, layer: Layer) -> "SoilProfile":
        """Insert a ``Layer`` at *index*.

        All layers from *index* onward have their depth coordinates updated
        automatically.

        :param index: Position in the layer list.
        :type index: int
        :param layer: The layer to insert.
        :type layer: Layer
        :returns: ``self``, for chaining.
        :rtype: SoilProfile
        :raises TypeError: If *layer* is not a ``Layer`` instance.
        """
        if not isinstance(layer, Layer):
            raise TypeError(f"Expected a Layer instance, got {type(layer).__name__!r}.")
        self._layers.insert(index, layer)
        self._assign_depths()
        self.split_at_water_table()
        return self

    def split_at_water_table(self) -> "SoilProfile":
        """Physically split the layer that straddles the water table.

        If the water table depth coincides exactly with an existing layer
        boundary, the profile is already correctly divided and the method
        returns ``self`` unchanged.

        If the water table falls *inside* a layer, that layer is replaced by
        two new ``Layer`` objects that share the same ``Soil``:

        * **upper sub-layer** — thickness from ``layer.top`` to
          ``water_table_depth``
        * **lower sub-layer** — thickness from ``water_table_depth`` to
          ``layer.bottom``

        Depth coordinates are reassigned after the split.

        :returns: ``self``, for chaining.
        :rtype: SoilProfile
        """
        wt = self.water_table_depth

        # WT above the surface or below the deepest layer — nothing to split
        if wt <= 0.0 or wt >= self.total_depth:
            return self

        # Find the layer that contains the WT (ignore exact boundary matches within 1e-3 tolerance)
        split_idx: Optional[int] = None
        tolerance = 1e-3
        for idx, layer in enumerate(self._layers):
            if (layer.top + tolerance) < wt < (layer.bottom - tolerance):
                split_idx = idx
                break

        # WT already falls on or extremely close to a layer boundary — no split needed
        if split_idx is None:
            return self

        target = self._layers[split_idx]

        upper = Layer(soil=target.soil, thickness=wt - target.top)
        lower = Layer(soil=target.soil, thickness=target.bottom - wt)

        # Replace the straddling layer with the two sub-layers
        self._layers[split_idx : split_idx + 1] = [upper, lower]
        self._assign_depths()
        return self

    def _total_stress_at(self, depth: float, idx: Optional[int] = None) -> float:
        """Total vertical stress σ_v at *depth*.

        Integrates unit weight from 0 to *depth* efficiently utilizing
        precomputed layer boundary stresses.

        :param depth: Absolute depth from ground surface.
        :param idx: Optional layer index if already known from caller.
        """
        if idx is None:
            layer, idx = self.layer_at(depth)
        else:
            layer = self._layers[idx]
            
        sigma_v = self._cum_stress[idx]
        
        seg_top = layer.top
        seg_bot = depth
        wt = self.water_table_depth

        if layer.bottom <= wt:
            # Entirely above or at WT
            sigma_v += layer.soil.unit_weight * (seg_bot - seg_top)
        elif layer.top >= wt:
            # Entirely below WT
            sigma_v += layer.soil.unit_weight_sat * (seg_bot - seg_top)  # type: ignore[operator]
        else:
            # Safely handles unsplit layers correctly despite guaranteed splits
            if seg_bot <= wt:
                sigma_v += layer.soil.unit_weight * (seg_bot - seg_top)
            else:
                sigma_v += layer.soil.unit_weight * (wt - seg_top)
                sigma_v += layer.soil.unit_weight_sat * (seg_bot - wt)  # type: ignore[operator]

        return sigma_v

    def _pore_pressure_at(self, depth: float) -> float:
        """Hydrostatic pore water pressure u at *depth*.

        Positive below the water table; zero (not negative) above it.

        When the water table is above the ground surface (``water_table_depth``
        < 0), the phreatic head extends down from that elevation, so pore
        pressure at depth z is ``gamma_w × (z − water_table_depth)``.  The
        standing-water total stress (``standing_water_pressure``) is carried
        separately in ``total_surcharge`` — there is no double-counting because
        total stress and pore pressure are both referenced to the same datum.
        """
        if depth <= self.water_table_depth:
            return 0.0
        return self.gamma_w * (depth - self.water_table_depth)

    def compute_point(self, point: Point) -> Point:
        """Compute and populate stresses on *point* in place.

        :param point: The point to compute stresses for.
        :type point: Point
        :returns: The same object with ``total_stress``, ``pore_pressure``,
            ``effective_stress``, and ``layer_index`` set.
        :rtype: Point
        """
        self._assert_in_profile(point.depth)
        _, idx = self.layer_at(point.depth)
        return self._compute_point_at_index(point, idx)

    def _compute_point_at_index(self, point: Point, idx: int) -> Point:
        """Compute stresses for a point with a known layer index.

        Internal optimization used by methods that already step through layers.
        Full vertical and lateral stresses are populated.
        """
        depth = point.depth
        point.total_stress = round(self._total_stress_at(depth, idx=idx), 6)
        point.pore_pressure = round(self._pore_pressure_at(depth), 6)
        point.effective_stress = round(point.total_stress - point.pore_pressure, 6)
        point.layer_index = idx

        # ── Lateral Stresses ──────────────────────────────────────
        soil = self._layers[idx].soil
        phi = getattr(soil, "phi", 0.0) or 0.0
        c = getattr(soil, "c", 0.0) or 0.0

        point.active_earth_pressure = round(
            calc_active_earth_pressure(point.effective_stress, phi, c), 6
        )
        point.passive_earth_pressure = round(
            calc_passive_earth_pressure(point.effective_stress, phi, c), 6
        )

        return point

    def compute_points(self, points: Iterable[Union[Point, float]]) -> List[Point]:
        """Compute stresses for multiple depths or Points.

        Accepts raw floats (auto-wrapped into ``Point``) or existing ``Point``
        objects.  Returns a list sorted by ascending depth.

        :param points: Depths or Point objects to compute.
        :type points: iterable of Point or float
        :returns: Computed points sorted by ascending depth.
        :rtype: list of Point
        """
        result: List[Point] = []
        for p in points:
            if isinstance(p, (int, float)):
                p = Point(depth=float(p))
            result.append(self.compute_point(p))
        result.sort()
        return result


    def vertical_stress_profile(self) -> List[Point]:
        """Return computed :class:`Point` objects at every structural depth boundary.

        Places one point at the ground surface (depth 0), one at the bottom of
        each layer, and therefore one at the base of the profile.  For an
        *n*-layer profile this always returns exactly *n* + 1 points.

        If the water table falls inside a layer, ``split_at_water_table`` will
        have already divided that layer into two, so the water-table depth
        appears as an ordinary layer boundary here without any special handling.

        :returns: Computed points in ascending depth order.
        :rtype: list of Point
        """
        depths = [0.0] + [layer.bottom for layer in self._layers]
        return self.compute_points(depths)



    def _lateral_stress_profile_points(self) -> List[Point]:
        """Return Point objects for every structural depth boundary.

        Produces one point at the ground surface, one or two points at every
        internal layer boundary (one if it's a WT split, two otherwise),
        and one point at the profile base. Points have depth and layer_index set.
        """
        wt = self.water_table_depth
        pts: List[Point] = []

        # ── Surface ───────────────────────────────────────────────
        pts.append(Point(depth=0.0, layer_index=0))

        # ── Internal boundaries ───────────────────────────────────
        for i in range(self.num_layers - 1):
            depth = self._layers[i].bottom
            wt_split = (
                abs(depth - wt) < 1e-9
                and self._layers[i].soil is self._layers[i + 1].soil
            )
            if wt_split:
                # One point; the upper layer i is used
                pts.append(Point(depth=depth, layer_index=i))
            else:
                # Two points: upper layer i followed by lower layer i+1
                pts.append(Point(depth=depth, layer_index=i))
                pts.append(Point(depth=depth, layer_index=i + 1))

        # ── Profile base ──────────────────────────────────────────
        pts.append(Point(depth=self.total_depth, layer_index=self.num_layers - 1))
        return pts

    def lateral_stress_profile(self) -> List[Point]:
        """Compute structural lateral earth pressures across the profile bounds.

        Generates and maps both active and passive Rankine earth pressures at every
        lateral analysis depth boundary.

        :returns: List of ``Point`` objects with all stress fields populated.
        :rtype: list[Point]
        """
        pts = self._lateral_stress_profile_points()
        for pt in pts:
            # pt already has depth and layer_index set
            self._compute_point_at_index(pt, pt.layer_index) # type: ignore
        return pts

    def regular_points(self, l: float, dl: float, at_midpoint: bool = True) -> List[Point]:
        """Return computed :class:`Point` objects sampled within every element.

        Divides the depth range ``[0, l]`` into elements of length *dl*,
        layer by layer.  The last element in each layer is capped at the layer
        boundary, so it may be shorter than *dl*.  A new element then starts
        fresh from the top of the next layer.

        :param l: Total depth to divide.  Must satisfy ``0 < l ≤ total_depth``.
        :type l: float
        :param dl: Target element length.  Must be positive.
        :type dl: float
        :param at_midpoint: If ``True`` (default), return a point at the
            **midpoint** of every element — the conventional integration point
            for beam-on-spring and finite-element analyses.  If ``False``,
            return a point at the **bottom** of every element instead.
        :type at_midpoint: bool
        :returns: Computed points in ascending depth order, one per element.
        :rtype: list of Point
        :raises ValueError: If *dl* ≤ 0 or *l* ≤ 0 or *l* > ``total_depth``.
        """
        if dl <= 0:
            raise ValueError(f"dl must be positive, got {dl!r}.")
        if l <= 0:
            raise ValueError(f"l must be positive, got {l!r}.")
        if l > self.total_depth:
            raise ValueError(
                f"l ({l!r}) exceeds the profile depth ({self.total_depth!r})."
            )

        points: List[Point] = []

        for idx, layer in enumerate(self._layers):
            if layer.top >= l:
                break
            layer_end = min(layer.bottom, l)
            cursor = layer.top
            while cursor < layer_end:
                elem_end = min(cursor + dl, layer_end)
                depth = cursor + (elem_end - cursor) / 2.0 if at_midpoint else elem_end
                # Optimized: bypass binary search since we know the layer index
                pt = self._compute_point_at_index(Point(depth=depth), idx)
                points.append(pt)
                cursor = elem_end

        return points


    def stresses_at_depths(
        self, depths: Union[float, Iterable[float]]
    ) -> Union[Point, List[Point]]:
        """Return Point or list of Points with both vertical and lateral stresses.

        Handles both single depth values and multiple depths.

        :param depths: A single depth (float) or a sequence of depths.
        :returns: A single Point or a list of Points, depending on input.
        :rtype: Point or list[Point]
        """
        is_single = isinstance(depths, (int, float))
        depth_list = [float(depths)] if is_single else list(depths)

        results: List[Point] = []
        for d in depth_list:
            pt = Point(depth=d)
            self.compute_point(pt)
            results.append(pt)

        return results[0] if is_single else results


    def layer_params(self, point: Point, keys: Sequence[str]) -> dict:
        """Return a dict of soil attributes from the layer containing *point*.

        Validates that keys are public and that every layer in the profile
        has all requested attributes set.
        """
        if point.layer_index is None:
            raise ValueError("point.layer_index is None — run compute_point() first.")

        # ── Validate keys ──────────────────────────────────────────
        for key in keys:
            if key.startswith("_"):
                raise AttributeError(
                    f"Access to private attribute {key!r} is not permitted."
                )

        # ── Check every layer has every requested attribute set ────
        missing = [
            (layer.soil.name, key)
            for layer in self._layers
            for key in keys
            if getattr(layer.soil, key, None) is None
        ]
        if missing:
            details = ", ".join(f"'{soil}'.{attr}" for soil, attr in missing)
            raise AttributeError(
                f"The following soil attributes are missing or unset: {details}"
            )

        soil = self._layers[point.layer_index].soil
        result = {}
        for key in keys:
            # We already know it exists from the check above
            result[key] = getattr(soil, key)
        return result

    def __len__(self) -> int:
        return self.num_layers

    def __getitem__(self, index: int) -> Layer:
        return self._layers[index]

    def __iter__(self):
        return iter(self._layers)

    def __repr__(self) -> str:
        return (
            f"SoilProfile({self.num_layers} layers, "
            f"depth={self.total_depth}, "
            f"WT={self.water_table_depth}, "
            f"q={self.surcharge})"
        )


# ── Internal Earth Pressure Helpers ────────────────────────────────


def _rankine_k(phi: Union[float, np.ndarray], active: bool) -> Union[float, np.ndarray]:
    """Rankine earth-pressure coefficient for active or passive conditions.

    Ka = tan²(45 − φ/2),  Kp = tan²(45 + φ/2).
    """
    angle = 45.0 - phi / 2.0 if active else 45.0 + phi / 2.0
    return np.tan(np.radians(angle)) ** 2


def _calc_earth_pressure(
    sigma: Union[float, np.ndarray],
    phi: Union[float, np.ndarray],
    c: Union[float, np.ndarray],
    active: bool,
) -> Union[float, np.ndarray]:
    """Rankine lateral earth pressure (Mohr-Coulomb c–φ soil)."""
    k = _rankine_k(phi, active)
    sign = -1.0 if active else 1.0
    return sigma * k + sign * 2.0 * c * np.sqrt(k)


def calc_active_earth_pressure(
    sigma: Union[float, np.ndarray],
    phi: Union[float, np.ndarray],
    c: Union[float, np.ndarray],
) -> Union[float, np.ndarray]:
    """Active lateral earth pressure: σ_a = σ · Ka − 2c√Ka."""
    return _calc_earth_pressure(sigma, phi, c, active=True)


def calc_passive_earth_pressure(
    sigma: Union[float, np.ndarray],
    phi: Union[float, np.ndarray],
    c: Union[float, np.ndarray],
) -> Union[float, np.ndarray]:
    """Passive lateral earth pressure: σ_p = σ · Kp + 2c√Kp."""
    return _calc_earth_pressure(sigma, phi, c, active=False)
