%matplotlib inline
import os
import string
import numpy as np
import seaborn as sns
import pandas as pd
from matplotlib import pyplot as plt, ticker, patches
from mintpy.objects import timeseries, sensor
from mintpy.utils import readfile, utils as ut
from mintpy import add
from ipynb.fs.full import utils
plt.rcParams.update({'font.size': 12})

work_dir = os.path.expanduser('~/Papers/2022_Geolocation/figs_src/topTEC')
os.chdir(work_dir)
print('Go to directory:', work_dir)

proj_dirs = [os.path.expanduser('~/data/geolocation/ChileSenAT149/mintpy_offset'),
             os.path.expanduser('~/data/geolocation/ChileSenDT156/mintpy_offset'),
             os.path.expanduser('~/data/geolocation/KyushuAlos2DT23/mintpy_offset')]

proj_names = []
rDicts = []
for proj_dir in proj_dirs:
    suffix = '' if 'Alos2' in proj_dir else '_S1Bias'
    infile = os.path.join(proj_dir, f'timeseriesRg{suffix}_SET_ERA5.h5')
    fbase = os.path.splitext(infile)[0]
    fnames = [infile] + [f'{fbase}_{x}.h5' for x in tbases + ['TECsub']]
    proj_name, tsDict, rDict = utils.read_ts_files(fnames, print_msg=True, print_max=False)
    proj_names.append(proj_name)
    rDicts.append(rDict)

x = [0,1,2]
fig, axs = plt.subplots(nrows=3, ncols=3, figsize=[6, 5.5], sharex=True)
for i in range(3):
    tecs = list(rDicts[i].values())
    clr, jlr, jhr = tecs[1:4], tecs[4:7], tecs[7:10]
    for j, tec in enumerate([tecs[1:4], tecs[4:7], tecs[7:10]]):
        axs[i, j].plot(x, tec, 'o-', color='k')
# axis format
nums = ['(a)', '(b)', '(c)',
        '(d)', '(e)', '(f)',
        '(g)', '(h)', '(i)']
gnames = ['CLR GIM', 'JLR GIM', 'JHR GIM']
for i, ax in enumerate(axs.flatten()):
    ax.tick_params(which='both', direction='in', top=True, bottom=True, left=True, right=True)
    ax.set_xlim(-0.4, 2.4)
    ax.annotate(f'{nums[i]}', xy=(0.08, 0.8), xycoords='axes fraction', ha='left')
for ax in axs[0,:2]:  ax.set_ylim(15.5, 23.5); ax.yaxis.set_major_locator(ticker.MultipleLocator(4));  ax.yaxis.set_minor_locator(ticker.MultipleLocator(1))
for ax in axs[0,2:]:  ax.set_ylim(5.5 , 9.5 ); ax.yaxis.set_major_locator(ticker.MultipleLocator(2));  ax.yaxis.set_minor_locator(ticker.MultipleLocator(0.5))
for ax in axs[1,: ]:  ax.set_ylim(4.5,  6.5 ); ax.yaxis.set_major_locator(ticker.MultipleLocator(1));  ax.yaxis.set_minor_locator(ticker.MultipleLocator(0.25))
for ax in axs[2,:2]:  ax.set_ylim(120 , 175 ); ax.yaxis.set_major_locator(ticker.MultipleLocator(20)); ax.yaxis.set_minor_locator(ticker.MultipleLocator(10))
for ax in axs[2,2:]:  ax.set_ylim(47  , 102 ); ax.yaxis.set_major_locator(ticker.MultipleLocator(20)); ax.yaxis.set_minor_locator(ticker.MultipleLocator(10))
for i in range(3):    axs[0,i].set_title(gnames[i])
axs[1,0].set_ylabel('RMSE [cm]')
#plt.xticks(x, ['no\nscale', 'scale\nw/\nfixed\nratio', 'scale\nw/\nadap.\nratio'])
plt.xticks(x, ['No', 'Fix', 'Adap.'])
axs[2,1].set_xlabel('Scaling method')
fig.tight_layout()
fig.subplots_adjust(wspace=0.3)
axs[0,0].annotate('Sentinel-1 asc.',  fontsize=11, xy=(0.52, 0.15), xycoords='axes fraction', ha='center')
axs[1,0].annotate('Sentinel-1 desc.', fontsize=11, xy=(0.52, 0.55), xycoords='axes fraction', ha='center')
axs[2,0].annotate('ALOS-2 desc.',     fontsize=11, xy=(0.52, 0.13), xycoords='axes fraction', ha='center')

# outputs
out_fig = os.path.join(work_dir, 'topTEC_scale_rmse.pdf')
print('save figure to file', out_fig)
plt.savefig(out_fig, bbox_inches='tight', transparent=True, dpi=300)
plt.show()