#
# Project NuriKit - Copyright 2024 SNU Compbio Lab.
# SPDX-License-Identifier: Apache-2.0
#

from pathlib import Path
from typing import List

import numpy as np
import pytest

import nuri
from nuri.core import Molecule, SubstructureCategory
from nuri.fmt import pdb

pdb_data = """\
MODEL        1
ATOM      7  N   MET A   1     -13.991  -6.903  33.129  1.00 22.97           N
ATOM      8  CA  MET A   1     -13.215  -8.093  33.479  1.00 22.29           C
ATOM      9  C   MET A   1     -13.314  -8.351  34.974  1.00 22.05           C
ATOM     10  O   MET A   1     -12.547  -9.129  35.537  1.00 22.25           O
ATOM     11  CB  MET A   1     -11.754  -7.911  33.072  1.00 22.24           C
ATOM     12  CG  MET A   1     -11.575  -7.558  31.609  1.00 22.06           C
ATOM     13  SD  MET A   1     -12.824  -8.351  30.571  1.00 21.21           S
ATOM     14  CE  MET A   1     -12.200 -10.034  30.491  1.00 21.07           C
ENDMDL
MODEL        2
ATOM     15  N   GLU A   2     -14.276  -7.680  35.597  1.00 21.61           N
ATOM     16  CA  GLU A   2     -14.539  -7.754  37.021  1.00 21.37           C
ATOM     17  C   GLU A   2     -14.838  -9.192  37.467  1.00 20.56           C
ATOM     18  O   GLU A   2     -14.382  -9.623  38.530  1.00 20.29           O
ATOM     19  CB  GLU A   2     -15.725  -6.832  37.337  1.00 21.91           C
ATOM     20  CG  GLU A   2     -16.178  -6.790  38.801  1.00 23.74           C
ATOM     21  CD  GLU A   2     -17.703  -6.770  38.937  1.00 26.25           C
ATOM     22  OE1 GLU A   2     -18.397  -7.133  37.957  1.00 26.86           O
ATOM     23  OE2 GLU A   2     -18.208  -6.402  40.025  1.00 27.14           O
ATOM     24  N   ASN B   3A    -15.590  -9.921  36.639  1.00 19.52           N
ATOM     25  CA  ASN B   3A    -16.096 -11.249  36.992  1.00 18.45           C
ATOM     26  C   ASN B   3A    -15.210 -12.437  36.598  1.00 17.91           C
ATOM     27  O   ASN B   3A    -15.565 -13.584  36.850  1.00 17.64           O
ATOM     28  CB  ASN B   3A    -17.500 -11.436  36.413  1.00 18.28           C
ATOM     29  CG  ASN B   3A    -18.490 -10.459  36.982  1.00 17.92           C
ATOM     30  OD1 ASN B   3A    -18.464 -10.165  38.175  1.00 18.06           O
ATOM     31  ND2 ASN B   3A    -19.375  -9.947  36.135  1.00 17.48           N
ENDMDL
"""


def _verify_mols(mols: List[Molecule]):
    assert len(mols) == 2

    assert len(mols[0]) == 8

    mol = mols[1]
    assert len(mol.subs) == 4
    for sub in mol.subs:
        if sub.category == SubstructureCategory.Chain:
            assert sub.name == "A" or sub.name == "B"
        elif sub.category == SubstructureCategory.Residue:
            if sub.id == 2:
                assert sub.name == "GLU"
                assert sub.props["chain"] == "A"
            elif sub.id == 3:
                assert sub.name == "ASN"
                assert sub.props["chain"] == "B"
            else:
                pytest.fail("Invalid residue ID")
        else:
            pytest.fail("Invalid substructure category")


def test_pdb_file(tmp_path: Path):
    file = tmp_path / "test.pdb"
    file.write_text(pdb_data)

    mols = list(nuri.readfile("pdb", file))
    _verify_mols(mols)


def test_pdb_str():
    mols = list(nuri.readstring("pdb", pdb_data))
    _verify_mols(mols)

    mols_re = [
        mol
        for pdb_re in map(nuri.to_pdb, mols)
        for mol in nuri.readstring("pdb", pdb_re)
    ]
    _verify_mols(mols_re)


