#! /usr/bin/env python3
# -*- coding:utf8 -*-

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

# data source: PRIS, OurWorldInData
years_reactors_construction_Eprod_Pcap = [
[1955, 1, 15, 'nan', 0.005],
[1956, 2, 27, 'nan', 0.055],
[1957, 5, 30, 'nan', 0.105],
[1958, 6, 35, 'nan', 0.155],
[1959, 11, 39, 'nan', 0.393],
[1960, 15, 42, 'nan', 0.845],
[1961, 16, 48, 'nan', 0.86],
[1962, 25, 43, 'nan', 1.75],
[1963, 33, 43, 'nan', 2.10],
[1964, 40, 44, 'nan', 3.06],
[1965, 48, 51, 26, 4.74],
[1966, 55, 68, 34, 6.11],
[1967, 64, 92, 41, 8.12],
[1968, 69, 102, 52, 9.14],
[1969, 78, 130, 61, 12.7],
[1970, 84, 138, 79, 16.3],
[1971, 99, 149, 110, 24.0],
[1972, 113, 160, 152, 33.0],
[1973, 132, 165, 204, 45.5],
[1974, 154, 176, 267, 62.5],
[1975, 169, 191, 370, 72.9],
[1976, 186, 187, 433, 87.2],
[1977, 200, 185, 539, 99.7],
[1978, 219, 189, 626, 115.8],
[1979, 225, 202, 651, 121.6],
[1980, 245, 196, 712, 135.0],
[1981, 267, 187, 841, 156.3],
[1982, 284, 178, 912, 170.2],
[1983, 306, 160, 1034, 189.8],
[1984, 336, 140, 1255, 221.2],
[1985, 363, 119, 1489, 253.0],
[1986, 389, 97, 1595, 277.9],
[1987, 407, 79, 1735, 299.3],
[1988, 416, 70, 1891, 311.6],
[1989, 420, 56, 1945, 320.1],
[1990, 416, 56, 2000, 325.1],
[1991, 415, 57, 2096, 324.7],
[1992, 418, 58, 2112, 329.3],
[1993, 427, 48, 2185, 336.3],
[1994, 429, 45, 2226, 339.4],
[1995, 434, 39, 2322, 342.5],
[1996, 438, 33, 2407, 349.2],
[1997, 434, 35, 2390, 349.3],
[1998, 430, 41, 2431, 346.2],
[1999, 432, 42, 2524, 348.9],
[2000, 435, 40, 2581, 351.5],
[2001, 438, 38, 2654, 354.5],
[2002, 444, 27, 2596, 358.6],
[2003, 443, 25, 2642, 360.8],
[2004, 443, 25, 2762, 366.1],
[2005, 442, 22, 2769, 369.0],
[2006, 442, 30, 2803, 371.7],
[2007, 438, 32, 2746, 371.6],
[2008, 436, 36, 2738, 369.7],
[2009, 436, 49, 2699, 369.9],
[2010, 438, 59, 2769, 372.1],
[2011, 443, 61, 2653, 375.9],
[2012, 429, 63, 2471, 364.6],
[2013, 428, 68, 2491, 365.7],
[2014, 426, 72, 2541, 364.9],
[2015, 435, 67, 2576, 374.7],
[2016, 438, 62, 2614, 380.6],
[2017, 438, 59, 2637, 382.1],
[2018, 444, 55, 2700, 390.1],
[2019, 443, 46, 2796, 390.3],
[2020, 435, 52, 2694, 385.8],
[2021, 435, 56, 2800, 386.3] ]
# please update in future.

years, reactors, construction, Eprod, Pcap = map(list, zip(*years_reactors_construction_Eprod_Pcap))
Eprod = np.array([float(x) for x in Eprod])
Pcap = np.array([float(x) for x in Pcap])
reactors = np.array(reactors)
construction = np.array(construction)
Pprod = Eprod * (1e3 / (24 * 365.25))
Nsum = reactors + construction

