# coding=utf-8

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

from netCDF4 import Dataset

from livvkit.util import elements as el

describe_ts = """
Time series of the ({model_color}) CESM 2-meter air temperature (°C) for the 
entire {model_start:.0f}--{model_end:.0f} simulation compared to the ({obs_color}) 
{obs_start:.0f}--{obs_end:.0f} RACMO 2.3 2-meter air temperature.   
"""

describe_bx = """
Box plot of the CESM 2-meter air temperature (°C) for the entire 
{model_start:.0f}--{model_end:.0f} simulation compared to the {obs_start:.0f}--{obs_end:.0f} 
RACMO 2.3 2-meter air temperature. Here, the colored lines represent the median 
values, the boxes are drawn from the first quartile to the third quartile 
(covering the interqartile range, or IQR), and the whiskers cover 1.5x the IQR.
Diamonds outside the whiskers represent suspected outliers. 
"""


def make_plot(config, out_path='.'):
    # ---------------- Data source in rhea, the two files save the same data ------------------------
    f_cesm_atm_ts_trefht = os.path.join(config['cesm_atm_climos'],
                                        'b.e10.BG20TRCN.f09_g16.002.cam2.aavg.h0.yrly.TREFHT.nc')
    f_racmo_ts_t2m_aavg = os.path.join(config['racmo_data'], 'racmo23_GRN_monthly.t2m.ann_tseries_aavg.nc')
    # --------------------------------------------------------------

    img_list = []

    # read f_cesm_atm_ts_trefht
    # time from 100-255, both time and tre have 156 data
    ncid1 = Dataset(f_cesm_atm_ts_trefht, mode='r')
    time = ncid1.variables['time'][:]
    tre = ncid1.variables['TREFHT'][:]  # K
    ncid1.set_auto_scale(True)  # put the scale-factor in the time-series plot

    # unit transfer
    tre = tre - 273.15  # convert from K to C
    time = (time - time[1]) / 365 + 1851  # time from 1851-2006

    # read f_racmo_ts_t2m_aavg
    ncid2 = Dataset(f_racmo_ts_t2m_aavg, mode='r')
    time2 = ncid2.variables['time'][:]
    t2m = ncid2.variables['t2m'][:]  # K
    ncid2.set_auto_scale(True)  # put the scale-factor in the time-series plot

    t2m = t2m - 273.15  # unit transfer

    # NOTE: add months to make the full year, divide by months to get years, add to get 1958
    time2 = (time2 + 6.5) / 12 + 1957

    # min_temp = (np.min(tre))
    # max_temp = (np.max(tre))
    # max_t2m = (np.max(t2m))
    # min_t2m = (np.min(t2m))
    #
    # print("Min TREFHT: {}".format(min_temp))
    # print("Max TREFHT: {}".format(max_temp))
    # print("Min t2m: {}".format(min_t2m))
    # print("Max t2m: {}".format(max_t2m))

    # ------- Time Series PLOT --------
    fig, ax = plt.subplots(1, 1, dpi=300, figsize=(10, 6))

    ax.plot(time, tre)
    ax.plot(time2, t2m, linestyle='--')

    ax.set_ylim(-24, -14)
    ax.set_xlim(1851, 2013)
    ax.xaxis.set_ticks(np.arange(1851, 2013, 26))
    ax.set_xticklabels(['1851', '1877', '1903', '1929', '1955', '1981', '2006'])
    ax.set_ylabel('2-meter air temperature ($^{\circ}$C)', fontsize=14)
    ax.set_xlabel('time', fontsize=14)

    ts_img_path = os.path.join(out_path, 'TimeSeries_T2M.png')
    plt.savefig(ts_img_path)
    plt.close()

    img_link = os.path.join(os.path.basename(out_path),
                            os.path.basename(ts_img_path))
    img_elem = el.image('TimeSeries_T2M',
                        ' '.join(describe_ts.format(model_color='blue',
                                                    model_start=time[0],
                                                    model_end=time[-1],
                                                    obs_color='orange',
                                                    obs_start=time2[0],
                                                    obs_end=time2[-1],
                                                    ).split()),
                        img_link)
    if config:
        img_elem['Height'] = config['image_height']
    img_list.append(img_elem)

    # ------ Box plot of statistics ---------
    # If I have 3 sets of data, put them a matrix with 3 columns,
    # then boxplot(matrix) in one figure data=[a1,a2,a3]

    fig, ax = plt.subplots(1, 1, dpi=300, figsize=(8, 6))

    ln = ax.boxplot([tre, t2m])

    ln['medians'][0].set_color('C0')
    ln['fliers'][0].set_marker('d')
    ln['fliers'][0].set_markeredgecolor('C0')

    ln['medians'][1].set_color('C1')
    ln['fliers'][1].set_marker('d')
    ln['fliers'][1].set_markeredgecolor('C1')

    ax.set_xticklabels(['CESM', 'RACMO 2.3'])
    ax.set_ylabel('2-meter air temperature ($^{\circ}$C)', fontsize=14)
    ax.set_ylim(-24, -14)
    ax.grid(False)

    plt.tight_layout()
    bx_img_path = os.path.join(out_path, 'BoxPlot_T2M.png')
    plt.savefig(bx_img_path)
    plt.close()

    img_link = os.path.join(os.path.basename(out_path),
                            os.path.basename(bx_img_path))
    img_elem = el.image('BoxPlot_T2M',
                        ' '.join(describe_bx.format(model_start=time[0],
                                                    model_end=time[-1],
                                                    obs_start=time2[0],
                                                    obs_end=time2[-1],
                                                    ).split()),
                        img_link)
    img_elem['Height'] = config['image_height']
    img_list.append(img_elem)

    return img_list