#!/usr/bin/env python
# coding: utf-8

import numpy as np
import netCDF4 as nc
import pylab as plt
import matplotlib
import numpy.ma as ma
from matplotlib import rcParams
import datetime
import cartopy.crs as ccrs
from cartopy.util import add_cyclic_point
import matplotlib.dates as mdates
import pandas as pd
from scipy.optimize import curve_fit
import cftime
import nc_time_axis
import xarray as xr

rcParams['font.size'] = 7

# In[2]:
var_lab = 'CH$_4$'
#output for plots
output_write = '/home/users/zosiast/jasmin_plots/no_anthro_ensemble_ssp370/output_plots_jan/'

#constants
mr_ch4 = 16.
mr_co = 28.
mr_o3 = 48.
per_sec_to_per_yr = 60*60*24*360
g_to_Tg = 1e12
n_a = 6.022e23

#stash codes
stash_ch4 = 'mass_fraction_of_methane_in_air'
stash_co = 'mass_fraction_of_carbon_monoxide_in_air'
stash_ch4_oh = 'm01s50i041'
stash_oh = 'mass_fraction_of_hydroxyl_radical_in_air'
stash_o3 = 'mass_fraction_of_ozone_in_air'
stash_trop = 'm01s50i062'
stash_mass = 'm01s50i063'
stash_temp = 'surface_temperature'

# # 2015-2050: net zero anthro

# ### u-by186

suite_id = 'u-by186'
suite_lab = 'uby186'
dates = '2015_2050'

data_1 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_{dates}.nc')
#extract variables to arrays
ch4_1 = data_1.variables[stash_ch4][:]*28.97/mr_ch4*1e9
lat_1 = data_1.variables['latitude'][:]
lon_1 = data_1.variables['longitude'][:]

time_1 = data_1.variables['time']
cf_dtime_1 = cftime.num2date(time_1[:],time_1.units, time_1.calendar)
dtime_1 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_1])

#tropospheric mask
data_trop_1 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_trop_mask_{dates}.nc')
trop_1 = data_trop_1.variables[stash_trop][:]

#air mass
data_mass_1 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_airmass_{dates}.nc')
mass_1 = data_mass_1.variables[stash_mass][:]

#temperature starts in 2014
data_temp_1 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_surf_temp_2015_2050.nc')
temp_1 = data_temp_1.variables[stash_temp][:]

#ch4_oh data
data_ch4_oh = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_oh_flux_{dates}.nc')
ch4_oh_1 = data_ch4_oh.variables[stash_ch4_oh][:]

#o3_conc
data_o3 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_o3_{dates}.nc')
o3_1 = data_o3.variables[stash_o3][:]

#oh_conc
data_oh = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_oh_{dates}.nc')
oh_1 = data_oh.variables[stash_oh][:]

# ### u-bz146

suite_id = 'u-bz146'
suite_lab = 'ubz146'
dates = '2015_2050'

data_146 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_{dates}.nc')
#extract variables to arrays
ch4_146 = data_146.variables[stash_ch4][:]*28.97/mr_ch4*1e9
lat_146 = data_146.variables['latitude'][:]
lon_146 = data_146.variables['longitude'][:]

time_146 = data_146.variables['time']
cf_dtime_146 = cftime.num2date(time_146[:],time_146.units, time_146.calendar)
dtime_146 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_146])

#tropospheric mask
data_trop_146 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_trop_mask_{dates}.nc')
trop_146 = data_trop_146.variables[stash_trop][:]

#air mass
data_mass_146 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_airmass_{dates}.nc')
mass_146 = data_mass_146.variables[stash_mass][:]

#temperature starts in 2014
data_temp_146 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_surf_temp_{dates}.nc')
temp_146 = data_temp_146.variables[stash_temp][:]

#ch4_oh
data_ch4_oh_146 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_oh_flux_{dates}.nc')
ch4_oh_146 = data_ch4_oh_146.variables[stash_ch4_oh][:]

#o3 conc
data_o3_146 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_o3_{dates}.nc')
o3_146 = data_o3_146.variables[stash_o3][:]

#oh_conc
data_oh_146 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_oh_{dates}.nc')
oh_146 = data_oh_146.variables[stash_oh][:]

#### u-bz473

suite_id = 'u-bz473'
suite_lab = 'ubz473'
dates = '2015_2050'

data_473 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_{dates}.nc')
#extract variables to arrays
ch4_473 = data_473.variables[stash_ch4][:]*28.97/mr_ch4*1e9
lat_473 = data_473.variables['latitude'][:]
lon_473 = data_473.variables['longitude'][:]

time_473 = data_473.variables['time']
cf_dtime_473 = cftime.num2date(time_473[:],time_473.units, time_473.calendar)
dtime_473 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_473])

#tropospheric mask
data_trop_473 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_trop_mask_{dates}.nc')
trop_473 = data_trop_473.variables[stash_trop][:]

#air mass
data_mass_473 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_airmass_{dates}.nc')
mass_473 = data_mass_473.variables[stash_mass][:]

#temperature starts in 2014
data_temp_473 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_surf_temp_{dates}.nc')
temp_473 = data_temp_473.variables[stash_temp][:]

#ch4_oh
data_ch4_oh_473 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_oh_flux_{dates}.nc')
ch4_oh_473 = data_ch4_oh_473.variables[stash_ch4_oh][:]

#o3_conc
data_o3_473 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_o3_{dates}.nc')
o3_473 = data_o3_473.variables[stash_o3][:]

#oh_conc
data_oh_473 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_oh_{dates}.nc')
oh_473 = data_oh_473.variables[stash_oh][:]

# ## SSP370 counterfactual: u-bo797

suite_id = 'u-bo797'
suite_lab = 'ubo797'
dates = '2015_2050'

#ssp370 u-bo797
data_3 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_{dates}.nc')
#extract variables to arrays
ch4_3 = data_3.variables['mass_fraction_of_methane_in_air'][:-1,:,:,:]*28.97/mr_ch4*1e9
lat_3 = data_3.variables['latitude'][:]
lon_3 = data_3.variables['longitude'][:]

time_3 = data_3.variables['time']
cf_dtime_3 = cftime.num2date(time_3[:-1],time_3.units, time_3.calendar)
dtime_3 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_3])

#surface temperature
data_temp_3 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ann_temp_2015_2050.nc')
temp_3 = data_temp_3.variables[stash_temp][:]

#tropospheric mask
data_trop_3 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_trop_mask_{dates}.nc')
trop_3 = data_trop_3.variables[stash_trop][:-1,:,:,:]

#air mass
data_mass_3 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_airmass_{dates}.nc')
mass_3 = data_mass_3.variables[stash_mass][:-1,:,:,:]

#ch4_oh
data_ch4_oh_3 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_oh_flux_{dates}.nc')
ch4_oh_3 = data_ch4_oh_3.variables[stash_ch4_oh][:-1,:,:,:]

#o3_conc
data_o3_3 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_o3_{dates}.nc')
o3_3 = data_o3_3.variables[stash_o3][:-1,:,:,:]

#oh_conc
data_oh_3 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_oh_{dates}.nc')
oh_3 = data_oh_3.variables[stash_oh][:-1,:,:,:]

# for some reason some of them are up to 2051 so the last year is removed

# ### u-ca723 SSP370

suite_id = 'u-ca723'
suite_lab = 'uca723'
dates = '2015_2050'

#ssp370 u-bo797
data_4 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_{dates}.nc')
#extract variables to arrays
ch4_4 = data_4.variables['mass_fraction_of_methane_in_air'][:,:,:,:]*28.97/mr_ch4*1e9
lat_4 = data_4.variables['latitude'][:]
lon_4 = data_4.variables['longitude'][:]

time_4 = data_4.variables['time']
cf_dtime_4 = cftime.num2date(time_4[:],time_4.units, time_4.calendar)
dtime_4 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_4])

#surface temperature
data_temp_4 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_surf_temp_2015_2050.nc')
temp_4 = data_temp_4.variables[stash_temp][:]

#tropospheric mask
data_trop_4 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_trop_mask_{dates}.nc')
trop_4 = data_trop_4.variables[stash_trop][:,:,:,:]

#air mass
data_mass_4 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_airmass_{dates}.nc')
mass_4 = data_mass_4.variables[stash_mass][:,:,:,:]

#ch4_oh
data_ch4_oh_4 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_oh_flux_{dates}.nc')
ch4_oh_4 = data_ch4_oh_4.variables[stash_ch4_oh][:,:,:,:]

#o3_conc
data_o3_4 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_o3_{dates}.nc')
o3_4 = data_o3_4.variables[stash_o3][:]

#oh_conc
data_oh_4 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_oh_{dates}.nc')
oh_4 = data_oh_4.variables[stash_oh][:]

# ### u-cb039 SSP370

suite_id = 'u-cb039'
suite_lab = 'ucb039'
dates = '2015_2050'

#ssp370
data_5 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_{dates}.nc')
#extract variables to arrays
ch4_5 = data_5.variables['mass_fraction_of_methane_in_air'][:,:,:,:]*28.97/mr_ch4*1e9
lat_5 = data_5.variables['latitude'][:]
lon_5 = data_5.variables['longitude'][:]

time_5 = data_5.variables['time']
cf_dtime_5 = cftime.num2date(time_5[:],time_5.units, time_5.calendar)
dtime_5 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_5])

#surface temperature
data_temp_5 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_surf_temp_2015_2050.nc')
temp_5 = data_temp_5.variables[stash_temp][:]

#tropospheric mask
data_trop_5 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_trop_mask_{dates}.nc')
trop_5 = data_trop_5.variables[stash_trop][:,:,:,:]

#air mass
data_mass_5 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_airmass_{dates}.nc')
mass_5 = data_mass_5.variables[stash_mass][:,:,:,:]

#ch4_oh
data_ch4_oh_5 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_oh_flux_{dates}.nc')
ch4_oh_5 = data_ch4_oh_5.variables[stash_ch4_oh][:,:,:,:]

#o3_conc
data_o3_5 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_o3_{dates}.nc')
o3_5 = data_o3_5.variables[stash_o3][:]

#oh_conc
data_oh_5 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_oh_{dates}.nc')
oh_5 = data_oh_5.variables[stash_oh][:]

#SSP126: u-bo812

suite_id = 'u-bo812'
suite_lab = 'ubo812'
dates = '2015_2050'

data_6 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_{dates}.nc')
#extract variables to arrays
ch4_6 = data_6.variables['mass_fraction_of_methane_in_air'][:-1,:,:,:]*28.97/mr_ch4*1e9
lat_6 = data_6.variables['latitude'][:]
lon_6 = data_6.variables['longitude'][:]

time_6 = data_6.variables['time']
cf_dtime_6 = cftime.num2date(time_6[:-1],time_6.units, time_6.calendar)
dtime_6 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_6])

#surface temperature
data_temp_6 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_surf_temp_{dates}.nc')
temp_6 = data_temp_6.variables[stash_temp][:-1,:,:]

#tropospheric mask
data_trop_6 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_trop_mask_{dates}.nc')
trop_6 = data_trop_6.variables[stash_trop][:-1,:,:,:]

#air mass
data_mass_6 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_airmass_{dates}.nc')
mass_6 = data_mass_6.variables[stash_mass][:-1,:,:,:]

#ch4_oh
data_ch4_oh_6 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_ch4_oh_flux_{dates}.nc')
ch4_oh_6 = data_ch4_oh_6.variables[stash_ch4_oh][:-1,:,:,:]

#o3_conc
data_o3_6 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_o3_{dates}.nc')
o3_6 = data_o3_6.variables[stash_o3][:-1,:,:,:]

#oh_conc
data_oh_6 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/{suite_id}/netcdf/{suite_lab}_oh_{dates}.nc')
oh_6 = data_oh_6.variables[stash_oh][:-1,:,:,:]

# bl593 data
ch4_dataset_int = xr.open_dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/u-bl593/netcdf/ubl593_ch4_1850_2014.nc')
nc_ch4_data_int = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/u-bl593/netcdf/ubl593_ch4_1850_2014.nc')

data_hist_bl593 = ch4_dataset_int.sel(time=slice('1984-01-01', '2015-12-30'))
stash_name_bl = 'UM_m01s34i009_vn1102'
ch4_hist_bl593 = data_hist_bl593.variables[stash_name_bl]*28.97/mr_ch4*1e9

time_bl593 = nc_ch4_data_int.variables['time']
cf_dtime_bl593 = cftime.num2date(time_bl593[-32:],time_bl593.units, time_bl593.calendar)
dtime_bl593 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_bl593])

# bn213 dataset
ch4_dataset_bn213 = xr.open_dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/u-bn213/netcdf/ubn213_ch4_ann_1985_2013.nc')
nc_ch4_data_bn213 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/u-bn213/netcdf/ubn213_ch4_ann_1985_2013.nc')

data_hist_bn213 = ch4_dataset_bn213.sel(time=slice('1984-01-01', '2015-12-30'))

ch4_hist_bn213 = data_hist_bn213.variables[stash_ch4]*28.97/mr_ch4*1e9

time_bn213 = nc_ch4_data_int.variables['time']
cf_dtime_bn213 = cftime.num2date(time_bn213[-30:],time_bn213.units, time_bn213.calendar)
dtime_bn213 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_bn213])

# bl998 dataset
ch4_dataset_bl998 = xr.open_dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/u-bl998/netcdf/ubl998_ch4_ann_1985_2013.nc')
nc_ch4_data_bl998 = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/znjs2/u-bl998/netcdf/ubl998_ch4_ann_1985_2013.nc')

