import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

from backend.latmathcher import PipelineLatMatch
from backend.db_utils.parse_c2db import from_c2db_structure, read_c2db_json
from backend.db_utils.utils import structure_ato_list
from backend.latmathcher.plots import plot_atom_list
from backend.latmathcher import atoms_to_greed
import matplotlib.pyplot as plt
import matplotlib.patches as patches



filepathA = "EXAMPLES/WSe2-1cfbe6183886.json"
filepathB = "EXAMPLES/C2-a6735a4a3797.json"


A_structure = from_c2db_structure(read_c2db_json(filepathA))
B_structure = from_c2db_structure(read_c2db_json(filepathB))

min_supercel = PipelineLatMatch(A_structure["cell"],B_structure["cell"],
                                Aatoms3D=structure_ato_list(A_structure),
                                Batoms3D=structure_ato_list(B_structure),
                                dim=(10, 10), optimize_angle=True, optimize_strain=True)


print(min_supercel.Alat)
print(min_supercel.Blat)
new_structure = min_supercel.get_new_structure()



super_xyz = structure_ato_list(new_structure)

atoms = atoms_to_greed(super_xyz, lat_v=min_supercel.sc_vec3, dim=(10,10,0))
print(atoms)
plot_atom_list(atoms, marker=".")
x = [0.0,min_supercel.sc_vec3[0][0],min_supercel.sc_vec3[0][0]+min_supercel.sc_vec3[0][1],min_supercel.sc_vec3[0][1],0.0]
y = [0.0,min_supercel.sc_vec3[1][0],min_supercel.sc_vec3[1][0]+min_supercel.sc_vec3[1][1],min_supercel.sc_vec3[1][1],0.0]
plt.gca().add_patch(patches.Polygon(xy=list(zip(x,y)), fill=True, alpha=0.4,color="green"))
plt.ylim(min(y) - 10, max(y) + 10)
plt.xlim(min(x) - 10, max(x) + 10)

super_a = [super_xyz[i] for i in range(len(new_structure["host_guest"])) if new_structure["host_guest"][i] == "host"]
super_b = [super_xyz[i] for i in range(len(new_structure["host_guest"])) if new_structure["host_guest"][i] == "guest"]

plt.show()
print(min_supercel.sc_vec)
print(min_supercel.Alat)
print(min_supercel.Blat)
print(new_structure)
plt.show()