"""This script calculates optical properties for various hydrometeors at mm and sub-mm
wavelengths for various space-borne passive applications""" 

from PyARTS import *
from scipy import *
from scipy import special
import copy

f_grid=[19.5e9]
#f_grid=[200.5e9]

def rain_pnds(RR=25,T=0):
    Lambda_r=8.2/(RR**0.21)
    N_0r=16.0*1e3
    A=special.laguerre(n=5)
    #abscissa
    x=real(A.weights[:,0])
    #weights
    w=real(A.weights[:,1])
    r=x/Lambda_r*1e3
    pnds=w*N_0r/Lambda_r
    return r,pnds

def ice_pnds(IWC=0.05,T=255):
    a=0.02#this can be tuned
    #First get Laguerre Gauss weights and abscissa
    A=special.laguerre(n=10)
    x=real(A.weights[:,0])
    w=real(A.weights[:,1])
    r=x/2/a
    n=clouds.mh97(IWC,r,T)[0]
    pnd=n*exp(2*a*r)/2/a*w
    return r,pnd

def liquid_pnds(LWC=0.2,T=280):
    c1=6
    c2=1
    rc=10
    alpha=(c1+c2+1)/c2
    A=special.genlaguerre(n=10,alpha=alpha)
    x=real(A.weights[:,0])
    w=real(A.weights[:,1])
    #turn the abscissa into radii
    B=float(c1)/(c2*rc**c2)
    r=(x/B)**(1/c2)
    A=3*c2*B**((c1+4.0)/c2)*1e12/(4*pi*x**alpha*special.gamma((c1+4.0)/c2))
    pnd_vec=A*w*r**(c1+1-c2)/c2/B
    return r,pnd_vec


def optprops(key,param1,param2):
    r,pnd=hydrometeors[key]['size_dist'](param1,param2)
    K_vec=[]
    Kabs_vec=[]
    for ri in r:
        params=copy.deepcopy(hydrometeors[key]['params'])
        params['equiv_radius']=ri
        scat_data=arts_types.SingleScatteringData(params)
        scat_data.calc()
        K_vec.append(squeeze(scat_data.ext_mat_data))
        Kabs_vec.append(squeeze(scat_data.abs_vec_data))
    K=sum(pnd*K_vec)
    Kabs=sum(pnd*Kabs_vec)
    return K,(K-Kabs)/K
        
        
                    
hydrometeors={
    'rain':{'title':"Typical raindrop at 37GHz (TMI)",
            'params':{'equiv_radius':700,
                      'T_grid':[283],
                      'ptype':20,
                      'NP':-1,
                      'f_grid':f_grid,
                      'phase':'liquid',
                      'aspect_ratio':1.000001},
            'size_dist':rain_pnds,
            'typical_size_dist_params':(25,None)
            },
    'liquid':{'title':"Typical liquid cloud droplet at 200.5 GHz (MLS)",
              'params':{'equiv_radius':15,
                        'T_grid':[275],
                        'ptype':20,
                        'NP':-1,
                        'f_grid':f_grid,
                        'phase':'liquid',
                        'aspect_ratio':1.000001},
            'size_dist':liquid_pnds,
            'typical_size_dist_params':(0.2,280)
              },
    'ice':{'title':"Typical ice cloud crystal at 200.5 GHz (MLS)",
           'params':{'equiv_radius':50,
                     'T_grid':[255],
                     'ptype':20,
                     'NP':-1,
                     'f_grid':f_grid,
                      'phase':'ice',
                     'aspect_ratio':1.000001},
            'size_dist':ice_pnds,
            'typical_size_dist_params':(0.05,255)
           },
    }

for k in ['rain','liquid','ice']:
    scat_data=arts_types.SingleScatteringData(hydrometeors[k]['params'])
    scat_data.calc()
    print
    print "=============="+k+"================"
    print hydrometeors[k]['title']
    # Note GH 2011-05-18: Next line cannot possibly ever have worked
    #print "Refractive index: (%.3f,%.3f)" % (squeeze(scat_data.mrr),squeeze(scat_data.mri))
    print "Extinction cross-section (m^2): %3e" % squeeze(scat_data.ext_mat_data)
    print "Single Scattering Albedo: %.3f" % squeeze((scat_data.ext_mat_data-scat_data.abs_vec_data)/ scat_data.ext_mat_data)
    print
    print "For typical polydispersion"
    r,pnd=hydrometeors[k]['size_dist'](hydrometeors[k]['typical_size_dist_params'][0],
                      hydrometeors[k]['typical_size_dist_params'][1])
    K,albedo=optprops(k,hydrometeors[k]['typical_size_dist_params'][0],
                      hydrometeors[k]['typical_size_dist_params'][1])
    print "Effective radius (microns): %.2f" % (sum(pnd*r**3)/sum(pnd*r**2))
    print "Extinction coefficient (m^-1): %3e" % K
    print "Single Scattering Albedo: %.3f" % albedo
    print