import torch
import unittest
import numpy as np
from torch_friction import FrictionLayer, RateStateLayer

class TestEdgeCases(unittest.TestCase):
    def setUp(self):
        torch.set_default_dtype(torch.float64)
        
    def tearDown(self):
        torch.set_default_dtype(torch.float32)

    def test_zero_velocity(self):
        """Test that zero velocity is handled gracefully by soft_clamp."""
        layer = FrictionLayer(law='aging')
        velocity = torch.tensor([0.0])
        state = torch.tensor([0.1])
        dt = 0.1
        
        friction, next_state = layer(velocity, state, dt)
        
        # Should not produce NaN or Inf
        self.assertTrue(torch.isfinite(friction))
        self.assertTrue(torch.isfinite(next_state))
        # Zero velocity should be clamped to minimum value
        self.assertTrue(friction.item() > 0)

    def test_very_small_dt(self):
        """Test that very small time steps work correctly."""
        layer = FrictionLayer(law='aging')
        velocity = torch.tensor([1e-6])
        state = torch.tensor([0.1])
        dt = 1e-10  # Very small time step
        
        friction, next_state = layer(velocity, state, dt)
        
        # Should produce finite results
        self.assertTrue(torch.isfinite(friction))
        self.assertTrue(torch.isfinite(next_state))
        # State should change very little with tiny dt
        self.assertAlmostEqual(state.item(), next_state.item(), places=5)

    def test_negative_dt_raises_error(self):
        """Test that negative dt raises ValueError."""
        layer = FrictionLayer(law='aging')
        velocity = torch.tensor([1e-6])
        state = torch.tensor([0.1])
        dt = -0.1
        
        with self.assertRaises(ValueError):
            layer(velocity, state, dt)

    def test_zero_dt_raises_error(self):
        """Test that zero dt raises ValueError."""
        layer = FrictionLayer(law='aging')
        velocity = torch.tensor([1e-6])
        state = torch.tensor([0.1])
        dt = 0.0
        
        with self.assertRaises(ValueError):
            layer(velocity, state, dt)

    def test_empty_sequence(self):
        """Test that empty velocity sequence is handled."""
        layer = FrictionLayer(law='aging')
        velocity_sequence = torch.empty((1, 0), dtype=torch.float64)
        dt = 0.1
        
        frictions, states = layer.forward_trajectory(velocity_sequence, dt)
        
        # Should return empty tensors with correct shape
        self.assertEqual(frictions.shape, (1, 0))
        self.assertEqual(states.shape, (1, 0))

    def test_1d_input_trajectory(self):
        """Test that 1D input to forward_trajectory works."""
        layer = FrictionLayer(law='aging')
        velocity_sequence = torch.ones(100, dtype=torch.float64) * 1e-6
        dt = 0.1
        
        frictions, states = layer.forward_trajectory(velocity_sequence, dt)
        
        # Should return 1D tensors
        self.assertEqual(frictions.ndim, 1)
        self.assertEqual(states.ndim, 1)
        self.assertEqual(frictions.shape[0], 100)
        self.assertEqual(states.shape[0], 100)

    def test_device_placement(self):
        """Test that inputs are automatically moved to model device."""
        layer = FrictionLayer(law='aging')
        
        # Model is on CPU by default
        velocity = torch.tensor([1e-6])
        state = torch.tensor([0.1])
        dt = 0.1
        
        # Should work without explicit device placement
        friction, next_state = layer(velocity, state, dt)
        
        # Results should be on same device as model
        self.assertEqual(friction.device, layer.mu0.device)
        self.assertEqual(next_state.device, layer.mu0.device)

    def test_ruina_law_comprehensive(self):
        """Test RuinaLaw with velocity step to verify behavior."""
        layer = FrictionLayer(law='slip', a=0.015, b=0.02, Dc=0.01)
        
        # Create velocity step
        steps = 1000
        velocities = torch.ones(1, steps, dtype=torch.float64) * 1e-6
        velocities[0, 200:] = 1e-5
        dt = 0.1
        
        frictions, states = layer.forward_trajectory(velocities, dt)
        
        # Should produce finite results
        self.assertTrue(torch.all(torch.isfinite(frictions)))
        self.assertTrue(torch.all(torch.isfinite(states)))
        
        # Friction should change at jump point
        mu_pre = frictions[0, 199].item()
        mu_post = frictions[0, 200].item()
        self.assertNotAlmostEqual(mu_pre, mu_post, places=6)

    def test_negative_state_handled(self):
        """Test that negative state values are clamped."""
        layer = FrictionLayer(law='aging')
        velocity = torch.tensor([1e-6])
        state = torch.tensor([-0.1])  # Negative state
        dt = 0.1
        
        friction, next_state = layer(velocity, state, dt)
        
        # Should produce finite results (soft_clamp handles it)
        self.assertTrue(torch.isfinite(friction))
        self.assertTrue(torch.isfinite(next_state))
        # Next state should be positive (clamped)
        self.assertTrue(next_state.item() > 0)

    def test_batch_dimension_mismatch_raises_error(self):
        """Test that mismatched batch dimensions raise error."""
        layer = FrictionLayer(law='aging')
        velocity = torch.tensor([1e-6, 2e-6])  # Batch size 2
        state = torch.tensor([0.1])  # Batch size 1
        dt = 0.1
        
        with self.assertRaises(ValueError):
            layer(velocity, state, dt)

    def test_initial_state_shape_validation(self):
        """Test initial_state shape validation in forward_trajectory."""
        layer = FrictionLayer(law='aging')
        velocity_sequence = torch.ones((3, 100), dtype=torch.float64) * 1e-6
        dt = 0.1
        
        # Valid: scalar
        initial_state = torch.tensor(0.1)
        frictions, states = layer.forward_trajectory(velocity_sequence, dt, initial_state)
        self.assertEqual(frictions.shape, (3, 100))
        
        # Valid: [1] - should broadcast
        initial_state = torch.tensor([0.1])
        frictions, states = layer.forward_trajectory(velocity_sequence, dt, initial_state)
        self.assertEqual(frictions.shape, (3, 100))
        
        # Valid: [batch_size]
        initial_state = torch.tensor([0.1, 0.2, 0.3])
        frictions, states = layer.forward_trajectory(velocity_sequence, dt, initial_state)
        self.assertEqual(frictions.shape, (3, 100))
        
        # Invalid: wrong length
        initial_state = torch.tensor([0.1, 0.2])  # Length 2, batch is 3
        with self.assertRaises(ValueError):
            layer.forward_trajectory(velocity_sequence, dt, initial_state)
        
        # Invalid: 2D tensor
        initial_state = torch.tensor([[0.1]])
        with self.assertRaises(ValueError):
            layer.forward_trajectory(velocity_sequence, dt, initial_state)

    def test_large_batch(self):
        """Test that large batches work correctly."""
        layer = FrictionLayer(law='aging')
        batch_size = 100
        time_steps = 1000
        velocity_sequence = torch.ones((batch_size, time_steps), dtype=torch.float64) * 1e-6
        dt = 0.1
        
        frictions, states = layer.forward_trajectory(velocity_sequence, dt)
        
        self.assertEqual(frictions.shape, (batch_size, time_steps))
        self.assertEqual(states.shape, (batch_size, time_steps))
        self.assertTrue(torch.all(torch.isfinite(frictions)))
        self.assertTrue(torch.all(torch.isfinite(states)))

if __name__ == '__main__':
    unittest.main()
