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

from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument("--plot", action='store_true')
parser.add_argument("--crystalsystem", default='tetragonal',
                    choices=["monoclinic", "tetragonal"])
parser.add_argument("--curvatures", action='store_true')
parser.add_argument("--kokkos", action="store_true")
args = parser.parse_args()

if args.kokkos:
    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 _:

    import numpy as np
    from scipy.spatial.transform import Rotation
    from scipy.stats import pearsonr
    import pylab as plt
    from scipy.stats import linregress
    from simtbx.diffBragg.hopper_utils import UCELL_ID_OFFSET
    from scitbx.matrix import sqr
    from scitbx.matrix import rec
    from simtbx.nanoBragg.nanoBragg_crystal import NBcrystal
    from cctbx import uctbx
    from simtbx.nanoBragg.sim_data import SimData
    from simtbx.diffBragg import utils
    from dxtbx.model.crystal import Crystal


    # STEP 1:
    # make a crystal and orient it randomly
    if args.crystalsystem=="tetragonal":
        ucell = (55, 55, 77, 90, 90, 90)
        symbol = "P43212"
    else:  # args.crystalsystem == "monoclinic"
        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)

    # random raotation
    rotation = Rotation.random(num=1, random_state=101)[0]
    Q = rec(rotation.as_quat(), n=(4, 1))
    rot_ang, rot_axis = Q.unit_quaternion_as_axis_and_angle()
    C.rotate_around_origin(rot_axis, rot_ang)

    # STEP3:
    # create the unit cell parameter manager
    UcellMan = utils.manager_from_params(ucell)
    n_ucell_params = len(UcellMan.variables)
    assert np.allclose(UcellMan.B_recipspace, C.get_B())

    # STEP4:
    # make a nanoBragg crystal to pass to diffBragg
    nbcryst = NBcrystal(init_defaults=True)
    nbcryst.dxtbx_crystal = C
    nbcryst.n_mos_domains = 1
    nbcryst.thick_mm = 0.01
    nbcryst.Ncells_abc = (7, 7, 7)

    # STEP5: make an instance of diffBRagg, use the simData wrapper
    SIM = SimData()
    # overwrite the default detector to use smaller pixels
    img_sh = 700,700
    SIM.detector = SimData.simple_detector(300, 0.1, img_sh)
    SIM.crystal = nbcryst
    SIM.instantiate_diffBragg(oversample=0, verbose=0, auto_set_spotscale=True)
    # D is an instance of diffBragg with sensible parameters
    # and our dxtbx crystal created above
    D = SIM.D
    D.progress_meter = True

    # STEP6:
    # initialize the derivative managers for the unit cell parameters
    for i_param in range(n_ucell_params):
        D.refine(UCELL_ID_OFFSET+i_param)
    for i in range(n_ucell_params):
        D.set_ucell_derivative_matrix(UCELL_ID_OFFSET+i, UcellMan.derivative_matrices[i])
        if args.curvatures:
            D.set_ucell_second_derivative_matrix(UCELL_ID_OFFSET + i, UcellMan.second_derivative_matrices[i])
    D.initialize_managers()


    # STEP7:
    # compute the scattering and its derivative
    print("Adding diffBragg spots")
    #D.printout_pixel_fastslow =150, 351
    D.add_diffBragg_spots()
    print("Done!")
    img = D.raw_pixels_roi.as_numpy_array()
    # reset all pixel values
    D.raw_pixels *= 0
    D.raw_pixels_roi *= 0

    derivs = []
    second_derivs = []
    for i_param in range(n_ucell_params):
        analy_deriv = D.get_derivative_pixels(UCELL_ID_OFFSET+i_param).as_numpy_array()
        derivs.append(analy_deriv)
        if args.curvatures:
            second_derivs.append(D.get_second_derivative_pixels(UCELL_ID_OFFSET+i_param).as_numpy_array())

    # STEP8
    # iterate over the parameters and do a finite difference test for each one
    # parameter shifts:
    shifts = [1e-4*(2*i) for i in range(1, 12, 2)]

    import copy
    starting_var = copy.copy(UcellMan.variables)
    for i_param in range(n_ucell_params):
        analy_deriv = derivs[i_param]
        cc_vals = []
        cc_vals2 = []
        error = []
        error2 = []
        h_vals = []
        for i_shift, percent_shift in enumerate(shifts):

            var = copy.copy(starting_var)

            param_shift = var[i_param] * percent_shift

            var[i_param] += param_shift
            UcellMan.variables = var

            D.Bmatrix = UcellMan.B_recipspace
            D.add_diffBragg_spots()

            img_forward = D.raw_pixels_roi.as_numpy_array()

            # reset for next computation
            D.raw_pixels_roi *= 0
            D.raw_pixels *= 0

            if args.curvatures:
                # estimate the second derivative
                var[i_param] = var[i_param] - 2*param_shift  # do the backwards finite deriv
                UcellMan.variables = var

                D.Bmatrix = UcellMan.B_recipspace
                D.add_diffBragg_spots()

                img_backward = D.raw_pixels_roi.as_numpy_array()

                # reset for next computation
                D.raw_pixels_roi *= 0
                D.raw_pixels *= 0

            finite_deriv = (img_forward-img) / param_shift

            if second_derivs:
                finite_second_deriv = (img_forward - 2*img + img_backward) / param_shift / param_shift

            bragg = img > 0.5

            ave_error = np.abs(finite_deriv[bragg] - analy_deriv[bragg]).mean()

            r = pearsonr(analy_deriv[bragg].ravel(), finite_deriv[bragg].ravel())[0]
            cc_vals.append(r)

            error.append(ave_error)
            h_vals.append( param_shift)
            print ("\tAverage error=%f; parameter shift h=%f" % (ave_error, abs(param_shift)))
            if args.curvatures:
                ave_error2 = np.abs(finite_second_deriv[bragg] - second_derivs[i_param][bragg]).mean()
                print("\tsecond derivative Average error=%f; parameter shift squared h^2=%f"
                      % (ave_error2, abs(param_shift)**2))
                r2 = pearsonr(second_derivs[i_param][bragg].ravel(), finite_second_deriv[bragg].ravel())[0]

                error2.append(ave_error2)
                cc_vals2.append(r2)
            if args.plot:
                plt.subplot(121)
                plt.imshow(finite_deriv)
                plt.title("finite diff")
                plt.subplot(122)
                plt.imshow(analy_deriv)
                plt.title("analytical")
                plt.draw()
                plt.suptitle("Shift %d / %d"
                             % (i_shift+1, len(shifts)))
                plt.pause(0.8)
                if args.curvatures:
                    plt.subplot(121)
                    plt.imshow(finite_second_deriv)
                    plt.title("finite second diff")
                    plt.subplot(122)
                    plt.imshow(second_derivs[i_param])
                    plt.title("analytical")
                    plt.draw()
                    plt.suptitle("Shift %d / %d"
                                 % (i_shift + 1, len(shifts)))
                    plt.pause(0.8)

        l = linregress(h_vals, error)

        print ("finite diff l.rvalue=%10.7g" % l.rvalue)
        assert l.rvalue > .99
        assert l.slope > 0
        assert l.pvalue < 1e-6

        if args.curvatures:
            l2 = linregress(np.array(h_vals)**2, error2)
            print ("finite 2nd diff l.rvalue=%10.7g" % l2.rvalue)
            assert l2.rvalue > .99
            assert l2.slope > 0
            assert l2.pvalue < 1e-6

        if args.plot:
            plt.close()
            plt.plot(shifts, cc_vals, 'o')
            title = "Unit cell parameter %d / %d" % (i_param+1, n_ucell_params)
            plt.title(title + "\nPearson corr between finite deriv and analytical")
            plt.xlabel("unit cell shifts")
            plt.ylabel("Pearson corr")
            plt.show()
            if args.curvatures:
                plt.close()
                plt.plot(np.array(shifts)**2, cc_vals2, 'o')
                title = "Unit cell parameter %d / %d" % (i_param + 1, n_ucell_params)
                plt.title(title + "\nPearson corr between finite second deriv and analytical")
                plt.xlabel("unit cell shifts")
                plt.ylabel("Pearson corr")
                plt.show()

        # verify a high correlation for the smallest parameter shift
        print("Check high pearson R between analytical and finite diff")
        print("Pearson correlection at smallest parameter shift=%f" % cc_vals[0])
        assert(cc_vals[0] > .98), "%f" % cc_vals[0]
        # check monotonic decrease
        print("Fit polynomial and check monotonic decrease")
        trend = np.polyval(np.polyfit(shifts, cc_vals, 2), shifts)
        assert np.all(np.diff(list(zip(trend[:-1], trend[1:])), axis=1) <= 0)
        if args.curvatures:
            assert (cc_vals2[0] > .99), "%f" % cc_vals2[0]
    print("OK!")
    for name in find_diffBragg_instances(globals()): del globals()[name]
