"""
Comprehensive end-to-end integration tests for LLRQ package.

Tests complete workflows, complex reaction networks, different equilibrium modes,
and integration between all components of the package.
"""

import numpy as np
import pytest
import sys
import os
from unittest.mock import patch
import warnings

# Add source directory to path for testing
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

from llrq.reaction_network import ReactionNetwork
from llrq.llrq_dynamics import LLRQDynamics
from llrq.solver import LLRQSolver


class TestBasicWorkflows:
    """Test basic end-to-end workflows."""

    def test_simple_reversible_reaction_workflow(self):
        """Test complete workflow for A ⇌ B reaction."""
        # Step 1: Create network
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        
        # Step 2: Create dynamics from mass action
        c_star = np.array([1.0, 2.0])
        k_plus = np.array([2.0])
        k_minus = np.array([1.0])
        
        dynamics = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        # Step 3: Create solver and solve
        solver = LLRQSolver(dynamics)
        initial_concentrations = np.array([0.5, 1.0])
        t_span = (0.0, 5.0)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        # Step 4: Verify results
        assert result['success']
        assert len(result['time']) > 0
        assert result['concentrations'].shape[1] == 2
        assert result['reaction_quotients'].shape[1] == 1
        
        # Should approach equilibrium
        final_concentrations = result['concentrations'][-1]
        final_Q = result['reaction_quotients'][-1, 0]
        expected_Keq = k_plus[0] / k_minus[0]
        
        # Check that we're at or approaching equilibrium
        initial_Q = result['reaction_quotients'][0, 0]
        initial_deviation = abs(initial_Q - expected_Keq)
        final_deviation = abs(final_Q - expected_Keq)
        
        # Either we started at equilibrium, or we approached it
        assert final_deviation <= initial_deviation
        
        # Mass should be conserved
        total_mass = result['concentrations'].sum(axis=1)
        initial_total = initial_concentrations.sum()
        assert np.allclose(total_mass, initial_total, rtol=1e-3)

    def test_two_reaction_chain_workflow(self):
        """Test workflow for A → B → C chain."""
        # Step 1: Create network A ⇌ B ⇌ C
        network = ReactionNetwork(['A', 'B', 'C'], ['R1', 'R2'],
                                np.array([[-1, 0], [1, -1], [0, 1]]))
        
        # Step 2: Create dynamics
        c_star = np.array([1.0, 1.5, 0.5])
        k_plus = np.array([2.0, 1.0])
        k_minus = np.array([1.0, 2.0])
        
        dynamics = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        # Step 3: Solve
        solver = LLRQSolver(dynamics)
        initial_concentrations = np.array([3.0, 0.5, 0.5])
        t_span = (0.0, 10.0)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        # Step 4: Verify
        assert result['success']
        assert result['concentrations'].shape[1] == 3
        assert result['reaction_quotients'].shape[1] == 2
        
        # Check mass conservation
        total_mass = result['concentrations'].sum(axis=1)
        initial_total = initial_concentrations.sum()
        assert np.allclose(total_mass, initial_total, rtol=1e-3)

    def test_bimolecular_reaction_workflow(self):
        """Test workflow for A + B ⇌ C reaction."""
        # Step 1: Create network
        network = ReactionNetwork(['A', 'B', 'C'], ['R1'],
                                np.array([[-1], [-1], [1]]))
        
        # Step 2: Create dynamics
        c_star = np.array([1.0, 1.0, 1.0])
        k_plus = np.array([1.0])
        k_minus = np.array([2.0])
        
        dynamics = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        # Step 3: Solve
        solver = LLRQSolver(dynamics)
        initial_concentrations = np.array([2.0, 2.0, 0.1])
        t_span = (0.0, 5.0)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        # Step 4: Verify
        assert result['success']
        assert result['concentrations'].shape[1] == 3
        
        # Reaction quotient should be Q = [C]/([A][B])
        final_c = result['concentrations'][-1]
        expected_Q = final_c[2] / (final_c[0] * final_c[1])
        final_Q = result['reaction_quotients'][-1, 0]
        assert np.isclose(final_Q, expected_Q, rtol=1e-6)

    def test_workflow_with_external_drive(self):
        """Test workflow with external drive."""
        # Step 1: Create network
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        
        # Step 2: Create dynamics with external drive
        def drive_func(t):
            return np.array([0.1 * np.sin(t)])
        
        K = np.array([[1.0]])
        Keq = np.array([2.0])
        dynamics = LLRQDynamics(network, Keq, K, drive_func)
        
        # Step 3: Solve
        solver = LLRQSolver(dynamics)
        initial_concentrations = np.array([1.0, 1.0])
        t_span = (0.0, 2*np.pi)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        # Step 4: Verify
        assert result['success']
        
        # With oscillatory drive, should see oscillations
        Q_values = result['reaction_quotients'][:, 0]
        assert np.std(Q_values) > 0  # Should have variation