labels = {"y1": {"en":"power [GW]", "de":"Leistung [GW]", "it":"Potenza [GW]"},
    "y2":{"en":"number of\npower reactors", "de":"Anzahl der\nKernreaktoren", "it":"Numero di\nReattori di potenza"},
    "installed":{"en":"installed", "de":"installiert", "it":"installata"},
    "realized":{"en":"realized", "de":"genutzt", "it":"realizzata"},
    "construction":{"en":"Under Construction", "de":"Im Bau", "it":"In Construzione"},
    "active":{"en":"Active Reactors", "de":"Aktive Reaktoren", "it":"Reattori attivi"},
    "threemile":{"en":"Three Mile Island", "de":"Three Mile Island", "it":"Three Mile Island"},
    "chernobyl":{"en":"Chernobyl", "de":"Tschernobyl", "it":"Chernobyl"},
    "fukushima":{"en":"Fukushima", "de":"Fukushima", "it":"Fukushima"} }

plt.rcParams.update({'text.usetex':False, "svg.fonttype":"none", 'font.size':14})
fig = plt.figure(figsize=(600 / 90.0, 480 / 90.0), dpi=72)
ax1 = plt.subplot(211)
plt.plot(years, Pcap, color='blue', lw=2, label=labels['installed']['en'])
plt.plot(years, Pprod, color='red', zorder=-1, lw=2, label=labels['realized']['en'])
plt.ylabel(labels['y1']['en'])
plt.xlim(years[0], years[-1])
plt.ylim(0, 400)
plt.grid(True)
plt.setp(ax1.get_xticklabels(), visible=False)
plt.legend(loc='upper left', framealpha=1, borderpad=0.8)

ax2 = plt.subplot(212, sharex=ax1)
plt.fill_between(years, Nsum, reactors, color='#d0d0d0', zorder=-5)
plt.plot(years, Nsum, color='#808080')
plt.fill_between(years, reactors, color='blue')
plt.plot(years, reactors, color='black')
ax2.yaxis.set_major_locator(ticker.MultipleLocator(100))
plt.ylim(0, 550)
#plt.xlabel('year')
plt.ylabel(labels['y2']['en'])
plt.grid(True)
ax2.text(1976, 270, labels['construction']['en'], ha="center", va="center", rotation=45)
ax2.text(2000, 230, labels['active']['en'], ha="center", va="center", bbox={'facecolor': 'white', 'edgecolor': 'black', 'boxstyle': 'square', 'pad':0.6})
arrowprops={'arrowstyle':'->', 'ec':'red', 'connectionstyle':'angle,angleA=0,angleB=270'}
ax2.annotate(labels['threemile']['en'], (1979.3, 427), (1977, 470), size=11, ha='right', va='center', arrowprops=arrowprops)
ax2.annotate(labels['chernobyl']['en'], (1986.3, 482), (1984, 520), size=11, ha='right', va='center', arrowprops=arrowprops)
ax2.annotate(labels['fukushima']['en'], (2011.2, 492), (2009, 520), size=11, ha='right', va='center', arrowprops=arrowprops)

plt.tight_layout()
fname = 'Nuclear-energy-timeline.svg'
plt.savefig(fname)

if True:
    # create multilingual labels with the svg switch element
    from copy import deepcopy
    from lxml import etree
    for k, v in list(labels.items()):
        # fix multi-line labels
        if '\n' in v['en']:
            for i, line in enumerate(v['en'].split('\n')):
                labels[k+str(i)] = dict([(k2, v2.split('\n')[i]) for k2, v2 in v.items()])
    
    with open(fname, "r") as svgfile:
        svg = etree.parse(svgfile, etree.XMLParser()).getroot()
    for label, values in labels.items():
        for el in svg.findall(".//{" + svg.nsmap[None] + "}text"):
            if el.text == values["en"]:
                switch = etree.SubElement(el.getparent(), "switch")
                for lang, text in values.items():
                    el2 = deepcopy(el)
                    el2.set("systemLanguage", lang)
                    el2.text = text
                    switch.append(el2)
                switch.append(el)
    with open(fname, "w") as svgfile:
        svgfile.write(etree.tostring(svg, pretty_print=True, encoding="unicode"))