# 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 surface mass balance (Gt a^-1) for the 
entire {model_start:.0f}--{model_end:.0f} simulation compared to the ({obs_color}) 
{obs_start:.0f}--{obs_end:.0f} RACMO 2.3 surface mass balance.   
"""

describe_bx = """
Box plot of the CESM surface mass balance (Gt a^-1) for the entire 
{model_start:.0f}--{model_end:.0f} simulation compared to the {obs_start:.0f}--{obs_end:.0f} 
RACMO 2.3 surface mass balance. 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_lnd_ts_qice = os.path.join(config['cesm_lnd_climos'],
                                      'b.e10.BG20TRCN.f09_g16.002.clm2.aavg.h0.yrly_ai.QICE.nc')
    f_racmo_ts_smb_aavg = os.path.join(config['racmo_data'], 'racmo23_GRN_monthly.smb.ann_tseries_aavg.nc')
    # --------------------------------------------------------------

    img_list = []

    # read f_cesm_lnd_ts_qice
    # time from 100-255, both time and acab have 156 data
    ncid1 = Dataset(f_cesm_lnd_ts_qice, mode='r')
    time = ncid1.variables['time'][:]
    acab = ncid1.variables['QICE'][:]  # m/yr

    # unit transfer
    acab = acab * 3.156e7  # convert from mm/s to kg/m^2/yr
    acab = acab * 1.7  # Convert to Gt/yr, for validation, compare to V13
    time = (time - time[1]) / 365 + 1851  # time from 1851-2006

    # read f_racmo_ts_smb_aavg
    ncid2 = Dataset(f_racmo_ts_smb_aavg, mode='r')
    time2 = ncid2.variables['time'][:]
    smb = ncid2.variables['smb'][:]  # mmWE = kg/m2
    smb = smb * 1.7  # unit Transfer to Gt from kg/m2

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

    # maxacab = (np.max(acab))
    # minacab = (np.min(acab))
    # maxtime = (np.max(time))
    # mintime = (np.min(time))
    # maxsmb = (np.max(smb))
    # minsmb = (np.min(smb))
    # maxtime2 = (np.max(time2))
    # mintime2 = (np.min(time2))

    # print("Max QICE: {}".format(maxacab))
    # print("Min QICE: {}".format(minacab))
    # print("Max time: {}".format(maxtime))
    # print("Min time: {}".format(mintime))
    # print("Max smb Gt/yr: {}".format(maxsmb))
    # print("Min smb Gt/yr: {}".format(minsmb))
    # print("Max time2: {}".format(maxtime2))
    # print("Min time2: {}".format(mintime2))

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

    ax.plot(time, acab)
    ax.plot(time2, smb, linestyle='--')

    ax.set_ylim(-50, 750)
    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('Surface mass balance (Gt a$^{-1}$)', fontsize=14)
    ax.set_xlabel('time', fontsize=14)

    plt.tight_layout()
    ts_img_path = os.path.join(out_path, 'TimeSeries_QICE_GT.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_QICE_GT',
                        ' '.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)

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

    ln = ax.boxplot([acab, smb])

    [box.set_color('k') for box in ln['boxes']]
    [box.set_color('k') for box in ln['whiskers']]

    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('Surface mass balance (Gt a$^{-1}$)', fontsize=14)
    ax.set_ylim(-50, 750)
    ax.grid(False)

    plt.tight_layout()
    bx_img_path = os.path.join(out_path, 'BoxPlot_QICE_GT.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_QICE_GT',
                        ' '.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