# read / save displacement data from time-series file
%matplotlib inline
import os
import numpy as np
import matplotlib.pyplot as plt
from mintpy.utils import readfile, utils as ut
plt.rcParams.update({'font.size': 12})

work_dir = os.path.expanduser('~/data/archives/Galapagos/AlcedoSenDT128/mintpy')
os.chdir(work_dir)
print('Go to directory:', work_dir)

# read displacement data between two dates in meters from time series file
dis, atr = readfile.read(ts_file, datasetName=date2)
dis -= readfile.read(ts_file, datasetName=date1)[0]
# OR read displacement data in meters per year from velocity file
# dis, atr = readfile.read(vel_file)

# read geometry data: latitude, longitude, incidence angle and azimuth angle
lat, lon = ut.get_lat_lon(atr)
inc_angle = readfile.read(geom_file, datasetName='incidenceAngle')[0] * np.pi / 180.
az_angle  = readfile.read(geom_file, datasetName='azimuthAngle')[0] * np.pi / 180.

# mask out invalid / unreliable pixels
mask = readfile.read(mask_file)[0]
mask *= ~np.isnan(inc_angle)
dis[mask == 0] = np.nan
lat[mask == 0] = np.nan
lon[mask == 0] = np.nan
inc_angle[mask == 0] = np.nan
az_angle[mask == 0] = np.nan

# calculate the unit vector for InSAR displacement [positive for motion toward from satellite]
ve = np.sin(inc_angle) * np.sin(az_angle) * -1
vn = np.sin(inc_angle) * np.cos(az_angle)
vu = np.cos(inc_angle)

fig, axs = plt.subplots(nrows=1, ncols=6, figsize=[18, 4], sharey=True)
titles = ['Latitude', 'Longitude', 'Displacement', 'V_east', 'V_north', 'V_up']
for ax, data, title in zip(axs, [lat, lon, dis, ve, vn, vu], titles):
    im = ax.imshow(data, interpolation='nearest')
    fig.colorbar(im, ax=ax, location='bottom')
    ax.set_title(title)
fig.tight_layout()
plt.show()

out_file = os.path.abspath(f'dis_{date1}_{date2}.txt')
header =  'number of points: {}\n'.format(np.sum(mask))
header += 'reference pixel in in (lat, lon): ({}, {})\n'.format(atr['REF_LAT'], atr['REF_LON'])
header += 'latitude       longitude    displacement[m]        Vz              Vn              Ve'
data = np.hstack((
    lat[mask].reshape(-1, 1),
    lon[mask].reshape(-1, 1),
    dis[mask].reshape(-1, 1),
    ve[mask].reshape(-1, 1),
    vn[mask].reshape(-1, 1),
    vu[mask].reshape(-1, 1),
))
print('writing to text file: {}'.format(out_file))
np.savetxt(out_file, data, fmt='%10.6f', delimiter='\t', header=header)
print('finished writing.')

# The saved text file is shown as below:
with open(out_file, 'r') as f:
    head = ''.join([next(f) for x in range(15)])
print(head)