data_hist_bl998 = ch4_dataset_bl998.sel(time=slice('1984-01-01', '2015-12-30'))

ch4_hist_bl998 = data_hist_bl998.variables[stash_ch4]*28.97/mr_ch4*1e9

time_bl998 = nc_ch4_data_int.variables['time']
cf_dtime_bl998 = cftime.num2date(time_bl998[-30:],time_bl998.units, time_bl998.calendar)
dtime_bl998 = np.array([nc_time_axis.CalendarDateTime(item, "360_day") for item in cf_dtime_bl998])

# ## Area and volume datasets

area = nc.Dataset(f'/gws/nopw/j04/ukca_vol1/jmw240/covid_crunch/area/areacella_fx_UKESM1-0-LL_piControl_r1i1p1f2_gn.nc')
box_area = area.variables['areacella'][:]

total_area = np.sum(box_area)
area_scaled = box_area/total_area
lat_area_scaled = area_scaled[:,0]#1D array of latitude scaling values

# Altitude data

alt_data = nc.Dataset(f'/home/users/zosiast/vol_n96.nc')
alt = alt_data.variables['level_height'][:]
vol = alt_data.variables['grid_cell_volume'][:]

# Population data
pop_data = nc.Dataset(f'/home/users/zosiast/ssp_population/pop_ssp3_n96_2006_2100.nc')
pop_2006_2100 = pop_data.variables['population'][:]

pop_xr = xr.open_dataset(f'/home/users/zosiast/ssp_population/pop_ssp3_n96_2006_2100.nc', decode_times=False)

pop_xr_time = pop_xr.variables['time']
year = np.array(pop_xr_time + 1661,dtype='int')

#select relevant years for population data
pop_data_2015_2050 = pop_2006_2100[9:45,:,:]

#ssp126 population data
pop_data_ssp1 = nc.Dataset(f'/home/users/zosiast/ssp_population/pop_ssp1_n96_2006_2100.nc')
pop_2006_2100_ssp1 = pop_data_ssp1.variables['population'][:]

pop_xr_1 = xr.open_dataset(f'/home/users/zosiast/ssp_population/pop_ssp1_n96_2006_2100.nc', decode_times=False)
#test = xr.decode_cf(pop_xr)
pop_xr_time_1 = pop_xr.variables['time']
year_1 = np.array(pop_xr_time + 1661,dtype='int')
pop_data_2015_2050_ssp1 = pop_2006_2100_ssp1[9:45,:,:]

print('All data read in')

#define functions
def molec_cm3(conc_kg_kg, mass, vol, mr):
    molec_box = conc_kg_kg*mass/mr*1000*n_a #molecules per box
    molec_cm3 = molec_box/(vol*1e6) #molec per cm3
    return molec_cm3

# ## CH4 burden

#sum of all CH4 including stratosphere
#ens 1
ch4_bl593_kg = np.multiply(mass_1,ch4_1/28*mr_ch4/1e9) #ch4 in kg
ch4_bur_bl593 = ch4_bl593_kg.sum(axis=(1,2,3))/1e9 # sum over lat, lon, alt, in Tg

#ens 2
ch4_bz146_kg = np.multiply(mass_146,ch4_146/28*mr_ch4/1e9) #ch4 in kg
ch4_bur_bz146 = ch4_bz146_kg.sum(axis=(1,2,3))/1e9 # sum over lat, lon, alt, in Tg

#ens 3
ch4_bz473_kg = np.multiply(mass_473,ch4_473/28*mr_ch4/1e9) #ch4 in kg
ch4_bur_bz473 = ch4_bz473_kg.sum(axis=(1,2,3))/1e9 # sum over lat, lon, alt, in Tg

ch4_bur_ens_mean = np.mean([ch4_bur_bz473,ch4_bur_bz146,ch4_bur_bl593],axis = 0)

#ssp370
ch4_bo797_kg = np.multiply(mass_3,ch4_3/28*mr_ch4/1e9) #ch4 in kg
ch4_bur_bo797 = ch4_bo797_kg.sum(axis=(1,2,3))/1e9 # sum over lat, lon, alt, in Tg

ch4_ca723_kg = np.multiply(mass_4,ch4_4/28*mr_ch4/1e9) #ch4 in kg
ch4_bur_ca723 = ch4_ca723_kg.sum(axis=(1,2,3))/1e9 # sum over lat, lon, alt, in Tg

ch4_cb039_kg = np.multiply(mass_5,ch4_5/28*mr_ch4/1e9) #ch4 in kg
ch4_bur_cb039 = ch4_cb039_kg.sum(axis=(1,2,3))/1e9 # sum over lat, lon, alt, in Tg

ch4_bur_ssp_mean = np.mean([ch4_bur_bo797,ch4_bur_ca723,ch4_bur_cb039],axis = 0)

ch4_bo812_kg = np.multiply(mass_6,ch4_6/28*mr_ch4/1e9) #ch4 in kg
ch4_bur_bo812 = ch4_bo812_kg.sum(axis=(1,2,3))/1e9 # sum over lat, lon, alt, in Tg

#methane lifetime calc

def ch4_lifetime_calc(ch4_bur, oh_ch4_flux, trop_nc):#ch4 in kg #oh in ppb

    #make trop mask
    trop_mask = ma.masked_where(trop_nc < 0.9999999, trop_nc)

    #flux calcs
    ch4_oh_tg_yr = oh_ch4_flux*mr_ch4*per_sec_to_per_yr/g_to_Tg
    ch4_oh_trop = ma.masked_where(trop_mask.mask,ch4_oh_tg_yr)
    #sum over lat lon alt
    flux_sum = np.sum(ch4_oh_trop, axis=(1,2,3))
    #flux_sum_all = np.sum(ch4_oh_tg_yr, axis=(1,2,3))

    #sum of only tropospheric methane
    #ch4_kg_trop = ma.masked_where(trop_mask.mask, ch4_kg)
    #bur_trop = ch4_kg_trop.sum(axis=(1,2,3))/1e9 # sum over lat, lon, alt, in Tg

    #lifetime calc
    t_ch4_strat_trop = ch4_bur/flux_sum

    return t_ch4_strat_trop, flux_sum

# calculate methane lifetime
t_ch4_by186, flux_by186 = ch4_lifetime_calc(ch4_bur_bl593,ch4_oh_1, trop_1)
t_ch4_bz146, flux_bz146 = ch4_lifetime_calc(ch4_bur_bz146,ch4_oh_146, trop_146)
t_ch4_bz473, flux_bz473 = ch4_lifetime_calc(ch4_bur_bz473,ch4_oh_473, trop_473)

t_ch4_bo797, flux_bo797 = ch4_lifetime_calc(ch4_bur_bo797,ch4_oh_3, trop_3)
t_ch4_ca723, flux_ca723 = ch4_lifetime_calc(ch4_bur_ca723,ch4_oh_4, trop_4)
t_ch4_cb039, flux_cb039 = ch4_lifetime_calc(ch4_bur_cb039,ch4_oh_5, trop_5)

t_ch4_bo812, flux_bo812 = ch4_lifetime_calc(ch4_bur_bo812,ch4_oh_6, trop_6)

# methane lifetime ens mean
t_ch4_ens_mean = np.mean([t_ch4_by186,t_ch4_bz146,t_ch4_bz473],axis=0)
t_ch4_ssp_mean = np.mean([t_ch4_bo797,t_ch4_ca723,t_ch4_cb039],axis=0)

#Fig: methane lifetime
fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax.plot(dtime_1, t_ch4_by186, c='skyblue',linewidth=1)
ax.plot(dtime_146, t_ch4_bz146, c='skyblue',linewidth=1)
ax.plot(dtime_473, t_ch4_bz473, c='skyblue',linewidth=1)

ax.plot(dtime_473, t_ch4_ens_mean, label = 'ZAME',c='C0')

ax.plot(dtime_6, t_ch4_bo812, c='orange',linewidth=1, label = 'SSP1-2.6')

ax.plot(dtime_3, t_ch4_bo797, c='pink')
ax.plot(dtime_4, t_ch4_ca723, c='pink')
ax.plot(dtime_5, t_ch4_cb039, c='pink')
ax.plot(dtime_3, t_ch4_ssp_mean, label = 'SSP370', c='#f11111')
ax.plot([dtime_1[0],dtime_1[-1]], [9.823, 9.823], 'grey',linestyle=':', label = 'PI level')

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

plt.text(dtime_1[2], 9.4 , 'b', weight='bold', fontsize=8)

plt.legend(loc=3)
plt.xlabel('Time')
plt.ylabel('CH$_4$ lifetime (years)')
#plt.title('CH$_4$ lifetime 2015-2050')
plt.savefig(f'{output_write}ch4_lifetime_2015_2050.eps', format = 'eps', bbox_inches='tight')

print('Lifetime calcs and plots done')

# ## NOAA obs data

#1984-2019 obs data
noaa_ch4 = pd.read_csv(f'/home/users/zosiast/scripts/noaa_annual_average_methane.csv',header=0,skiprows=1,usecols=[1,2,3],index_col=0)
noaa_years = np.array(noaa_ch4.index, dtype='str')
noaa_mean_conc = np.array(noaa_ch4['mean'])

# ## 2014 data points, from historical runs
# in ppb
ch4_2014_bl593 = 1633.55
ch4_2014_bl998 = 1617.53
ch4_2014_bn213 = 1620.09
ch4_2014_ens_mean = np.mean([ch4_2014_bl593,ch4_2014_bl998,ch4_2014_bn213])

# In[29]:
d_time_2014 = cftime.datetime(year=2014, month=6, day=1)
nc_dtime_2014 = nc_time_axis.CalendarDateTime(d_time_2014, "360_day")
dtime_2014_2050 = np.insert(dtime_1,0, nc_dtime_2014)

d_time_2015 = cftime.datetime(year=2015, month=6, day=1)
nc_dtime_2015 = nc_time_axis.CalendarDateTime(d_time_2015, "360_day")

noaa_time_1985_2015 = np.append(dtime_bn213,nc_dtime_2015)
noaa_dtime_1985_2020 = np.concatenate((noaa_time_1985_2015,dtime_2014_2050[2:7]),axis=0)

# ## CH4 surface conc

# mean surface conc
#area weighted by latitude and mean over lon
ch4_lat_weighted_by186 = np.mean(np.average(ch4_1[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)
ch4_lat_weighted_bz146 = np.mean(np.average(ch4_146[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)
ch4_lat_weighted_bz473 = np.mean(np.average(ch4_473[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)

ch4_lat_weighted_ens = np.mean([ch4_lat_weighted_by186,ch4_lat_weighted_bz146,ch4_lat_weighted_bz473],axis=0)

ch4_lat_weighted_bo797 = np.mean(np.average(ch4_3[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)
ch4_lat_weighted_ca723 = np.mean(np.average(ch4_4[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)
ch4_lat_weighted_cb039 = np.mean(np.average(ch4_5[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)

ch4_lat_weighted_ssp = np.mean([ch4_lat_weighted_bo797,ch4_lat_weighted_ca723,ch4_lat_weighted_cb039],axis=0)

ch4_lat_weighted_bo812 = np.mean(np.average(ch4_6[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)

# surface conc for hist ensemble members
ch4_lat_weighted_bl593 = np.mean(np.average(ch4_hist_bl593[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)

ch4_lat_weighted_bn213 = np.mean(np.average(ch4_hist_bn213[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)

ch4_lat_weighted_bl998 = np.mean(np.average(ch4_hist_bl998[:,0,:,:],axis=(1),weights = lat_area_scaled), axis = 1)

# hist ensemble mean
hist_mean_ch4 = np.mean([ch4_lat_weighted_bl998,ch4_lat_weighted_bn213,ch4_lat_weighted_bl593[-30:]],axis=0)

# 2014-2050 area weighted CH4 surface conc with obs

fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_by186,0, ch4_2014_bl593),  c='skyblue',linewidth=1)
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_bz146,0,ch4_2014_bn213),  c='skyblue',linewidth=1)
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_bz473,0, ch4_2014_bl998),  c='skyblue',linewidth=1)

ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_ens,0,ch4_2014_ens_mean), label = 'ZAME', c='C0',linewidth=1)

ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_bo797,0, ch4_2014_bl593), c='pink')
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_ca723,0, ch4_2014_bl998), c='pink')
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_cb039,0,ch4_2014_bn213), c='pink')
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_ssp,0,ch4_2014_ens_mean), label = 'SSP3-7.0', c='#f11111')

ax.scatter(dtime_2014_2050[:6], noaa_mean_conc[-6:], c='darkslategray', label = 'NOAA obs', marker='+')

ax.plot([dtime_2014_2050[0],dtime_2014_2050[-1]], [776, 776], 'grey',linestyle=':', label = 'PI level')

ax.spines['top'].set_visible(False)
#ax2.spines['top'].set_visible(False)

plt.legend()
plt.xlabel('Time')
plt.ylabel(f'Mean surface [{var_lab}] (ppb)')
#plt.title(f'Mean surface {var_lab} 2015-2050')
#plt.savefig(f'{output_write}ch4_surface_conc_obs_2015_2050.eps', format = 'eps', bbox_inches='tight')

# Normalise to year 2000 values

# year 2000 conc from bl593 hist run
ch4_2000_bl593 = 1537.89
ch4_2000_noaa = noaa_mean_conc[-20]
pi_ratio_2000 = 776./ch4_2000_bl593

# Normalised surface conc
fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_by186,0, ch4_2014_bl593)/ch4_2000_bl593,  c='skyblue',linewidth=1)
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_bz146,0,ch4_2014_bn213)/ch4_2000_bl593,  c='skyblue',linewidth=1)
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_bz473,0, ch4_2014_bl998)/ch4_2000_bl593,  c='skyblue',linewidth=1)

ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_ens,0,ch4_2014_ens_mean)/ch4_2000_bl593, label = 'NZAME', c='C0',linewidth=1)

ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_bo797,0, ch4_2014_bl593)/ch4_2000_bl593, c='pink')
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_ca723,0, ch4_2014_bl998)/ch4_2000_bl593, c='pink')
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_cb039,0,ch4_2014_bn213)/ch4_2000_bl593, c='pink')
ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_ssp,0,ch4_2014_ens_mean)/ch4_2000_bl593, label = 'SSP3-7.0', c='#f11111')

