from colorama import Fore, Back, Style
from blessings import Terminal
import PIL
from PIL import Image
import subprocess
import ffmpy
import json
import os

term = Terminal()

def append_to_name(file_name, prefix=None):
    if prefix is None:
        return file_name

    file_base = os.path.dirname(file_name)
    file_name, ext = os.path.splitext(os.path.basename(file_name))
    out_file_name = os.path.join(file_base, file_name + prefix + ext)
    return out_file_name

def prepare_image_for_video(in_file_name: str, video_info: dict) -> str:
    out_file_name = append_to_name(in_file_name, '_forVideo')

    assert os.path.isfile(in_file_name)
    assert 'height' in video_info
    assert 'width' in video_info

    W = video_info['width']
    H = video_info['height']

    in_img = Image.open(in_file_name)
    in_img = in_img.resize((W,H))
    in_img.save(out_file_name)

    return out_file_name

def get_video_meta_data(filename: str):

    ff = ffmpy.FFprobe(
        inputs={filename: None},
        global_options=[
            '-v', 'quiet',
            '-print_format', 'json',
            '-show_format',
            '-show_streams']).run(stdout=subprocess.PIPE)

    meta = json.loads(ff[0].decode('utf-8'))
    return meta

def extract_video_info(meta: dict, verbose = False):
    info = {}

    assert len(meta['streams']) == 2
    assert meta['streams'][0]['codec_name'] == 'h264'

    keys_fun = {
        'height': lambda x: int(x),
        'width': lambda x: int(x),
        'r_frame_rate': lambda x: x.split('/')[0]
    }

    if verbose:
        print(Fore.BLUE + \
              term.underline + term.bold + \
              "Video stream information:" + \
              term.normal)

    for k,f in keys_fun.items():
        info[k] = f(meta['streams'][0][k])
        if verbose:
            print(term.blue + \
                  "{:20s}: {}".format(k, str(info[k])) + \
                  term.normal)

    return info

def extract_audio_info(meta: dict, verbose = False):

    info = {}

    assert len(meta['streams']) == 2
    assert meta['streams'][1]['codec_name'] == 'aac'

    keys_fun = {
        'sample_rate': lambda x: int(x),
        'channel_layout': lambda x: x,
    }

    if verbose:
        print(term.blue + \
              term.underline + term.bold + \
              "Audio stream information:" + \
              term.normal)

    for k,f in keys_fun.items():
        info[k] = f(meta['streams'][1][k])
        if verbose:
            print(term.blue + \
                  "{:20s}: {}".format(k, str(info[k])) + \
                  term.normal)
    return info


import os
import json
import requests
import numpy as np
import base64
import progressbar
import time
import warnings
import random

from PyPDF2 import PdfFileMerger, PdfFileReader, PdfFileWriter
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5
from requests.packages.urllib3.exceptions import InsecureRequestWarning
warnings.simplefilter('ignore',InsecureRequestWarning)

key = json.load(open('key.json', 'r', encoding='utf-8'))
authorize = json.load(open('authorize.json', 'r', encoding='utf-8'))
file_id = str(authorize['Data']['FileId'])
book_urls = authorize['Data']['SplitFileUrls']

rsa_private = key['privateKey']
rsa_public = key['publicKey']
pdf_key = authorize['Data']['Key']

# RSA解密出AES的Key
data = base64.b64decode(pdf_key)
key = RSA.importKey(rsa_private)
cipher = PKCS1_v1_5.new(key)
decrypted = cipher.decrypt(data,1).decode()
#print(decrypted)
aes_key = ''.join([hex(ord(x))[2:] for x in decrypted])
print('AES Key:' + str(aes_key))
headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/79.0.3945.130 Safari/537.36',
    'Accept': '*/*',
    'Accept-Encoding': 'gzip, deflate',
    'Accept-Language': 'zh-TW,zh;q=0.9,en-US;q=0.8,en;q=0.7',
    'Origin': 'http://yd.51zhy.cn/',
    'Referer': 'https://yd.51zhy.cn/ebook/reader/index.html',
    'Sec-Fetch-Site': 'cross-site',
    'Sec-Fetch-Mode': 'cors'
}

# 保存PDF
p = progressbar.ProgressBar()

if not os.path.isdir(file_id):
    os.makedirs(file_id)

for url in p(book_urls):
    page = book_urls.index(url)+1
    file = f"x-{page}"
    if not os.path.isfile(file_id + '/' + file + '.pdf'):
        r = requests.get(url=url, headers=headers, stream=True, verify=False)
        with open(file_id + '/' + file + '.aes','wb') as f:
            f.write(r.content)
        os.system(f'openssl enc -d -aes-128-ecb -K {aes_key} -in {file_id}/{file}.aes -out {file_id}/{file}.pdf -iv 0')
        #os.remove(file_id + '/' + file + '.aes')
        time.sleep(random.uniform(3,5))

print("正在合成PDF文件.........")
files = os.listdir(file_id)
pdf_files = []
for file in files:
    if '.pdf' in file:
        pdf_files.append(file)
pdf_files.sort(key=lambda x: int(x[x.rfind('-') + 1:][:-4]))

merger = PdfFileMerger(strict=False)
for file_name in pdf_files:
    # print(file_name)
    merger.append(PdfFileReader(open(file_id + '/' + file_name, 'rb')))
merger.write(file_id + '.pdf')
print('PDF合成完成!!')


import requests
from bs4 import BeautifulSoup
from urllib.parse import quote

def search_books(book_name):
    urls = ["https://www.yes24.com/Product/Search", "https://www.aladin.co.kr/search/wsearchresult.aspx", "https://search.kyobobook.co.kr/web/search"]
    encoded_book_name = quote(book_name)

    for url in urls:
        if 'yes24' in url:
            search_url = url + '?domain=ALL&query=' + encoded_book_name
        elif 'aladin' in url:
            search_url = url + '?SearchTarget=All&SearchWord=' + encoded_book_name + '&x=0&y=0'
        elif 'kyobobook' in url:
            search_url = url + '?keyword=' + encoded_book_name + '&gbCode=TOT&target=total'

        response = requests.get(search_url)
        soup = BeautifulSoup(response.text, 'html.parser')

        if 'yes24' in url:
            book_list = soup.find('ul', {'id': 'yesSchList'})
            books = book_list.find_all('li')[:10]

            for book in books:
                title_tag = book.find('a', {'class': 'gd_name'})
                price_tag = book.find('em', {'class': 'yes_b'})

                title = title_tag.text.strip()
                price = price_tag.text.strip()

                print('Title:', title)
                print('Price:', price)
                print()

        elif 'aladin' in url:
            book_list = soup.find('div', {'id': 'Search3_Result'})
            books = book_list.find_all('div', {'class': 'ss_book_box'})[:10]

            for book in books:
                title_tag = book.find('a', {'class': 'bo3'})
                price_tag = book.find('span', {'class': 'ss_p2'})

                title = title_tag.text.strip()
                price = price_tag.text.strip()

                print('Title:', title)
                print('Price:', price)
                print()

        elif 'kyobobook' in url:
            book_list = soup.find('div', {'id': 'shopData_list'})
            books = book_list.find_all('li', {'class': 'prod_item'})[:10]

            for book in books:
                title_tag = book.find('a', {'class': 'prod_info'})
                price_tag = book.find('span', {'class': 'val'})

                title = title_tag.text.strip()
                price = price_tag.text.strip()

                print('Title:', title)
                print('Price:', price)
                print()