class TestComplexNetworks:
    """Test complex reaction networks."""

    def test_branched_network(self):
        """Test branched network: A ⇌ B, B ⇌ C, B ⇌ D."""
        # Network structure:
        #     C
        #     ↑
        # A ⇌ B ⇌ D
        network = ReactionNetwork(['A', 'B', 'C', 'D'], ['R1', 'R2', 'R3'],
                                np.array([
                                    [-1, 0, 0],   # A
                                    [1, -1, -1],  # B
                                    [0, 1, 0],    # C
                                    [0, 0, 1]     # D
                                ]))
        
        c_star = np.array([1.0, 2.0, 0.5, 0.5])
        k_plus = np.array([2.0, 1.0, 1.0])
        k_minus = np.array([1.0, 2.0, 2.0])
        
        dynamics = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        solver = LLRQSolver(dynamics)
        initial_concentrations = np.array([4.0, 0.5, 0.25, 0.25])
        t_span = (0.0, 10.0)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        assert result['success']
        assert result['concentrations'].shape[1] == 4
        assert result['reaction_quotients'].shape[1] == 3
        
        # Mass conservation
        total_mass = result['concentrations'].sum(axis=1)
        assert np.allclose(total_mass, total_mass[0], rtol=1e-3)

    def test_cyclic_network(self):
        """Test cyclic network: A → B → C → A."""
        network = ReactionNetwork(['A', 'B', 'C'], ['R1', 'R2', 'R3'],
                                np.array([
                                    [-1, 0, 1],   # A
                                    [1, -1, 0],   # B
                                    [0, 1, -1]    # C
                                ]))
        
        c_star = np.array([1.0, 1.0, 1.0])
        k_plus = np.array([1.0, 1.0, 1.0])
        k_minus = np.array([0.1, 0.1, 0.1])  # Strongly forward bias
        
        dynamics = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        solver = LLRQSolver(dynamics)
        initial_concentrations = np.array([3.0, 0.5, 0.5])
        t_span = (0.0, 5.0)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        assert result['success']
        
        # Mass should be conserved in cycle
        total_mass = result['concentrations'].sum(axis=1)
        assert np.allclose(total_mass, total_mass[0], rtol=1e-3)

    def test_coupled_reactions_with_conservation(self):
        """Test coupled reactions with multiple conservation laws."""
        # Two separate subsystems: A ⇌ B and C ⇌ D
        network = ReactionNetwork(['A', 'B', 'C', 'D'], ['R1', 'R2'],
                                np.array([
                                    [-1, 0],   # A
                                    [1, 0],    # B
                                    [0, -1],   # C
                                    [0, 1]     # D
                                ]))
        
        # Two separate equilibria
        c_star = np.array([1.0, 2.0, 3.0, 1.0])
        k_plus = np.array([2.0, 1.5])
        k_minus = np.array([1.0, 0.5])
        
        dynamics = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        solver = LLRQSolver(dynamics)
        initial_concentrations = np.array([0.5, 1.5, 2.0, 2.0])
        t_span = (0.0, 10.0)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        assert result['success']
        
        # Two separate conservation laws: A+B = const, C+D = const
        AB_total = result['concentrations'][:, 0] + result['concentrations'][:, 1]
        CD_total = result['concentrations'][:, 2] + result['concentrations'][:, 3]
        
        assert np.allclose(AB_total, AB_total[0], rtol=1e-3)
        assert np.allclose(CD_total, CD_total[0], rtol=1e-3)

    def test_large_network(self):
        """Test larger network with many species and reactions."""
        # Create A1 ⇌ A2 ⇌ A3 ⇌ A4 ⇌ A5 chain
        n_species = 5
        species_ids = [f'A{i+1}' for i in range(n_species)]
        reaction_ids = [f'R{i+1}' for i in range(n_species-1)]
        
        # Create stoichiometric matrix
        S = np.zeros((n_species, n_species-1))
        for j in range(n_species-1):
            S[j, j] = -1      # Reactant
            S[j+1, j] = 1     # Product
        
        network = ReactionNetwork(species_ids, reaction_ids, S)
        
        # Random but reasonable equilibrium
        c_star = np.array([2.0, 1.5, 1.0, 1.2, 0.8])
        k_plus = np.array([1.0, 1.2, 0.8, 1.1])
        k_minus = np.array([0.5, 0.6, 1.0, 0.9])
        
        dynamics = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        solver = LLRQSolver(dynamics)
        initial_concentrations = np.array([5.0, 0.5, 0.5, 0.5, 0.5])
        t_span = (0.0, 10.0)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        assert result['success']
        assert result['concentrations'].shape[1] == n_species
        assert result['reaction_quotients'].shape[1] == n_species-1
        
        # Mass conservation
        total_mass = result['concentrations'].sum(axis=1)
        assert np.allclose(total_mass, total_mass[0], rtol=1e-3)