ax.plot([dtime_bl593[0],dtime_2014_2050[-1]], [pi_ratio_2000, pi_ratio_2000], 'grey',linestyle=':', label = 'PI level')

ax.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_bo812,0, ch4_2014_bl998)/ch4_2000_bl593,  label='SSP1-2.6', c='orange',linewidth=1)

ax.plot(dtime_bl593, ch4_lat_weighted_bl593/ch4_2000_bl593,  c='lightgrey', linewidth=1,zorder=1)
ax.plot(dtime_bn213, ch4_lat_weighted_bn213/ch4_2000_bl593,  c='lightgrey',linewidth=1,zorder=1)
ax.plot(dtime_bl998, ch4_lat_weighted_bl998/ch4_2000_bl593,  c='lightgrey',linewidth=1,zorder=1)

ax.plot(dtime_bl998, hist_mean_ch4/ch4_2000_bl593,  c='grey',label='historical',zorder=2)

ax.plot(noaa_dtime_1985_2020, noaa_mean_conc[1:]/ch4_2000_noaa, c='k', label = 'NOAA obs', marker='+',zorder=3, ms=3, alpha = 0.7, ls="none")

ax.text(dtime_1[-15], 1.4 , 'b', weight='bold', fontsize=8)

ax2=ax.twinx()
#ax2.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_ens,0,ch4_2014_ens_mean), c='black',linewidth=1, linestyle='--')
#ax2.plot(dtime_2014_2050, np.insert(ch4_lat_weighted_ssp,0,ch4_2014_ens_mean), c='black',linestyle='--')
ax2.set_ylim(545, 2260)

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax2.spines['top'].set_visible(False)

ax.legend(fontsize = 7,ncol=2, loc=2)
plt.xlabel('Time')
ax.set_ylabel(f'Mean surface [{var_lab}] / year 2000 [CH$_4$]')
ax2.set_ylabel(f'Mean surface model [{var_lab}] (ppb)')
#plt.title(f'Mean surface {var_lab} 2015-2050 relative to year 2000')
plt.savefig(f'{output_write}ch4_relative_surface_conc_2015_2050.eps', format = 'eps', bbox_inches='tight')

# GMST plots

# ## GMST calculation

#by186 temp meaning
temp_1_lon_mean = np.mean(temp_1, axis = 2)
mean_sur_temp_186 = np.average(temp_1_lon_mean,weights = lat_area_scaled, axis = (1))

temp_146_lon_mean = np.mean(temp_146, axis = 2)
mean_sur_temp_146 = np.average(temp_146_lon_mean,weights = lat_area_scaled, axis = (1))

temp_473_lon_mean = np.mean(temp_473, axis = 2)
mean_sur_temp_473 = np.average(temp_473_lon_mean,weights = lat_area_scaled, axis = (1))

gmst_ens_mean = np.mean([mean_sur_temp_473,mean_sur_temp_146,mean_sur_temp_186],axis = 0)

# In[13]:

#bo797 global mean, area weighted
temp_3_lon_mean = np.mean(temp_3, axis = 2)
mean_sur_temp_bo797 = np.average(temp_3_lon_mean,weights = lat_area_scaled, axis = (1))

#ca723 global mean, area weighted
temp_4_lon_mean = np.mean(temp_4, axis = 2)
mean_sur_temp_ca723 = np.average(temp_4_lon_mean,weights = lat_area_scaled, axis = (1))

temp_5_lon_mean = np.mean(temp_5, axis = 2)
mean_sur_temp_cb039 = np.average(temp_5_lon_mean,weights = lat_area_scaled, axis = (1))

gmst_ens_mean_ssp = np.mean([mean_sur_temp_bo797,mean_sur_temp_ca723,mean_sur_temp_cb039],axis = 0)

#bo812
temp_6_lon_mean = np.mean(temp_6, axis = 2)
mean_sur_temp_bo812 = np.average(temp_6_lon_mean,weights = lat_area_scaled, axis = (1))

# GMST vs time

fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax.plot(dtime_1, mean_sur_temp_186, c='skyblue',linewidth=1) #temp for by186 is from 2014(?)
ax.plot(dtime_146, mean_sur_temp_146,  c='skyblue',linewidth=1)
ax.plot(dtime_473, mean_sur_temp_473, c='skyblue',linewidth=1)

ax.plot(dtime_473, gmst_ens_mean, label = 'ZAME',c='C0')

ax.plot(dtime_3, mean_sur_temp_bo797, c='pink')
ax.plot(dtime_4, mean_sur_temp_ca723, c='pink')
ax.plot(dtime_5, mean_sur_temp_cb039, c='pink')
#ax.plot(dtime_3, gmst_ens_mean_ssp, label = 'SSP3-7.0',c='#f11111')

ax.plot(dtime_6, mean_sur_temp_bo812, c='orange')

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

plt.legend()
plt.xlabel('Time')
plt.ylabel(f'GMST (K)')
#plt.title('GMST 2015-2050')
plt.savefig(f'{output_write}gmst_2015_2050.eps', format = 'eps', bbox_inches='tight')

# GMST anomaly wrt 2015

fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax.plot(dtime_1, mean_sur_temp_186 - mean_sur_temp_186[0], c='skyblue',linewidth=1) #temp for by186 is from 2014(?)
ax.plot(dtime_146, mean_sur_temp_146 - mean_sur_temp_146[0],  c='skyblue',linewidth=1)
ax.plot(dtime_473, mean_sur_temp_473 - mean_sur_temp_473[0], c='skyblue',linewidth=1)

ax.plot(dtime_473, gmst_ens_mean - gmst_ens_mean[0], label = 'ZAME',c='C0')

ax.plot(dtime_3, mean_sur_temp_bo797 - mean_sur_temp_bo797[0], c='pink')
ax.plot(dtime_4, mean_sur_temp_ca723 - mean_sur_temp_ca723[0], c='pink')
ax.plot(dtime_5, mean_sur_temp_cb039 - mean_sur_temp_cb039[0], c='pink')

ax.plot(dtime_6, mean_sur_temp_bo812 - mean_sur_temp_bo812[0], c='orange', label = 'SSP1-2.6')

ax.plot(dtime_3, gmst_ens_mean_ssp - gmst_ens_mean_ssp[0], label = 'SSP3-7.0',c='#f11111')

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

plt.text(dtime_1[2], 1.9 , 'a', weight='bold', fontsize=8)

plt.legend(bbox_to_anchor=(0.5,0.75))
plt.xlabel('Time')
plt.ylabel(f'GMST anomaly wrt 2015 (K)')
#plt.title('GMST anomaly 2015-2050')
plt.savefig(f'{output_write}gmst_anomaly_2015_2050_ssp126.eps', format = 'eps', bbox_inches='tight')

#gmst map plot
gmst_mean_map_nzame = np.mean([temp_1,temp_146,temp_473],axis = 0)
gmst_mean_map_ssp = np.mean([temp_3,temp_4,temp_5],axis = 0)

from matplotlib.colors import LinearSegmentedColormap,ListedColormap
from matplotlib import cm

top = cm.get_cmap('Blues_r', 128)
bottom = cm.get_cmap('Reds', 128)

newcolors = np.vstack((top(np.linspace(0, 1, 90)),
                       bottom(np.linspace(0, 0.1, 10))))
newcmp = ListedColormap(newcolors, name='Red_blue')

fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax = plt.axes(projection=ccrs.Robinson(central_longitude=0, globe=None))
ax.set_global()
ax.coastlines(linewidth=0.5)

var_name = 'Surface Temperature$_{ZAME - SSP37.0}$'

ch4_diff_cyclic, lon_plot = add_cyclic_point(np.mean(gmst_mean_map_nzame[-10:,:,:] - gmst_mean_map_ssp[-10:,:,:],
                                                     axis=0), coord=lon_473)

plt.pcolormesh(lon_plot,lat_473,ch4_diff_cyclic, cmap=newcmp,transform=ccrs.PlateCarree(central_longitude=0), vmin=-9, vmax=1)

plt.annotate('b', (0.1,0.85), xycoords='figure fraction',weight='bold', fontsize=8)

plt.colorbar(label = f'{var_name} / K',orientation='horizontal',pad=0.05)
#plt.title(f'Surface temperature 2040-2050: ZAME - SSP3-7.0')
plt.savefig(f'{output_write}gmst_surf_map_2040_2050.eps', format = 'eps', bbox_inches='tight')

#Trop mask based on trop mask variable

#tropospheric mask
trop_mask_186 = ma.masked_where(trop_1 < 1, trop_1)
trop_mask_146 = ma.masked_where(trop_146 < 1, trop_146)
trop_mask_473 = ma.masked_where(trop_473 < 1, trop_473)

trop_mask_3 = ma.masked_where(trop_3 < 1, trop_3)

#ozone in ppb
o3_186_ppb = o3_1*28.97/mr_o3*1e9
o3_146_ppb = o3_146*28.97/mr_o3*1e9
o3_473_ppb = o3_473*28.97/mr_o3*1e9

o3_bo797 = o3_3*28.97/mr_o3*1e9
o3_ca723 = o3_4*28.97/mr_o3*1e9
o3_cb039 = o3_5*28.97/mr_o3*1e9

o3_bo812 = o3_6*28.97/mr_o3*1e9

#surface selected
surf_o3_186 = o3_186_ppb[:,0,:,:]
surf_o3_146 = o3_146_ppb[:,0,:,:]
surf_o3_473 = o3_473_ppb[:,0,:,:]

surf_o3_ens_mean = np.mean([surf_o3_186,surf_o3_146,surf_o3_473],axis=0)

surf_o3_bo797 = o3_bo797[:,0,:,:]
surf_o3_ca723 = o3_ca723[:,0,:,:]
surf_o3_cb039 = o3_cb039[:,0,:,:]

surf_o3_ens_mean_ssp = np.mean([surf_o3_bo797[:36,:,:],surf_o3_ca723,surf_o3_cb039],axis=0)

surf_o3_ssp126 = o3_bo812[:,0,:,:]

#area weighted by latitude and mean over lon
o3_lat_weighted_186 = np.mean(np.average(surf_o3_186,axis=(1),weights = lat_area_scaled), axis = 1)
o3_lat_weighted_146 = np.mean(np.average(surf_o3_146,axis=(1),weights = lat_area_scaled), axis = 1)
o3_lat_weighted_473 = np.mean(np.average(surf_o3_473,axis=(1),weights = lat_area_scaled), axis = 1)

o3_lat_weighted_ens_mean = np.mean([o3_lat_weighted_186,o3_lat_weighted_146,o3_lat_weighted_473],axis=0)

o3_lat_weighted_bo797 = np.mean(np.average(surf_o3_bo797,axis=(1),weights = lat_area_scaled), axis = 1)
o3_lat_weighted_ca723 = np.mean(np.average(surf_o3_ca723,axis=(1),weights = lat_area_scaled), axis = 1)
o3_lat_weighted_cb039 = np.mean(np.average(surf_o3_ca723,axis=(1),weights = lat_area_scaled), axis = 1)

o3_lat_weighted_ens_mean_ssp = np.mean([o3_lat_weighted_bo797[:36],o3_lat_weighted_ca723,o3_lat_weighted_cb039],axis=0)

# o3 surface conc
fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax.plot(dtime_1, o3_lat_weighted_186, c='skyblue',linewidth=1)
ax.plot(dtime_146, o3_lat_weighted_146,  c='skyblue',linewidth=1)
ax.plot(dtime_473, o3_lat_weighted_473,  c='skyblue',linewidth=1)

ax.plot(dtime_473, o3_lat_weighted_ens_mean, label = 'ZAME', c='C0')

ax.plot(dtime_3,  o3_lat_weighted_bo797[:36], c='pink')
ax.plot(dtime_4, o3_lat_weighted_ca723, c='pink')
ax.plot(dtime_5, o3_lat_weighted_cb039, c='pink')

ax.plot(dtime_4,  o3_lat_weighted_ens_mean_ssp, label = 'SSP370', c='#f11111')
ax.plot([dtime_1[0],dtime_1[-1]], [20, 20], 'grey',linestyle=':', label = 'PI level')

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

plt.legend()
plt.xlabel('Time')
plt.ylabel(f'O$_3$ (ppb)')
#plt.title('O$_3$ mean surface conc 2015-2050')
plt.savefig(f'{output_write}o3_surface_area_weighted_2015_2050.eps', format = 'eps', bbox_inches='tight')

