from pathlib import Path
from unittest import TestCase
from pytest import mark, exit
from MetAromatic import MetAromatic, MetAromaticLocal

TEST_PARAMETERS = {
    "cutoff_distance": 4.9,
    "cutoff_angle": 109.5,
    "chain": "A",
    "model": "cp",
}

VALID_RESULTS_1RCY = [
    {
        "aromatic_residue": "TYR",
        "aromatic_position": 122,
        "methionine_position": 18,
        "norm": 4.211,
        "met_theta_angle": 75.766,
        "met_phi_angle": 64.317,
    },
    {
        "aromatic_residue": "TYR",
        "aromatic_position": 122,
        "methionine_position": 18,
        "norm": 3.954,
        "met_theta_angle": 60.145,
        "met_phi_angle": 68.352,
    },
    {
        "aromatic_residue": "TYR",
        "aromatic_position": 122,
        "methionine_position": 18,
        "norm": 4.051,
        "met_theta_angle": 47.198,
        "met_phi_angle": 85.151,
    },
    {
        "aromatic_residue": "TYR",
        "aromatic_position": 122,
        "methionine_position": 18,
        "norm": 4.39,
        "met_theta_angle": 53.4,
        "met_phi_angle": 95.487,
    },
    {
        "aromatic_residue": "TYR",
        "aromatic_position": 122,
        "methionine_position": 18,
        "norm": 4.62,
        "met_theta_angle": 68.452,
        "met_phi_angle": 90.771,
    },
    {
        "aromatic_residue": "TYR",
        "aromatic_position": 122,
        "methionine_position": 18,
        "norm": 4.537,
        "met_theta_angle": 78.568,
        "met_phi_angle": 76.406,
    },
    {
        "aromatic_residue": "PHE",
        "aromatic_position": 54,
        "methionine_position": 148,
        "norm": 4.777,
        "met_theta_angle": 105.947,
        "met_phi_angle": 143.022,
    },
    {
        "aromatic_residue": "PHE",
        "aromatic_position": 54,
        "methionine_position": 148,
        "norm": 4.61,
        "met_theta_angle": 93.382,
        "met_phi_angle": 156.922,
    },
    {
        "aromatic_residue": "PHE",
        "aromatic_position": 54,
        "methionine_position": 148,
        "norm": 4.756,
        "met_theta_angle": 93.287,
        "met_phi_angle": 154.63,
    },
]


def test_pair_1rcy_valid_results():
    results = MetAromatic(TEST_PARAMETERS).get_met_aromatic_interactions("1rcy")

    tc = TestCase()
    tc.maxDiff = None
    tc.assertCountEqual(results["interactions"], VALID_RESULTS_1RCY)


def test_pair_1rcy_valid_results_use_local():
    # File downloaded from RSCB PDB
    path_pdb_file = Path(__file__).resolve().parent / "data_1rcy.pdb"

    if not path_pdb_file.exists():
        exit(f"File {path_pdb_file} is missing")

    results = MetAromaticLocal(TEST_PARAMETERS).get_met_aromatic_interactions(
        path_pdb_file
    )

    tc = TestCase()
    tc.maxDiff = None
    tc.assertCountEqual(results["interactions"], VALID_RESULTS_1RCY)


def test_pair_1rcy_valid_results_use_local_invalid_file():
    # Simulating someone passing a non-PDB formatted file into program
    path_pdb_file = Path(__file__).resolve().parent / "data_lorem_ipsum.pdb"

    if not path_pdb_file.exists():
        exit(f"File {path_pdb_file} is missing")

    results = MetAromaticLocal(TEST_PARAMETERS).get_met_aromatic_interactions(
        path_pdb_file
    )

    assert results["status"] != "Success"
    assert not results["OK"]


@mark.parametrize(
    "code, cutoff_distance, cutoff_angle, model, status",
    [
        ("1rcy", -0.01, 109.5, "cp", "Invalid cutoff distance"),
        ("1rcy", 4.95, -60.0, "cp", "Invalid cutoff angle"),
        ("1rcy", 4.95, 720.0, "cp", "Invalid cutoff angle"),
        ("2rcy", 4.95, 109.5, "cp", "No MET residues"),
        ("3nir", 4.95, 109.5, "cp", "No MET residues"),
        ("6mwm", 4.95, 109.5, "cp", "No PHE/TYR/TRP residues"),
        ("abcd", 4.95, 109.5, "cp", "Invalid PDB entry"),
        ("1rcy", 4.95, 109.5, "pc", "Invalid model"),
        ("1rcy", "4.95", 109.5, "cp", "Cutoff distance must be a valid float"),
        ("1rcy", 4.95, "109.5", "cp", "Cutoff angle must be a valid float"),
        ("1rcy", 4.95, 109.5, 25, "Model must be a valid string"),
    ],
)
def test_pair_invalid_inputs(code, cutoff_distance, cutoff_angle, model, status):
    params = {
        "cutoff_angle": cutoff_angle,
        "cutoff_distance": cutoff_distance,
        "chain": "A",
        "model": model,
    }

    results = MetAromatic(params).get_met_aromatic_interactions(code)
    assert not results["OK"]
    assert results["status"] == status


def test_pair_no_results_error():
    results = MetAromatic(TEST_PARAMETERS).get_met_aromatic_interactions("1a5r")

    assert results["status"] == "No interactions"
    assert results["OK"]