def test_pdb_models(tmp_path: Path):
    file = tmp_path / "test.pdb"
    file.write_text(pdb_data)

    models = pdb.read_models(file)
    assert len(models) == 2

    model = models[0]
    assert len(model.atoms) == 8
    for atom in model.atoms:
        assert len(atom.sites) == 1
        assert atom.sites[0].altloc == " "
        assert atom.element.symbol in atom.name
        assert atom.formal_charge == 0
        assert not atom.hetero

    atom = model.atoms[0]
    assert atom.res_id.chain_id == "A"
    assert atom.res_id.res_seq == 1
    assert atom.res_id.ins_code == ""

    site = model.atoms[0].sites[0]
    assert site.occupancy == pytest.approx(1.00)
    assert site.tempfactor == pytest.approx(22.97)

    np.testing.assert_array_equal(
        model.major_conf,
        np.stack([atom.sites[0].pos for atom in model.atoms]),
    )
    np.testing.assert_array_equal(
        model.major_conf,
        [
            [-13.991, -6.903, 33.129],
            [-13.215, -8.093, 33.479],
            [-13.314, -8.351, 34.974],
            [-12.547, -9.129, 35.537],
            [-11.754, -7.911, 33.072],
            [-11.575, -7.558, 31.609],
            [-12.824, -8.351, 30.571],
            [-12.200, -10.034, 30.491],
        ],
    )

    model = models[1]
    assert len(model.residues) == 2
    for res in model.residues:
        if res.id.res_seq == 2:
            assert res.id.chain_id == "A"
            assert res.id.ins_code == ""
            assert res.name == "GLU"
            assert len(res.atom_idxs) == 9
        elif res.id.res_seq == 3:
            assert res.id.chain_id == "B"
            assert res.id.ins_code == "A"
            assert res.name == "ASN"
            assert len(res.atom_idxs) == 8
        else:
            pytest.fail(f"Invalid residue ID {res.id.res_seq}")

    assert len(model.chains) == 2
    for chain in model.chains:
        if chain.id == "A":
            assert len(chain.res_idxs) == 1
        elif chain.id == "B":
            assert len(chain.res_idxs) == 1
        else:
            pytest.fail(f"Invalid chain ID {chain.id}")


def test_pdb_model_asdict(tmp_path: Path):
    file = tmp_path / "test.pdb"
    file.write_text(pdb_data)

    models = pdb.read_models(file)
    assert len(models) == 2

    model = models[0]

    md = model.as_dict()
    for ma, da in zip(model.atoms, md["atoms"]):
        for ms, ds in zip(ma.sites, da["sites"]):
            np.testing.assert_array_equal(ms.pos, ds["pos"])
            del ds["pos"]

    for mr, dr in zip(model.residues, md["residues"]):
        np.testing.assert_array_equal(mr.atom_idxs, dr["atom_idxs"])
        del dr["atom_idxs"]

    for mc, dc in zip(model.chains, md["chains"]):
        np.testing.assert_array_equal(mc.res_idxs, dc["res_idxs"])
        del dc["res_idxs"]

    np.testing.assert_array_equal(model.major_conf, md["major_conf"])
    del md["major_conf"]

    assert md == {
        "atoms": [
            {
                "atomic_number": atom.element.atomic_number,
                "formal_charge": atom.formal_charge,
                "hetero": atom.hetero,
                "name": atom.name,
                "res_id": {
                    "chain_id": atom.res_id.chain_id,
                    "res_seq": atom.res_id.res_seq,
                    "ins_code": atom.res_id.ins_code,
                },
                "sites": [
                    {
                        "altloc": site.altloc,
                        "occupancy": site.occupancy,
                        "tempfactor": site.tempfactor,
                    }
                    for site in atom.sites
                ],
            }
            for atom in model.atoms
        ],
        "residues": [
            {
                "id": {
                    "chain_id": res.id.chain_id,
                    "res_seq": res.id.res_seq,
                    "ins_code": res.id.ins_code,
                },
                "name": res.name,
            }
            for res in model.residues
        ],
        "chains": [{"id": chain.id} for chain in model.chains],
        "props": dict(model.props),
    }
