#!/usr/bin/env python
# -*- coding: UTF-8 -*-

'''
Plotting tool for BeCOOL datasets 1.2 & 2.2

reference python environment :
- python        3.9
- numpy         1.20
- xarray        2022.3
- matplotlib    3.4

module "python/meso-3.9" on Mesocentre ESPRI IPSL contains these packages
'''

__project__   = 'plotting tools for BeCOOL'
__author__    = 'Thomas LESIGNE (thomas.lesigne@latmos.ipsl.fr)'
__date__      = '2023/07/06'
__version__   = '1.0'

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors


def becool_curtain(ds, ax=None, var='bt2', ymin=-1, ymax=21, colorbar=True,
                   concat_mode='normal', title=None):
        '''
        plotting 2D lidar curtain (time Vs altitude)
        
        Parameters
        - ds (xarray dataset)
        - ax (matplotlib.axes, default:None) - matplotlib axes where to plot the data. If None a new figure will be created
        - var ('bt2', 'beta_part' or 'scattering_ratio' ) - variable to plot
        - ymin, ymax (float or int) : altitude boundaries in km
        - colorbar (bool, default:True) - draw colorbar
        - concat_mode ('normal' or 'night') - temporal concatenationn of profiles
        - title (string) - title of the plot
        '''
        
        # initialize figure if necessary
        if ax is None:
                fig, ax = plt.subplots(figsize=(18,9))
        
        # set dynamic bounds and colormap
        if var =='bt2':
                vmin=1e-6
                vmax=1e-2
                cmap = 'inferno'
        elif var == 'beta_part':
                vmin = 1e-10
                vmax = 5e-3
                cmap='turbo'
        elif var == 'scattering_ratio':
                vmin = 1
                vmax = 2e4
                cmap='jet'
        
        # build 2D altitudes variable
        ds['altitudes'] = ds.alti - ds.range
        
        if concat_mode == 'normal':
                data = ds[var].values
                altitudes = ds.altitudes.values
                
                # start time of measurements
                t_start = ds.time_start.values
                
                # delta_t between start of consecutive measurements
                delta = list(ds.time_start.diff('time', label='lower').values)
                delta.append(min(delta))
                delta = np.array(delta)
                
                # end time of measurements
                t_stop = ds.time_stop.values
                
                # if delta[i,i+1] < 2min, t_stop[i] = t-start[i+1]
                t_start_plus = list(t_start[1:])
                t_start_plus.append(t_stop[-1])
                t_start_plus = np.array(t_start_plus)
                t_stop = np.where(delta < np.timedelta64(2,'m'), t_start_plus, t_stop)
        
        elif concat_mode == 'night':
                # add sequence_idx as coordinate
                ds.coords['sequence_idx'] = ds.sequence_idx
                
                # start time of measurements
                t_start = ds.groupby('sequence_idx').map(lambda x: x.time_start[0]).values
                
                # stop time of measurements
                t_stop = ds.groupby('sequence_idx').map(lambda x: x.time_stop[-1]).values
                
                # delta_t between start of consecutive measurements
                delta = list(np.diff(t_start))
                delta.append(min(delta))
                delta = np.array(delta)
                
                # if delta[i,i+1] < 25 min, t_stop[i] = t-start[i+1]
                t_start_plus = list(t_start[1:])
                t_start_plus.append(t_stop[-1])
                t_start_plus = np.array(t_start_plus)
                t_stop = np.where(delta < np.timedelta64(25,'m').astype('<m8[ns]'), t_start_plus, t_stop)

                # compute 10_minutes mins
                data = ds[var].groupby('sequence_idx').mean('time').values
                altitudes = ds.altitudes.groupby('sequence_idx').mean('time').values
                
        # switch to kilometers
        altitudes /= 1000
        if var in ['bt2','beta_part']:
                data *= 1000
        
        # array of start and stop times
        time = [None]*len(t_start)*2
        time[::2] = t_start
        time[1::2] = t_stop
        
        # increasing time index for start/stop
        t_full = np.sort(np.concatenate((t_start,t_start + np.timedelta64(30,'s'))))
        
        # adding nan profiles between measurements
        data = np.repeat(data,2,axis=0)
        data[1::2,:] = np.nan
        altitudes = np.repeat(altitudes,2,axis=0)
        
        # building new dataarray
        a_tracer = xr.DataArray(data,
                                coords={'t':t_full,
                                        'z':ds.range.values,
                                        'time':('t',time),
                                        'altitude':(['t','z'],altitudes)},
                                dims=['t','z'])
        
        # add name and units attributes
        a_tracer.altitude.attrs['units'] = 'km'
        if var == 'bt2':
                a_tracer.attrs['units'] = 'sr-1.km-1'
                a_tracer.attrs['long_name'] = 'Attenuated Backscatter'
        elif var == 'beta_part':
                a_tracer.attrs['units'] = 'sr-1.km-1'
                a_tracer.attrs['long_name'] = 'Particular Backscatter'
        elif var == 'scattering_ratio':
                a_tracer.attrs['units'] = 'no_unit'
                a_tracer.attrs['long_name'] = 'Scattering Ratio'
        
        # colorbar properties
        if colorbar:
                cbar_kwargs = {'location':'right'}
        else:
                cbar_kwargs = {}
        
        # draw curtain with log dynamic
        norm = mcolors.LogNorm(vmin=vmin, vmax=vmax)
        a_tracer.plot(x='time',y='altitude',cmap=cmap,ax=ax,
                norm=norm,
                infer_intervals=False,
                cbar_kwargs=cbar_kwargs,
                add_colorbar=colorbar)
                
        # some final adjustments
        ax.set_facecolor('k')
        ax.set_ylim([ymin,ymax])
        if title is not None:
                ax.set_title(title)

        return ax








if __name__ == '__main__':
        '''
        Examples
        '''
        # flight name
        flight = 'ST2_C1_02_STR1'
        
        # level 1 file
        L1_file = f"{flight}_BECOOL_1.2.nc"
        
        # level 2 file
        L2_file = f"{flight}_BECOOL_2.2.nc"
        
        # reading files
        ds1 = xr.open_dataset(L1_file)
        ds2 = xr.open_dataset(L2_file)
        
        
        ### first example : 5 nights, concatenated 10-minutes measurements
        # temporal restriction
        ds1 = ds1.where(ds1.night_idx < 5, drop=True)
        ds2 = ds2.where(ds2.night_idx < 5, drop=True)
        
        fig, ax = plt.subplots(2,1, sharex=True, sharey=True, figsize = (12,12))
        becool_curtain(ds1, ax=ax[0], concat_mode='night', 
                       title='Level 1 : Attenuated Backscatter')
        becool_curtain(ds2, ax=ax[1], concat_mode='night', var='scattering_ratio',
                       title='Level 2 : Scattering Ratio')
        plt.show()
         
        
        ### second example : 1 night, real 1-minute resolution for attenuated backscatter
        # new temporal restriction
        ds1 = ds1.where(ds1.night_idx == 2, drop=True)
        ds2 = ds2.where(ds2.night_idx == 2, drop=True)
        
        fig, ax = plt.subplots(2,1, sharex=True, sharey=True, figsize = (12,12))
        becool_curtain(ds1, ax=ax[0], concat_mode='normal', 
                       title='Level 1 : Attenuated Backscatter')
        becool_curtain(ds2, ax=ax[1], concat_mode='normal', var='scattering_ratio',
                          title='Level 2 : Scattering Ratio')
        plt.show()