#population weighting
#area weighted by population
o3_pop_weighted_186 = np.average(surf_o3_186,axis=(1,2),weights = pop_data_2015_2050)
o3_pop_weighted_146 = np.average(surf_o3_146,axis=(1,2),weights = pop_data_2015_2050)
o3_pop_weighted_473 = np.average(surf_o3_473,axis=(1,2),weights = pop_data_2015_2050)

o3_pop_weighted_ens_mean = np.mean([o3_pop_weighted_186,o3_pop_weighted_146,o3_pop_weighted_473],axis=0)

o3_pop_weighted_bo797 = np.average(surf_o3_bo797[:36,:,:],axis=(1,2),weights = pop_data_2015_2050)
o3_pop_weighted_ca723 = np.average(surf_o3_ca723,axis=(1,2),weights = pop_data_2015_2050)
o3_pop_weighted_cb039 = np.average(surf_o3_cb039,axis=(1,2),weights = pop_data_2015_2050)

o3_pop_weighted_ens_mean_ssp = np.mean([o3_pop_weighted_bo797,o3_pop_weighted_ca723,o3_pop_weighted_cb039],axis=0)

o3_pop_weighted_bo812 = np.average(surf_o3_ssp126,axis=(1,2),weights = pop_data_2015_2050_ssp1)

#ozone surface population weighted
fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax.plot(dtime_1, o3_pop_weighted_186, c='skyblue',linewidth=1)
ax.plot(dtime_146, o3_pop_weighted_146, c='skyblue',linewidth=1)
ax.plot(dtime_473, o3_pop_weighted_473, c='skyblue',linewidth=1)

ax.plot(dtime_473, o3_pop_weighted_ens_mean, label = 'ZAME', c='C0')

ax.plot(dtime_3,  o3_pop_weighted_bo797, c='pink')
ax.plot(dtime_3,  o3_pop_weighted_ca723, c='pink')
ax.plot(dtime_3,  o3_pop_weighted_cb039, c='pink')
ax.plot(dtime_3,  o3_pop_weighted_bo812, label = 'SSP1-2.6', c='orange')

plt.text(dtime_1[2], 34.5 , 'd', weight='bold', fontsize=8)

ax.plot(dtime_3,  o3_pop_weighted_ens_mean_ssp, label = 'SSP3-7.0', c='#f11111')

#ax.plot([dtime_1[0],dtime_1[-1]], [20, 20], 'grey',linestyle=':', label = 'PI level')

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

plt.legend(loc=3)
plt.xlabel('Time')
plt.ylabel(f'Pop. wt. O$_3$ (ppb)')
#plt.title('O$_3$ mean surface conc 2015-2050: population weighted')
plt.savefig(f'{output_write}o3_surface_pop_weighted_2015_2050.eps', format = 'eps', bbox_inches='tight')

# ozone surface concentration plot

fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax = plt.axes(projection=ccrs.Robinson(central_longitude=0.0))
ax.coastlines(linewidth=0.3)

var_name = 'Surface [O$_3$]$_{ZAME - SSP37.0}$'

ch4_diff_cyclic, lon_plot = add_cyclic_point(np.mean(surf_o3_ens_mean[-10:,:,:] - surf_o3_ens_mean_ssp[-10:,:,:],axis=0),
                                             coord=lon_473)

plt.pcolormesh(lon_plot,lat_473,ch4_diff_cyclic, cmap='Oranges_r',vmin = -18, vmax = 0,
               transform=ccrs.PlateCarree(central_longitude=0))#, cmap='bwr', vmin = -8, vmax = 8)

plt.annotate('c', (0.1,0.85), xycoords='figure fraction',weight='bold', fontsize=8)

plt.colorbar(label = f'{var_name} (ppb)',orientation='horizontal',pad=0.05)
#plt.title(f'Surface ozone 2040-2050: ZAME - SSP3-7.0')
plt.savefig(f'{output_write}o3_surf_map_2040_2050.eps', format = 'eps', bbox_inches='tight')

#OH concentration

# convert to molec cm -3
oh_molec_cm3_186 = molec_cm3(oh_1, mass_1, vol, 17)
oh_molec_cm3_146 = molec_cm3(oh_146, mass_146, vol, 17)
oh_molec_cm3_473 = molec_cm3(oh_473, mass_473, vol, 17)

oh_molec_cm3_bo797 = molec_cm3(oh_3[:36,:,:,:], mass_3, vol, 17)
oh_molec_cm3_ca723 = molec_cm3(oh_4, mass_4, vol, 17)
oh_molec_cm3_cb039 = molec_cm3(oh_5, mass_5, vol, 17)

oh_molec_cm3_bo812 = molec_cm3(oh_6, mass_6, vol, 17)

#by186 oh
oh_trop_186 = ma.masked_where(o3_186_ppb > 125,oh_molec_cm3_186)
oh_trop_mean_186 = np.average(oh_trop_186, weights=mass_1, axis = (1,2,3))

#by146 oh
oh_trop_146 = ma.masked_where(o3_146_ppb > 125,oh_molec_cm3_146)
oh_trop_mean_146 = np.average(oh_trop_146, weights=mass_146, axis = (1,2,3))

#by473 oh
oh_trop_473 = ma.masked_where(o3_473_ppb > 125,oh_molec_cm3_473)
oh_trop_mean_473 = np.average(oh_trop_473, weights=mass_473, axis = (1,2,3))

#bo797 oh
oh_trop_bo797 = ma.masked_where(o3_bo797 > 125,oh_molec_cm3_bo797)
oh_trop_mean_bo797 = np.average(oh_trop_bo797, weights=mass_3, axis = (1,2,3))

#ca723 oh
oh_trop_ca723 = ma.masked_where(o3_ca723 > 125,oh_molec_cm3_ca723)
oh_trop_mean_ca723 = np.average(oh_trop_ca723, weights=mass_4, axis = (1,2,3))

#cb039 oh
oh_trop_cb039 = ma.masked_where(o3_cb039 > 125,oh_molec_cm3_cb039)
oh_trop_mean_cb039 = np.average(oh_trop_cb039, weights=mass_5, axis = (1,2,3))

#bo812 oh
oh_trop_bo812 = ma.masked_where(o3_bo812 > 125,oh_molec_cm3_bo812)
oh_trop_mean_bo812 = np.average(oh_trop_bo812, weights=mass_5, axis = (1,2,3))

#ensemble mean
oh_trop_ens_mean = np.mean([oh_trop_mean_186,oh_trop_mean_146,oh_trop_mean_473], axis = 0)
oh_trop_ens_mean_ssp = np.mean([oh_trop_mean_bo797,oh_trop_mean_ca723,oh_trop_mean_cb039], axis = 0)

# OH trop mean conc

fig = plt.figure(figsize=(3.5,2.625), dpi=300)
ax = plt.axes()

ax.plot(dtime_1, oh_trop_mean_186/1e6, c='skyblue',linewidth=1)
ax.plot(dtime_146, oh_trop_mean_146/1e6,  c='skyblue',linewidth=1)
ax.plot(dtime_473, oh_trop_mean_473/1e6,  c='skyblue',linewidth=1)
ax.plot(dtime_473, oh_trop_ens_mean/1e6, label = 'ZAME', c='C0')

ax.plot(dtime_3, oh_trop_mean_bo797/1e6, c='pink')
ax.plot(dtime_4, oh_trop_mean_ca723/1e6, c='pink')
ax.plot(dtime_5, oh_trop_mean_cb039/1e6, c='pink')

ax.plot(dtime_6, oh_trop_mean_bo812/1e6, label='SSP1-2.6', c='orange',linewidth=1)

ax.plot(dtime_4, oh_trop_ens_mean_ssp/1e6, label = 'SSP3-7.0', c='#f11111')
ax.plot([dtime_1[0],dtime_1[-1]], [0.825, 0.825], 'grey',linestyle=':', label = 'PI level')

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
#ax.set_xticks(labels)

plt.text(dtime_1[2], 1.3 , 'a', weight='bold', fontsize=8)

plt.legend(loc=3)
plt.xlabel('Time')
plt.ylabel('OH ($10^6$ molec cm$^{-3}$)')
#plt.title('OH trop mean conc 2015-2025')
plt.savefig(f'{output_write}oh_trop_mean_2015_2050.eps', format = 'eps', bbox_inches='tight')

## Quantitative print data for 2040-2050

def mean_std_2040_2050(a,b,c,a_ssp,b_ssp,c_ssp,var,units):
    mean_2040_2050_nzame = np.mean([a[-10],b[-10],c[-10]])
    std_2040_2050_nzame = np.mean(np.std([a[-10],b[-10],b[-10]],axis=0))

    mean_2040_2050_ssp = np.mean([a_ssp[-10],b_ssp[-10],c_ssp[-10]])
    std_2040_2050_ssp = np.mean(np.std([a_ssp[-10],b_ssp[-10],c_ssp[-10]],axis=0))

    diff_2040_2050 = mean_2040_2050_nzame - mean_2040_2050_ssp
    diff_2040_2050_err = np.sqrt(std_2040_2050_nzame**2 + std_2040_2050_ssp**2)

    print(f'{var} 2040-2050 ZAME: {mean_2040_2050_nzame:.2f} ± {std_2040_2050_nzame:.2f} {units}')
    print(f'{var} 2040-2050 SSP: {mean_2040_2050_ssp:.2f} ± {std_2040_2050_ssp:.2f} {units}')
    print('SSP increases over 2040-2050')
    print(f'{var} difference 2040-50: {diff_2040_2050:.2f} ± {diff_2040_2050_err:.2f} {units}')

def mean_std_2050(a,b,c,a_ssp,b_ssp,c_ssp,var,units):
    mean_2050_nzame = np.mean([a[-1],b[-1],c[-1]])
    std_2050_nzame = np.mean(np.std([a[-1],b[-1],b[-1]],axis=0))

    mean_2050_ssp = np.mean([a_ssp[-1],b_ssp[-1],c_ssp[-1]])
    std_2050_ssp = np.mean(np.std([a_ssp[-1],b_ssp[-1],c_ssp[-1]],axis=0))

    diff_2050 = mean_2050_nzame - mean_2050_ssp
    diff_2050_err = np.sqrt(std_2050_nzame**2 + std_2050_ssp**2)

    print(f'{var} 2050 ZAME: {mean_2050_nzame:.3f} ± {std_2050_nzame:.3f} {units}')
    print(f'{var} 2050 SSP: {mean_2050_ssp:.3f} ± {std_2050_ssp:.3f} {units}')
    print(f'{var} difference 2050: {diff_2050:.3f} ± {diff_2050_err:.3f} {units}')

def mean_std_2015(a,b,c,a_ssp,b_ssp,c_ssp,var,units):
    mean_2015_nzame = np.mean([a[0],b[0],c[0]])
    std_2015_nzame = np.mean(np.std([a[0],b[0],b[0]],axis=0))

    mean_2015_ssp = np.mean([a_ssp[0],b_ssp[0],c_ssp[0]])
    std_2015_ssp = np.mean(np.std([a_ssp[0],b_ssp[0],c_ssp[0]],axis=0))

    diff_2015 = mean_2015_nzame - mean_2015_ssp
    diff_2015_err = np.sqrt(std_2015_nzame**2 + std_2015_ssp**2)

    print(f'{var} 2015 ZAME: {mean_2015_nzame:.3f} ± {std_2015_nzame:.3f} {units}')
    print(f'{var} 2015 SSP: {mean_2015_ssp:.3f} ± {std_2015_ssp:.3f} {units}')
    print(f'{var} difference 2015: {diff_2015:.3f} ± {diff_2015_err:.3f} {units}')

# prints output file with 2040-2050 data
import sys
stdoutOrigin=sys.stdout
sys.stdout = open("conc_data_nzame_final.txt", "w")

print('Conc data for 2040-2050:')
print()
# OH conc
mean_std_2050(oh_trop_mean_186/1e6,oh_trop_mean_146/1e6,oh_trop_mean_473/1e6,
              oh_trop_mean_bo797/1e6, oh_trop_mean_ca723/1e6,oh_trop_mean_cb039/1e6,'OH surface conc','x10^6 molec cm-3')
print()
mean_std_2040_2050(oh_trop_mean_186/1e6,oh_trop_mean_146/1e6,oh_trop_mean_473/1e6,
              oh_trop_mean_bo797/1e6, oh_trop_mean_ca723/1e6,oh_trop_mean_cb039/1e6,'OH surface conc','x10^6 molec cm-3')

#pop weighted ozone
mean_std_2050(o3_pop_weighted_186,o3_pop_weighted_146,o3_pop_weighted_473,o3_pop_weighted_bo797,
              o3_pop_weighted_ca723,o3_pop_weighted_cb039,'O3 (pop)','ppb')
print()
mean_std_2040_2050(o3_pop_weighted_186,o3_pop_weighted_146,o3_pop_weighted_473,o3_pop_weighted_bo797,
              o3_pop_weighted_ca723,o3_pop_weighted_cb039,'O3 (pop)','ppb')
#area weighted ozone
print()
mean_std_2050(o3_lat_weighted_186,o3_lat_weighted_146,o3_lat_weighted_473,o3_lat_weighted_bo797,
              o3_lat_weighted_ca723,o3_lat_weighted_cb039,'O3 (area)','ppb')
