#!/usr/bin/env python
'''
Diffraction functions require for fitting and analyzing data.

mkak 2017.02.06 (originally written spring 2016)
'''


import numpy as np
from numpy import cos, sin, arcsin,  degrees
from .physical_constants import PLANCK_HC, TAU, DEG2RAD, RAD2DEG


def wavelength2energy(wavelength, E_units='keV', lambda_units='A'):
    '''
    Converts lambda into energy
    E = hf ; E = hc/lambda

    E_units      : default keV; can convert to 'eV' if given
    lambda_units : default 'A'; can convert from 'm' or 'nm' if given
    '''
    if lambda_units == 'm':
        wavelength = wavelength*1e10
    elif lambda_units == 'nm':
        wavelength = wavelength*10
    if E_units.lower() == 'kev':
        return PLANCK_HC/wavelength*1e-3 # keV
    else:
        return (PLANCK_HC/wavelength)    # eV


def hkl2d(hkl, a, b, c, alpha, beta, gamma, **kws):
    "convert hkl to d"
    h, k, l = hkl[:, 0], hkl[:, 1], hkl[:, 2]
    alpha, beta, gamma = DEG2RAD*alpha, DEG2RAD*beta, DEG2RAD*gamma
    cos_a, cos_b, cos_g = cos(alpha), cos(beta), cos(gamma)
    x = 1-cos_a**2 - cos_b**2 - cos_g**2 + 2*cos_a*cos_b*cos_g
    y = ((h*sin(alpha)/a)**2 + 2*k*l*(cos_b*cos_g-cos_a)/(b*c) +
         (k*sin(beta)/b)**2  + 2*l*h*(cos_g*cos_a-cos_b)/(c*a) +
         (l*sin(gamma)/c)**2 + 2*h*k*(cos_a*cos_b-cos_g)/(a*b))
    d = np.sqrt(x/y)
    return d


def q2twotheta(q, wavelength, units='degrees'):
    '''
    Converts q axis into 2th (q and wavelength will have inverse units)
    2th = 2*sin^-1(lambda/[2*d])

    ang_unit : default in degrees; will convert to 'rad' if given
    '''
    twoth = 2*arcsin((q*wavelength)/(2*TAU))
    if units.startswith('rad'):
        return twoth
    else:
        return RAD2DEG*twoth


def generate_hkl(hmax=15, kmax=15, lmax=15, positive_only=True):
    if positive_only:
        hklall = np.mgrid[0:hmax+1, 0:kmax+1, 0:lmax+1].reshape(3, -1).T
    else:
        hklall = np.mgrid[-hmax:hmax+1, -kmax:kmax+1, -lmax:lmax+1].reshape(3, -1).T
    return np.array([hkl for hkl in hklall if hkl[0]**2 + hkl[1]**2 + hkl[2]**2 > 0])

#
# def d_from_q(q):
#     '''
#     Converts q axis into d (returned units inverse of provided units)
#     d = 2*PI/q
#     '''
#     return TAU/q
#
# def d_from_twth(twth,wavelength,ang_units='degrees'):
#     '''
#     Converts 2th axis into d (returned units same as wavelength units)
#     d = lambda/[2*sin(2th/2)]
#
#     ang_unit : default in degrees; will convert from 'rad' if given
#     '''
#     if not ang_units.startswith('rad'):
#         twth = DEG2RAD*twth
#     return wavelength/(2*sin(twth/2.))
#
#
# def twth_from_d(d,wavelength,ang_units='degrees'):
#     '''
#     Converts d axis into 2th (d and wavelength must be in same units)
#     2th = 2*sin^-1(lambda/[2*d])
#
#     ang_unit : default in degrees; will convert to 'rad' if given
#     '''
#     twth = 2*arcsin(wavelength/(2.*d))
#     if ang_units.startswith('rad'):
#         return twth
#     else:
#         return RAD2DEG*twth
#
#
# def q_from_d(d):
#     '''
#     Converts d axis into q (returned units inverse of provided units)
#     q = 2*PI/d
#     '''
#     return TAU/d
#
#
# def q_from_twth(twth, wavelength, ang_units='degrees'):
#     '''
#     Converts 2th axis into q (q returned in inverse units of wavelength)
#     q = [(4*PI)/lamda]*sin(2th/2)
#
#     ang_unit : default in degrees; will convert from 'rad' if given
#     '''
#     if not ang_units.startswith('rad'):
#         twth = DEG2RAD*twth
#     return ((2*TAU)/wavelength)*sin(twth/2.)
#
# def qv_from_hkl(hklall, a, b, c, alpha, beta, gamma):
#
#     qv = np.zeros(np.shape(hklall))
#     uvol = unit_cell_volume(a,b,c,alpha,beta,gamma)
#     alpha, beta, gamma = DEG2RAD*alpha, DEG2RAD*beta, DEG2RAD*gamma
#     q0 = [(b*c*sin(alpha))/uvol,(c*a*sin(beta))/uvol,(a*b*sin(gamma))/uvol]
#
#     for i, hkl in enumerate(hklall):
#         qv[i] = [TAU*hkl[0]*q0[0], TAU*hkl[1]*q0[1], TAU*hkl[2]*q0[2]]
#     return qv
#
#
# def unit_cell_volume(a, b, c, alpha, beta, gamma):
#     alpha, beta, gamma = DEG2RAD*alpha, DEG2RAD*beta, DEG2RAD*gamma
#     return a*b*c*(1-cos(alpha)**2-cos(beta)**2-cos(gamma)**2+
#                   2*cos(alpha)*cos(beta)*cos(gamma))**0.5
#
#
#
# def lambda_from_E(E, E_units='keV', lambda_units='A'):
#     '''
#     Converts lambda into energy
#     E = hf ; E = hc/lambda
#
#     E_units      : default keV; can convert from 'eV' if given
#     lambda_units : default 'A'; can convert to 'm' or 'nm' if given
#     '''
#     if E_units.lower() != 'kev':
#         E = E*1e-3 # keV
#     scale = 1.e-3
#     if lambda_units == 'm':
#         scale = 1.e-13
#     elif lambda_units == 'nm':
#         scale = 1e-4
#     return scale*PLANCK_HC/E
#
