from __future__ import division
##from simtbx.kokkos import gpu_instance
#kokkos_run = gpu_instance(deviceId = 0)

from simtbx.nanoBragg import sim_data
from scitbx.matrix import sqr
from cctbx import uctbx
from dxtbx.model import Crystal

from simtbx.nanoBragg.nanoBragg_crystal import NBcrystal
import numpy as np
import sys
if "--kokkos" in sys.argv:
    import os
    os.environ["DIFFBRAGG_USE_KOKKOS"]="1"
from simtbx.diffBragg.utils import find_diffBragg_instances
from simtbx.diffBragg.device import DeviceWrapper
with DeviceWrapper(0) as _:
    ucell = (70, 60, 50, 90.0, 110, 90.0)
    symbol = "C121"
    a_real, b_real, c_real = sqr(uctbx.unit_cell(ucell).orthogonalization_matrix()).transpose().as_list_of_lists()
    C = Crystal(a_real, b_real, c_real, symbol)

    nbr = NBcrystal()
    nbr.dxtbx_crystal = C

    S = sim_data.SimData(use_default_crystal=True)
    S.crystal = nbr
    S.instantiate_diffBragg(auto_set_spotscale=True)
    S.D.add_diffBragg_spots()
    img = S.D.raw_pixels.as_numpy_array()

    # simulate the primitive cell directly
    to_p1 = C.get_space_group().info().change_of_basis_op_to_primitive_setting()
    Cp1 = C.change_basis(to_p1)
    nbr2 = NBcrystal()
    nbr2.dxtbx_crystal = Cp1

    S2 = sim_data.SimData()
    S2.crystal = nbr2
    S2.instantiate_diffBragg(auto_set_spotscale=True)
    S2.D.add_diffBragg_spots()
    img2 = S2.D.raw_pixels.as_numpy_array()

    # rescale because currently volume is computed incorrectly
    img2 = img2 * S.D.spot_scale / S2.D.spot_scale

    assert S.D.Omatrix == tuple(to_p1.c_inv().r().transpose().as_double())
    assert S2.D.Omatrix == (1, 0, 0, 0, 1, 0, 0, 0, 1)
    assert np.allclose(img, img2)

    print("OK")
    for name in find_diffBragg_instances(globals()): del globals()[name]