print()
mean_std_2040_2050(o3_lat_weighted_186,o3_lat_weighted_146,o3_lat_weighted_473,o3_lat_weighted_bo797,
              o3_lat_weighted_ca723,o3_lat_weighted_cb039,'O3 (area)','ppb')
# CH4 burden
print()
mean_std_2040_2050(ch4_bur_bl593,ch4_bur_bz146,ch4_bur_bz473,ch4_bur_bo797,
              ch4_bur_ca723,ch4_bur_cb039,'CH4 burden','Tg')
print()
mean_std_2050(ch4_bur_bl593,ch4_bur_bz146,ch4_bur_bz473,ch4_bur_bo797,
              ch4_bur_ca723,ch4_bur_cb039,'CH4 burden','Tg')
# CH4 surface conc
print()
mean_std_2040_2050(ch4_lat_weighted_by186,ch4_lat_weighted_bz146,ch4_lat_weighted_bz473,ch4_lat_weighted_bo797,
              ch4_lat_weighted_ca723,ch4_lat_weighted_cb039,'CH4 (area)','ppb')
print()
mean_std_2050(ch4_lat_weighted_by186,ch4_lat_weighted_bz146,ch4_lat_weighted_bz473,ch4_lat_weighted_bo797,
              ch4_lat_weighted_ca723,ch4_lat_weighted_cb039,'CH4 (area)','ppb')
print()
mean_std_2050(mean_sur_temp_473,mean_sur_temp_146,mean_sur_temp_186,mean_sur_temp_bo797,
              mean_sur_temp_ca723,mean_sur_temp_cb039,'GMST','K')
print()
mean_std_2015(mean_sur_temp_473,mean_sur_temp_146,mean_sur_temp_186,mean_sur_temp_bo797,
              mean_sur_temp_ca723,mean_sur_temp_cb039,'GMST','K')

sys.stdout.close()
sys.stdout=stdoutOrigin
print('end')


'''
Given an integer array nums, find the contiguous subarray (containing at least one number)
which has the largest sum and return its sum.

A subarray is a contiguous part of an array.

Example 1:

Input: nums = [-2,1,-3,4,-1,2,1,-5,4]
Output: 6
Explanation: [4,-1,2,1] has the largest sum = 6.
Example 2:

Input: nums = [1]
Output: 1
Example 3:

Input: nums = [5,4,-1,7,8]
Output: 23
'''
class MaxSubArray:
    def maxSubArray(self, nums):
        currentSum = 0
        # assign large negative value
        maxSum = -9999999
        for index, value in enumerate(nums): #enumerate gives the array index and value
            currentSum += value
            if currentSum > maxSum: #current sum is compared with maxSum
                maxSum = currentSum #re setting

            if value > maxSum: #individual number is compared with the maxSum
                maxSum = value #re setting
                currentSum = value #re setting

            if value > currentSum: #individual number should be greater currentSum is checked
                currentSum = value #re setting

        return maxSum

if __name__ == '__main__':
    maxSubArray = MaxSubArray()
    nums = [-2, 1, -3, 4, -1, 2, 1, -5, 4]
    maxSubArraySum = maxSubArray.maxSubArray(nums)
    print(f'Max sub array is {maxSubArraySum}')


import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Embedding, GlobalAveragePooling1D, Dense
from cowait.tasks import Task

class ImdbTask(Task):
    async def run(self):
        # get the pre-processed data
        data_train, data_test = self.get_data()

        # train and evaluate model
        model = self.fit(data_train, data_test)

        # test model on example data point
        test_index = 500
        test_example = data_test[0][test_index]
        pred = model.predict(test_example.reshape(1, len(test_example)))
        pred = 1 if pred[0][0] > 0.5 else 0
        print("##### Test example #####")
        print("Correct class:", data_test[1][test_index])    # gt=1
        print("Predicted class:", pred)

        return int(pred)

    def get_data(self):
        # load the data
        # the text reviews have already been converted to numerical features
        (x_train, y_train), (x_test, y_test) = tf.keras.datasets.imdb.load_data()

        # pad all sequences to the same length
        pad_length = 1000
        x_train = pad_sequences(x_train, pad_length)
        x_test = pad_sequences(x_test, pad_length)

        return (x_train, y_train), (x_test, y_test)

    def create_model(self):
        # size of the vocabulary
        vocab_size = len(tf.keras.datasets.imdb.get_word_index().keys()) + 3
        # size of the embedding vectors learned by the model
        embedding_dim = 16

        # set up the model structure
        model = Sequential([
            Embedding(input_dim=vocab_size, output_dim=embedding_dim),    # embedding layer
            GlobalAveragePooling1D(),                                     # pooling layer
            Dense(1, activation='sigmoid')                                # single sigmoid output node
        ])

        # compile model with optimizer, loss function and evaluation metrics
        model.compile(
            optimizer='adam',
            loss='binary_crossentropy',
            metrics=['accuracy']
        )

        return model

    def fit(self, data_train, data_test):
        # get the data
        (x_train, y_train), (x_test, y_test) = data_train, data_test

        # create the model
        model = self.create_model()

        # fit the model to the train set
        print("##### Train model #####")
        model.fit(
            x=x_train,
            y=y_train,
            epochs=2
        )

        # evaluate the model on the test set
        print("##### Evaluate model #####")
        model.evaluate(
            x=x_test,
            y=y_test
        )

        return model


import requests
import time
import sys
import subprocess

import redis
from redis import Redis

redis_db = redis.Redis(host='localhost', charset="utf-8", decode_responses=True)

id = input('ID: ').upper()
input = input('Text: ').lower()

letter = 'a'

