/*
* Copyright 2021 Gao's lab, Peking University, CCME. All rights reserved.
*
* NOTICE TO LICENSEE:
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
* http://www.apache.org/licenses/LICENSE-2.0
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

#ifndef XRD3D_H
#define XRD3D_H
#include "../common.cuh"
#include "../control.cuh"
#include <cufft.h>

//相关拓展阅读
//https://en.wikipedia.org/wiki/X-ray_crystallography
//https://en.wikipedia.org/wiki/Powder_diffraction

/*
*采用的公式定义为I=1/N*sum(sum(sin(k0*r)/k0/r,N),N),
*其中sin(k0*r)/k0/r按无限周期求和会得到狄拉克delta函数
*因此实际操作是加上了有限的截断，为了连续性和方便公式推导
*有限截断反应为衰减因子s，sin(k0*r)/k0/r*exp(-s*r*r)
*由于sin(k0*r)/k0/r形式的势场不好像PME一样分解成近程和远程两部分
*因此整个作用计算只是在倒空间进行。由于sin(k0*r)/k0/r*exp(-s*r*r)
*的傅里叶变换结果Pi^1.5/k/k0/sqrt(s)*(exp(-(k-k0)^2./4./s)-exp(-(k+k0)^2./4./s))
*在k空间收敛迅速（在远离k0处高斯衰减），因此只要格点间距小于0.25*Pi/k0对应的长度
*就能获得比较准确的结果。根据一般k0的选取，格点间距在0.5 埃左右合适。
*/
struct XRD3D
{
	char module_name[CHAR_LENGTH_MAX];
	int is_initialized = 0;
	int is_controller_printf_initialized = 0;
	int last_modify_date = 20210531;

	//XRD3D计算的物理参数
	float theta = 0.;//衍射角,角度制
	float lambda = 0.;//X-ray波长
	float k0 = 0.;//4*PI/lambda * sinf(theta * PI/180);
	float s = 1.;//物理上XRD3D的衰减因子exp(-s*r*r)

	//注意，由于XRD3D计算很多时候只需要部分原子参与，
	//因此atom_numbers不等于体系真正的atom_numbers
	//在不特别指出s或system时，所有atom均代表xrd3d atom，xatom
	int system_atom_numbers=0;//整个分子模拟系统的所有原子数目
	int xatom_numbers = 0;//需要进行XRD3D计算的原子数目
	int *h_atom_serial = NULL;//需要进行XRD3D原子的原子序号，长度xatom_numbers
	int *d_atom_serial = NULL;
	float *h_fi = NULL;//需要进行XRD3D原子的原子形状因子
	float *d_fi = NULL;
	int *h_satom2xatom_serial = NULL;//体系所有原子向XRD3D原子的映射序号表，等于-1的体系原子不做XRD3D，长度satom_numbers=system_atom_numbers
	int *d_satom2xatom_serial = NULL;
	UNSIGNED_INT_VECTOR *atom_uint_crd = NULL;//只针对XRD3D原子的坐标和力（Cv梯度）记录，长度xatom_numbers
	VECTOR *atom_frc = NULL;


	//体系周期性相关记录
	VECTOR box_length;
	float volume = 0.;
	INT_VECTOR grid_dimension;
	int grid_numbers = 0;
	int layer_numbers = 0;
	
	//能量相关
	float *d_sum_of_reciprocal_energy = NULL;
	float h_sum_of_reciprocal_energy = 0.;
	//能量对system原子坐标的梯度
	VECTOR *system_atom_crd_grad = NULL;

	struct FFT//使用倒空间求和所需要的变量和参数
	{
		XRD3D *xrd3d = NULL;
		cufftResult cuda_result;//在初始fft plan时监控初始化是否成功
		cufftHandle plan_3d_r2c;
		cufftHandle plan_3d_c2r;
		int NzNy_Nx21_numbers=0;
		INT_VECTOR grid_dimension_lower_half_plus_one;
		INT_VECTOR grid_dimension_upper_half;
		VECTOR box_length_inverse;
		VECTOR crd_to_grid_serial_factor;
		VECTOR uint_crd_to_grid_serial_factor;
		float scalor=0.;//Pi^1.5/V/k0/sqrt(s)/N,构建的fft_fijk列表乘上此系数可保证
						//最终的结果为直接求和的结果

		float *fft_fijk = NULL;//用于与fft_rho对应相乘再fft获得phi，NzNy(Nx/2+1)

		float *h_B_spline_grid = NULL;//四阶B样条线插值卷积时的格点权重分布：{...，0.，1./6.，2./3.，1./3.，0.，...}->{...,-2,-1,0,1,2,...}
		float *d_B_spline_grid = NULL;//NzNyNx
		cufftComplex *fft_B_spline = NULL;//上一变量的三维傅里叶变换：NzNy(Nx/2+1)

		float *rho = NULL;//xrd原子在空间中的分布离散化情况，进行过四阶B样条线插值，NzNyNx
		cufftComplex *fft_rho = NULL;//上一变量的三维傅里叶变换，NzNy(nX/2+1)
		float *phi = NULL;//xrd原子在空间中产生的场的情况，已经被四阶B样条线插值过的分布产生的，NzNyNx


		//四阶B样条线插值所需变量
		dim3 thread_XRD3D;//计算四阶B样条线插值时GPUkernel参数
		int **XRD3D_atom_near;
		UNSIGNED_INT_VECTOR *XRD3D_kxyz=NULL;//初始化后不再改变
		UNSIGNED_INT_VECTOR *XRD3D_uxyz = NULL;
		VECTOR *XRD3D_frxyz = NULL;

		//fft内部函数，不要随意调用
		void Initial(CONTROLLER *controller, XRD3D *xrd3d);
		void FFT_One_Time();//在刷新需要xrd的atom_uint_crd列表后
							//运行此函数，可获得Cv：d_sum_of_reciprocal_energy
							//和对应的对体系全原子坐标的Cv梯度：system_atom_crd_grad
		void Refresh_FFT_Volume();//在更新了box_length后，进行fft内部相关体积更新
	}fft;


	void Initial(CONTROLLER *controller, const int system_atom_numbers, const VECTOR boxlength, char *module_name = NULL);
	void Calculate_XRD3D_And_Grad(const UNSIGNED_INT_VECTOR *uint_crd);
	float XRD3D_Cv(const UNSIGNED_INT_VECTOR *uint_crd,int is_download=1);
	void Refresh_Volume(const VECTOR boxlength);
	//一些小函数
	float Get_k(float theta, float lambda);
};

#endif //XRD3D_H(XRD3D.cuh)