import sys
from pathlib import Path
import tempfile

import numpy as np

sys.path.append(str(Path(__file__).resolve().parents[1] / "src"))

from knot_recognition.protein import extract_ca_polyline


PDB_SAMPLE = """\
ATOM      1  CA  GLY A   1      10.000  11.000  12.000  1.00 20.00           C
ATOM      2  CA  GLY A   2      13.000  14.000  15.000  1.00 20.00           C
ATOM      3  CA  GLY A   3      16.000  17.000  18.000  1.00 20.00           C
ATOM      4  CA  GLY B   1      20.000  21.000  22.000  1.00 20.00           C
ATOM      5  CA  GLY B   2      23.000  24.000  25.000  1.00 20.00           C
"""


def test_extract_ca_polyline_default_chain():
    with tempfile.NamedTemporaryFile(mode="w", suffix=".pdb", delete=False) as f:
        f.write(PDB_SAMPLE)
        path = f.name
    result = extract_ca_polyline(path)
    assert result.chain_id == "A"
    assert result.ca_coords.shape == (3, 3)
    assert np.allclose(result.ca_coords[0], [10.0, 11.0, 12.0])


def test_extract_ca_polyline_chain_b():
    with tempfile.NamedTemporaryFile(mode="w", suffix=".pdb", delete=False) as f:
        f.write(PDB_SAMPLE)
        path = f.name
    result = extract_ca_polyline(path, chain_id="B")
    assert result.chain_id == "B"
    assert result.ca_coords.shape == (2, 3)
    assert np.allclose(result.ca_coords[-1], [23.0, 24.0, 25.0])
