import pygedm
import numpy as np

skip_ddm = False 
skip_dmd = False
skip_xyz = False

gl = np.linspace(-180, 180, 360*2+1)
gb = np.linspace(-90, 90, 180*2+1)
dist = np.arange(0, 30, 1)
dist[0] = 0.1
dml   = np.array((1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000))

np.save('datadim_gl.npy', gl)
np.save('datadim_gb.npy', gb)   
np.save('datadim_dist.npy', dist)
np.save('datadim_dm.npy', dml)   

ddm_data_ymw = np.zeros((len(dist), len(gb), len(gl)))
ddm_data_ne = np.zeros((len(dist), len(gb), len(gl)))
ddm_data_ymw_tau = np.zeros((len(dist), len(gb), len(gl)))
ddm_data_ne_tau = np.zeros((len(dist), len(gb), len(gl)))

dmd_data_ymw = np.zeros((len(dist), len(gb), len(gl)))
dmd_data_ne = np.zeros((len(dist), len(gb), len(gl)))
dmd_data_ymw_tau = np.zeros((len(dist), len(gb), len(gl)))
dmd_data_ne_tau = np.zeros((len(dist), len(gb), len(gl)))

if not skip_xyz:
    # Create empty arrays for NE2001 and YMW16 model output
    d_2016 = np.zeros((100, 100, 100))
    d_2001 = np.zeros((100, 100, 100))

    # Loop through pixels 
    for ii in range(100):
        if ii % 10 == 0: print(ii)
        for jj in range(100):
            for kk in range(100):
                # generate X, Y, Z in PC from indexes
                x, y, z = (ii - 50) * 400, (jj - 50) * 400, (kk - 50) * 400

                # Compute EDM
                d_2016[ii, jj, kk] = pygedm.calculate_electron_density_xyz(x, y, z, method='ymw16').value
                d_2001[ii, jj, kk] = pygedm.calculate_electron_density_xyz(x, y, z, method='ne2001').value

    np.save('ne2001_xyz.npy', d_2001)
    np.save('ymw16_xyz.npy', d_2016)

    ii = np.array(range(100))
    jj = np.array(range(100))
    kk = np.array(range(100))

    x, y, z = (ii - 50) * 400, (jj - 50) * 400, (kk - 50) * 400
    np.save('xyz_x.npy', x)
    np.save('xyz_y.npy', y)
    np.save('xyz_z.npy', z)


if not skip_ddm:
    for igl, ll in enumerate(gl):
        print(ll)
        for igb, bb in enumerate(gb):
            for idist, dd in enumerate(dist):
                dd *= 1000
                dm, tau = pygedm.dist_to_dm(ll, bb, dd, method='ymw16')
                ddm_data_ymw[idist, igb, igl] = dm.value
                ddm_data_ymw_tau[idist, igb, igl] = tau.value
                dm, tau = pygedm.dist_to_dm(ll, bb, dd, method='ne2001')
                ddm_data_ne[idist, igb, igl] = dm.value
                ddm_data_ne_tau[idist, igb, igl] = tau.value

    np.save('datacube_ddm_ymw16.npy', ddm_data_ymw)
    np.save('datacube_ddm_ne2001.npy', ddm_data_ne)   
    np.save('datacube_ddm_tau_ymw16.npy', ddm_data_ymw_tau)
    np.save('datacube_ddm_tau_ne2001.npy', ddm_data_ne_tau)   

    
if not skip_dmd:
    for igl, ll in enumerate(gl):
        print(ll)
        for igb, bb in enumerate(gb):
            for idm, dd in enumerate(dml):
                dist, tau =  pygedm.dm_to_dist(ll, bb, dd, method='ymw16')
                dmd_data_ymw[idm, igb, igl] = dist.value
                dmd_data_ymw_tau[idm, igb, igl] = tau.value

                dist, tau =  pygedm.dm_to_dist(ll, bb, dd, method='ne2001')
                dmd_data_ne[idm, igb, igl] = dist.value
                dmd_data_ne_tau[idm, igb, igl] = tau.value
                
    np.save('datacube_dmd_ymw16.npy', dmd_data_ymw)
    np.save('datacube_dmd_ne2001.npy', dmd_data_ne)   
    np.save('datacube_dmd_tau_ymw16.npy', dmd_data_ymw_tau)
    np.save('datacube_dmd_tau_ne2001.npy', dmd_data_ne_tau)