# ==============================================================================
# Authors: ralf mueller, stephan siemen
#
#
# Plan is to create a plot similar to the scatter plot for co2 concentration and
# september minimum of sea ice extent
#
# reference:
#   https://www.mpg.de/10579957/W004_Environment_climate_062-069.pdf, p. 7
#
# ==============================================================================
import os
from ecmwfapi import ECMWFDataServer
from cdo import Cdo
from multiprocessing import Pool
from tarfile import TarFile
import matplotlib.pyplot as plt
import matplotlib.transforms as mtrans

# basic setup    {{{ ===========================================================
server = ECMWFDataServer()
cdo = Cdo()
cdo.debug = True
tasks = 4
startYear = 1980
endYear = 2014
# }}} ==========================================================================
# helper methods {{{ ===========================================================
def getDataFromTarfile(tarfile):
    tf = TarFile(tarfile)
    members = [m.name for m in tf.getmembers()]
    if list(set([os.path.exists(x) for x in members])) != [True]:
        tf.extractall()
    tf.close
    return members


def computeTimeSeries(file, varname, useCellArea=False):
    if useCellArea:
        ofile = cdo.mul(
            input="-fldmean -selname,%s %s -fldsum -gridarea %s"
            % (varname, file, file),
            options="-b F32",
            output="_" + os.path.basename(file),
            force=False,
        )
    else:
        ofile = cdo.fldmean(
            input="-selname,%s %s" % (varname, file),
            options="-b F32",
            output="_" + os.path.basename(file),
            force=False,
        )
    return ofile


def computeTimeSeriesOfFilelist(pool, files, varname, ofile, useCellArea=False):
    results = dict()
    for file in files:
        rfile = pool.apply_async(computeTimeSeries, (file, varname, False))
        results[file] = rfile
    pool.close()
    pool.join()

    for k, v in results.items():
        results[k] = v.get()

    cdo.yearmean(
        input="-cat %s" % (" ".join([results[x] for x in files])),
        output=ofile,
        force=False,
        options="-f nc",
    )
    return ofile


# }}} ==========================================================================
# Sea Ice Cover retrival + processing {{{
iceCover_file = "ci_interim_%s-%s-NH.grb" % (startYear, endYear)
if not os.path.exists(iceCover_file):
    server.retrieve(
        {
            "stream": "oper",
            "levtype": "sfc",
            "param": "31.128",
            "dataset": "interim",
            "step": "0",
            "grid": "0.5/0.5",
            "time": "12",
            "date": "%s-01-01/to/%s-01-01" % (startYear, endYear),
            "type": "an",
            "class": "ei",
            "area": "90/-180/0/180",
            "target": iceCover_file,
        }
    )
else:
    print("use existing file '%s'" % (iceCover_file))
# compute the nh ice extent: minimum usually happens in September
iceExtent = "ice_extent_%s-%s-daymean-SeptMin.nc" % (startYear, endYear)
cdo.setattribute(
    "sea_ice_extent@unit=m2,sea_ice_extent@standard_name=sea_ice_extent",
    input="-setname,sea_ice_extent -yearmin -fldsum -mul -selmon,9 %s -gridarea %s"
    % (iceCover_file, iceCover_file),
    output=iceExtent,
    force=False,
    options="-f nc",
)
iceExtent_ds = cdo.readXDataset(iceExtent)
# }}} ==========================================================================
# {{{ CO2 retrieval + processing ===========================================================
# cams return tarballs of netcdf files
co2_tarball = "co2_totalColumn_%s-%s.tar" % (startYear, endYear)
if not os.path.exists(co2_tarball):
    server.retrieve(
        {  # CO2
            "dataset": "cams_ghg_inversions",
            "datatype": "ra",
            "date": "%s-01-01/to/%s-01-01" % (startYear, endYear),
            "frequency": "3h",
            "param": "co2",
            "quantity": "total_column",
            "version": "v16r2",
            "target": co2_tarball,
        }
    )
else:
    print("use existing file '%s'" % (co2_tarball))
co2_files = getDataFromTarfile(co2_tarball)
co2_timeSeries = "co2_timeseries_%s-%s.nc" % (startYear, endYear)
computeTimeSeriesOfFilelist(Pool(tasks), co2_files, "XCO2", co2_timeSeries, False)
co2_ds = cdo.readXDataset(co2_timeSeries)

# }}} ==========================================================================
# scatter plot {{{ =============================================================
# some debugging output
iceExtent_ds.info()
co2_ds.info()
# shaping the data for plotting it
xSelection = co2_ds.sel(time=slice("%s-01-01" % (startYear), "%s-01-01" % (endYear)))
ySelection = iceExtent_ds.sel(
    time=slice("%s-01-01" % (startYear), "%s-01-01" % (endYear))
)
# create the final scatter plot
fig = plt.figure(figsize=(10, 7))
ax = plt.subplot(1, 1, 1)
trans_offset = mtrans.offset_copy(
    ax.transData, fig=fig, x=0.05, y=-0.20, units="inches"
)  # inches because we are in UK

x = xSelection.to_array()[1, :, 0, 0, 0]
y = ySelection.to_array()[1, :, 0, 0, 0]
plt.scatter(x, y)
# put years as labels
years = xSelection.time.dt.year
for _x, _y, _year in zip(x, y, years):
    plt.text(_x, _y, "%d" % (_year), transform=trans_offset)

plt.grid(True)
plt.ylabel("sea ice extent [m2]")
plt.xlabel("co2 concentration [ppm]")
plt.title("Correlation of NH Sea Ice extent minimum and CO2 emissions")
plt.savefig("seaIce_CO2_correlation.png")
# }}} ==========================================================================
# vim:fdm=marker