book_name = input("Enter the book name: ")
search_books(book_name)


import argparse

SEC_FORM_TYPE_10K = "10-K"
SEC_FORM_TYPE_10Q = "10-Q"

parser = argparse.ArgumentParser(description='EDGAR program')
parser.add_argument(
    '-f', '--form-types', type=str.upper, nargs='+', required=False,
    default=["10-K", "10-Q"],
    help='specify the form types to select e.g. 10-Q, 10-K'
)
parser.add_argument(
    '-t', '--test-mode', action="store_true",
    help='specify to use the test mode'
)

args = vars(parser.parse_args())
print(args)

# Validations
assert all([
    form in [SEC_FORM_TYPE_10Q, SEC_FORM_TYPE_10K]
    for form in args['form_types']
]), "Invalid form type(s) in [%s]" % args['form_types']

assert isinstance(args["test_mode"], bool), f"Invalid type {type(args['test_mode'])}"


import pytest

from Tree.Tree import Node, Tree

@pytest.fixture
def bTree():
    t = Tree()
    Node(
        TREE=t,
        NAME='Root',
        options=['finance', 'housing', 'healthcare'],
        fields={'label': 'Select the category of assistance the client requires.'}
    )


from io import StringIO

from django.conf import settings
from django.http import JsonResponse, HttpResponse, QueryDict
from django.shortcuts import render
import pandas as pd
import json
import os
from django.core.files.storage import default_storage
from django.core.files.base import ContentFile

from django.views.decorators.csrf import csrf_exempt
from django.core.files.uploadedfile import InMemoryUploadedFile

@csrf_exempt
def jsonTest2(request):
    address_data = json.loads(request.POST['JsonData'])
    member_name = json.loads(request.POST['member_name'])
    print(address_data['name'])
    print(member_name)
    imgName = address_data['name'] + ".png"#snapShot name
    data = request.FILES['file']#request imgFile

    if os.path.exists(settings.MEDIA_ROOT+member_name) :
        print("이미 폴더가 존재합니")
    else:
        print("폴더를 생성합니다")
        os.mkdir(settings.MEDIA_ROOT+member_name)

    #SnapShot Save
    if os.path.exists(settings.MEDIA_ROOT+member_name+"/"+imgName):
        print('파일이 이미 존재합니다')
        return JsonResponse("exist", json_dumps_params={'ensure_ascii': False}, safe=False)
    else:
        print('파일을 생성합니다')
        path = default_storage.save(member_name+"/"+imgName, ContentFile(data.read()))
        tmp_file = os.path.join(settings.MEDIA_ROOT, path)#setting MEDIA_ROOT

    #Json Save
    imsiDic = dict()
    imsiList = list()
    imsiList.append(address_data)
    fileName = address_data['name']+".json"#이거 아이디로 바꿔야됨
    if os.path.exists(settings.MEDIA_ROOT+member_name+"/"+member_name+".json"):#Json file check
        with open(settings.MEDIA_ROOT+member_name+"/"+member_name+".json",'r',encoding='UTF-8') as file:
            after = json.load(file)
            print(after)
        after.append(address_data)
        with open(settings.MEDIA_ROOT+member_name+"/"+member_name+".json", 'w', encoding='UTF-8') as file:
            json.dump(after, file, ensure_ascii=False, indent=2)  # new file
    else:
        with open(settings.MEDIA_ROOT+member_name+"/"+member_name+".json", 'w', encoding='UTF-8') as file:
            json.dump(imsiList, file, ensure_ascii=False, indent=2) #new file
    return JsonResponse("success", json_dumps_params={'ensure_ascii':False},safe=False)

@csrf_exempt
def myList(request):
    if request.method=='POST':
        member_name = json.loads(request.body.decode('utf-8'))
        print(member_name)
        #member_name =  json.loads(request.body.decode('utf-8'))#member name
        path = settings.MEDIA_ROOT+member_name+"/"
    else:
        member_name = request.GET.get('name')
        path = settings.MEDIA_ROOT + member_name + "/"
        print(member_name)

    print(path)

    fileName = path + member_name + ".json"
    if os.path.exists(fileName):  # Json file check
        f = open(fileName)
        jsonResponse = json.load(f)
    else:
        jsonResponse = "None"

    print(jsonResponse)

    return JsonResponse(jsonResponse, json_dumps_params={'ensure_ascii':False},safe=False)


# -*- coding: utf-8 -*-
"""
Created on Thu Dec 28 09:55:42 2017

@author: zmzhang
"""

from rdkit import Chem
from rdkit.Chem.rdMolDescriptors import CalcMolFormula, CalcExactMolWt
from rdkit.Chem.rdmolops import GetFormalCharge
import gzip
from collections import Counter
import json, glob, os, itertools
from pathos.multiprocessing import ProcessingPool

from IsoSpecPy import IsoThreshold
from pylab import vlines
import numpy as np

def filter_pubchem(ms):
    ms_filtered = []
    elements = set(['C', 'H', 'O', 'N', 'S', 'P', 'Cl', 'B','Br','Se'])
    for m in ms:
        mw = CalcExactMolWt(m)
        if mw < 100 or mw> 1500:
            continue

        if GetFormalCharge(m) != 0:
            continue

        atoms = [a.GetSymbol() for a in m.GetAtoms()]
        c = Counter(atoms)
        if 'C' in c and 'H' in c:
            if 'S' in c and c['S'] > 5:
                continue
            if 'Cl' in c and c['Cl'] > 5:
                continue
            if 'Br' in c and c['Br'] > 5:
                continue
            if 'B' in c and c['B'] > 5:
                continue
            if set(c.keys()).issubset(elements):
                ms_filtered.append(CalcMolFormula(m))
    return ms_filtered

def write_json(ms, fname):
    f = open(fname, 'w')
    json.dump(ms, f)
    f.close()

def process_sdf(sdf, i, total):
    print( '{:2.4f}%  {}'.format((100 * float(i)/total), sdf))
    if not(os.path.exists(sdf[0:-6]+'json')):
        sdf_gz = gzip.open(sdf)
        gzsuppl = Chem.ForwardSDMolSupplier(sdf_gz)
        ms = [Chem.AddHs(x) for x in gzsuppl if x is not None]
        ms_filtered = filter_pubchem(ms)
        write_json(ms_filtered, sdf[0:-6]+'json')

def load_jsons(json_path):
    mols = []
    jsons =  glob.glob(sdf_path + '*.json')
    for i, json_file in enumerate(jsons):
        print( '{:2.4f}%'.format((100 * float(i)/len(jsons))))
        f = open(json_file)
        mol = json.load(f)
        f.close()
        mols.append(mol)
    return list(itertools.chain(*mols))