while True:
    temp = subprocess.Popen(['/opt/vc/bin/vcgencmd', 'measure_temp'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    stdout, stderr = temp.communicate()
    temp1 = str(stdout.decode('utf-8')).split('=')
    print(temp1)
    if len(input) > 0:
        tx_string = '6{}T{}:{}[{}]'.format(letter, temp1[1][:4], input, id)
    else:
        tx_string = '6{}T{}[{}]'.format(letter, temp1[1][:4], id)

    print(tx_string)
    redis_db.set((int(time.time()) - 1000), tx_string)
    print('Loaded into the broadcast queue')
    letter = chr(ord(letter) + 1)
    if letter == '{':
        letter = 'b'

    time.sleep(300)


import numpy as np
from matplotlib import pyplot as plt
import C_Re
import airprops
import RK

class Particles:
    def __init__(self, dp, rop):
        self.dp = dp
        self.name = int(dp*1000000)
        self.rop = rop

    def __repr__(self):
        return "[dp = " + str(self.name) + ", rop = " + str(self.rop) + "]"

class Gas:
    def __init__(self, R, P, T):
        self.R = R
        self.P = P
        self.T = T
        self.muf = airprops.mu_air(self.T)
        self.rof = airprops.density_air(self.R, self.P, self.T)

    def __repr__(self):
        return "[R = " + str(self.R) + ", P = " + str(self.P) + " , T = " + str(self.T) + "]"

class Flow:
    def __init__(self, l, u, gas, particle):
        self.u = u
        self.gas = gas
        self.particles = particle
        self.l = l
        self.w_set_c = C_Re.w_settling_corrected(self.particles.rop,
                                                 self.particles.dp,
                                                 self.gas.muf,
                                                 self.gas.rof,
                                                 self.u,
                                                 9.81)

    def __repr__(self):
        return "u = " + str(self.u) + ", particles: " + repr(self.particles)

    def caaasmp(self):
        k = 80001
        i = 1
        dt = 0.00005
        self.v = np.zeros(k)
        self.Re_p = np.zeros(k - 1)
        self.cd = np.zeros(k - 1)
        self.time = np.zeros(k)
        self.x = np.zeros(k)
        while i < k:
            self.Re_p[i - 1] = C_Re.reynolds_particle(self.u,
                                                      self.v[i - 1],
                                                      self.particles.dp,
                                                      self.gas.muf,
                                                      self.gas.rof)
            self.cd[i - 1] = C_Re.c_drag(self.Re_p[i - 1])
            self.v[i] = RK.f_plus(dt,
                                  9.81,
                                  self.cd[i - 1],
                                  self.u,
                                  self.v[i - 1],
                                  self.gas.rof,
                                  self.particles.rop,
                                  self.particles.dp)
            self.x[i] = self.x[i - 1] + (self.v[i] + self.v[i - 1]) / 2.0 * dt
            self.time[i] = self.time[i - 1] + dt
            i += 1
    """def caaasmp_height(self):
        k = 10001
        i = 1
        dt = 0.00005
        #self.vh = - self.v[-1] * np.ones(k)
        self.vh = - 1 * self.v[self.l_element] * np.ones(k)
        #self.vh = np.zeros(k)
        #self.v = -np.ones(k)*31
        self.Re_ph = np.zeros(k - 1)
        self.cdh = np.zeros(k - 1)
        self.timeh = np.zeros(k)
        self.xh = np.zeros(k)
        while self.vh[i-1] < 0:
            self.Re_ph[i - 1] = C_Re.reynolds_particle(self.u,
                                                      self.vh[i - 1],
                                                      self.particles.dp,
                                                      self.gas.muf,
                                                      self.gas.rof)
            self.cdh[i - 1] = C_Re.c_drag(self.Re_ph[i - 1])
            self.vh[i] = RK.f_plus(dt,
                                  9.81,
                                  self.cdh[i - 1],
                                  self.u,
                                  self.vh[i - 1],
                                  self.gas.rof,
                                  self.particles.rop,
                                  self.particles.dp)
            self.xh[i] = self.xh[i - 1] + (self.vh[i] + self.vh[i - 1]) / 2.0 * dt
            self.timeh[i] = self.timeh[i - 1] + dt
            self.height = self.xh[i]
            i += 1

    def velocity_lenght(self):
        self.l_element = np.argmin(np.abs(self.x - self.l))"""

if __name__ == '__main__':
    air = Gas(287.0, 101325.0, 293.15)

    u = np.linspace(-20,0.0001,5)
    partlist = [Particles(160*1e-6, 2550.0), Particles(8000*1e-6, 1000.0)]

    flows = []

    for velocity in u:
        for p in partlist:
            #flows[velocity, p] = Flow(3.0, velocity, air, p)
            flows.append(Flow(3.0, velocity, air, p))

    for flow in flows:
        flow.caaasmp()

    for flow in flows:
        if flow.particles.name == 160:
            plt.plot(flow.x, flow.v, linestyle="--")
        else:
            plt.plot(flow.x, flow.v)
    plt.grid()
    plt.show()
    #for flow in flows:
    #    print(flow, ":")
    #    flow.caaasmp()
    #    print("-------------------------")
    """for i in partlist:
        f = {}
        for j in u:
            f[j] = Flow(3.0, j, air, i)
        flows[i.dp] = f
    for i in flows:
        for j in flows[i]:
            print('u = ', flows[i][j].u)
            flows[i][j].caaasmp()
            flows[i][j].velocity_lenght()
            flows[i][j].caaasmp_height()
    for i in flows:
        v = []
        u = []
        h = []
        for j in flows[i]:
            indx = flows[i][j].l_element
            #v.append(flows[i][j].v[indx]-flows[i][j].u)
            v.append(flows[i][j].v[indx]-flows[i][j].u)
            u.append(flows[i][j].u)
            h.append(-flows[i][j].height)
        plt.plot(u,v)
    #plt.axhline(1.069)
    plt.grid(True)
    plt.show()
    plt.plot(u, h)
    plt.show()"""


# encoding:utf-8
import copy
import random
import bisect #bisect_left　これで二部探索の大小検索が行える
import fractions #最小公倍数などはこっち
import math
import sys
import bisect

mod = 10**9+7
sys.setrecursionlimit(mod) # 再帰回数上限はでdefault1000

def LI(): return list(map(int, sys.stdin.readline().split()))

N = int(input())
A = LI()

position = {}
for i in range(1, N+1):
    position[i] = []

for i, a in enumerate(A):
    position[a].append(i)
# print(position)

diff = [0 for i in range(2*N - 1)]
for key in position.keys():
    # print(position[key])
    l, r = position[key]
    diff[abs(l - r)] += 1

ans = 0
for i in range(1, 2*N-1):
    ans += diff[i]
    print(ans)

print(ans)


from fractions import Fraction
from random import randint, choice

# val is the secret (the value when x = 0).
# n is the number of points to distribute.
# k is the number of points needed to reconstruct the polynomial.
# Return a list of n random points such that x != 0 for every point.
def split(val, n, k):
	r = 1000
	neg = [1, -1]
	coefs = [choice(neg) * Fraction(randint(1, r), randint(1, r)) for i in range(0, k)]
	coefs.insert(0, val)

	xs = []
	ys = []
	for i in range(n):
		p = 0
		while p == 0 or p in xs:
			p = randint(-r, r)

		xs.append(p)
		ys.append(polyVal(coefs, p))

	keys = [(xs[i], ys[i]) for i in range(len(xs))]
	return keys

#given list of coefficents for p, find p(x)
def polyVal(coefs, x):
	val = 0
	i = 0
	for c in coefs:
		val += c * (x**i)
		i += 1
	return val

# points is a list of shares.
# x is the x-coordinate in which to compute the secret at.
# Return the computed secret value.
def interpolate(points, x): #using lagrange interpolation from wiki
	val = 0
	for j in range(len(points)):
		y = points[j][1]
		l = basis(points, x, j)
		val += y*l
	return val

def basis(points, x, j):
	val = 1
	for m in range(len(points)):
		if not m == j:
			a = (x - points[m][0])
			b = (points[j][0] - points[m][0])
			v=Fraction(a)/Fraction(b)
			val *= v
	return val

if __name__ == "__main__":
	pass


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from refinery.units import Arg, Unit

class netbios(Unit):
    """
    Encodes and decodes strings using the same algorithm that is used for NetBIOS
    labels. Each byte 0xUL is encoded as two bytes, which are the sum of 0xU and
    0xL with an offset character, respectively. The default offset is the capital
    letter A.
    """

    def __init__(self, key: Arg(help="Provide a single letter to use as the offset.") = B'A'):
        if len(key) != 1:
            raise ValueError("The key must be a binary string of length exactly 1")
        super().__init__(key=key[0])

    def reverse(self, data):
        result = bytearray(2 * len(data))
        for k, byte in enumerate(data):
            hi, lo = byte >> 4, byte & 15
            result[2 * k + 0] = hi + self.args.key
            result[2 * k + 1] = lo + self.args.key
        return result

    def process(self, data):
        def merge(it):
            while True:
                try:
                    hi = next(it) - self.args.key
                    lo = next(it) - self.args.key
                    if hi not in range(16) or lo not in range(16):
                        raise ValueError(F'Invalid character encoding detected: hi={hi:X}, lo={lo:X}.')
                    yield (hi << 4) | lo
                except StopIteration:
                    break
        return bytearray(merge(iter(data)))


import json
import os
import requests

class Region():

    root = os.path.dirname(os.path.abspath(__file__))

    def files(self, fileName):
        with open(self.root+"/"+fileName+".json", encoding="utf8") as jsonData:
            data = jsonData.read()

        return data

    def callCountry(self, country=None):
        countries = json.loads(self.files('countries'))

        countriesDict = {}
        for i in range(len(countries)):
            countriesDict[countries[i]['iso2']] = countries[i]['name']

        return countriesDict

    def callState(self, countryID):
        states = json.loads(self.files('states'))
        statesDict = {}
        for i in range(len(states)):
            if (states[i]['country_code'].upper() == countryID.upper()):
                statesDict[states[i]['state_code']] = states[i]['name']

        return statesDict

    def callCity(self, stateID):
        cities = json.loads(self.files('cities'))

        citiesDict = {}
        for i in range(len(cities)):
            if (cities[i]['state_code'] == stateID.upper()):
                citiesDict[cities[i]['name']] = cities[i]['name']

        return citiesDict

    # def remoteCities(self, hi):
    #     data = requests.get('https://raw.githubusercontent.com/dr5hn/countries-states-cities-database/master/cities.json').json()
    #     with open(self.root+"/cities.json", 'w', encoding='utf-8') as f:
    #         json.dump(data, f, ensure_ascii=False, indent=4)


#!/usr/bin/env python
# encoding: utf-8
#@contact: 罗湘飞
#@time: 2019/4/19 17:55
import random
# from pip._internal import main as install_requirements
# install_requirements(['install', '-r', 'requirements.txt'])
from requestProt.K8S_call.K8S_call_log_save import K8sCallLogSave
from common.timeDecorator import my_time
import threading   #多线程模块

a=(i for i in range(1,100000000))
def port_fun():
    phone = random.randint(10000000, 99999999)
    data = '{"caller":"136%s","callTime":"2019-04-16 00:00:00","region":"020","time":1,"type":"1"}' % phone
    try:
        result=K8sCallLogSave().call_log_dave_port(data)
        print(result+"次数为%s"%next(a))
    except Exception as e:
        print("连接失败")

def run(bingfa):
    '''并发测试--多线程'''
    a=[]
    for i in range(int(bingfa)):                    #设置并发数量
        b=threading.Thread(target=port_fun)
        a.append(b)
        b.start()
    for c in a:
        c.join()
@my_time
def run_num(bingfa,xunhuan):
    '''测试循环次数'''
    for i in range(int(xunhuan)):                      #设置循环次数
        run(bingfa)
if __name__ == '__main__':
    bingfa=1000
    xunhuan=1000
    run_num(bingfa,xunhuan)
    CGL = (next(a) - 1) / (bingfa * xunhuan)
    print("成功为【{}】次，执行请求【{}】次，成功率为【{}】".format((next(a) - 2), (bingfa * xunhuan), str(CGL * 100) + '%'))


import os
from datetime import datetime as dt
counter = 0
flag = 0
pull_counter=0
date_Time = dt.now()
#print(date_Time)
#print("Previous working Directory : {0}".format(os.getcwd()))
#os.chdir('d:')
f = open("gitlogs.txt","w+")
log_File_Path = "gitlogs.txt"
print("GitLogs file is created in: {0}".format(os.getcwd()))
#os.chdir('UIProject_OH')
print("Present working directory : {0}".format(os.getcwd()))
###check for status###
status = os.popen("git status")

f.write("----------------------------------------------------------\n")
f.write("GIT STATUS:\n")
f.write(date_Time.strftime("%m/%d/%Y, %H:%M:%S:%f")+"\n\n")
status = os.popen('git status')
f.write(status.read())
f.write("\n----------------------------------------------------------\n")
f.close()
Text = "Changes not staged for commit:"
fileR = open(log_File_Path,"r")
status_data = fileR.read()
if Text in status_data:
    print("\nThere are pending changes, Now going to stash the changes\n")
    flag = 1
else:
    print("There are no pending changes, Ready to Pull...")
    flag = 0
fileR.close()
if flag == 1:
    ###Stash the changes###
    stash = os.popen("git stash")
    stash_data = stash.read()
    stashed_Text = "Saved working directory"
    if stashed_Text in stash_data:
        print("Stashed\n")
        file_Stash = open(log_File_Path,"a+")
        file_Stash.write("\nGIT STASH:\n")
        #print("\nSTASH: \n {0}".format(file_Stash.read()))
        file_Stash.write(stash.read())
        file_Stash.write("\n----------------------------------------------------------\n")
        file_Stash.close()
        ###check for status again###
        status1 = os.popen('git status')
        status1_data = status1.read()
        #print (status1.read())
        status1_Text = "working tree clean"
        if status1_Text in status1_data:
            print("Working tree is clean, \nReady to Pull...\n")
            GitPull()
        else :
            print("Error, Please proceed manually\n")
    else :
        print("Not Stashed, Please proceed masnually\n")
else :
    print("Ready to Pull...\n")
    GitPull()

def GitPull():
    file_pull = open(log_File_Path,"a+")
    pull = os.popen('git pull')
    file_pull.write("\n----------------------------------------------------------\n")
    file_pull.write("GIT PULL:\n")
    file_pull.write(pull.read())
    print("Pull Successfull\n")
    print("Restoring Stashed Data\n")
    os.popen('git stash pop')
    print("Stashed data restored, You can now proceed...!\n")
    exit()

exit()

#os.system('cmd /k "git status"')


from ipaddress import NetmaskValueError
import socket
import pickle

class Network:
    def __init__(self):
        self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.host = "localhost"
        self.port = 5555
        self.addr = (self.host, self.port)
        self.board = self.connect()
        self.board = pickle.loads(self.board)

    def connect(self):
        self.client.connect(self.addr)
        return self.client.recv(4096*8)

    def disconnect(self):
        self.client.close()

    def send(self, message):
        """
        :param data: str
        :return: str
        """
        self.client.send(message.encode())
        data = self.client.recv(1024).decode()  # receive response
        print('Received from server: ' + data)  # show in terminal


""""
Performs the PCA to find the most relevant axis and compute the features of UFSR
on the points (000),mean(100,-100), mean(010,0-10), mean(001,00-1)

Compute distances
"""

import scipy as np
import os
import csv
from sklearn.decomposition import PCA
import pandas as pd
import time

# returns the centroid of an array
def centroidnp(arr):
    length, dim = arr.shape
    sums = np.array([np.sum(arr[:, i]) for i in range(dim)])
    center = sums / length
    return center

# takes an array of positions and returns the positions of the ufsr pivots
# distance_of_points represents the distance in Angstroms of the pivots to the centroids
def find_pivots(X, distance_of_points=4):
    center = centroidnp(X)
    # make PCA
    n_components = X.shape[1]
    pca = PCA(n_components=n_components)
    pca.fit(X)
    pca_vectors = pca.components_

    # find usfr normalized eigenvectors and add them to the centroid to find the USFR starting point
    full_vectors = np.concatenate((pca_vectors, -1 * pca_vectors))
    # print(full_vectors)
    full_points = distance_of_points * full_vectors + center
    return full_points

# random distributions of 30 nodes, 1000 examples
# find pivots : 0.28 seconds ---
# dist to pivots : 0.15 seconds ---
# moments : 1.26 (6 moments), 0.55 (3), 0.33 (2) seconds ---
# encode a matrix of positions with the distribution of its moments
# number_moments is the number of moments computed, for instance, [mu, sigma] means number_moments = 2
def encode(X, number_of_moments=3):
    n_dimension = X.shape[1]
    center = centroidnp(X)
    full_points = find_pivots(X)

    # get distance matrix where rows are pivot points and columns are data points
    dist_to_centroid = np.array([[np.linalg.norm(X[j] - center) for j in range(X.shape[0])]])
    dist_matrix = np.spatial.distance_matrix(full_points, X)
    # code à la main, 4 fois plus lent
    # dist_matrix = np.array(
    #     [[np.linalg.norm(X[j] - full_point) for j in range(X.shape[0])] for full_point in full_points])

    # aggregate the symetric pivots
    dist_ufsr = np.array([(dist_matrix[i] + dist_matrix[n_dimension + i]) * 0.5 for i in range(n_dimension)])

    # add the distance to center of mass, dist_ufsr is now a matrix with distribution of distances wrt
    # (000),mean(100,-100), mean(010,0-10), mean(001,00-1) as rows
    dist_ufsr = np.concatenate((dist_to_centroid, dist_ufsr))

    # get the features
    means = np.array([[np.mean(dist_ufsr[i]) for i in range(n_dimension + 1)]])
    means = np.transpose(means)
    moments = np.array(
        [[np.stats.moment(dist_ufsr[i], j) for j in range(2, number_of_moments + 1)] for i in range(n_dimension + 1)])
    ufsr_feature = np.concatenate((means, moments), axis=1)

    # mean, moment2, moment3..., moment6, mean, moment1... for each pivot
    ufsr_feature = ufsr_feature.ravel()

    return ufsr_feature

# returns the similarity measure of two arrays
def usfr_similarity(x, y):
    x = np.asarray(x)
    y = np.asarray(y)
    try:
        similarity_measure = 1 / (1 + np.sum(abs(x - y)))
    except:
        return None
    return similarity_measure

# copied from scipy documentation
def distance_matrix(x, y, output_name=0):
    x = np.asarray(x)
    m, k = x.shape
    y = np.asarray(y)
    n, kk = y.shape
    result = np.empty((m, n), dtype=float)  # FIXME: figure out the best dtype
    for i in range(m):
        result[i, :] = [usfr_similarity(x[i], y_j) for y_j in y]
        print(result[i, :])
        if output_name:
            with open(os.path.join('../data/output_csv', output_name + '.csv'), "a", newline='') as fp:
                wr = csv.writer(fp, dialect='excel')
                wr.writerow(result[i, :])
    return result

'''
apply these computations to data
'''

def build_colnames(moments_kept, centroids):
    # build the names of the columns
    lists = [['m' + str(i + 1) + 'c' + str(j + 1) for i in range(moments_kept)] for j in range(centroids)]
    L = []
    for list in lists:
        L += list
    colnames = ['pdb', 'ligand'] + L
    return colnames

def csv_reader(input_name):
    df = pd.read_csv(os.path.join('../data/output_csv', input_name + '.csv'))
    return df

# keep only a few moments from a pandas df
# returns a df with name, ligand and moments compressed
def reduce_number_of_moment(df, moments_after, moments_before=3, centroids=4):
    # keep the name and ligand, then select only the first moments for each feature
    criteria = [True, True] + [(i % moments_before < moments_after) for i in range(moments_before * centroids)]
    reduced_df = df[df.columns[criteria]]
    return reduced_df

def DM_from_csv(input_name, output_name, moments_kept=0):
    df = csv_reader(input_name)
    if moments_kept:
        df = reduce_number_of_moment(df, moments_kept)
    values = df[df.columns[2:]]
    DM = distance_matrix(values, values, output_name)

if __name__ == "__main__":
    import matplotlib.pyplot as plt

    X = np.array([[2, 1], [5, 2], [1, 1], [2, 2], [3, 2], [5, 5], [6, 6.57]])

    X_x = X[:, 0]
    X_y = X[:, 1]
    center = centroidnp(X)
    full_points = find_pivots(X)
    usfr_x = full_points[:, 0]
    usfr_y = full_points[:, 1]

    plt.scatter(X_x, X_y, alpha=0.5)
    plt.scatter(center[0], center[1], alpha=0.5)
    plt.scatter(usfr_x, usfr_y, alpha=0.5)
    plt.show()

'''
plot

x = X[:, 0]
y = X[:, 1]

fig, ax = plt.subplots(figsize=(10, 5), ncols=3, nrows=2)

left = 0.125  # the left side of the subplots of the figure
right = 0.9  # the right side of the subplots of the figure
bottom = 0.1  # the bottom of the subplots of the figure
top = 0.9  # the top of the subplots of the figure
wspace = .5  # the amount of width reserved for blank space between subplots
hspace = 1.1  # the amount of height reserved for white space between subplots

# This function actually adjusts the sub plots using the above paramters
plt.subplots_adjust(
    left=left,
    bottom=bottom,
    right=right,
    top=top,
    wspace=wspace,
    hspace=hspace
)

y_title_margin = 0.8
ax[0][0].set_title("zero", y=y_title_margin)
ax[0][1].set_title("first", y=y_title_margin)
ax[0][2].set_title("second", y=y_title_margin)

sns.distplot(dist_ufsr[0], ax=ax[0][0], kde=False, rug=True)
sns.distplot(dist_ufsr[1], ax=ax[0][1], kde=False, rug=True)
sns.distplot(dist_ufsr[2], ax=ax[0][2], kde=False, rug=True)

plt.show()

usfr_x = full_points[:, 0]
usfr_y = full_points[:, 1]
plt.scatter(x, y, alpha=0.5)
plt.scatter(center[0], center[1], alpha=0.5)
plt.scatter(usfr_x, usfr_y, alpha=0.5)
plt.show()

# print(preprocessing.normalize(pca.components_, norm='l2'))
'''


from itertools import product
from pygame import Rect, RLEACCEL
import pygame

__all__ = ('GraphicBox', 'draw_text')

class GraphicBox(object):
    """Generic class for drawing graphical boxes

    box = GraphicBox('border.png')
    box.draw(rect)

    The border graphic must contain 9 tiles laid out in a box.
    Each tile must be the same size
    """
    def __init__(self, image, hollow=False):
        iw, ih = image.get_size()
        self.tw = int(iw / 3)
        self.th = int(ih / 3)
        self.hollow = hollow

        tiles = [image.subsurface((x, y, self.tw, self.th))
                 for x, y in product(range(0, iw, self.tw),
                                     range(0, ih, self.th))]

        self.tiles = dict(zip("nw w sw n c s ne e se".split(), tiles))
        self.background = self.tiles['c'].get_at((0, 0))

        if self.hollow:
            ck = self.tiles['c'].get_at((0, 0))
            [t.set_colorkey(ck, RLEACCEL) for t in self.tiles.values()]

    def draw(self, surface, rect=None):
        if rect is None:
            rect = surface.get_rect()

        ox, oy, w, h = Rect(rect)
        surface_blit = surface.blit
        tiles = self.tiles
        tw = self.tw
        th = self.th

        if not self.hollow:
            p = product(range(tw + ox, w - tw + ox, tw),
                        range(th + oy, h - th + oy, th))
            [surface_blit(tiles['c'], (x, y)) for x, y in p]

        for x in range(tw + ox, w - tw + ox, tw):
            surface_blit(tiles['n'], (x, oy))
            surface_blit(tiles['s'], (x, h - th + oy))

        for y in range(th + oy, h - th + oy, th):
            surface_blit(tiles['e'], (w - tw + ox, y))
            surface_blit(tiles['w'], (ox, y))

        surface_blit(tiles['nw'], (ox, oy))
        surface_blit(tiles['ne'], (w - tw + ox, oy))
        surface_blit(tiles['sw'], (ox, h - th + oy))
        surface_blit(tiles['se'], (w - tw + ox, h - th + oy))

def draw_text(surface, text, rect, font=None, fg_color=None, bg_color=None, aa=False):
    """ draw some text into an area of a surface

    automatically wraps words
    returns size and any text that didn't get blit
    passing None as the surface is ok
    """
    if fg_color is None:
        fg_color = (0, 0, 0)

    total_width = 0
    rect = Rect(rect)
    y = rect.top
    line_spacing = -2

    if font is None:
        full_path = pygame.font.get_default_font()
        font = pygame.font.Font(full_path, 16)

    # get the height of the font
    font_height = font.size("Tg")[1]

    # for very small fonts, turn off antialiasing
    if font_height < 16:
        aa = 0
        bg_color = None

    while text:
        char_index = 1

        # determine if the row of text will be outside our area
        if y + font_height > rect.bottom:
            break

        # determine maximum width of line
        line_width = font.size(text[:char_index])[0]
        while line_width < rect.width and char_index < len(text):
            if text[char_index] == "\n":
                text = text[:char_index] + text[char_index + 1:]
                break

            char_index += 1
            line_width = font.size(text[:char_index])[0]
            total_width = max(total_width, line_width)
        else:
            # if we've wrapped the text, then adjust the wrap to the last word
            if char_index < len(text):
                char_index = text.rfind(" ", 0, char_index) + 1

        if surface:
            # render the line and blit it to the surface
            if bg_color:
                image = font.render(text[:char_index], 1, fg_color, bg_color)
                image.set_colorkey(bg_color)
            else:
                image = font.render(text[:char_index], aa, fg_color)

            surface.blit(image, (rect.left, y))

        y += font_height + line_spacing

        # remove the text we just blitted
        text = text[char_index:]

    return total_width, text


from PIL import Image
import os

def resize_image(image_path, size):
    original_image = Image.open(image_path)
    resized_image = original_image.resize(size)
    return resized_image

def resize_images_in_directory(directory, size):
    for filename in os.listdir(directory):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            image_path = os.path.join(directory, filename)
            resized_image = resize_image(image_path, size)
            resized_image.save(f"resized_{filename}")

directory = input("Enter the path of the directory containing the images: ")
width = int(input("Enter the desired width for the resized images: "))
height = int(input("Enter the desired height for the resized images: "))
size = (width, height)
resize_images_in_directory(directory, size)


from load_dataset import ImageDataset
import tensorflow as tf

imd = ImageDataset('icons/')
images = imd.get_image_arrays()
labels = imd.get_image_lables()

layers = [
    tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), padding="same", activation=tf.nn.relu,
                                                                  input_shape=images.shape[1:]),
    tf.keras.layers.MaxPool2D(pool_size=(2,2), strides=(2,2)),
    tf.keras.layers.Conv2D(filters=32, kernel_size=(3,3), padding="same", activation=tf.nn.relu),
    tf.keras.layers.MaxPool2D(pool_size=(2,2), strides=(2,2)),
    tf.keras.layers.Conv2D(filters=64, kernel_size=(3,3), padding="same", activation=tf.nn.relu),
    tf.keras.layers.MaxPool2D(pool_size=(2,2), strides=(2,2)),
    tf.keras.layers.Conv2D(filters=128, kernel_size=(3,3), padding="same", activation=tf.nn.relu),
    tf.keras.layers.MaxPool2D(pool_size=(2,2), strides=(2,2)),
    tf.keras.layers.Conv2D(filters=256, kernel_size=(3,3), padding="same", activation=tf.nn.relu),
    tf.keras.layers.MaxPool2D(pool_size=(2,2), strides=(2,2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(units=512, activation=tf.nn.relu),
    tf.keras.layers.Dense(units=256, activation=tf.nn.relu),
    tf.keras.layers.Dense(1)
]
model = tf.keras.Sequential(layers)
model.compile(optimizer=tf.keras.optimizers.Adam(), loss='mse',metrics=['mae', 'mse'])
model.fit(images, labels, epochs=20, batch_size=20)
model.save_weights('model.tf')

eval_model = tf.keras.Sequential(layers)
eval_model.load_weights('model.tf')

print(eval_model.predict(images[:10]))
print(labels[:10])


import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import streamlit as st
import numpy as np

@ticker.FuncFormatter
def major_formatter(x, pos):
    return "d%d" % (x / 24 / 3600)

def plot(days, sources, color="#E28BFD"):
    fig, ax = plt.subplots(figsize=(15, 5), nrows=1, ncols=1)
    for s in sources:
        data = s[0].xt()
        ax.plot(data[1], data[0], color, label=s[1])
    ax.legend(loc="lower center", fontsize=18.0)
    ax.xaxis.set_ticks(np.arange(0, 24 * 3600 * (days + 1), 24 * 3600))
    ax.xaxis.set_tick_params(labelrotation=45)
    ax.xaxis.set_major_formatter(major_formatter)

    st.write(fig)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Microsoft Crypto API structures
"""
import enum

from Cryptodome.PublicKey import RSA
from Cryptodome.Math.Numbers import Integer

from refinery.lib.structures import Struct, StructReader

class _ENUM(enum.IntEnum):
    def __str__(self): return self.name
    def __repr__(self): return self.name

class TYPES(_ENUM):
    KEYSTATEBLOB               = 0xC # noqa
    OPAQUEKEYBLOB              = 0x9 # noqa
    PLAINTEXTKEYBLOB           = 0x8 # noqa
    PRIVATEKEYBLOB             = 0x7 # noqa
    PUBLICKEYBLOB              = 0x6 # noqa
    PUBLICKEYBLOBEX            = 0xA # noqa
    SIMPLEBLOB                 = 0x1 # noqa
    SYMMETRICWRAPKEYBLOB       = 0xB # noqa

class ALGORITHMS(_ENUM):
    CALG_3DES                  = 0x00006603 # noqa
    CALG_3DES_112              = 0x00006609 # noqa
    CALG_AES                   = 0x00006611 # noqa
    CALG_AES_128               = 0x0000660e # noqa
    CALG_AES_192               = 0x0000660f # noqa
    CALG_AES_256               = 0x00006610 # noqa
    CALG_AGREEDKEY_ANY         = 0x0000aa03 # noqa
    CALG_CYLINK_MEK            = 0x0000660c # noqa
    CALG_DES                   = 0x00006601 # noqa
    CALG_DESX                  = 0x00006604 # noqa
    CALG_DH_EPHEM              = 0x0000aa02 # noqa
    CALG_DH_SF                 = 0x0000aa01 # noqa
    CALG_DSS_SIGN              = 0x00002200 # noqa
    CALG_ECDH                  = 0x0000aa05 # noqa
    CALG_ECDH_EPHEM            = 0x0000ae06 # noqa
    CALG_ECDSA                 = 0x00002203 # noqa
    CALG_ECMQV                 = 0x0000a001 # noqa
    CALG_HASH_REPLACE_OWF      = 0x0000800b # noqa
    CALG_HUGHES_MD5            = 0x0000a003 # noqa
    CALG_HMAC                  = 0x00008009 # noqa
    CALG_KEA_KEYX              = 0x0000aa04 # noqa
    CALG_MAC                   = 0x00008005 # noqa
    CALG_MD2                   = 0x00008001 # noqa
    CALG_MD4                   = 0x00008002 # noqa
    CALG_MD5                   = 0x00008003 # noqa
    CALG_NO_SIGN               = 0x00002000 # noqa
    CALG_OID_INFO_CNG_ONLY     = 0xffffffff # noqa
    CALG_OID_INFO_PARAMETERS   = 0xfffffffe # noqa
    CALG_PCT1_MASTER           = 0x00004c04 # noqa
    CALG_RC2                   = 0x00006602 # noqa
    CALG_RC4                   = 0x00006801 # noqa
    CALG_RC5                   = 0x0000660d # noqa
    CALG_RSA_KEYX              = 0x0000a400 # noqa
    CALG_RSA_SIGN              = 0x00002400 # noqa
    CALG_SCHANNEL_ENC_KEY      = 0x00004c07 # noqa
    CALG_SCHANNEL_MAC_KEY      = 0x00004c03 # noqa
    CALG_SCHANNEL_MASTER_HASH  = 0x00004c02 # noqa
    CALG_SEAL                  = 0x00006802 # noqa
    CALG_SHA1                  = 0x00008004 # noqa
    CALG_SHA_256               = 0x0000800c # noqa
    CALG_SHA_384               = 0x0000800d # noqa
    CALG_SHA_512               = 0x0000800e # noqa
    CALG_SKIPJACK              = 0x0000660a # noqa
    CALG_SSL2_MASTER           = 0x00004c05 # noqa
    CALG_SSL3_MASTER           = 0x00004c01 # noqa
    CALG_SSL3_SHAMD5           = 0x00008008 # noqa
    CALG_TEK                   = 0x0000660b # noqa
    CALG_TLS1_MASTER           = 0x00004c06 # noqa
    CALG_TLS1PRF               = 0x0000800a # noqa

class BCRYPT_MAGIC(_ENUM):
    BCRYPT_RSAPUBLIC_MAGIC      = 0x31415352 # noqa
    BCRYPT_RSAPRIVATE_MAGIC     = 0x32415352 # noqa
    BCRYPT_RSAFULLPRIVATE_MAGIC = 0x33415352 # noqa

class BLOBHEADER(Struct):
    def __init__(self, reader: StructReader):
        t, self.version, self.reserved, a = reader.read_struct('BBHI')
        self.type = TYPES(t)
        self.algorithm = ALGORITHMS(a)

class PLAINTEXTKEYBLOB(Struct):
    def __bytes__(self): return bytes(self.data)

    def __init__(self, reader: StructReader):
        self.size = reader.u32()
        self.data = reader.read(self.size)

class SIMPLEBLOB(Struct):
    def __bytes__(self): return bytes(self.data)

    def __init__(self, reader: StructReader):
        self.magic = reader.read(4)
        if self.magic != B'\0\xA4\0\0':
            raise ValueError(F'Invalid magic bytes: {self.magic.hex(":").upper()}')
        self.data = reader.read(0x100)

class BCRYPT_RSAKEY_BLOB(Struct):
    def __init__(self, reader: StructReader):
        magic, bits, e_size, n_size, p_size, q_size = reader.read_struct('<6L')
        e_size *= 8
        n_size *= 8
        self.magic = BCRYPT_MAGIC(magic)
        reader.bigendian = True
        self.exponent = reader.read_integer(e_size)
        self.modulus = reader.read_integer(n_size)
        self.bit_size = bits
        if self.has_private_key:
            p_size *= 8
            q_size *= 8
            self.prime1 = reader.read_integer(p_size)
            self.prime2 = reader.read_integer(q_size)
            if self.magic is BCRYPT_MAGIC.BCRYPT_RSAFULLPRIVATE_MAGIC:
                self.exp1 = reader.read_integer(p_size)
                self.exp2 = reader.read_integer(q_size)
                self.coefficient = reader.read_integer(p_size)
                self.exp_private = reader.read_integer(n_size)
            else:
                self.exp1 = None
                self.exp2 = None
                self.coefficient = None
                self.exp_private = None
        else:
            self.prime1 = None
            self.prime2 = None

    @property
    def has_private_key(self):
        return self.magic in (
            BCRYPT_MAGIC.BCRYPT_RSAPRIVATE_MAGIC,
            BCRYPT_MAGIC.BCRYPT_RSAFULLPRIVATE_MAGIC
        )

    def convert(self, force_public=False):
        components = self.modulus, self.exponent
        if not force_public and self.has_private_key:
            components += self.exp_private, self.prime1, self.prime2
        return RSA.construct(components)

class RSAPUBKEY(Struct):
    def __init__(self, reader: StructReader):
        self.magic = reader.read(4)
        if self.magic not in (B'RSA2', B'RSA1'):
            raise ValueError(F'Invalid signature: {self.magic.hex()}')
        self.size, self.exponent = reader.read_struct('II')
        if self.size % 8 != 0:
            raise ValueError(F'The bitlength {self.size} is not a multiple of 8.')
        self.modulus = reader.read_integer(self.size)

    def convert(self):
        return RSA.construct((self.modulus, self.exponent))

    def __str__(self):
        return self.key().export_key(format='PEM')

    def __bytes__(self):
        return str(self).encode('ascii')

class PRIVATEKEYBLOB(Struct):
    def __init__(self, reader: StructReader):
        self.pub = RSAPUBKEY(reader)
        halfsize = self.pub.size // 2
        self.prime1 = reader.read_integer(halfsize)
        self.prime2 = reader.read_integer(halfsize)
        self.exp1 = reader.read_integer(halfsize)
        self.exp2 = reader.read_integer(halfsize)
        self.coefficient = reader.read_integer(halfsize)
        self.exponent = reader.read_integer(self.pub.size)
        self._check()

    def _check(self):
        if self.pub.modulus // self.prime1 != self.prime2:
            raise ValueError('Product of primes does not equal the modulus.')
        from math import gcd
        a = self.prime1 - 1
        b = self.prime2 - 1
        totient = (a * b) // gcd(a, b)
        if self.pub.exponent * self.exponent % totient != 1:
            raise ValueError('Public exponent is not a modular inverse of private exponent.')

    def convert(self):
        parameters = (
            self.pub.modulus,
            self.pub.exponent,
            self.exponent,
            self.prime1,
            self.prime2,
            self.coefficient
        )
        try:
            return RSA.construct(parameters, consistency_check=True)
        except ValueError as V:
            try:
                return RSA.RsaKey(
                    n=Integer(self.pub.modulus),
                    e=Integer(self.pub.exponent),
                    d=Integer(self.exponent),
                    p=Integer(self.prime1),
                    q=Integer(self.prime2),
                    u=Integer(self.coefficient),
                )
            except Exception as E:
                raise E from V

    def __str__(self):
        return self.key().export_key(format='PEM')

    def __bytes__(self):
        return str(self).encode('ascii')

class DHPUBKEY(Struct):
    def __init__(self, reader: StructReader):
        self.magic, self.size = reader.read_struct('4sI')
        if self.magic not in (B'\0DH1', B'\0DH2'):
            raise ValueError(F'Invalid magic bytes: {self.magic.hex(":").upper()}')
        if self.size % 8 != 0:
            raise ValueError('Bit length is not a multiple of 8.')
        self.public = reader.read_integer(self.size)
        self.prime = reader.read_integer(self.size)
        self.generator = reader.read_integer(self.size)

    def __bytes__(self):
        raise NotImplementedError

class CRYPTOKEY(Struct):
    def __init__(self, reader: StructReader):
        self.header = BLOBHEADER(reader)
        if self.header.type in {
            TYPES.KEYSTATEBLOB,
            TYPES.OPAQUEKEYBLOB,
            TYPES.SYMMETRICWRAPKEYBLOB
        }:
            raise ValueError(F'Unsupported type: {self.header.type}')
        elif self.header.type == TYPES.PLAINTEXTKEYBLOB:
            self.key = PLAINTEXTKEYBLOB(reader)
        elif self.header.type == TYPES.SIMPLEBLOB:
            self.key = SIMPLEBLOB(reader)
        else:
            if self.header.algorithm not in {
                ALGORITHMS.CALG_RSA_KEYX,
                ALGORITHMS.CALG_RSA_SIGN
            }:
                raise ValueError(F'Unknown algorithm for {self.header.type}: {self.header.algorithm}')
            elif self.header.type == TYPES.PRIVATEKEYBLOB:
                self.key = PRIVATEKEYBLOB(reader)
            elif self.header.type == TYPES.PUBLICKEYBLOB:
                self.key = RSAPUBKEY(reader)
            elif self.header.type == TYPES.PUBLICKEYBLOBEX:
                self.key = DHPUBKEY(reader)


from decimal import Decimal
from twython import TwythonStreamer
import json
import os
import requests
import re

settings = {
    #'twitter_userid': "12137672", # gorillamania (for testing)
    'twitter_userid': "961445378", # officialmcafee
    'exchange_key': os.getenv("EXCHANGE_KEY"), # your Bittrex API key
    'exchange_secret': os.getenv("EXCHANGE_SECRET"), # your Bittrex API secret
    'buy_amount_btc': Decimal(2), # Amount to buy
    'limit_multiplier': Decimal(1.075), # How much over the current ticker to set the limit order
    'twitter_app_key': os.getenv("TWITTER_APP_KEY"),
    'twitter_app_secret': os.getenv("TWITTER_APP_SECRET"),
    'twitter_oauth_token': os.getenv("TWITTER_OAUTH_TOKEN"),
    'twitter_oauth_secret': os.getenv("TWITTER_OAUTH_SECRET"),
}

"""
Flow
1. Set up the streamer to listen to all tweets from twitter_username
2. When a tweet is found, parse for symbol
3. Is it a valid symbol for Bittrex?
4. BUY!
"""

def get_markets():
    # Fetch the current markets so we know if it is a valid symbol
    response = requests.get("https://bittrex.com/api/v1.1/public/getmarkets")
    assert response.status_code == 200
    markets = []
    for m in response.json()["result"]:
        markets.append(m["MarketCurrency"])
    return markets
MARKETS = get_markets()

def handle_tweet(tweet):
    match = re.search("[A-Z]{3,5}", tweet)
    if not match:
        print("No symbol found in tweet: '%s'" % tweet)
        return

    symbol = match.group(0)
    if symbol not in MARKETS:
        print("%s is not tradeable on Bittrex" % symbol)
        return

    buy_crypto(symbol)

def buy_crypto(symbol):
    from bitex import Bittrex
    e = Bittrex(settings["exchange_key"], settings["exchange_secret"])

    pair = "BTC-" + symbol
    print("Fetching current quote for %s" % pair)
    quote = e.ticker(pair).formatted
    current_price = Decimal(quote[6])

    price = current_price * settings["limit_multiplier"]
    amount = settings["buy_amount_btc"] / price
    print("Buying %s of %s at %s" % (amount, symbol, price))
    response = e.bid(pair, price, amount, False)
    print(response.json())

## Set up the streamer
# https://twython.readthedocs.io/en/latest/usage/streaming_api.html
class OurStreamer(TwythonStreamer):

    def on_success(self, data):
        if "user" in data and data["user"]["id_str"] == settings["twitter_userid"]:
            print(json.dumps(data, indent=2))
            if 'text' in data:
                handle_tweet(data["text"])

    def on_error(self, status_code, data):
        print("Error:", status_code, data)
        self.disconnect()

print("Setting up Twitter feed listener")
stream = OurStreamer(
        settings["twitter_app_key"],
        settings["twitter_app_secret"],
        settings["twitter_oauth_token"],
        settings["twitter_oauth_secret"])
stream.statuses.filter(follow=settings["twitter_userid"])
print("Listening...")


import numpy as np

class Markov_Chain:
    def __init__(self, state, initial, transition, step):
        # state_name_list is 1 by n array of int
        # initial_prob_dist is 1 by n array of floats
        # transition_matrix is n by n array of floats
        # step is the number of simulated state in a markov chain
        self.initial = initial
        self.transition = transition
        self.step = step
        self.state = state

    def simulator(self):
        # array index starts from ZERO!!

        # how many states are there
        # P is m by n array, len(P)=m
        num_of_state = len(self.initial)

        # initial an empty array to fill state index, 1 by step
        # range(3)=0,1,2
        final = [0 for i in range(self.step)]

        # array of cumulative initial probabilities, 1 by n+1
        c_initial = [0 for i in range(num_of_state + 1)]

        for i in range(num_of_state):
            # P[0][0:1]=P[0][0];
            c_initial[i + 1] = sum(self.initial[0:i + 1])

        # simulate the first state
        first = np.random.uniform(0, 1)
        for i in range(num_of_state):
            if (first >= c_initial[i]) & (first <= c_initial[i + 1]):
                temp = i

        # now using transition matrix to simulate the rest of the chain

        # matrix of cumulative transition matrix
        c_transition = [[0 for i in range(num_of_state + 1)]
                        for j in range(num_of_state)]

        for i in range(num_of_state):
            for j in range(num_of_state):
                c_transition[i][j + 1] = sum(self.transition[i][0:j + 1])

        # loop to update each state in the chain
        for i in range(self.step):
            final[i] = temp
            prob = c_transition[temp]
            first = np.random.uniform(0, 1)
            for j in range(num_of_state):
                if (first >= prob[j]) & (first <= prob[j + 1]):
                    temp = j

        # translate "final"(the state index) to state
        chain = [0 for i in range(self.step)]
        for i in range(self.step):
            chain[i] = self.state[final[i]]

        return chain

    def invariant_dist(self):
        # P=[1,2,3], R=asmatrix(P) is a 1 by 3 matrix,len(P)=3, len(R)=1
        # Q=[[1],[2],[3]], S=asmatrix(Q) is a 3 by 1 matrix, len(Q)=3, len(S)=3
        # matrix index can only use [i,j] instead of [i][j] for array

        # transfer array into matrix
        m_initial = np.asmatrix(self.initial)
        m_transition = np.asmatrix(self.transition)

        # matrix multiplication notice:
        # P*Q^2=(P*Q)^2 != P*Q*Q

        # track the number of m_transition in product_old
        mtply = 1
        product_old = m_initial * m_transition
        product_new = product_old * m_transition

        while (np.any(abs(product_new - product_old) > 0.000000000001)):
            product_old = product_new
            mtply = mtply + 1
            product_new = product_new * m_transition

        print(mtply)
        return product_old


from django.db import models

class Recipient(models.Model):
    name = models.CharField(max_length=250)
    description = models.TextField(null=True, blank=True)
    recipient_type = models.ForeignKey(to='RecipientType', on_delete=models.SET_NULL, null=True)
    iban = models.CharField(max_length=50, null=True, blank=True)
    street = models.CharField(max_length=100, null=True, blank=True)
    housenumber = models.CharField(max_length=50, null=True, blank=True)
    postcode = models.CharField(max_length=5, null=True, blank=True)
    city = models.CharField(max_length=50, null=True, blank=True)
    phonenumber = models.CharField(max_length=20, null=True, blank=True)

    active = models.BooleanField(default=True)
    user = models.ForeignKey(to='users.CustomUser', on_delete=models.SET_NULL, null=True)

    def __str__(self):
        return f"{self.name}"

class RecipientFavorit(models.Model):
    user = models.ForeignKey(to='users.CustomUser', on_delete=models.CASCADE, null=True)
    recipient = models.ForeignKey(to='Recipient', on_delete=models.CASCADE)

    def __str__(self):
        return f"{self.id}"

class RecipientType(models.Model):
    name = models.CharField(max_length=250, unique=True)
    description = models.TextField(null=True, blank=True)

    def __str__(self):
        return f"{self.name}"

class OpeningHours(models.Model):
    class Weekdays(models.IntegerChoices):
        Montag = 1
        Dienstag = 2
        Mittwoch = 3
        Donnerstag = 4
        Freitag = 5
        Samstag = 6
        Sonntag = 7

    recipient = models.ForeignKey(to='Recipient', on_delete=models.CASCADE)
    day = models.IntegerField(choices=Weekdays.choices)
    beginn_time = models.TimeField()
    end_time = models.TimeField()

    def __str__(self):
        return f"{self.recipient}"


# hangman game
import random

print("Welcome to hangman guess the right word before it's to late :) ")

myWord = ("Yeet", "Reet", "Bruh", "Potato", "Buddy")

word = random.choice(myWord)

choice = input("Type a word: ")

if choice == myWord:
	print("It's a match")
else:
	print("Not a match")

letter = input("Type a letter")
if letter in myWord:
	print("Letter is in the word")
else:
	print("Letter is not in the word")

count = 0
for s in myWord:
	if s == letter:
		print(count)
	count += 1