from typing import Tuple, List

from numpy.typing import NDArray
from chemsh import *
from chemsh.hybrid.Layer import Layer
from chemsh.utils import testutils
from chemsh.hybrid.NLayerSubtractive import NLayerSubtractive
import numpy as np


# Define the fragment
butanol = Fragment(coords="butanol.cjson")

# Define layer 1 (inner/high)
atoms_layer_1 = np.array([3, 4], dtype=np.int32)
method_1 = Gaussian(method="HF", basis="6-31+G(d,p)", maxiter=10000)
layer_1 = Layer(method_1, atoms_layer_1)

# Define Layer 2 (middle)
atoms_layer_2 = np.array([0, 1, 2], dtype=np.int32)
method_2 = Gaussian(method="HF", basis="6-31G")
layer_2 = Layer(method_2, atoms_layer_2)

# Define layer 3 (outer/low)
atoms_layer_3 = np.array([5, 6, 7, 8, 9, 10, 11, 12, 13, 14], np.int32)
method_3 = Gaussian(method="HF", basis="3-21G", maxiter=10000)
layer_3 = Layer(method_3, atoms_layer_3)


# Define the layers
layers: Tuple[Layer, Layer, Layer] = (layer_1, layer_2, layer_3)

nlayers = NLayerSubtractive(frag=butanol, layers=layers, embedding="mechanical", coupling="covalent")


sp = SP(theory=nlayers)
sp.run()

# 3 Layer embedding from Gaussian
g16_energy = -231.494258144610

# test validation
testutils.validate(sp.result.energy, g16_energy)