if __name__ == '__main__':
    sdf_path = 'F:/resources/isotope/pubchem/ftp.ncbi.nlm.nih.gov/pubchem/Compound/Monthly/2016-12-01/SDF/'
#    sdfs =  glob.glob(sdf_path + '*.sdf.gz')
#    pool = ProcessingPool(os.cpu_count())
#    pool.map(process_sdf, sdfs, range(len(sdfs)), [len(sdfs)]*len(sdfs))
    mols = load_jsons(sdf_path)
    mols_rd = list(set(mols))

    iso = IsoThreshold(formula="C254H377N65O75S6")
    mass = np.array([i for i in iso.masses])
    probs = np.array([i for i in iso.probs])
    idx = np.argsort(probs)[::-1]

    for i in idx[0:10]:
        print(mass[i], "\t", probs[i])

    vlines(mass, np.zeros(probs.shape), probs)


#!/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')


from urllib.request import urlopen
from bs4 import BeautifulSoup

url = input("Enter URL to parse (ex: https://0xskar.github.io): ")
extract_tag = input("Enter the tag to extract (ex: p for <p>): ")

page = urlopen(url)
html_bytes = page.read()
html = html_bytes.decode("utf-8")

# create a BeautifulSoup object to parse the HTML
soup = BeautifulSoup(html, 'html.parser')

# find all the tags
tag_outputs = soup.find_all(extract_tag)

# loop through the list of tags and print the text
for tag_output in tag_outputs:
    print(tag_output.get_text())


import pickle
import logging
import itertools
from io import StringIO
from fnmatch import fnmatch

from collections import OrderedDict, defaultdict

import colorama
from colorama import Fore, Back, Style

from stringparser import Parser

from lantz.log import ColorizingFormatter, SocketListener

def queue_from_record(rec):
    # TODO
    return ('a', 'b')
    #return rec.lantz_driver, rec.lantz_name

try:
    from msvcrt import getch
except ImportError:
    import sys, tty, termios
    def getch():
        fd = sys.stdin.fileno()
        old_settings = termios.tcgetattr(fd)
        try:
            tty.setraw(sys.stdin.fileno())
            ch = sys.stdin.read(1)
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
        return ch

COLS = 80

def restyle(text, new):
    """Replace style with a new one.

    :param text: the styled text.
    :param new: the new style, a value from colorame.Style.
    """
    for style in (Style.DIM, Style.BRIGHT, Style.NORMAL):
        text = text.replace(style, new)
    return text

def pos(row, column):
    """Move cursor to a given row and column.
    """
    return '\x1b[%d;%dH' % (row, column)

def clear():
    """Clear screen.
    """
    return '\x1b[2J'

RETITLE = logging.LogRecord('', 0, '', '', '', '', '', '')
REDRAW = logging.LogRecord('', 0, '', '', '', '', '', '')

def level_name(level):
    """For a level defined in the logging module, return its name
    """
    if level == logging.DEBUG:
        return 'DEBUG'
    elif level == logging.INFO:
        return 'INFO'
    elif level == logging.WARNING:
        return 'WARNING'
    elif level == logging.CRITICAL:
        return 'CRITICAL'
    elif level == logging.ERROR:
        return 'ERROR'
    return str(level)

class MessageListFormatter(ColorizingFormatter):
    """Log formatter for Message List View.

    :param queues: The list of queues with sorted logging messages. Each queue
                   corresponds to one instrument.
    """

    FMT = Style.NORMAL + '%(asctime)s <color>%(levelname)-8s</color>' + Style.RESET_ALL + ' %(message)s'

    FIRST_ROW = 2
    ROWS = 24

    def __init__(self, queues, *args, **kwargs):
        super().__init__(fmt=self.FMT, datefmt='%H:%M:%S', *args, **kwargs)
        self.queues = queues
        self.redraw = True
        self._title = ''
        self._current = ''
        self.last_records = []

    @property
    def current(self):
        return self._current

    @current.setter
    def current(self, value):
        self._current = value
        self.redraw = True

    @property
    def title(self):
        return self._title

    @title.setter
    def title(self, value):
        self._title = value
        self.redraw = True

    SEL = Style.BRIGHT + Back.RED + Fore.WHITE

    def format(self, record):
        if record is RETITLE:
            return pos(0, 0) + self.title

        if not record is REDRAW:
            self.last_records.append(record)
            if len(self.last_records) > self.ROWS:
                self.last_records = self.last_records[-self.ROWS:]

        buf = StringIO()
        self.redraw = False
        buf.write(clear())
        buf.write(pos(0, 0))
        buf.write(self.title)
        for row, saved_record in zip(range(self.ROWS), self.last_records):
            buf.write(pos(row + self.FIRST_ROW, 0))
            queue = queue_from_record(saved_record)
            if queue == self._current:
                buf.write(self.SEL + ' ' + Style.RESET_ALL)
            else:
                buf.write(' ')
            buf.write(queue[1].ljust(30))

            buf.write(super().format(saved_record))

        return buf.getvalue()

class DriverListFormatter(ColorizingFormatter):
    """Log formatter for Driver List View.

    :param queues: The list of queues with sorted logging messages. Each queue
                   corresponds to one instrument.
    """

    FMT = Style.NORMAL + '%(asctime)s <color>%(levelname)-8s</color>' + Style.NORMAL + ' %(message)s'

    FIRST_ROW = 2
    LOG_HEIGHT = 4

    def __init__(self, queues, *args, **kwargs):
        super().__init__(fmt=self.FMT, datefmt='%H:%M:%S', *args, **kwargs)
        self.queues = queues
        self.redraw = True
        self._title = ''
        self._current = ''

    @property
    def current(self):
        return self._current

    @current.setter
    def current(self, value):
        self._current = value
        self.redraw = True

    @property
    def title(self):
        return self._title

    @title.setter
    def title(self, value):
        self._title = value
        self.redraw = True

    def format(self, record):
        if record is RETITLE:
            return pos(0, 0) + self.title

        buf = StringIO()
        if self.redraw or record is REDRAW:
            self.redraw = False
            buf.write(clear())
            buf.write(pos(0, 0))
            buf.write(self.title)
            self._write_view(buf)
        else:
            self._write_record(record, buf)
        return buf.getvalue()

    def _write_view(self, buf):
        for num, (name, state) in enumerate(self.queues.items()):
            self._write_record_helper(num, name, state, buf)

    def _write_record(self, record, buf):
        queue = queue_from_record(record)
        num = list(self.queues.keys()).index(queue)
        state = self.queues[queue]
        self._write_record_helper(num, queue, state, buf)

    def _write_record_helper(self, num, name, state, buf):
        row = num * self.LOG_HEIGHT + self.FIRST_ROW
        buf.write(pos(row, 0))

        if name == self._current:
            buf.write(Back.RED + Fore.WHITE + Style.BRIGHT)
            buf.write((' ' + '.'.join(name)).ljust(COLS))
        else:
            buf.write(Back.WHITE + Fore.BLACK)
            buf.write((' ' + '.'.join(name)).ljust(COLS))
        buf.write(Style.RESET_ALL)
        for subrow, saved_record in zip(range(1, self.LOG_HEIGHT), state.last_records):
            buf.write(pos(row + subrow, 0))
            buf.write('  ' + super().format(saved_record))

