import numpy
import random 
#from netCDF4 import Dataset
from scipy.io import netcdf

# Read peak stratospheric volcanic aerosol mass loads 
#   These are derived from sulphate depositions of Sigl et al. 2015 and
#   converted to stratospheric aerosol mass using the Gao et al. 2007 scaling
volc_data = numpy.genfromtxt('nature14565-s6_stratAerosolMass.txt')
year = volc_data[:,0]   # extract years
flag = volc_data[:,1]   # extract location flag (1=tropical, 2=NH, 3=SH)
mass = volc_data[:,2]   # extract stratospheric volcanic aerosol mass (Tg)

# Load the CCSM data for horizontal shape 
ncin = netcdf.netcdf_file('volc_samples.nc','r')     
month = ncin.variables['month'][:]
St = ncin.variables['Tropical_Volcanoe'][:,:]
Sh = ncin.variables['Hemisphere_Volcanoe'][:,:]
lat = ncin.variables['lat'][:]
ncin.close()

# Normalise horizontal shapes by time maximum of global integral
area_earth = 4*numpy.pi*6371000**2
dlat = 2.79;
weights = numpy.cos(lat/180.*numpy.pi)
weights = weights/sum(weights)*area_earth
St = St/max(sum(numpy.transpose(St*weights)))*1e9 # convert from Tg to kg
Sh = Sh/max(sum(numpy.transpose(Sh*weights)))*1e9 # convert from Tg to kg

# Load the CCSM data for vertical shape
ncin = netcdf.netcdf_file('CCSM4_volcanic_1850-2008_prototype1.nc','r')     
time = ncin.variables['time'][:]
date = ncin.variables['date'][:]
lev = ncin.variables['lev'][:]
lat = ncin.variables['lat'][:]
date_sec = ncin.variables['datesec'][:]
ccsm_volc = ncin.variables['MMRVOLC'][:,:,:]
ccsm_colmass = ncin.variables['colmass'][:,:]
ncin.close()
 
# Create vertical profile (from matlab code of EXPLAIN)
profile_lat = numpy.zeros((64,8))
for i in range(8):
    zlev = ccsm_volc[:,i,:]/ccsm_colmass
    profile_lat[:,i] = numpy.nanmean(zlev, axis=0)

################## Begin constructing artificial forcing ##################
# Create time variable
start_year = 2006
new_time = time[0:(2100-start_year)*12+2] + (start_year - 1850)
new_date = date[0:(2100-start_year)*12+2] + ((start_year - 1850) * 10000)
new_len = new_time.shape[0]
new_datesec = numpy.int32(numpy.zeros(new_len))

# define ensemble size and time period for synthetic forcing
NMEM = 60
YEAR1 = 2006
YEARN = 2099

# Create time variable
new_time = time[0:(YEARN-YEAR1+1)*12+2] + (YEAR1 - 1850)
new_date = date[0:(YEARN-YEAR1+1)*12+2] + ((YEAR1 - 1850) * 10000)
NMONTH = new_time.shape[0]
new_datesec = numpy.int32(numpy.zeros(NMONTH))

# re-sample eruption events from reconstruction

# loop over members
numpy.random.seed(0)
P=1./2500./12.
for member in range(NMEM):

    # Create empty col_mass variable
    new_colmass = numpy.zeros((NMONTH,ccsm_colmass.shape[1]))

    # loop over months 
    for m in range(NMONTH):

        # loop over eruptions in Sigl el al. 2015 
        for eruption in range(len(mass)):

            # check if eruption is triggered and add to column mass if yes
            if ( numpy.random.rand() < P ):
                maxlen=min(St.shape[1],NMONTH-m)
               
                if flag[eruption] == 1: # tropical eruption 
                    new_colmass[m:m+maxlen,:] = new_colmass[m:m+maxlen,:] + St[0:maxlen,:] * mass[eruption]
                elif flag[eruption] == 2: # NH eruption      
                    new_colmass[m:m+maxlen,:] = new_colmass[m:m+maxlen,:] + Sh[0:maxlen,:] * mass[eruption]
                elif flag[eruption] == 3: # SH eruption      
                    new_colmass[m:m+maxlen,:] = new_colmass[m:m+maxlen,:] + numpy.fliplr(Sh[0:maxlen,:]) * mass[eruption]

    # Create new_mmrvolc from the new_colmn mass as it was created in Explain
    new_mmrvolc = numpy.zeros((new_time.shape[0],lev.shape[0],lat.shape[0]))
    for i in range(8):
        prof_temp = profile_lat[:,i]
        prof_temp = numpy.tile(prof_temp,(new_time.shape[0],1))
        new_mmrvolc[:,i,:] = new_colmass * prof_temp

    # Output new NorESM forcing file
    # Determine size of dimensions
    ntime = new_time.shape[0]
    nlev = lev.shape[0]
    nlat = lat.shape[0]

    # open a new netCDF file for writing.
    fnout = 'Data/v4f_volcanic_v5_mem' + str(member+1).zfill(2) + '.nc'
    ncfile = netcdf.netcdf_file(fnout,'w')

    # create the lat and lon dimensions.
    ncfile.createDimension('time',None)
    ncfile.createDimension('date',ntime)
    ncfile.createDimension('lev',nlev)
    ncfile.createDimension('lat',nlat)

    # Define the coordinate variables.
    time_out = ncfile.createVariable('time',numpy.dtype('float64').char,('time',))
    date_out = ncfile.createVariable('date',numpy.dtype('float64').char,('date',))
    lev_out = ncfile.createVariable('lev',numpy.dtype('float64').char,('lev',))
    lat_out = ncfile.createVariable('lat',numpy.dtype('float64').char,('lat',))

    # Assign units attributes to coordinate variable data.
    date_out.units = 'yyyymmdd'
    date_out.FillValue = -9999.
    lev_out.units = 'hPa'
    lat_out.units = 'degrees_north'

    # write data to coordinate vars.
    time_out[:] = new_time
    date_out[:] = new_date
    lev_out[:] = lev
    lat_out[:] = lat

    # create main variables
    mmrvolc_out = ncfile.createVariable('MMRVOLC',numpy.dtype('float64').char,('time','lev','lat'))
    colmass_out = ncfile.createVariable('colmass',numpy.dtype('float64').char,('time','lat'))
    datesec_out = ncfile.createVariable('datesec',numpy.dtype('int32').char,('time',))


    # set the units attribute.
    mmrvolc_out.units = 'kg kg-1'
    mmrvolc_out.FillValue = -9999.
    mmrvolc_out.long_name = 'layer volcanic aerosol mass mixing ratio'
    colmass_out.FillValue = -999.
    datesec_out.FillValue = -999

    # write data to variables along record (unlimited) dimension.
    mmrvolc_out[:,::] = new_mmrvolc[:,:]
    colmass_out[:,::] = new_colmass[:,:]
    datesec_out[:] = new_datesec[:]

    # close the file.
    ncfile.close()