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

class TestRateState(unittest.TestCase):
    def setUp(self):
        # Use Double precision for verification as enforced by the library
        torch.set_default_dtype(torch.float64)
        self.dt = 0.1
        self.steps = 1000
        self.jump_step = 200
        
    def tearDown(self):
        torch.set_default_dtype(torch.float32)

    def test_velocity_step_aging(self):
        """Verify standard Aging Law response to velocity step."""
        layer = RateStateLayer(law='aging')
        
        # Create Data
        velocities = torch.ones(1, self.steps, dtype=torch.float64) * 1e-6
        velocities[0, self.jump_step:] = 1e-5
        # The layer enforces proper types now
        
        # Use new trajectory method
        # Signature: velocity_sequence, dt, initial_state=None -> (frictions, states)
        frictions, _ = layer.forward_trajectory(velocities, self.dt)
        frictions = frictions.detach().numpy()[0] # Batch 0
        
        # There is no state return in the new forward_trajectory, only friction?
        # Checked my implementation: return torch.stack(frictions, dim=1)
        # Ah, the user plan showed `states.append(mu)`... wait.
        # User Plan:
        # states = []
        # for t...: mu, next_state = self.step(...); states.append(mu)
        # return torch.stack(states, dim=1)
        # So "states" list actually held "mu"s?
        # And the method returns just a tensor of frictions.
        # Previous test check: `frictions, states = layer.forward_trajectory...`
        # I should probably return (frictions, states) to be useful for analysis.
        # But my current implementation of `forward_trajectory` returns ONLY `torch.stack(frictions, dim=1)`.
        # I will update the test to expect just frictions for now, OR update the layer to return both.
        # Given "Analysis" usually needs state, I should probably return Tuple.
        # Use simple return for now to match the code I wrote.
        
        # 1. Direct Effect Check
        # Calculate theoretical jump on fixed state
        # We don't have the state history easily if I don't return it.
        # However, I can check friction values.
        
        # Direct effect: Instantaneous change in friction when V jumps.
        # Friction at jump_step should be higher than jump_step-1
        
        # V changes at `jump_step`.
        # Friction depends on V(t) and Theta(t).
        # Theta(t) is evolved from previous step.
        # At t=jump_step, V is high. Theta is still "old" (from previous V low).
        # So Friction should jump UP.
        
        mu_pre = frictions[self.jump_step-1]
        mu_post = frictions[self.jump_step]
        
        actual_jump = mu_post - mu_pre
        
        # Theoretical jump: a * ln(V_new/V_old)
        V_old = 1e-6
        V_new = 1e-5
        expected_jump = layer.a.item() * np.log(V_new / V_old)
        
        self.assertTrue(np.isclose(actual_jump, expected_jump, rtol=0.1),
                        f"Direct effect mismatch. Expected {expected_jump}, Got {actual_jump}")

    def test_slip_law_existence(self):
        """Ensure Slip law runs without error."""
        layer = RateStateLayer(law='slip')
        v = torch.tensor([1e-6], dtype=torch.float64)
        theta = torch.tensor([0.1], dtype=torch.float64)
        f, next_theta = layer(v, theta, 0.1)
        self.assertTrue(torch.isfinite(f))
        self.assertTrue(torch.isfinite(next_theta))

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