class TestEquilibriumModes:
    """Test different equilibrium modes."""

    def test_equilibrium_vs_nonequilibrium_modes(self):
        """Compare equilibrium and nonequilibrium modes."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        
        c_star = np.array([1.0, 2.0])
        k_plus = np.array([2.0])
        k_minus = np.array([1.0])
        
        # Test equilibrium mode
        dynamics_eq = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        # Test nonequilibrium mode
        dynamics_neq = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='nonequilibrium'
        )
        
        # Both should create valid dynamics
        assert dynamics_eq.K.shape == (1, 1)
        assert dynamics_neq.K.shape == (1, 1)
        
        # Both should have same equilibrium constants
        assert np.allclose(dynamics_eq.Keq, dynamics_neq.Keq)
        
        # K matrices might differ
        # This is expected as they use different algorithms

    def test_basis_reduction_option(self):
        """Test dynamics computation with and without basis reduction."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        
        c_star = np.array([1.0, 2.0])
        k_plus = np.array([2.0])
        k_minus = np.array([1.0])
        
        # Compute dynamics matrix directly
        result_with_reduction = network.compute_dynamics_matrix(
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium',
            reduce_to_image=True
        )
        
        result_without_reduction = network.compute_dynamics_matrix(
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium',
            reduce_to_image=False
        )
        
        # Both should be stable
        assert result_with_reduction['eigenanalysis']['is_stable']
        assert result_without_reduction['eigenanalysis']['is_stable']
        
        # With reduction should have additional fields
        assert 'K_reduced' in result_with_reduction
        assert 'basis' in result_with_reduction
        assert 'K_reduced' not in result_without_reduction

    def test_symmetry_enforcement(self):
        """Test symmetry enforcement in dynamics matrix."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        
        c_star = np.array([1.0, 1.0])  # Symmetric equilibrium
        k_plus = np.array([1.0])
        k_minus = np.array([1.0])  # Symmetric rates
        
        result_sym = network.compute_dynamics_matrix(
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            enforce_symmetry=True
        )
        
        result_no_sym = network.compute_dynamics_matrix(
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            enforce_symmetry=False
        )
        
        # Symmetric version should be exactly symmetric
        K_sym = result_sym['K']
        assert np.allclose(K_sym, K_sym.T)
        
        # Both should be stable
        assert result_sym['eigenanalysis']['is_stable']
        assert result_no_sym['eigenanalysis']['is_stable']


class TestSolutionMethods:
    """Test different solution methods."""

    def test_analytical_vs_numerical_comparison(self):
        """Compare analytical and numerical solutions."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        K = np.array([[2.0]])
        dynamics = LLRQDynamics(network, relaxation_matrix=K)
        solver = LLRQSolver(dynamics)
        
        initial_concentrations = np.array([2.0, 0.5])
        t_span = (0.0, 2.0)
        
        # Solve analytically
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")  # Ignore warnings about method choice
            result_analytical = solver.solve(
                initial_concentrations, t_span, method='analytical'
            )
        
        # Solve numerically
        result_numerical = solver.solve(
            initial_concentrations, t_span, method='numerical'
        )
        
        # Both should succeed (or analytical might fallback to numerical)
        assert result_numerical['success']
        
        if result_analytical['method'] == 'analytical':
            # If analytical worked, compare solutions
            t_common = result_analytical['time']
            
            # Interpolate numerical solution to analytical time points
            Q_analytical = result_analytical['reaction_quotients'][:, 0]
            Q_numerical = np.interp(t_common, result_numerical['time'], 
                                  result_numerical['reaction_quotients'][:, 0])
            
            # Should be close
            assert np.allclose(Q_analytical, Q_numerical, rtol=1e-3)

    def test_solver_with_different_tolerances(self):
        """Test numerical solver with different tolerances."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        dynamics = LLRQDynamics(network)
        solver = LLRQSolver(dynamics)
        
        initial_concentrations = np.array([2.0, 0.5])
        t_span = (0.0, 1.0)
        
        # Loose tolerance
        result_loose = solver.solve(
            initial_concentrations, t_span, method='numerical',
            rtol=1e-3, atol=1e-6
        )
        
        # Tight tolerance
        result_tight = solver.solve(
            initial_concentrations, t_span, method='numerical',
            rtol=1e-8, atol=1e-10
        )
        
        # Both should succeed
        assert result_loose['success']
        assert result_tight['success']
        
        # Tight tolerance should be more accurate (but we can't easily test this)
        # Just verify both give reasonable results
        assert np.isfinite(result_loose['reaction_quotients']).all()
        assert np.isfinite(result_tight['reaction_quotients']).all()


class TestStabilityAndConvergence:
    """Test stability and convergence properties."""

    def test_convergence_to_equilibrium(self):
        """Test that solutions converge to equilibrium."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        
        c_star = np.array([1.0, 3.0])
        k_plus = np.array([3.0])
        k_minus = np.array([1.0])
        
        dynamics = LLRQDynamics.from_mass_action(
            network=network,
            equilibrium_point=c_star,
            forward_rates=k_plus,
            backward_rates=k_minus,
            mode='equilibrium'
        )
        
        solver = LLRQSolver(dynamics)
        
        # Start far from equilibrium
        initial_concentrations = np.array([4.0, 0.1])
        t_span = (0.0, 20.0)  # Long time
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        assert result['success']
        
        # Should approach equilibrium
        expected_Keq = k_plus[0] / k_minus[0]
        final_Q = result['reaction_quotients'][-1, 0]
        initial_Q = result['reaction_quotients'][0, 0]
        
        # Final should be closer to equilibrium than initial
        final_error = abs(final_Q - expected_Keq)
        initial_error = abs(initial_Q - expected_Keq)
        assert final_error < 0.5 * initial_error

    def test_stability_for_stable_system(self):
        """Test that stable systems remain stable."""
        network = ReactionNetwork(['A', 'B', 'C'], ['R1', 'R2'],
                                np.array([[-1, 0], [1, -1], [0, 1]]))
        
        # Create stable K matrix (positive definite)
        K = np.array([[2.0, 0.1], [0.1, 1.5]])
        eigenvals = np.linalg.eigvals(K)
        assert np.all(eigenvals > 0)  # Should be stable
        
        dynamics = LLRQDynamics(network, relaxation_matrix=K)
        solver = LLRQSolver(dynamics)
        
        initial_concentrations = np.array([2.0, 1.0, 1.0])
        t_span = (0.0, 10.0)
        
        result = solver.solve(initial_concentrations, t_span, method='numerical')
        
        assert result['success']
        
        # Solution should not blow up
        assert np.isfinite(result['concentrations']).all()
        assert np.isfinite(result['reaction_quotients']).all()
        
        # Should not grow without bound
        max_concentration = np.max(result['concentrations'])
        assert max_concentration < 100.0  # Reasonable bound

    def test_conservation_law_preservation(self):
        """Test that conservation laws are preserved throughout integration."""
        # A ⇌ B ⇌ C system with total mass conservation
        network = ReactionNetwork(['A', 'B', 'C'], ['R1', 'R2'],
                                np.array([[-1, 0], [1, -1], [0, 1]]))
        
        dynamics = LLRQDynamics(network)
        solver = LLRQSolver(dynamics)
        
        initial_concentrations = np.array([3.0, 1.0, 1.0])
        t_span = (0.0, 5.0)
        
        result = solver.solve(initial_concentrations, t_span, 
                            enforce_conservation=True, method='numerical')
        
        assert result['success']
        
        # Total mass should be conserved at all times
        total_mass = result['concentrations'].sum(axis=1)
        expected_total = initial_concentrations.sum()
        
        assert np.allclose(total_mass, expected_total, rtol=1e-6)
        
        # Should not have negative concentrations
        assert np.all(result['concentrations'] >= 0)


