import torch
import torch.nn as nn
import torch.nn.functional as F

from .utils import soft_clamp

class RateStateLaw(nn.Module):
    """Base class for Rate-and-State friction laws."""
    def __init__(self):
        super().__init__()

    def forward(self, _t: float, theta: torch.Tensor, velocity: torch.Tensor, Dc: torch.Tensor) -> torch.Tensor:
        """Calculate d_theta / dt"""
        raise NotImplementedError

class DieterichLaw(RateStateLaw):
    """
    Dieterich (Aging) Law:
    d_theta / dt = 1 - (V * theta / Dc)
    """
    def forward(self, _t: float, theta: torch.Tensor, velocity: torch.Tensor, Dc: torch.Tensor) -> torch.Tensor:
        # d_theta / dt = 1 - (V * theta / Dc)
        return 1.0 - (velocity * theta / (Dc + 1e-12))

class RuinaLaw(RateStateLaw):
    """
    Ruina (Slip) Law:
    d_theta / dt = - (V * theta / Dc) * ln(V * theta / Dc)
    """
    def forward(self, _t: float, theta: torch.Tensor, velocity: torch.Tensor, Dc: torch.Tensor) -> torch.Tensor:
        epsilon = 1e-12
        term = (velocity * theta / Dc) + epsilon
        # Safety: Soft clamp the term to avoid log(negative) if used standalone
        # Also, relies on inputs being clamped > 0 to avoid 0*log(0) NaN.
        term = soft_clamp(term) 
        return -1.0 * term * torch.log(term)
