import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import warnings
from typing import Tuple, Optional, Union, List
from .laws import DieterichLaw, RuinaLaw
from .utils import soft_clamp, robust_inv_softplus

# -----------------------------------------------------------------------------
# Phase 1: Numerical Fortification Utils
# -----------------------------------------------------------------------------
# Moved to unique utils.py file

class FrictionLayer(nn.Module):
    def __init__(
        self, 
        mu0: float = 0.6, 
        a: float = 0.015, 
        b: float = 0.02, 
        Dc: float = 0.01, 
        V0: float = 1e-6,
        law: str = 'aging',
        seed: Optional[int] = None
    ):
        super().__init__()
        
        # ---------------------------------------------------------------------
        # Phase 1.1: Enforce Double Precision
        # ---------------------------------------------------------------------
        # We handle this by casting parameters to float64 by default, but allow flexible dtype later.
        # Original strict enforcement removed to allow float32 if desired.

        
        # ---------------------------------------------------------------------
        # Phase 4.2: Deterministic Initialization
        # ---------------------------------------------------------------------
        if seed is not None:
            torch.manual_seed(seed)
            # For specific parameter initialization if random, use generator. 
            # Here params are scalars, but good practice.
        
        self.law_name = law
        
        # Physics Logic
        if law == 'aging':
            self.physics = DieterichLaw()
        elif law == 'slip':
            self.physics = RuinaLaw()
        else:
            raise ValueError(f"Unknown law: {law}")
            
        # Parameters (enforced float64)
        # Using softplus inverse for initialization to ensure positive start if we were training from scratch,
        # but here we just take the values. We will constrain them in property accessors.

        # Store as constraints? Typically we want a, b, Dc > 0.
        # We default to float64 for precision but allow Autocast/Float model behavior if user calls .float()
        # To do this, we just use the python floats for now, they become parameters.
        # If we want default double, we can cast them here.
        self._a = nn.Parameter(torch.tensor(robust_inv_softplus(a), dtype=torch.float64))
        self._b = nn.Parameter(torch.tensor(robust_inv_softplus(b), dtype=torch.float64))
        self._Dc = nn.Parameter(torch.tensor(robust_inv_softplus(Dc), dtype=torch.float64))
        self.mu0 = nn.Parameter(torch.tensor(mu0, dtype=torch.float64))
        
        self.register_buffer('V0', torch.tensor(V0, dtype=torch.float64), persistent=False)
        
    @property
    def a(self): return F.softplus(self._a)
    
    @property
    def b(self): return F.softplus(self._b)
    
    @property
    def Dc(self): return F.softplus(self._Dc)

    def forward(
        self, 
        velocity: torch.Tensor, 
        state: torch.Tensor, 
        dt: Union[float, torch.Tensor]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass for Rate-and-State friction.
        
        Args:
            velocity: Sliding velocity (m/s). Shape: [B] or [1] for batch size B.
            state: State variable theta (seconds). Shape: [B] or [1], must match velocity batch dimension.
            dt: Time step (seconds). Must be positive. All units must be consistent (e.g. SI).
        
        Returns:
            Tuple of (friction, next_state) tensors with same batch dimension as inputs.
        """
        
        # Validate dt
        dt_val = dt.item() if isinstance(dt, torch.Tensor) else dt
        if dt_val <= 0:
            raise ValueError(f"Time step dt must be positive, got {dt_val}")
        
        # ---------------------------------------------------------------------
        # Phase 1.1: Enforce Input Precision
        # ---------------------------------------------------------------------
        # Respect the model's current precision state.
        target_dtype = self.mu0.dtype
        target_device = self.mu0.device
        velocity = velocity.to(dtype=target_dtype, device=target_device)
        state = state.to(dtype=target_dtype, device=target_device)
        
        # Validate batch dimensions match
        if velocity.shape != state.shape:
            raise ValueError(f"Velocity and state must have matching shapes. Got velocity: {velocity.shape}, state: {state.shape}")
        
        # ---------------------------------------------------------------------
        # Phase 4.1: Input Guards
        # ---------------------------------------------------------------------
        if (velocity < 0).any():
             # In physical reality V must be positive for log(V). 
             # We can soft_clamp it, or error. Plan says Raise ValueError.
             raise ValueError("Physics violation: Negative velocity detected.")
        
        # Optional: Warn about negative state (soft_clamp handles it, but explicit is clearer)
        if (state < 0).any():
            warnings.warn(
                "Negative state values detected. These will be clamped to positive values via soft_clamp.",
                stacklevel=2
            )
        
        # ---------------------------------------------------------------------
        # Phase 1.3: Adaptive Stability Checks (CFL Condition)
        # ---------------------------------------------------------------------
        # critical_time = Dc / (velocity.max() + 1e-9) # User suggested
        # Actually for RK4 it's often stricter. User said: 
        # critical_time = self.Dc / (velocity.max() + 1e-9)
        # if dt > critical_time * 0.5: warn...
        
        Dc_val = self.Dc
        
        with torch.no_grad():
            v_max = velocity.max()
            critical_time = Dc_val / (v_max + 1e-9)
            critical_dt = critical_time * 0.5
            dt_val = torch.tensor(dt, dtype=self.mu0.dtype) # Ensure dt is a tensor for comparison
            if dt_val > critical_dt:
                warnings.warn(
                    f"RK4 Instability Risk: dt={dt_val.item()} > {critical_dt.item()}",
                    stacklevel=2
                )

        # ---------------------------------------------------------------------
        # Integration (RK4)
        # ---------------------------------------------------------------------
        # V must be positive for log law in friction, and usually for evolution
        # Use soft_clamp for safety/smoothness as Phase 1.2
        V = soft_clamp(velocity)
        theta = soft_clamp(state)
        
        # Ensure dt is float for physics signature consistency
        dt_float = float(dt) if not isinstance(dt, torch.Tensor) else dt.item()

        # Step
        next_theta = self.rk4_step(0.0, theta, V, dt_float, Dc_val)
        
        # Friction Check
        # mu = mu0 + a ln(V/V0) + b ln(theta * V0 / Dc)
        # Safe logs
        V_safe = soft_clamp(V)
        theta_safe = soft_clamp(next_theta)
        
        term1 = self.a * torch.log(V_safe / self.V0)
        term2 = self.b * torch.log(theta_safe * self.V0 / Dc_val)
        
        friction = self.mu0 + term1 + term2
        
        return friction, next_theta

    def rk4_step(self, t: float, theta: torch.Tensor, v: torch.Tensor, dt: float, Dc: torch.Tensor) -> torch.Tensor:
        """
        Perform a single Runge-Kutta 4 (RK4) integration step for state evolution.
        
        RK4 is a fourth-order numerical integration method that provides good accuracy
        and stability for solving the ODE: dtheta/dt = f(t, theta, v, Dc)
        
        Args:
            t: Current time (not used in rate-and-state laws, but kept for ODE interface).
            theta: Current state variable theta (seconds). Shape: [B] for batch size B.
            v: Current velocity (m/s). Shape: [B], must match theta batch dimension.
            dt: Time step (seconds). Must be positive.
            Dc: Characteristic slip distance (m). Shape: scalar or [B].
        
        Returns:
            Next state variable theta after one RK4 step. Shape: [B], same as input theta.
        
        Algorithm:
            1. k1 = f(t, theta, v, Dc)                    # Evaluate at current point
            2. k2 = f(t + dt/2, theta + k1*dt/2, v, Dc)   # Evaluate at midpoint using k1
            3. k3 = f(t + dt/2, theta + k2*dt/2, v, Dc)   # Evaluate at midpoint using k2
            4. k4 = f(t + dt, theta + k3*dt, v, Dc)       # Evaluate at endpoint using k3
            5. theta_new = theta + (dt/6) * (k1 + 2*k2 + 2*k3 + k4)  # Weighted average
        """
        k1 = self.physics(t, theta, v, Dc)
        
        k2 = self.physics(t + dt/2, theta + k1*dt/2, v, Dc)
        k3 = self.physics(t + dt/2, theta + k2*dt/2, v, Dc)
        k4 = self.physics(t + dt, theta + k3*dt, v, Dc)
        
        return theta + (dt/6.0)*(k1 + 2*k2 + 2*k3 + k4)

    # -------------------------------------------------------------------------
    # Phase 3: Performance & Scale
    # -------------------------------------------------------------------------
    def forward_trajectory(
        self, 
        velocity_sequence: torch.Tensor, 
        dt: Union[float, torch.Tensor],
        initial_state: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Unrolls the solver over the time dimension efficiently.
        
        Args:
            velocity_sequence: Tensor of shape (Batch, Time). Batch-first is enforced.
            dt: Time step (seconds). Must be positive.
            initial_state: Optional starting state (seconds). 
                - If scalar: broadcast to all batch elements
                - If 1D with length 1: broadcast to all batch elements
                - If 1D with length matching batch size: used as-is
                - Otherwise: raises ValueError

        Returns:
            Tuple of (frictions, states) tensors of shape (Batch, Time).

        Note: If initial_state is None, this assumes Steady State Initialization
        (theta = Dc / V_init), putting the system in equilibrium at t=0.
        """
        # Validate dt
        dt_val = dt.item() if isinstance(dt, torch.Tensor) else dt
        if dt_val <= 0:
            raise ValueError(f"Time step dt must be positive, got {dt_val}")
        
        # Type enforcement - Respect layer dtype and device
        target_dtype = self.mu0.dtype
        target_device = self.mu0.device
        velocity_sequence = velocity_sequence.to(dtype=target_dtype, device=target_device)
        if initial_state is not None:
            initial_state = initial_state.to(dtype=target_dtype, device=target_device)
            
        # Forensic Polish: Handle 1D input
        is_1d_input = velocity_sequence.ndim == 1
        if is_1d_input:
            velocity_sequence = velocity_sequence.unsqueeze(0)
            
        # Atomic Polish #2: Handle empty sequence
        steps = velocity_sequence.shape[1]
        if steps == 0:
            if is_1d_input:
                return torch.empty(0, dtype=target_dtype, device=velocity_sequence.device), torch.empty(0, dtype=target_dtype, device=velocity_sequence.device)
            return torch.empty_like(velocity_sequence), torch.empty_like(velocity_sequence)
        
        # Pre-allocate outputs (Zero-Copy Optimization)
        batch_size = velocity_sequence.shape[0]
        device = velocity_sequence.device
        
        # Shape: [Batch, Time]
        frictions = torch.empty((batch_size, steps), dtype=target_dtype, device=device)
        states = torch.empty((batch_size, steps), dtype=target_dtype, device=device)
        
        # Steady state initialization if None
        Dc = self.Dc
        if initial_state is None:
            # v_init = velocity_sequence[:, 0]
            v_init = velocity_sequence[:, 0]
            v_init_safe = soft_clamp(v_init)
            current_state = Dc / v_init_safe
        else:
            # Validate initial_state shape
            batch_size = velocity_sequence.shape[0]
            if initial_state.ndim == 0:
                # Scalar: broadcast to all batch elements
                if batch_size > 1:
                    current_state = initial_state.expand(batch_size)
                else:
                    current_state = initial_state
            elif initial_state.ndim == 1:
                if initial_state.shape[0] == 1:
                    # 1D with length 1: broadcast to all batch elements
                    if batch_size > 1:
                        current_state = initial_state.expand(batch_size)
                    else:
                        current_state = initial_state
                elif initial_state.shape[0] == batch_size:
                    # 1D with length matching batch: use as-is
                    current_state = initial_state
                else:
                    raise ValueError(
                        f"initial_state shape mismatch. Expected scalar, [1], or [{batch_size}], "
                        f"got shape {initial_state.shape}"
                    )
            else:
                raise ValueError(
                    f"initial_state must be scalar or 1D tensor, got shape {initial_state.shape}"
                )
            
        a = self.a
        b = self.b
        mu0 = self.mu0
        V0 = self.V0
        
        dt_float = float(dt) if not isinstance(dt, torch.Tensor) else dt.item()

        # Loop
        for t_idx in range(steps):
            v = velocity_sequence[:, t_idx]
            
            v_safe = soft_clamp(v)
            theta_safe = soft_clamp(current_state)
            
            # RK4
            next_state = self.rk4_step(0.0, theta_safe, v_safe, dt_float, Dc)
            
            # Friction (using next_state as per legacy implementation logic)
            theta_next_safe = soft_clamp(next_state)
            term1 = a * torch.log(v_safe / V0)
            term2 = b * torch.log(theta_next_safe * V0 / Dc)
            f_val = mu0 + term1 + term2
            
            # Direct assignment (No cat/stack overhead)
            frictions[:, t_idx] = f_val
            states[:, t_idx] = next_state
            
            current_state = next_state
            
        
        if is_1d_input:
            frictions = frictions.squeeze(0)
            states = states.squeeze(0)
            
        return frictions, states

    def extra_repr(self) -> str:
        """String representation of physical parameters."""
        return (f"law='{self.law_name}', "
                f"a={self.a.item():.4g}, b={self.b.item():.4g}, "
                f"Dc={self.Dc.item():.4g}, mu0={self.mu0.item():.4g}")

    def freeze_parameters(self, params: List[str]):
        """
        Helper to freeze specific physical parameters for inversion.
        Args:
            params: List of parameter names to freeze (e.g. ['a', 'Dc']).
                    Use public property names ('a', 'b', 'Dc', 'mu0').
                    
        WARNING: Call this BEFORE creating your optimizer. Optimizers typically
        store references to parameters. Freezing afterwards implies the optimizer
        might still track the param but gradients will be None or zero, which works,
        but it's cleaner to freeze before.
        """
        param_map = {
            'a': self._a,
            'b': self._b,
            'Dc': self._Dc,
            'mu0': self.mu0
        }
        
        for name in params:
            if name not in param_map:
                raise ValueError(f"Unknown parameter: {name}. Valid: {list(param_map.keys())}")
            param_map[name].requires_grad = False

# Backwards compatibility / Alias
RateStateLayer = FrictionLayer
