## nuclear weapon fallout v 0007
## using python glasstone library 2017

import matplotlib.pyplot as plt
from matplotlib import cm, colors, colorbar
import matplotlib.lines as mlines
import matplotlib.patches as mpatches
import numpy as np
from scipy.stats import lognorm
from glasstone.fallout import WSEG10
             
## 1 dose , 2 mortality , 2 dose 1, 2, ...
      
#picname="fallout_1hr_doserate_1mt_burst_15mph.svg"      
 
picname="fallout_30days_1mt_burst_15mph.svg"
 
tyyppi=1

mitta=1100.0

#mitta=360.0

### if pic not drawn, select this 4 or greater, not 1
## grid resolution
mdx=4

# yield in megatons
yld = 1

## wind speed m/s
ms=6.7

# fission fraction
ff = 1.0

# wind direction (in degrees with wind from north = 0)
wind_direction = 270
# wind shear (change in mph per kilofoot change in altitude)
wind_shear = 0.25


if tyyppi==1:
	paletti=['b','g','y','y', '#ff9900', '#ff9900','r','r']
	kontours=[0.0005,0.05,0.1,0.3,1.5,3,5,10]
	kontkoeff=1.0
	laabeli1="Sv"

if tyyppi==2:
	paletti=('g', '#ff9900','r','#7F0000')
	kontours=[0.00000001,0.05, 0.5,1.0]
	kontkoeff=100.0
	laabeli1="%" 

if tyyppi==3:
	paletti=['b','g','y','y', '#ff9900', '#ff9900','r','r']
	kontours=[0.0005,0.05,0.1,0.3,1.5,3,5,10]
	kontkoeff=1.0
	laabeli1="Sv/h"

if tyyppi==4:
	kontours=[0.25,0.5,1,2,4,8,16]
	paletti=['g','g','g','y','y','r','r','r']
	kontkoeff=1.0
	laabeli1="Sv"


kaptitail=" MT burst "+ str(round(ff*100))+ "% fission"

if tyyppi==1:
	kaptibase="WSEG-10 30-day total dose contours for "

if tyyppi==2:
	kaptibase="WSEG-10 30-day total death rate contours for "

if tyyppi==3:
	kaptibase="WSEG-10 1-hour dose rate contours for "
	
if tyyppi==4:
	kaptibase="WSEG-10 30-day total dose contours for "

kaptioni=kaptibase+str(yld)+ kaptitail 


wind_speed=ms/0.44704

maili=1.609

# ground zero x & y locations (st. mi)
gzx = 20
gzy=(mitta/2)
ymax = (mitta*maili)+10

## wind speed (mph)

shear_si=round((wind_shear*0.44707*3.2808),3)

leka=[]
pitu=len(kontours)
for n in range(0, pitu):
	print n, kontours[n], paletti[n]
	meme=mlines.Line2D([], [], color=paletti[n], linewidth=5, marker='*',markersize=1, label=str(kontours[n]*kontkoeff)+ " "+laabeli1)
	leka.append(meme)

#print "Koebreak"		
#exit(1)		
	

print "TYP ", tyyppi 
print "Wait ..."

x = np.arange(-1, mitta, mdx)
y = np.arange(-1, mitta, mdx)
X, Y = np.meshgrid(x, y)
# use WSEG10's native units

w = WSEG10(gzx, gzy, yld, ff, wind_speed, wind_direction, wind_shear, dunits='mi', wunits='mph', yunits='MT', shearunits='mph/kilofoot')
dose = np.vectorize(w.dose)
doserate = np.vectorize(w.D_Hplus1)

def fatality_fraction(x, y):
    erd = w.dose(x, y, dunits='mi', doseunits='Sv')
    if erd > 20.0:
        return 1.01
    else:
        return lognorm.cdf(erd, 0.42, scale=4.5)
    
deaths = np.vectorize(fatality_fraction)

if tyyppi==1:
	Z = dose(X, Y, dunits='mi', doseunits='Sv')

if tyyppi==2:
	Z = deaths(X, Y)

if tyyppi==3:
	Z = doserate(X, Y, dunits='mi', doseunits='Sv')
	
if tyyppi==4:
	Z = dose(X, Y, dunits='mi', doseunits='Sv')


fig = plt.figure()
ax1 = fig.add_axes([0.1, 0.1, 0.9, 0.9])

x = np.arange(-1, mitta, mdx)
y = np.arange(-1, mitta, mdx)

for n in range(1,len(x)):
	x[n]=x[n]*maili
	y[n]=y[n]*maili    

X, Y = np.meshgrid(x, y)

CS = ax1.contour(X, Y, Z, kontours, colors=paletti, linewidths=2.5)

ax1.grid(True)

ax1.text(15, 80, kaptioni, size=14)

outs1="  Wind "+str(ms)+ " m/s shear "+str(shear_si)+" m/s per km"

ax1.text(15, 20, outs1, size=14)

ax1.set_ylim([-0.5, ymax])
ax1.set_ylabel("km")
ax1.set_xlabel("km")
ax1.xaxis.label.set_size(14)
ax1.yaxis.label.set_size(14)

ax1.legend(handles=leka)
plt.show()
fig.savefig(picname, format="svg")