class DriverFormatter(ColorizingFormatter):
    """Log formatter for Driver View.

    :param queues: The list of queues with sorted logging messages. Each queue
                   corresponds to one instrument.
    """

    FIRST_ROW = 3
    COL_WIDTH = 40
    PROP_HEIGHT = 15
    LOG_HEIGHT = 4

    FMT = Style.NORMAL + '%(asctime)s <color>%(levelname)-8s</color>' + Style.NORMAL + ' %(message)s'

    def __init__(self, queues, *args, **kwargs):
        super().__init__(fmt=self.FMT, datefmt='%H:%M:%S', *args, **kwargs)
        self.queues = queues
        self.redraw = True
        self._title = ''
        self._current = ''

    @property
    def current(self):
        return self._current

    @current.setter
    def current(self, value):
        self._current = value
        self.redraw = True

    @property
    def title(self):
        return self._title

    @title.setter
    def title(self, value):
        self._title = value
        self.redraw = True

    def format(self, record):
        if record is RETITLE:
            return pos(0, 0) + self.title

        buf = StringIO()
        if self.redraw or record is REDRAW:
            self.REDRAW = False
            buf.write(clear())
            buf.write(pos(0, 0))
            buf.write(self.title)
            self._write_view(buf)
        else:
            self._write_record(record, buf)

        return buf.getvalue()

    def _write_view(self, buf):
        current = self._current
        buf.write(pos(2, 0))
        buf.write(Back.RED + Fore.WHITE + Style.BRIGHT)
        buf.write((' ' + '.'.join(current)).ljust(COLS))
        buf.write(pos(self.FIRST_ROW + self.PROP_HEIGHT, 0))
        buf.write(Style.RESET_ALL)
        buf.write(Back.WHITE + Fore.BLACK)
        buf.write((' Messages').ljust(COLS))
        buf.write(Style.RESET_ALL)
        if current not in self.queues:
            return
        state = self.queues[current]
        for num, (name, value) in enumerate(state.feats.items()):
            col = num // self.PROP_HEIGHT * self.COL_WIDTH
            row = num % self.PROP_HEIGHT + self.FIRST_ROW
            self._write_property(row, col, name, value, buf)

        self._write_log_helper(state, buf)

    def _write_log_helper(self, state, buf):
        row = self.FIRST_ROW + self.PROP_HEIGHT
        for subrow, saved_record in zip(range(1, self.PROP_HEIGHT), state.last_records):
            buf.write(pos(row + subrow, 0))
            buf.write('  ' + super().format(saved_record))

    def _write_record(self, record, buf):
        queue = queue_from_record(record)
        state = self.queues[queue]
        if hasattr(record, 'lantz_feat'):
            if queue not in self.queues:
                return
            name, value = record.lantz_feat
            num = list(state.feats.keys()).index()
            col = num // self.PROP_HEIGHT * self.COL_WIDTH
            row = num % self.PROP_HEIGHT + self.FIRST_ROW
            self._write_property(row, col, name, value, buf)

        self._write_record_helper(0, '.'.join(queue), state, buf)

    def _write_property(self, row, col, name, value, buf):
        if name:
            buf.write(pos(row, col + 0))
            buf.write(name)
        buf.write(pos(row, col + 20))
        buf.write(value)

#: Parser for set messages.
SETTER = Parser('{0:s} was set to {1:s}')

#: Parser for get messages.
GETTER = Parser('Got {1:s} for {0:s}')

class InstrumentState(object):
    """A class to build up the state of the instrument by parsing incoming
    log records. It also remember the last received log records.
    """

    #: Number of log records to remember.
    REMEMBER = 10

    def __init__(self):
        self.count = 0
        self.last_records = []
        self.feats = OrderedDict()

    def handle(self, record):
        self.count += 1
        self.last_records.append(record)
        if len(self.last_records) > self.REMEMBER:
            self.last_records = self.last_records[-self.REMEMBER:]
        if hasattr(record, 'lantz_feat'):
            self.feats[record.lantz_feat[0]] = record.lantz_feat[1]

class LantzMonitor(SocketListener):
    """Text UI to display log records incoming via tcp or udp.
    """

    def __init__(self, tcphost, udphost, level=logging.INFO,
                 filter_='lantz.*'):
        super().__init__(tcphost, udphost)

        self.queues = defaultdict(InstrumentState)
        self.filter = filter_ or '*'
        self._level = level
        self.formatters = [MessageListFormatter(self.queues),
                           DriverListFormatter(self.queues),
                           DriverFormatter(self.queues)]
        self.handler = logging.StreamHandler()
        self.handler.setFormatter(self.formatters[0])
        self.retitle('Welcome, press h for help.')

    @property
    def level(self):
        return self._level

    @level.setter
    def level(self, value):
        self._level = value
        self.retitle()

    def loop(self):
        view = 0
        queue = 0
        last_help = 0
        messages = itertools.cycle(('Change logging level: 1-5',
                                    'Change view: [ ]',
                                    'Select instrument: , .',
                                    'Quit: q',
                                    'Press h for help.'))
        while 1:
            command = getch()

            if command in ('1', '2', '3', '4', '5'):
                self.level = int(command) * 10

            elif command in ('[', ']'):
                view = (view + (-1 if command == '[' else 1)) % 3
                self.handler.setFormatter(self.formatters[view])
                self.handler.handle(REDRAW)

            elif command in (',', '.'):
                queue = (queue + (-1 if command == ',' else 1)) % len(self.queues)
                for formatter in self.formatters:
                    formatter.current = list(self.queues.keys())[queue]
                self.handler.handle(REDRAW)

            elif command in ('q', chr(3)):
                raise SystemExit

            elif command == 'h':
                self.retitle(next(messages))

            elif command in (chr(13), ):
                view = 2
                self.handler.setFormatter(self.formatters[view])
                self.handler.handle(REDRAW)

            else:
                self.retitle("Unknown command '{}' (ord {})".format(command, ord(command)))

    def retitle(self, message=''):
        title = '<> Lantz <> Queues: {:3d} <> Level: {:8s} <> {}'.format(len(self.queues),
                                                                    level_name(self._level),
                                                                    message)
        title = Back.BLUE + Fore.WHITE + title.ljust(COLS) + Style.RESET_ALL
        for formatter in self.formatters:
            formatter.title = title
        self.handler.handle(RETITLE)

    def on_record(self, record):
        if not fnmatch(record.name or '', self.filter):
            return
        if record.levelno < self.level:
            return
        with self._lock:
            queue = queue_from_record(record)
            if queue not in self.queues:
                self.queues[queue] = InstrumentState()
                self.retitle()
            if len(self.queues) == 1:
                for formatter in self.formatters:
                    formatter.current = queue
            self.queues[queue].handle(record)
            self.handler.handle(record)

###########################
# standalone main
###########################

