import numpy as np
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 = np.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')     
month = ncin.variables['month'].data.copy()
St = ncin.variables['Tropical_Volcanoe'].data.copy()
Sh = ncin.variables['Hemisphere_Volcanoe'].data.copy()
lat = ncin.variables['lat'].data.copy()
ncin.close()

# Normalise horizontal shapes by time maximum of global integral
area = np.zeros(len(lat))
for n in range(len(lat)):
    if n==0:
       lats = -90
       latn = 0.5*(lat[n]+lat[n+1])
    elif n==len(lat)-1:  
       lats = 0.5*(lat[n]+lat[n-1])
       latn = 90
    else:
       lats = 0.5*(lat[n]+lat[n-1])
       latn = 0.5*(lat[n]+lat[n+1])    
    area[n] = 2*np.pi*6371000**2*(np.sin(latn/180*np.pi)-np.sin(lats/180*np.pi))
St = St/max(sum(np.transpose(St*area)))*1e9 # convert from Tg to kg
Sh = Sh/max(sum(np.transpose(Sh*area)))*1e9 # convert from Tg to kg

# Load the CCSM data for vertical shape
ncin = netcdf.netcdf_file('CCSM4_volcanic_1850-2008_prototype1.nc')     
time = ncin.variables['time'].data.copy()
date = ncin.variables['date'].data.copy()
lev = ncin.variables['lev'].data.copy()
lat = ncin.variables['lat'].data.copy()
date_sec = ncin.variables['datesec'].data.copy()
ccsm_volc = ncin.variables['MMRVOLC'].data.copy()
ccsm_colmass = ncin.variables['colmass'].data.copy()
ncin.close()
 
# Create vertical profile (from matlab code of EXPLAIN)
profile_lat = np.zeros((64,8))
for i in range(8):
    zlev = ccsm_volc[:,i,:]/ccsm_colmass
    profile_lat[:,i] = np.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 = np.int32(np.zeros(new_len))

# define ensemble size and time period for synthetic forcing
NMEM = 60
YEAR1 = 2006
YEARN = 2099
EFREQ = 1. / 2500. / 12. # frequency for eruption event to occur in specific month

# 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 = np.int32(np.zeros(NMONTH))

# re-sample eruption events from reconstruction

# seed random generator 
#   default seed uses system clock but we want results to be reproducible
np.random.seed(0)

# loop over members
for member in range(NMEM):

    # Create empty col_mass variable
    new_colmass = np.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 (np.random.random() <= EFREQ):
                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,:] + np.fliplr(Sh[0:maxlen,:]) * mass[eruption]
    # Create new_mmrvolc from the new_colmn mass as it was created in Explain
    new_mmrvolc = np.zeros((new_time.shape[0],lev.shape[0],lat.shape[0]))
    for i in range(8):
        prof_temp = profile_lat[:,i]
        prof_temp = np.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/volcSynth_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',np.dtype('float64').char,('time',))
    date_out = ncfile.createVariable('date',np.dtype('float64').char,('date',))
    lev_out = ncfile.createVariable('lev',np.dtype('float64').char,('lev',))
    lat_out = ncfile.createVariable('lat',np.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',np.dtype('float64').char,('time','lev','lat'))
    colmass_out = ncfile.createVariable('colmass',np.dtype('float64').char,('time','lat'))
    datesec_out = ncfile.createVariable('datesec',np.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()