"""Build aligned TPSA and Morgan arrays with published chematic 1.1.4."""
import csv
from pathlib import Path

import chematic
import numpy as np

assert chematic.__version__ == "1.1.4"
records = [
    ("ethanol", "CCO"),
    ("aspirin", "CC(=O)Oc1ccccc1C(=O)O"),
    ("caffeine", "Cn1cnc2c1c(=O)n(C)c(=O)n2C"),
]
molecules = [chematic.from_smiles(smiles) for _, smiles in records]
features = np.stack([mol.rdkit_ecfp_config_numpy(2, 2048) for mol in molecules])
tpsa = np.array([mol.rdkit_tpsa for mol in molecules])
assert features.shape == (3, 2048)
assert features.dtype == np.uint8
assert np.isin(features, [0, 1]).all()
for mol, row in zip(molecules, features):
    packed = np.frombuffer(bytes(mol.rdkit_ecfp_config(2, 2048)), dtype=np.uint8)
    np.testing.assert_array_equal(row, np.unpackbits(packed, bitorder="little"))

np.save("morgan.npy", features)
np.save("tpsa.npy", tpsa)
with Path("features.csv").open("w", newline="", encoding="utf-8") as handle:
    writer = csv.writer(handle)
    writer.writerow(["row", "compound_id", "source_smiles", "tpsa"])
    for index, ((compound_id, smiles), value) in enumerate(zip(records, tpsa)):
        writer.writerow([index, compound_id, smiles, f"{value:.2f}"])

print(features.shape, features.dtype)
print(tpsa)