import os
import numpy as np
from matplotlib import pyplot as plt, ticker
from scipy import stats
from mintpy.objects import timeseries, sensor
from mintpy.utils import readfile, utils as ut, plot as pp


##---------- plot range delay time series
def adjust_ts_axis_format(ax, x, vstep=0.3, xlim=None, ylim=None, ylabel='cumulative\nslant range offset [m]', xticklabels=True):
    ax.tick_params(which='both', direction='in', bottom=True, top=True, left=True, right=True)
    ax.grid('on', linestyle='--')
    ax.yaxis.set_major_locator(ticker.MultipleLocator(vstep))
    pp.auto_adjust_xaxis_date(ax, x, every_year=1, buffer_year=None)
    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)
    if ylabel:
        ax.set_ylabel(ylabel)
    if not xticklabels:
        ax.set_xticklabels([])
    return

def plot_dot_figure(ax, x, y, vlim, vstep=0.2, xname=None, yname=None, fbase='offset'):
    # omit nan value
    flag = (~np.isnan(x)) * (~np.isnan(y))
    x = x[flag]
    y = y[flag]

    ax.plot(x, y, 'k.')
    ax.plot(vlim, vlim, 'k--', lw=1)

    # axis format
    ax.tick_params(which='both', direction='in', bottom=True, top=True, left=True, right=True)
    ax.xaxis.set_major_locator(ticker.MultipleLocator(vstep))
    ax.yaxis.set_major_locator(ticker.MultipleLocator(vstep))
    #ax.set_yticklabels([])
    ax.yaxis.tick_right()
    ax.yaxis.set_label_position("right")
    ax.set_xlim(vlim)
    ax.set_ylim(vlim)
    #ax.set_xlabel('${{{b}}}_{{{n}}}$ [m]'.format(b=fbase, n=xname), color='C0')
    #ax.set_ylabel('${{{b}}}_{{{n}}}$ [m]'.format(b=fbase, n=yname), color='C1', labelpad=ylabelpad)
    ax.set_aspect('equal', 'box')

    # stats
    rmse = np.sqrt(np.sum((y - x)**2) / (x.size - 1))
    r2 = stats.linregress(y, x)[2]
    msg  = r'$R^2$ = {:.2f}'.format(r2) + '\n'
    msg += r'$RMSE$ = {:.1f} cm'.format(rmse*100) + '\n'
    #ax.annotate(msg, xy=(0.05, 0.65), xycoords='axes fraction', color='k', ha='left')
    #ax.annotate(msg, xy=(0.95, 0.07), xycoords='axes fraction', color='k', ha='right')
    print(msg)

    return

##---------- read list of time series files for RMSE analysis
def read_ts_files(fnames, print_msg=True, print_max=True, rm_SenD_outlier=True):
    fnames = [x for x in fnames if os.path.isfile(x)]
    proj_dir = os.path.dirname(fnames[0])
    proj_name = os.path.basename(os.path.dirname(proj_dir))

    ### geometry files
    geom_file = os.path.join(proj_dir, 'inputs/geometryRadar.h5')
    mask_file = os.path.join(proj_dir, 'maskResInv.h5')

    # date info
    date_list = timeseries(fnames[0]).get_date_list()
    num_date = len(date_list)

    # for ChileSenDT156, remove the outlier at 2015-04-02
    flag = np.ones(num_date, dtype=np.bool_)
    if rm_SenD_outlier and proj_name == 'ChileSenDT156':
        flag[date_list.index('20150402')] = 0

    # lalo --> box
    lalo = [-21.30, -67.39] if proj_name.startswith('ChileSen') else None
    if lalo is None:
        box = None
    else:
        atr = readfile.read_attribute(fnames[0])
        coord = ut.coordinate(atr, lookup_file=geom_file)
        y, x = coord.geo2radar(lalo[0], lalo[1])[:2]
        win = 10
        box = (x-win, y-win, x+win+1, y+win+1)

    # read data into lists
    tsDict, rDict = dict(), dict()
    mask = readfile.read(mask_file, box=box)[0].flatten()
    for i, fname in enumerate(fnames):
        # label
        fbase = os.path.basename(fname).replace('timeseriesRg', 'SAR').replace('.h5', '')
        label = ' - '.join(fbase.split('_'))

        # read data
        ts_data = readfile.read(fname, box=box)[0].reshape(num_date, -1)[flag][:, mask]
        ts_med = np.nanmedian(ts_data, axis=-1)
        ts_med -= np.nanmedian(ts_med)

        # print out stats
        rmse = ut.root_mean_sq_error(ts_med) * 100
        dmax = np.nanmax(np.abs(ts_med)) * 100
        if print_msg:
            msg = ''
            if i == 0:
                msg += f'{proj_name}: RMSE / MAX\n' if print_max else f'{proj_name}: RMSE\n'
            #fdigit = 0 if 'Alos2' in proj_name else 1
            fdigit = 1
            msg += '    {l:<40}: {r:6.{f}f}'.format(l=label, r=rmse, f=fdigit)
            msg += ' / {:4.0f}'.format(dmax) if print_max else ''
            msg += ' cm'
            print(msg)

        # save for outputs
        tsDict[label] = ts_med
        rDict[label] = rmse

    return proj_name, tsDict, rDict