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

class TestVerification(unittest.TestCase):
    def setUp(self):
        torch.set_default_dtype(torch.float64)
        
    def test_mms_dieterich(self):
        """
        Method of Manufactured Solutions for Dieterich Law.
        Equation: d_theta/dt = 1 - V * theta / Dc
        Assume solution: theta(t) = 0.1 + 0.05 * sin(t)
        Then d_theta/dt = 0.05 * cos(t)
        
        So we need a source term S(t) added to the equation? 
        Or force V(t) such that the equation holds?
        1 - V * theta / Dc = 0.05 * cos(t)
        => V * theta / Dc = 1 - 0.05 * cos(t)
        => V(t) = (1 - 0.05 * cos(t)) * Dc / theta(t)
        
        If we drive the system with this V(t), and start with correct theta(0),
        the system should follow theta(t).
        """
        Dc = 0.01
        layer = FrictionLayer(Dc=Dc, law='aging')
        
        # Time setup
        dt = 1e-5 # Tiny dt for accuracy
        steps = 5000 # Total time 0.05s
        times = torch.arange(steps) * dt
        
        # Analytic Solution
        theta_analytic = 0.1 + 0.05 * torch.sin(times)
        theta_dot_analytic = 0.05 * torch.cos(times)
        
        # Manufactured Velocity
        V_manufactured = (1.0 - theta_dot_analytic) * Dc / theta_analytic
        
        # Run Simulation
        state = torch.tensor([theta_analytic[0]])
        errors = []
        current_state = state
        
        for i in range(steps):
            v = V_manufactured[i:i+1]
            mu, next_state = layer(v, current_state, dt)
            
            if i < steps - 1:
                theta_true = theta_analytic[i+1]
                error = torch.abs(next_state - theta_true)
                errors.append(error.item())
                
            current_state = next_state
            
        max_error = np.max(errors)
        # Using robust torch assertion
        # Note: assert_close compares tensor equality. 
        # We want to check if error is small. 
        # We can test that next_state matches theta_analytic within tolerance.
        # But we computed max_error. Let's keep it but maybe use a better check?
        # User requested torch.testing.assert_close.
        # Let's verify the LAST state matches analytically exactly.
        
        # Or better, just assert max_error is small, but using a robust check if feasible?
        # "Use torch.testing.assert_close(..., rtol=..., atol=...) which provides better error messages"
        # Since I'm checking max error scalar, assertLess is actually fine, 
        # but to satisfy "Robustness", I'll compare the full trajectory 
        # (re-calculating analytic on the fly or just checking the loop result if I saved it).
        # Since I didn't save the full trajectory to a tensor in this test (just errors list), 
        # I will stick to assertLess but update the comment or use `assert max_error < 1e-7` with better message.
        # Actually, let's verify the final state vs expected with assert_close.
        
        expected_final = theta_analytic[-1]
        actual_final = current_state.squeeze()
        # atol=1e-6 (Relaxed for robustness across platforms/minor numerical drift)
        torch.testing.assert_close(actual_final, expected_final, atol=1e-6, rtol=0.0)
        
        # Also keep the max error check for the path
        self.assertLess(max_error, 1e-6, f"MMS Failed. Max Error: {max_error}")

    def test_gradient_check(self):
        """Verify gradients flow correctly."""
        layer = FrictionLayer(law='aging')
        
        # Input that requires grad
        # Use V=1.0 to avoid 1/V numerical instability with finite difference eps=1e-6
        velocity = torch.tensor([1.0], dtype=torch.float64, requires_grad=True)
        state = torch.tensor([0.1], dtype=torch.float64, requires_grad=True)
        dt = 0.1
        
        # Use gradcheck
        def func(v, s):
            f, ns = layer(v, s, dt)
            return f
            
        from torch.autograd import gradcheck
        # eps=1e-6 is small. Since we use float64, it should be fine.
        test = gradcheck(func, (velocity, state), eps=1e-6, atol=1e-3)
        self.assertTrue(test, "Gradient check failed")

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