import dsd
from disser.units import to_dB, to_dBz, angle, exp_to_dB
from disser.scatter import bulk_scatter
import quantities as pq

d = (np.linspace(0.01, 8, 100).reshape(-1, 1, 1).astype(np.float32) 
        * pq.mm)
qr = (np.linspace(0, 15.0, 40).reshape(1, -1, 1).astype(np.float32)
        * pq.g / pq.m**3)
n = (np.logspace(-2, 6, 110).reshape(1, 1,-1).astype(np.float32)
        / pq.m**3)
dist = dsd.constrained_gamma_from_moments(n, qr, d)

qr_calc = dsd.lwc(d, dist).rescale('g/m**3')
print np.abs(qr - qr_calc).max()
rr = dsd.rainrate(d, dist, dsd.rain_fallspeed(d)).rescale('mm/hr')
print rr.min(), rr.max()

wavelength = 0.053 * pq.m
temp_colors = {0:'b', 10:'g', 20:'y', 30:'r'}
temps = np.array(temp_colors.keys())

d_plot = d.squeeze()

z = np.empty((temps.size,) + dist.shape[1:], dtype=np.float32)
zdr = np.empty_like(z)
atten = np.empty_like(z)
diff_atten = np.empty_like(z)
kdp = np.empty_like(z)
delta = np.empty_like(z)
d = (np.linspace(0.01, 8, 100).reshape(-1, 1, 1).astype(np.float32) 
        * pq.mm)
lam = (np.linspace(0.01, 15, 40).reshape(1, -1, 1).astype(np.float32)
        / pq.mm)
nr_lut = (np.logspace(-2, 6, 110).reshape(1, 1,-1).astype(np.float32)
        / (pq.m**3 * pq.cm))
d0 = (np.linspace(0.01, 4, 40).reshape(1, -1, 1).astype(np.float32)
        * pq.mm)
#dist = dsd.volume_gamma(d, d0, nr_lut * pq.cm, nu=-0.8)
#dist = dsd.modified_gamma(d, lam, nr_lut, dsd.constrained_gamma_shape(lam))
#dist = dsd.modified_gamma(d, lam, nr_lut, 0.8)
#dist = dsd.mp_from_lwc(d, pq.Quantity(10.0, 'g/m^3'))
qr = np.linspace(0.01, 15, 100).reshape(1, -1) * pq.g / pq.m**3
dist = dsd.mp_from_lwc(d, qr)
fallspeed = dsd.rain_fallspeed(d)
print fallspeed.min(), fallspeed.max()
rr = dsd.rainrate(d, dist, fallspeed).rescale('mm/hr')
print rr.min(), rr.max()

for title,wavelength in [('S-band', 10.8*pq.cm),('C-band', 5.3*pq.cm),
        ('X-band', 3.3*pq.cm)]:
    temps = np.arange(0, 35, 10)
    temp_colors = {0:'b', 10:'g', 20:'y', 30:'r'}

    d_plot = d.squeeze()

    z = np.empty((temps.size,) + dist.shape[1:], dtype=np.float32)
    zdr = np.empty_like(z)
    atten = np.empty_like(z)
    diff_atten = np.empty_like(z)