def main():
    import argparse

    colorama.init()
    parser = argparse.ArgumentParser(description='Text based tool to monitor Lantz messages logged to a socket..')
    parser.add_argument('-t', '--tcp', default='0.0.0.0', dest='tcphost',
                        help='Where to listen for TCP traffic (host[:port])')
    parser.add_argument('-u', '--udp', default='0.0.0.0', dest='udphost',
                        help='Where to listen for UDP traffic (host[:port])')
    parser.add_argument('--filter', default=None,
                        help='Filter to apply to incomming records')
    parser.add_argument('-l', '--level', choices=range(1, 6), default=2, type=int,
                        help='Logging level (debug, info, warn, error, critical')
    args = parser.parse_args()

    print(clear())
    monitor = LantzMonitor(args.tcphost, args.udphost,
                           args.level * 10, args.filter)

    try:
        monitor.loop()
    except (KeyboardInterrupt, SystemExit):
        monitor.stop()
    print(clear())
    print('Lantz Monitor stopped')

if __name__ == '__main__':
    main()


"""
A collection of useful function snippets that irrelevant with the main modules
or await for further refactor
"""
import logging
from typing import Any, Callable, Iterator, Optional, Sequence, Tuple

import numpy as np
import tensorflow as tf

logger = logging.getLogger(__name__)

try:
    import cirq
except ImportError as e:
    logger.warning(e)
    logger.warning("Therefore some functionality in %s may not work" % __name__)

from ..circuit import Circuit
from .. import gates as G
from ..gates import array_to_tensor, num_to_tensor

Array = Any
Tensor = Any
Graph = Any

class FakeModule:
    def __getattr__(self, name: str) -> str:
        return name

fake_module = FakeModule()

##############
## qml related functions to be refactored
##############

# tf specific

def amplitude_encoding(
    fig: Tensor,
    qubits: int,
    index: Optional[Sequence[int]] = None,
    index_func: Optional[Callable[[int, int], Sequence[int]]] = None,
) -> Tensor:
    if fig.shape[-1] == 1:
        fig = tf.reshape(fig, shape=fig.shape[:-1])
    if len(fig.shape) == 2:
        fig = fig[tf.newaxis, ...]
    fig = tf.reshape(fig, shape=(fig.shape[0], fig.shape[1] * fig.shape[2]))
    norm = tf.linalg.norm(fig, axis=1)
    norm = norm[..., tf.newaxis]
    fig = fig / norm
    if fig.shape[1] < 2**qubits:
        fig = tf.concat(
            [
                fig,
                tf.zeros([fig.shape[0], 2**qubits - fig.shape[1]], dtype=tf.float64),
            ],
            axis=1,
        )
    if index is None and index_func is not None:
        index = []
        for i in range(32):
            # magic number here, only valid for MNIST data
            for j in range(32):
                l = index_func(i, j)
                r = 0
                for p, q in enumerate(l):
                    r += q * 2 ** (9 - p)
                index.append(r)
    if index is not None:
        fig = tf.constant(fig.numpy()[:, index], dtype=fig.dtype)
    return fig

