"""
Tests for conformational diversity feature.
Following TDD methodology - these tests should fail initially.
"""
import pytest
from synth_pdb.generator import generate_pdb_content
from synth_pdb.validator import PDBValidator
from synth_pdb.data import RAMACHANDRAN_PRESETS
import numpy as np


class TestConformationalDiversity:
    """Test suite for conformational diversity feature."""
    
    def test_ramachandran_presets_exist(self):
        """Test that RAMACHANDRAN_PRESETS dictionary exists in data.py."""
        assert 'alpha' in RAMACHANDRAN_PRESETS
        assert 'beta' in RAMACHANDRAN_PRESETS
        assert 'extended' in RAMACHANDRAN_PRESETS
        assert 'ppii' in RAMACHANDRAN_PRESETS
    
    def test_ramachandran_preset_structure(self):
        """Test that each preset has correct keys."""
        for conformation, angles in RAMACHANDRAN_PRESETS.items():
            # Check for standard presets (phi/psi)
            if 'phi' in angles:
                assert 'psi' in angles, f"{conformation} missing psi"
                assert isinstance(angles['phi'], (int, float))
                assert isinstance(angles['psi'], (int, float))
            # Check for residue-specific presets (mean/std)
            elif 'phi_mean' in angles:
                assert 'phi_std' in angles
                assert 'psi_mean' in angles
                assert 'psi_std' in angles
            # Check for region-based presets (RAMACHANDRAN_REGIONS structure)
            elif 'favored' in angles:
                assert isinstance(angles['favored'], list)
            else:
                pytest.fail(f"Unknown preset format for {conformation}: {angles.keys()}")
    
    def test_alpha_helix_angles(self):
        """Test alpha helix preset has correct angles."""
        alpha = RAMACHANDRAN_PRESETS['alpha']
        assert alpha['phi'] == pytest.approx(-57.0, abs=1.0)
        assert alpha['psi'] == pytest.approx(-47.0, abs=1.0)
    
    def test_beta_sheet_angles(self):
        """Test beta sheet preset has correct angles."""
        beta = RAMACHANDRAN_PRESETS['beta']
        assert beta['phi'] == pytest.approx(-135.0, abs=5.0)
        assert beta['psi'] == pytest.approx(135.0, abs=5.0)
    
    def test_default_conformation_is_alpha(self):
        """Test that default conformation is alpha helix."""
        # Use a fixed seed for reproducibility
        random.seed(42)
        np.random.seed(42)
        
        generated_pdb = generate_pdb_content(length=10)  # Default is alpha
        
        # Verify basic structure
        lines = generated_pdb.splitlines()
        assert lines[0].startswith("HEADER")
        assert len([l for l in lines if l.startswith("ATOM")]) > 0
        assert lines[-1].strip() == "END"
    
    def test_generate_with_beta_conformation(self):
        """Test generating PDB with beta sheet conformation."""
        pdb_content = generate_pdb_content(
            length=5,
            sequence_str="AAAAA",
            conformation='beta'
        )
        
        # Should generate valid PDB
        assert "ATOM" in pdb_content
        assert "END" in pdb_content
        
        # Validate it parses correctly
        validator = PDBValidator(pdb_content)
        atoms = validator._parse_pdb_atoms(pdb_content)
        assert len(atoms) > 0
    
    def test_generate_with_extended_conformation(self):
        """Test generating PDB with extended conformation."""
        pdb_content = generate_pdb_content(
            length=5,
            sequence_str="AAAAA",
            conformation='extended'
        )
        
        # Should generate valid PDB
        assert "ATOM" in pdb_content
        assert "END" in pdb_content
    
    def test_generate_with_ppii_conformation(self):
        """Test generating PDB with polyproline II conformation."""
        pdb_content = generate_pdb_content(
            length=5,
            sequence_str="PPPPP",
            conformation='ppii'
        )
        
        # Should generate valid PDB
        assert "ATOM" in pdb_content
        assert "END" in pdb_content
    
    def test_generate_with_random_conformation(self):
        """Test generating PDB with random conformation sampling."""
        # Generate two structures with random conformations
        pdb1 = generate_pdb_content(length=10, conformation='random')
        pdb2 = generate_pdb_content(length=10, conformation='random')
        
        # Both should be valid
        assert "ATOM" in pdb1
        assert "ATOM" in pdb2
        
        # They should be different (with high probability)
        # This is a probabilistic test, but with 10 residues the chance of identical
        # random structures is vanishingly small
        assert pdb1 != pdb2, "Random conformations should produce different structures"
    
    def test_default_conformation_is_alpha(self):
        """Test that default conformation (when not specified) is alpha helix."""
        pdb_default = generate_pdb_content(length=5, sequence_str="AAAAA")
        
        # Verify it generates valid PDB content
        assert "ATOM" in pdb_default
        assert "HEADER" in pdb_default
        assert "END" in pdb_default
    
    def test_invalid_conformation_raises_error(self):
        """Test that invalid conformation name raises ValueError."""
        with pytest.raises(ValueError, match="Invalid conformation"):
            generate_pdb_content(length=5, conformation='invalid_name')
    
    def test_conformation_affects_backbone_geometry(self):
        """Test that different conformations produce different backbone geometries."""
        alpha_pdb = generate_pdb_content(length=10, sequence_str="A"*10, conformation='alpha')
        beta_pdb = generate_pdb_content(length=10, sequence_str="A"*10, conformation='beta')
        
        # Parse CA coordinates
        alpha_validator = PDBValidator(alpha_pdb)
        beta_validator = PDBValidator(beta_pdb)
        
        alpha_atoms = alpha_validator._parse_pdb_atoms(alpha_pdb)
        beta_atoms = beta_validator._parse_pdb_atoms(beta_pdb)
        
        # Get CA atoms
        alpha_ca = [a for a in alpha_atoms if a['atom_name'].strip() == 'CA']
        beta_ca = [a for a in beta_atoms if a['atom_name'].strip() == 'CA']
        
        # Calculate end-to-end distance
        alpha_distance = np.linalg.norm(
            np.array(alpha_ca[-1]['coords']) - np.array(alpha_ca[0]['coords'])
        )
        beta_distance = np.linalg.norm(
            np.array(beta_ca[-1]['coords']) - np.array(beta_ca[0]['coords'])
        )
        
        # Beta sheets are more extended than alpha helices
        # So end-to-end distance should be greater for beta
        assert beta_distance > alpha_distance, \
            f"Beta sheet ({beta_distance:.2f}Å) should be more extended than alpha helix ({alpha_distance:.2f}Å)"