class TestErrorRecoveryAndRobustness:
    """Test error recovery and robustness."""

    def test_recovery_from_analytical_failure(self):
        """Test recovery when analytical solution fails."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        
        # Create ill-conditioned K matrix
        K = np.array([[1e-12]])  # Very small, might cause numerical issues
        dynamics = LLRQDynamics(network, relaxation_matrix=K)
        solver = LLRQSolver(dynamics)
        
        initial_concentrations = np.array([1.0, 1.0])
        t_span = (0.0, 1.0)
        
        # Request analytical but expect fallback to numerical
        result = solver.solve(initial_concentrations, t_span, method='analytical')
        
        # Should succeed even if it falls back to numerical
        assert result['success']
        assert result['method'] in ['analytical', 'numerical']

    def test_robustness_to_extreme_parameters(self):
        """Test robustness to extreme parameter values."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        
        # Extreme equilibrium point
        c_star = np.array([1e-3, 1e3])
        k_plus = np.array([1e6])
        k_minus = np.array([1.0])
        
        try:
            dynamics = LLRQDynamics.from_mass_action(
                network=network,
                equilibrium_point=c_star,
                forward_rates=k_plus,
                backward_rates=k_minus,
                mode='equilibrium'
            )
            
            solver = LLRQSolver(dynamics)
            initial_concentrations = np.array([1.0, 1.0])
            t_span = (0.0, 1e-3)  # Short time for fast dynamics
            
            result = solver.solve(initial_concentrations, t_span, method='numerical')
            
            # Should either succeed or fail gracefully
            if result['success']:
                assert np.isfinite(result['concentrations']).all()
                
        except (ValueError, OverflowError, np.linalg.LinAlgError):
            # It's acceptable to raise errors for extreme parameters
            pass

    def test_handling_of_conservation_violations(self):
        """Test handling when conservation laws might be violated numerically."""
        network = ReactionNetwork(['A', 'B'], ['R1'], np.array([[-1], [1]]))
        dynamics = LLRQDynamics(network)
        solver = LLRQSolver(dynamics)
        
        initial_concentrations = np.array([1.0, 1.0])
        t_span = (0.0, 10.0)
        
        # Solve with loose tolerances that might violate conservation
        result = solver.solve(initial_concentrations, t_span, 
                            method='numerical', rtol=1e-2, atol=1e-4)
        
        assert result['success']
        
        # Even with loose tolerances, mass conservation should be approximately maintained
        if result['concentrations'] is not None:
            total_mass = result['concentrations'].sum(axis=1)
            expected_total = initial_concentrations.sum()
            
            # Allow larger tolerance for loose integration
            assert np.allclose(total_mass, expected_total, rtol=1e-1)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])