def recursive_index(x: int, y: int) -> Sequence[int]:
    rl = []
    for k in range(5):
        rl.append((x // (2 ** (4 - k))) % 2)
        rl.append((y // (2 ** (4 - k))) % 2)
    return rl

def mnist_amplitude_data(
    a: int,
    b: int,
    binarize: bool = False,
    index: Optional[Sequence[int]] = None,
    index_func: Optional[Callable[[int, int], Sequence[int]]] = None,
    loader: Any = None,
    threshold: float = 0.4,
) -> Tensor:
    def filter_pair(x: Tensor, y: Tensor, a: int, b: int) -> Tuple[Tensor, Tensor]:
        keep = (y == a) | (y == b)
        x, y = x[keep], y[keep]
        y = y == a
        return x, y

    if loader is None:
        loader = tf.keras.datasets.mnist
    (x_train, y_train), (x_test, y_test) = loader.load_data()
    x_train, x_test = x_train[..., np.newaxis] / 255.0, x_test[..., np.newaxis] / 255.0

    if binarize:
        x_train[x_train > threshold] = 1.0
        x_train[x_train <= threshold] = 0.0
        x_test[x_test > threshold] = 1.0
        x_test[x_test <= threshold] = 0.0
    x_train, y_train = filter_pair(x_train, y_train, a, b)
    x_test, y_test = filter_pair(x_test, y_test, a, b)
    x_train = tf.image.pad_to_bounding_box(x_train, 2, 2, 32, 32)
    x_test = tf.image.pad_to_bounding_box(x_test, 2, 2, 32, 32)
    x_train_q = amplitude_encoding(x_train, 10, index=index, index_func=index_func)
    x_test_q = amplitude_encoding(x_test, 10, index=index, index_func=index_func)
    return (x_train_q, y_train), (x_test_q, y_test)

def mnist_generator(
    x_train: Tensor, y_train: Tensor, batch: int = 1, random: bool = True
) -> Iterator[Any]:
    if isinstance(x_train, tf.Tensor):
        x_train = x_train.numpy()
    if isinstance(y_train, tf.Tensor):
        y_train = y_train.numpy()
    i = np.array(list(range(batch)))
    while True:
        if random:
            i = np.random.randint(low=0, high=x_train.shape[0], size=batch)
        else:
            i += batch
            i = i % x_train.shape[0]
        yield tf.constant(x_train[i]), tf.constant(y_train[i])

def generate_random_circuit(
    inputs: Tensor, nqubits: int = 10, epochs: int = 3, layouts: Tensor = None
) -> Circuit:
    inputs = tf.cast(inputs, dtype=tf.complex64)
    c = Circuit(nqubits, inputs=inputs)
    if layouts is None:
        layouts = np.random.choice([0, 1, 2], size=[epochs, nqubits])
    layouts = tf.reshape(layouts, shape=[epochs, nqubits])
    for epoch in range(epochs):
        for i in range(nqubits):
            flg = layouts[epoch, i]
            if flg == 0:
                c.rx(i, theta=num_to_tensor(np.pi / 2))  # type: ignore
            elif flg == 1:
                c.ry(i, theta=num_to_tensor(np.pi / 2))  # type: ignore
            else:
                c.wroot(i)  # type: ignore
        for i in range(nqubits):  # entangling layer
            c.swap(i, (i + 1) % nqubits)  # type: ignore
    return c

# the following vags are not the recommended fashion to write qml tasks in tensorcircuit now.
# also the hardcoded circuit structure is a bad practice, they are here for backward compatibility
# instead use vvag or keras paradigm now

def naive_qml_vag(
    gdata: Graph,
    nnp: Tensor,
    preset: Sequence[int],
    nqubits: int = 10,
    epochs: int = 3,
    target: int = 0,
) -> Tuple[Tensor, Tensor]:
    # used for structure searching?
    xs, ys = gdata
    loss = 0.0
    for x, y in zip(xs, ys):
        circuit = generate_random_circuit(
            x, nqubits=nqubits, epochs=epochs, layouts=preset
        )
        value = circuit.expectation(
            (
                G.z(),  # type: ignore
                [
                    target,
                ],
            )
        )
        y = tf.cast(y, dtype=tf.complex64)
        y = 2 * y - 1
        loss += (value - y) ** 2
    return tf.constant(tf.cast(loss, dtype=tf.float32)), tf.zeros_like(nnp)  # MSE loss

def train_qml_vag(
    gdata: Graph,
    nnp: Tensor,
    preset: Optional[Sequence[int]] = None,
    nqubits: int = 10,
    epochs: int = 3,
    batch: int = 64,
    validation: bool = False,
) -> Any:
    xs, ys = gdata
    with tf.GradientTape() as tape:
        tape.watch(nnp)
        cnnp = tf.cast(nnp, dtype=tf.complex64)
        c = Circuit(nqubits, inputs=np.ones([1024], dtype=np.complex64) / 2**5)
        for epoch in range(epochs):
            for i in range(nqubits):
                c.rz(i, theta=cnnp[3 * epoch, i])  # type: ignore
            for i in range(nqubits):
                c.ry(i, theta=cnnp[3 * epoch + 1, i])  # type: ignore
            for i in range(0, nqubits, 2):
                # c.swap(i, (i + 1) % 10)  # type: ignore
                c.exp(  # type: ignore
                    i,
                    (i + 1) % 10,
                    unitary=array_to_tensor(G._swap_matrix),
                    theta=cnnp[3 * epoch + 2, i],
                )
            for i in range(1, nqubits, 2):
                # c.swap(i, (i + 1) % 10)  # type: ignore
                c.exp(  # type: ignore
                    i,
                    (i + 1) % 10,
                    unitary=array_to_tensor(G._swap_matrix),
                    theta=cnnp[3 * epoch + 2, i],
                )
        for i in range(nqubits):
            c.rx(i, theta=cnnp[3 * epochs, i])  # type: ignore
        count = 0
        loss = 0
        for x, y in zip(xs, ys):
            y = tf.cast(y, dtype=tf.float32)  # [1, 0] binary
            c.replace_inputs(tf.cast(x, dtype=tf.complex64))
            flg = False
            yp = 0.0

            for i in range(nqubits):
                yp += cnnp[3 * epochs + 1, i] * c.expectation(
                    (
                        G.z(),  # type: ignore
                        [
                            i,
                        ],
                    ),
                    reuse=flg,
                )

                flg = True

            yp = tf.cast(yp, dtype=tf.float32)
            yp = tf.math.sigmoid(
                (yp + tf.math.real(cnnp[3 * epochs + 2, 0])) * 15
            )  ## [0,1] intepreted as probability
            if tf.math.abs(yp - y) < 0.5:
                count += 1
            loss += (yp - y) ** 2
        # print("accuracy:", count / batch)
        gr = tape.gradient(loss, cnnp)
        if not validation:
            return tf.constant(count / batch, dtype=tf.float32), gr
        else:
            return count / batch  # evaluate as validation kernel

def validate_qml_vag(
    gdata: Graph,
    nnp: Tensor,
    preset: Optional[Sequence[int]] = None,
    nqubits: int = 10,
    epochs: int = 3,
    batch: int = 64,
) -> Any:
    xs, ys = gdata
    cnnp = tf.cast(nnp, dtype=tf.complex64)
    c = Circuit(nqubits, inputs=np.ones([1024], dtype=np.complex64) / 2**5)
    for epoch in range(epochs):
        for i in range(nqubits):
            c.rz(i, theta=cnnp[3 * epoch, i])  # type: ignore
        for i in range(nqubits):
            c.ry(i, theta=cnnp[3 * epoch + 1, i])  # type: ignore
        for i in range(0, nqubits, 2):
            # c.swap(i, (i + 1) % 10)  # type: ignore
            c.exp(  # type: ignore
                i,
                (i + 1) % 10,
                unitary=array_to_tensor(G._swap_matrix),
                theta=cnnp[3 * epoch + 2, i],
            )
        for i in range(1, nqubits, 2):
            # c.swap(i, (i + 1) % 10)  # type: ignore
            c.exp(  # type: ignore
                i,
                (i + 1) % 10,
                unitary=array_to_tensor(G._swap_matrix),
                theta=cnnp[3 * epoch + 2, i],
            )
    for i in range(nqubits):
        c.rx(i, theta=cnnp[3 * epochs, i])  # type: ignore
    count = 0
    loss = 0
    for x, y in zip(xs, ys):
        y = tf.cast(y, dtype=tf.float32)  # [1, 0] binary
        c.replace_inputs(tf.cast(x, dtype=tf.complex64))
        flg = False
        yp = 0.0

        for i in range(nqubits):
            yp += cnnp[3 * epochs + 1, i] * c.expectation(
                (
                    G.z(),  # type: ignore
                    [
                        i,
                    ],
                ),
                reuse=flg,
            )

            flg = True

        yp = tf.cast(yp, dtype=tf.float32)
        yp = tf.math.sigmoid(
            (yp + tf.math.real(cnnp[3 * epochs + 2, 0])) * 15
        )  ## [0,1] intepreted as probability
        if tf.math.abs(yp - y) < 0.5:
            count += 1
        loss += (yp - y) ** 2

    return {
        "val_loss": loss / batch,
        "val_accuracy": count / batch,
    }  # evaluate as validation kernel

##############
## really small and irrelavant pieces below
##############

def color_svg(circuit: cirq.Circuit, *coords: Tuple[int, int]) -> Any:
    """
    color cirq circuit SVG for given gates,
    a small tool to hack the cirq SVG

    :param circuit:
    :param coords: integer coordinate which gate is colored
    :return:
    """
    import xml
    from cirq.contrib.svg import SVGCircuit

    svg_str = SVGCircuit(circuit)._repr_svg_()
    DOMTree = xml.dom.minidom.parseString(svg_str)  # type: ignore
    xpos = []
    ypos = []
    for r in DOMTree.getElementsByTagName("rect"):  # [0].setAttribute("fill", "gray")
        xpos.append(int(float(r.getAttribute("x"))))
        ypos.append(int(float(r.getAttribute("y"))))
    xpos = sorted(list(set(xpos)))
    ypos = sorted(list(set(ypos)))
    # xpos_dict = dict(zip(range(len(xpos)), xpos))
    # ypos_dict = dict(zip(range(len(ypos)), ypos))
    i_xpos_dict = dict(zip(xpos, range(len(xpos))))
    i_ypos_dict = dict(zip(ypos, range(len(ypos))))
    for r in DOMTree.getElementsByTagName("rect"):
        x, y = int(float(r.getAttribute("x"))), int(float(r.getAttribute("y")))
        if (i_xpos_dict[x], i_ypos_dict[y]) in coords:
            r.setAttribute("fill", "gray")
    return DOMTree.toxml()

def repr2array(inputs: str) -> Array:
    """
    transform repr form of an array to real numpy array

    :param inputs:
    :return:
    """
    # used to transform np.array in print form into live np.array
    inputs = inputs.split("]")  # type: ignore
    inputs = [l.strip().strip("[") for l in inputs if l.strip()]  # type: ignore
    outputs = []
    for l in inputs:
        o = [float(c.strip()) for c in l.split(" ") if c.strip()]
        outputs.append(o)
    return np.array(outputs)

# duplicated, will remove later
# use the version in applications.physics.baseline.py
def TFIM1Denergy(
    L: int, Jzz: float = 1.0, Jx: float = 1.0, Pauli: bool = True
) -> float:
    # PBC
    # nice tutorial: https://arxiv.org/pdf/2009.09208.pdf
    # further investigation on 1 TFIM solution structure is required
    # will fail on AFM phase Jzz>Jx and Jzz>0 and odd sites (frustration in boundary)
    e = 0
    if Pauli:
        Jx *= 2
        Jzz *= 4
    for i in range(L):
        q = np.pi * (2 * i - (1 + (-1) ** L) / 2) / L
        e -= np.abs(Jx) / 2 * np.sqrt(1 + Jzz**2 / 4 / Jx**2 - Jzz / Jx * np.cos(q))
    return e

def Heisenberg1Denergy(L: int, Pauli: bool = True, maxiters: int = 1000) -> float:
    # PBC
    error = 1e-15
    eps = 1e-20  # avoid zero division
    phi = np.zeros([L // 2, L // 2])
    phi2 = np.zeros([L // 2, L // 2])
    lamb = np.array([2 * i + 1 for i in range(L // 2)])
    for _ in range(maxiters):
        k = 1 / L * (2 * np.pi * lamb + np.sum(phi, axis=-1) - np.diag(phi))
        for i in range(L // 2):
            for j in range(L // 2):
                phi2[i, j] = (
                    np.arctan(
                        2
                        / (
                            1 / (np.tan(k[i] / 2) + eps)
                            - 1 / (np.tan(k[j] / 2) + eps)
                            + eps
                        )
                    )
                    * 2
                )
        if np.allclose(phi, phi2, rtol=error):  # converged
            break
        phi = phi2.copy()
    else:
        raise ValueError(
            "the maxiters %s is too small for bethe ansatz to converge" % maxiters
        )
    e = -np.sum(1 - np.cos(k)) + L / 4
    if Pauli is True:
        e *= 4
    return e  # type: ignore


import torch
import torch.nn as nn

from utils import get_gaussian_mean

class IMM(nn.Module):
    def __init__(self, dim=10, heatmap_std=0.1, in_channel=3, h_channel=32):
        """
        It should be noted all params has been fixed to Jakab 2018 paper.
        Goto the original class if params and layers need to be changed.
        Images should be rescaled to 128*128
        """
        super(IMM, self).__init__()
        self.content_encoder = ContentEncoder(in_channel, h_channel)
        self.pose_encoder = PoseEncoder(dim, heatmap_std, in_channel, h_channel)
        self.generator = Generator(channels=8*h_channel+dim, h_channel=h_channel)
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

    def forward(self, x, y):
        content_x = self.content_encoder(x)
        pose_y, pose_coord = self.pose_encoder(y)
        code = torch.cat((content_x, pose_y), dim=1)
        recovered_y = self.generator(code)
        return recovered_y, pose_coord

class ContentEncoder(nn.Module):
    def __init__(self, in_channel=3, h_channel=64):
        super(ContentEncoder, self).__init__()
        self.conv1_1 = Conv_Block(in_channel, h_channel, (3, 3))
        self.conv1_2 = Conv_Block(h_channel, h_channel, (3, 3))

        self.conv2_1 = Conv_Block(h_channel, 2 * h_channel, (3, 3), downsample=True)
        self.conv2_2 = Conv_Block(2 * h_channel, 2 * h_channel, (3, 3))

        self.conv3_1 = Conv_Block(2 * h_channel, 4 * h_channel, (3, 3), downsample=True)
        self.conv3_2 = Conv_Block(4 * h_channel, 4 * h_channel, (3, 3))

        self.conv4_1 = Conv_Block(4 * h_channel, 8 * h_channel, (3, 3), downsample=True)
        self.conv4_2 = Conv_Block(8 * h_channel, 8 * h_channel, (3, 3))

        # self.out_conv = Conv_Block(8 * h_channel, h_channel, (3, 3))
        self.conv_layers = nn.ModuleList([
            self.conv1_1,
            self.conv1_2,
            self.conv2_1,
            self.conv2_2,
            self.conv3_1,
            self.conv3_2,
            self.conv4_1,
            self.conv4_2
            # self.out_conv
        ])

    def forward(self, x):
        for layer in self.conv_layers:
            x = layer(x)
        return x

class PoseEncoder(nn.Module):
    def __init__(self, dim=10, heatmap_std=0.1, in_channel=3, h_channel=64, heatmap_size=16):
        """

        Args:
            dim (int): Num of keypoints
        """
        super(PoseEncoder, self).__init__()
        self.conv1_1 = Conv_Block(in_channel, h_channel, (3, 3))
        self.conv1_2 = Conv_Block(h_channel, h_channel, (3, 3))

        self.conv2_1 = Conv_Block(h_channel, 2 * h_channel, (3, 3), downsample=True)
        self.conv2_2 = Conv_Block(2 * h_channel, 2 * h_channel, (3, 3))

        self.conv3_1 = Conv_Block(2 * h_channel, 4 * h_channel, (3, 3), downsample=True)
        self.conv3_2 = Conv_Block(4 * h_channel, 4 * h_channel, (3, 3))

        self.conv4_1 = Conv_Block(4 * h_channel, 8 * h_channel, (3, 3), downsample=True)
        self.conv4_2 = Conv_Block(8 * h_channel, 8 * h_channel, (3, 3))

        self.out_conv = nn.Sequential(nn.Conv2d(8 * h_channel, dim, (1, 1)))
        self.conv_layers = nn.ModuleList([
            self.conv1_1,
            self.conv1_2,
            self.conv2_1,
            self.conv2_2,
            self.conv3_1,
            self.conv3_2,
            self.conv4_1,
            self.conv4_2,
            self.out_conv
        ])
        self.heatmap = HeatMap(heatmap_std, (heatmap_size, heatmap_size))

    def forward(self, x):
        for layer in self.conv_layers:
            x = layer(x)
        heatmap, coord = self.heatmap(x)
        return heatmap, coord

class Generator(nn.Module):
    """"""

    def __init__(self, channels=64 + 10, h_channel=64):
        super(Generator, self).__init__()
        self.conv1_1 = Conv_Block(channels, 8 * h_channel, (3, 3))
        self.conv1_2 = Conv_Block(8 * h_channel, 8 * h_channel, (3, 3), upsample=True)

        self.conv2_1 = Conv_Block(8 * h_channel, 4 * h_channel, (3, 3))
        self.conv2_2 = Conv_Block(4 * h_channel, 4 * h_channel, (3, 3), upsample=True)

        self.conv3_1 = Conv_Block(4 * h_channel, 2 * h_channel, (3, 3))
        self.conv3_2 = Conv_Block(2 * h_channel, 2 * h_channel, (3, 3), upsample=True)

        self.conv4_1 = Conv_Block(2 * h_channel, h_channel, (3, 3))
        self.conv4_2 = Conv_Block(h_channel, h_channel, (3, 3))

        self.final_conv = nn.Conv2d(h_channel, 3, (3, 3), padding=[1, 1])
        self.conv_layers = nn.ModuleList([
            self.conv1_1,
            self.conv1_2,
            self.conv2_1,
            self.conv2_2,
            self.conv3_1,
            self.conv3_2,
            self.conv4_1,
            self.conv4_2,
            self.final_conv
        ])

    def forward(self, x):
        for layer in self.conv_layers:
            x = layer(x)
        # return x
        #return (nn.functional.tanh(x) + 1) / 2.0
        return nn.functional.relu(x)

class HeatMap(nn.Module):
    """
    Refine the estimated pose map to be gaussian distributed heatmap.
    Calculate the gaussian mean value.
    Params:
    std: standard deviation of gaussian distribution
    output_size: output feature map size
    """

    def __init__(self, std, output_size):
        super(HeatMap, self).__init__()
        self.std = std
        self.out_h, self.out_w = output_size

    def forward(self, x, h_axis=2, w_axis=3):
        """
        x: feature map BxCxHxW
        h_axis: the axis of Height
        w_axis: the axis of width
        """
        # self.in_h, self.in_w = x.shape[h_axis:]
        batch, channel = x.shape[:h_axis]

        # Calculate weighted position of joint(0~1)
        x_mean = get_gaussian_mean(x, h_axis, w_axis)  # BxC
        y_mean = get_gaussian_mean(x, w_axis, h_axis)  # BxC

        coord = torch.stack([y_mean, x_mean], dim=2)

        x_mean = x_mean.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, self.out_h, self.out_w)
        y_mean = y_mean.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, self.out_h, self.out_w)

        x_ind = torch.tensor(torch.linspace(0, 1, self.out_h)).unsqueeze(-1).repeat(batch, channel, 1, self.out_w).to(
            x.device)
        y_ind = torch.tensor(torch.linspace(0, 1, self.out_w)).unsqueeze(0).repeat(batch, channel, self.out_h, 1).to(
            x.device)

        dist = (x_ind - x_mean) ** 2 + (y_ind - y_mean) ** 2

        res = torch.exp(-(dist + 1e-6).sqrt_() / (2 * self.std ** 2))
        return res, coord

class Conv_Block(nn.Module):
    def __init__(self, inc, outc, size, downsample=False, upsample=False):
        super(Conv_Block, self).__init__()
        block = [
            nn.ReplicationPad2d(1),
            nn.Conv2d(in_channels=inc, out_channels=outc, kernel_size=size),
            nn.BatchNorm2d(outc),
            nn.LeakyReLU()
        ]
        if downsample:
            block += [nn.MaxPool2d(kernel_size=2, stride=2)]
        if upsample:
            block += [nn.UpsamplingBilinear2d(scale_factor=2)]
        self.block = nn.Sequential(*block)

    def forward(self, x):
        return self.block(x)

if __name__ == '__main__':
    t = IMM(10, 0.1)
    x = torch.randn(10, 3, 128, 128)
    # x: the original image
    y = torch.randn(10, 3, 128, 128)
    # y: the warped image

    tmp = t(x, y)


# simple image alpha blending
import numpy as np
import cv2

dir1 = 'img1_directory'
dir2 = 'img2_directory'

blending_area = 15
img1 = cv2.imread(dir1)
img2 = cv2.imread(dir2)
result = np.zeros_like(img1)

height, width = img1.shape[:2]
middle = width // 2

result[:, :middle, :] = img1[:, :middle, :]
result[:, middle:, :] = img2[:, middle:, :]
cv2.imshow('img1', img1)
cv2.imshow('img2', img2)
cv2.imshow('half', result)
cv2.waitKey()
start = middle - (blending_area * width // 100) // 2

for i in range((middle - start) * 2):
    beta = (i) / 100
    alpha = 1 - beta
    result[:, start + i] = img1[:, start + i] * alpha + \
                                              img2[:, start + i] * beta

cv2.imshow("result", result)
cv2.waitKey()


#!/usr/bin/python3
"""script takes in URL, sends request to URL, displays value of X-Request-Id"""
import requests
import sys

if __name__ == '__main__':
    url = 'http://0.0.0.0:5000/search_user'
    if len(sys.argv) > 1:
        data = {'q': sys.argv[1]}
    else:
        data = {'q': ""}
    try:
        r = requests.post(url, data)
        answer = r.json()
        if answer:
            print('[{}] {}'.format(answer.get('id'), answer.get('name')))
        else:
            print('No result')
    except ValueError:
        print('Not a valid JSON')


import boto3

# Get the service resource.
dynamodb = boto3.resource('dynamodb',
                          region_name='eu-west-1',
                          endpoint_url='http://localhost:8000')

# Instantiate a table resource object without actually
# creating a DynamoDB table. Note that the attributes of this table
# are lazy-loaded: a request is not made nor are the attribute
# values populated until the attributes
# on the table resource are accessed or its load() method is called.
table = dynamodb.Table('users')

# Print out some data about the table.
# This will cause a request to be made to DynamoDB and its attribute
# values will be set based on the response.
print(table.creation_date_time)


from apscheduler.schedulers.blocking import BlockingScheduler
from flask import Blueprint
from sentry_sdk import set_tag

from pcapi.local_providers.provider_manager import synchronize_data_for_provider
from pcapi.models.feature import FeatureToggle
from pcapi.scheduled_tasks import utils
from pcapi.scheduled_tasks.decorators import cron_context
from pcapi.scheduled_tasks.decorators import cron_require_feature
from pcapi.scheduled_tasks.decorators import log_cron_with_transaction

blueprint = Blueprint(__name__, __name__)

@cron_context
@log_cron_with_transaction
@cron_require_feature(FeatureToggle.SYNCHRONIZE_TITELIVE_PRODUCTS)
def synchronize_titelive_things():
    synchronize_data_for_provider("TiteLiveThings")

@cron_context
@log_cron_with_transaction
@cron_require_feature(FeatureToggle.SYNCHRONIZE_TITELIVE_PRODUCTS_DESCRIPTION)
def synchronize_titelive_thing_descriptions():
    synchronize_data_for_provider("TiteLiveThingDescriptions")

@cron_context
@log_cron_with_transaction
@cron_require_feature(FeatureToggle.SYNCHRONIZE_TITELIVE_PRODUCTS_THUMBS)
def synchronize_titelive_thing_thumbs():
    synchronize_data_for_provider("TiteLiveThingThumbs")

@blueprint.cli.command("titelive_clock")
def titelive_clock():
    set_tag("pcapi.app_type", "titelive_clock")
    scheduler = BlockingScheduler()
    utils.activate_sentry(scheduler)

    scheduler.add_job(synchronize_titelive_things, "cron", day="*", hour="1")

    scheduler.add_job(synchronize_titelive_thing_descriptions, "cron", day="*", hour="2")

    scheduler.add_job(synchronize_titelive_thing_thumbs, "cron", day="*", hour="3")

    scheduler.start()