import copy

import numpy as np
import pytest

import qcelemental as qcel
from qcelemental.testing import compare_molrecs

_schema_prov_stamp = {"creator": "QCElemental", "version": "1.0", "routine": "qcelemental.molparse.from_schema"}


@pytest.mark.parametrize(
    "inp,expected",
    [
        (
            {"frag_pattern": [[0], [1]], "geom": [0.0, 0.0, 0.0, 1.0, 0.0, 0.0], "elbl": ["O", "H"]},
            {
                "fragment_separators": np.array([1]),
                "geom": np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]),
                "elbl": np.array(["O", "H"]),
            },
        ),
        (
            {
                "frag_pattern": [[2, 0], [1]],
                "geom": np.array([[0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 0.0]]),
                "elem": np.array(["Li", "H", "He"]),
            },
            {
                "fragment_separators": np.array([2]),
                "geom": np.array([0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0]),
                "elem": np.array(["He", "Li", "H"]),
            },
        ),
        (
            {"frag_pattern": [[2, 0], [1]], "elez": [3, 1, 2]},
            {"fragment_separators": np.array([2]), "elez": np.array([2, 3, 1])},
        ),
    ],
)
def test_contiguize_from_fragment_pattern(inp, expected):
    ans = qcel.molparse.contiguize_from_fragment_pattern(**inp)

    # compare_molrecs instead of compare_dicts handles some fragment_separators types issues
    assert compare_molrecs(expected, ans, atol=1.0e-6)


@pytest.mark.parametrize(
    "inp,expected",
    [
        (
            {
                "frag_pattern": [[2, 0], [1, 3]],
                "geom": np.array([[0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 0.0]]),
                "elem": np.array(["Li", "H", "He"]),
            },
            "dropped atoms",
        ),
        ({"frag_pattern": [[2, 0], [1, 4]]}, "Fragmentation pattern skips atoms"),
        (
            {
                "frag_pattern": [[2, 0], [1, 3]],
                "elem": np.array(["U", "Li", "H", "He"]),
                "elbl": np.array(["Li", "H", "He"]),
            },
            "wrong number of atoms in array",
        ),
        (
            {"frag_pattern": [[2, 0], [1]], "elez": [3, 1, 2], "throw_reorder": True},
            "reorder atoms to accommodate non-contiguous fragments",
        ),
    ],
)
def test_contiguize_from_fragment_pattern_error(inp, expected):
    with pytest.raises(qcel.ValidationError) as e:
        qcel.molparse.contiguize_from_fragment_pattern(**inp)

    assert expected in str(e.value)


schema14_1 = {
    "geometry": [0.0, 0.0, -5.0, 0.0, 0.0, 5.0],
    "symbols": ["He", "He"],
    "fragments": [[0], [1]],
    "fragment_charges": [0.0, 0.0],
    "fragment_multiplicities": [3, 1],
    "masses": [4.00260325413, 4.00260325413],
    "name": "He2",
    "fix_com": False,
    "fix_orientation": False,
    "molecular_charge": 0.0,
    "molecular_multiplicity": 3,
    "real": [True, False],
}

schema14_psi4_np = {
    "geom": np.array([0.0, 0.0, -5.0, 0.0, 0.0, 5.0]),
    "elem": np.array(["He", "He"]),
    "elea": np.array([4, 4]),
    "elez": np.array([2, 2]),
    "fragment_charges": [0.0, 0.0],
    "fragment_multiplicities": [3, 1],
    "mass": np.array([4.00260325413, 4.00260325413]),
    "name": "He2",
    "fix_com": False,
    "fix_orientation": False,
    "molecular_charge": 0.0,
    "molecular_multiplicity": 3,
    "units": "Bohr",
    "fragment_separators": [1],
    "elbl": np.array(["", ""]),
    "real": np.array([True, False]),
    "provenance": _schema_prov_stamp,
}


def test_from_schema_1_14e():
    schema = {"schema_name": "qc_schema", "schema_version": 1, "molecule": copy.deepcopy(schema14_1)}

    ans = qcel.molparse.from_schema(schema)
    assert compare_molrecs(schema14_psi4_np, ans, 4)


def test_from_schema_1p5_14e():
    # this is the oddball case where passing code has internally a dtype=2 molecule
    #   but it's still passing the outer data structure
    schema = {"schema_name": "qc_schema", "schema_version": 1, "molecule": copy.deepcopy(schema14_1)}
    schema["molecule"].update({"schema_name": "qcschema_molecule", "schema_version": 2})

    ans = qcel.molparse.from_schema(schema)
    assert compare_molrecs(schema14_psi4_np, ans, 4)


def test_from_schema_2_14e():
    schema = copy.deepcopy(schema14_1)
    schema.update({"schema_name": "qcschema_molecule", "schema_version": 2})

    ans = qcel.molparse.from_schema(schema)
    assert compare_molrecs(schema14_psi4_np, ans, 4)


def test_from_schema_error_f():
    schema = {"schema_name": "private_schema", "schema_version": 1, "molecule": copy.deepcopy(schema14_1)}

    with pytest.raises(qcel.ValidationError) as e:
        qcel.molparse.from_schema(schema)

    assert "Schema not recognized" in str(e.value)


def test_from_schema_1_nfr_error_14g():
    schema = {"schema_name": "qc_schema", "schema_version": 1, "molecule": copy.deepcopy(schema14_1)}
    schema["molecule"].pop("fragments")

    with pytest.raises(qcel.ValidationError) as e:
        qcel.molparse.from_schema(schema)

    assert "Dimension mismatch among fragment quantities: sep + 1 (1), chg (2), and mult(2)" in str(e.value)
