#!/usr/bin/env python
u"""
read_cryosat_L2.py
Written by Tyler Sutterley (10/2018)

Reads CryoSat Level-2 data products from baselines A, B and C
Supported CryoSat Modes: LRM, SAR, SARin, FDM, SID, GDR

INPUTS:
	full_filename: full path of CryoSat .DBL file

OUTPUTS:
	Data_1Hz: Time and Orbit Parameters
	Corrections: Elevation Corrections and Flags
	Data_20Hz: Geolocation and Elevation Measurements with Quality Parameters
	METADATA: MPH, SPH and DSD Header data

UPDATE HISTORY:
Updated 10/2018: updated header read functions for python3
Updated 11/2016: added Abs_Orbit and Ascending_Flg to Data_1Hz outputs
	Abs_Orbit should be same as in read_cryosat_ground_tracks.py
	Ascending_Flg can use in surface regression fits following McMillan (2014)
Updated 05/2016: using __future__ print and division functions
Written 03/2016
"""
from __future__ import print_function
from __future__ import division

import os
import re
import numpy as np

#-- PURPOSE: Initiate L2 MDS variables for CryoSat Baselines A and B
def cryosat_baseline_AB(fid,record_size,n_records):
	#-- CryoSat-2 1 Hz data fields (Location Group)
	#-- Time and Orbit Parameters plus Measurement Mode
	L2_1Hz_parameters = {}
	#-- Time: day part
	L2_1Hz_parameters['Day'] = np.zeros((n_records),dtype=np.int32)
	#-- Time: second part
	L2_1Hz_parameters['Second'] = np.zeros((n_records),dtype=np.int32)
	#-- Time: microsecond part
	L2_1Hz_parameters['Micsec'] = np.zeros((n_records),dtype=np.int32)
	#-- SIRAL mode
	L2_1Hz_parameters['Siral_mode'] = np.zeros((n_records),dtype=np.uint64)
	#-- Lat_1Hz: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_1Hz_parameters['Lat_1Hz'] = np.zeros((n_records),dtype=np.int32)
	#-- Lon_1Hz: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_1Hz_parameters['Lon_1Hz'] = np.zeros((n_records),dtype=np.int32)
	#-- Alt_1Hz: packed units (mm, 1e-3 m)
	#-- Altitude of COG above reference ellipsoid (interpolated value)
	L2_1Hz_parameters['Alt_1Hz'] = np.zeros((n_records),dtype=np.int32)
	#-- Mispointing: packed units (millidegrees, 1e-3 degrees)
	L2_1Hz_parameters['Mispointing'] = np.zeros((n_records),dtype=np.int16)
	#-- Number of valid records in the block of twenty that contain data
	#-- Last few records of the last block of a dataset may be blank blocks
	#-- inserted to bring the file up to a multiple of twenty.
	L2_1Hz_parameters['N_valid'] = np.zeros((n_records),dtype=np.int16)

	#-- CryoSat-2 geophysical corrections (External Corrections Group)
	L2_final_corrections = {}
	#-- Dry Tropospheric Correction packed units (mm, 1e-3 m)
	L2_final_corrections['dryTrop'] = np.zeros((n_records),dtype=np.int16)
	#-- Wet Tropospheric Correction packed units (mm, 1e-3 m)
	L2_final_corrections['wetTrop'] = np.zeros((n_records),dtype=np.int16)
	#-- Inverse Barometric Correction packed units (mm, 1e-3 m)
	L2_final_corrections['InvBar'] = np.zeros((n_records),dtype=np.int16)
	#-- Dynamic Atmosphere Correction packed units (mm, 1e-3 m)
	L2_final_corrections['DynAtm'] = np.zeros((n_records),dtype=np.int16)
	#-- Ionospheric Correction packed units (mm, 1e-3 m)
	L2_final_corrections['Iono'] = np.zeros((n_records),dtype=np.int16)
	#-- Sea State Bias Correction packed units (mm, 1e-3 m)
	L2_final_corrections['SSB'] = np.zeros((n_records),dtype=np.int16)
	#-- Ocean tide Correction packed units (mm, 1e-3 m)
	L2_final_corrections['ocTideElv'] = np.zeros((n_records),dtype=np.int16)
	#-- Long period equilibrium ocean tide Correction packed units (mm, 1e-3 m)
	L2_final_corrections['lpeTideElv'] = np.zeros((n_records),dtype=np.int16)
	#-- Ocean loading tide Correction packed units (mm, 1e-3 m)
	L2_final_corrections['olTideElv'] = np.zeros((n_records),dtype=np.int16)
	#-- Solid Earth tide Correction packed units (mm, 1e-3 m)
	L2_final_corrections['seTideElv'] = np.zeros((n_records),dtype=np.int16)
	#-- Geocentric Polar tide Correction packed units (mm, 1e-3 m)
	L2_final_corrections['gpTideElv'] = np.zeros((n_records),dtype=np.int16)
	L2_final_corrections['Spare1'] = np.zeros((n_records),dtype=np.int16)
	#-- Surface Type: Packed in groups of three bits for each of the 20 records
	L2_final_corrections['Surf_type'] = np.zeros((n_records),dtype=np.uint64)
	#-- Mean Sea Surface or Geoid packed units (mm, 1e-3 m)
	L2_final_corrections['MSS_Geoid'] = np.zeros((n_records),dtype=np.int32)
	#-- Ocean Depth/Land Elevation Model (ODLE) packed units (mm, 1e-3 m)
	L2_final_corrections['ODLE'] = np.zeros((n_records),dtype=np.int32)
	#-- Ice Concentration packed units (%/100)
	L2_final_corrections['Ice_conc'] = np.zeros((n_records),dtype=np.int16)
	#-- Snow Depth packed units (mm, 1e-3 m)
	L2_final_corrections['Snow_depth'] = np.zeros((n_records),dtype=np.int16)
	#-- Snow Density packed units (kg/m^3)
	L2_final_corrections['Snow_density'] = np.zeros((n_records),dtype=np.int16)
	L2_final_corrections['Spare2'] = np.zeros((n_records),dtype=np.int16)
	#-- Corrections Status Flag
	L2_final_corrections['C_status'] = np.zeros((n_records),dtype=np.uint32)
	#-- Significant Wave Height (SWH) packed units (mm, 1e-3)
	L2_final_corrections['SWH'] = np.zeros((n_records),dtype=np.int16)
	#-- Wind Speed packed units (mm/s, 1e-3 m/s)
	L2_final_corrections['Wind_speed'] = np.zeros((n_records),dtype=np.uint16)
	L2_final_corrections['Spare3'] = np.zeros((n_records),dtype=np.int16)
	L2_final_corrections['Spare4'] = np.zeros((n_records),dtype=np.int16)
	L2_final_corrections['Spare5'] = np.zeros((n_records),dtype=np.int16)
	L2_final_corrections['Spare6'] = np.zeros((n_records),dtype=np.int16)

	#-- CryoSat-2 20 Hz data fields (Measurement Group)
	#-- Derived from instrument measurement parameters
	n_blocks = 20
	L2_final_measurements = {}
	#-- Delta between the timestamps for 20Hz record and the 1Hz record
	#-- D_time_mics packed units (microseconds)
	L2_final_measurements['D_time_mics'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Lat: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_final_measurements['Lat'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Lon: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_final_measurements['Lon'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Measured elevation above ellipsoid from retracker: packed units (mm, 1e-3 m)
	L2_final_measurements['Elev'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Interpolated Sea Surface Height Anomaly: packed units (mm, 1e-3 m)
	L2_final_measurements['SSHA_interp'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Interpolated Sea Surface Height measurement count
	L2_final_measurements['SSHA_num'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Interpolation quality estimate RSS: packed units (mm, 1e-3 m)
	L2_final_measurements['SSHA_qual'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Sigma Zero Backscatter for retracker: packed units (1e-2 dB)
	L2_final_measurements['Sig0'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Peakiness: packed units (1e-2)
	L2_final_measurements['Peakiness'] = np.zeros((n_records,n_blocks),dtype=np.uint16)
	#-- Freeboard: packed units (mm, 1e-3 m)
	#-- -9999 default value indicates computation has not been performed
	L2_final_measurements['Freeboard'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Number of averaged echoes or beams
	L2_final_measurements['N_avg'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	L2_final_measurements['Spare1'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Quality flags
	L2_final_measurements['Quality_Flg'] = np.zeros((n_records,n_blocks),dtype=np.uint32)
	L2_final_measurements['Spare2'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	L2_final_measurements['Spare3'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	L2_final_measurements['Spare4'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	L2_final_measurements['Spare5'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- for each record in the CryoSat file
	for r in range(n_records):
		#-- CryoSat-2 Location Group for record r
		L2_1Hz_parameters['Day'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_1Hz_parameters['Second'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_1Hz_parameters['Micsec'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_1Hz_parameters['Siral_mode'][r] = np.fromfile(fid,dtype='>u8',count=1)
		L2_1Hz_parameters['Lat_1Hz'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_1Hz_parameters['Lon_1Hz'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_1Hz_parameters['Alt_1Hz'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_1Hz_parameters['Mispointing'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_1Hz_parameters['N_valid'][r] = np.fromfile(fid,dtype='>i2',count=1)
		#-- CryoSat-2 External Corrections Group for record r
		L2_final_corrections['dryTrop'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['wetTrop'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['InvBar'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['DynAtm'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Iono'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['SSB'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['ocTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['lpeTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['olTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['seTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['gpTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Spare1'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Surf_type'][r] = np.fromfile(fid,dtype='>u8',count=1)
		L2_final_corrections['MSS_Geoid'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_final_corrections['ODLE'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_final_corrections['Ice_conc'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Snow_depth'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Snow_density'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Spare2'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['C_status'][r] = np.fromfile(fid,dtype='>u4',count=1)
		L2_final_corrections['SWH'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Wind_speed'][r] = np.fromfile(fid,dtype='>u2',count=1)
		L2_final_corrections['Spare3'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Spare4'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Spare5'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_final_corrections['Spare6'][r] = np.fromfile(fid,dtype='>i2',count=1)
		#-- CryoSat-2 Measurements Group for record r and block b
		for b in range(n_blocks):
			L2_final_measurements['D_time_mics'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_final_measurements['Lat'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_final_measurements['Lon'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_final_measurements['Elev'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_final_measurements['SSHA_interp'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['SSHA_num'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['SSHA_qual'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['Sig0'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['Peakiness'][r,b] = np.fromfile(fid,dtype='>u2',count=1)
			L2_final_measurements['Freeboard'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['N_avg'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['Spare1'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['Quality_Flg'][r,b] = np.fromfile(fid,dtype='>u4',count=1)
			L2_final_measurements['Spare2'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['Spare3'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['Spare4'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_final_measurements['Spare5'][r,b] = np.fromfile(fid,dtype='>i2',count=1)

	#-- Bind all the bits of the l2_mds together into a single dictionary
	CS_l2_mds = {}
	CS_l2_mds['Data_1Hz'] = L2_1Hz_parameters
	CS_l2_mds['Corrections'] = L2_final_corrections
	CS_l2_mds['Data_20Hz'] = L2_final_measurements
	#-- return the output dictionary
	return CS_l2_mds

#-- PURPOSE: Initiate L2 MDS variables for CryoSat Baseline C
def cryosat_baseline_C(fid,record_size,n_records):
	#-- CryoSat-2 1 Hz data fields (Location Group)
	#-- Time and Orbit Parameters plus Measurement Mode
	L2_c_1Hz_parameters = {}
	#-- Time: day part
	L2_c_1Hz_parameters['Day'] = np.zeros((n_records),dtype=np.int32)
	#-- Time: second part
	L2_c_1Hz_parameters['Second'] = np.zeros((n_records),dtype=np.int32)
	#-- Time: microsecond part
	L2_c_1Hz_parameters['Micsec'] = np.zeros((n_records),dtype=np.int32)
	#-- SIRAL mode
	L2_c_1Hz_parameters['Siral_mode'] = np.zeros((n_records),dtype=np.uint64)
	#-- Lat_1Hz: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_c_1Hz_parameters['Lat_1Hz'] = np.zeros((n_records),dtype=np.int32)
	#-- Lon_1Hz: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_c_1Hz_parameters['Lon_1Hz'] = np.zeros((n_records),dtype=np.int32)
	#-- Alt_1Hz: packed units (mm, 1e-3 m)
	#-- Altitude of COG above reference ellipsoid (interpolated value)
	L2_c_1Hz_parameters['Alt_1Hz'] = np.zeros((n_records),dtype=np.int32)
	#-- Roll: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_c_1Hz_parameters['Roll'] = np.zeros((n_records),dtype=np.int32)
	#-- Pitch: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_c_1Hz_parameters['Pitch'] = np.zeros((n_records),dtype=np.int32)
	#-- Yaw: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_c_1Hz_parameters['Yaw'] = np.zeros((n_records),dtype=np.int32)
	L2_c_1Hz_parameters['Spare'] = np.zeros((n_records),dtype=np.int16)
	#-- Number of valid records in the block of twenty that contain data
	#-- Last few records of the last block of a dataset may be blank blocks
	#-- inserted to bring the file up to a multiple of twenty.
	L2_c_1Hz_parameters['N_valid'] = np.zeros((n_records),dtype=np.int16)

	#-- CryoSat-2 geophysical corrections (External Corrections Group)
	L2_c_final_corrections = {}
	#-- Dry Tropospheric Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['dryTrop'] = np.zeros((n_records),dtype=np.int16)
	#-- Wet Tropospheric Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['wetTrop'] = np.zeros((n_records),dtype=np.int16)
	#-- Inverse Barometric Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['InvBar'] = np.zeros((n_records),dtype=np.int16)
	#-- Dynamic Atmosphere Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['DynAtm'] = np.zeros((n_records),dtype=np.int16)
	#-- Ionospheric Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['Iono'] = np.zeros((n_records),dtype=np.int16)
	#-- Sea State Bias Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['SSB'] = np.zeros((n_records),dtype=np.int16)
	#-- Ocean tide Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['ocTideElv'] = np.zeros((n_records),dtype=np.int16)
	#-- Long period equilibrium ocean tide Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['lpeTideElv'] = np.zeros((n_records),dtype=np.int16)
	#-- Ocean loading tide Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['olTideElv'] = np.zeros((n_records),dtype=np.int16)
	#-- Solid Earth tide Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['seTideElv'] = np.zeros((n_records),dtype=np.int16)
	#-- Geocentric Polar tide Correction packed units (mm, 1e-3 m)
	L2_c_final_corrections['gpTideElv'] = np.zeros((n_records),dtype=np.int16)
	L2_c_final_corrections['Spare1'] = np.zeros((n_records),dtype=np.int16)
	#-- Surface Type: Packed in groups of three bits for each of the 20 records
	L2_c_final_corrections['Surf_type'] = np.zeros((n_records),dtype=np.uint64)
	#-- Mean Sea Surface or Geoid packed units (mm, 1e-3 m)
	L2_c_final_corrections['MSS_Geoid'] = np.zeros((n_records),dtype=np.int32)
	#-- Ocean Depth/Land Elevation Model (ODLE) packed units (mm, 1e-3 m)
	L2_c_final_corrections['ODLE'] = np.zeros((n_records),dtype=np.int32)
	#-- Ice Concentration packed units (%/100)
	L2_c_final_corrections['Ice_conc'] = np.zeros((n_records),dtype=np.int16)
	#-- Snow Depth packed units (mm, 1e-3 m)
	L2_c_final_corrections['Snow_depth'] = np.zeros((n_records),dtype=np.int16)
	#-- Snow Density packed units (kg/m^3)
	L2_c_final_corrections['Snow_density'] = np.zeros((n_records),dtype=np.int16)
	L2_c_final_corrections['Spare2'] = np.zeros((n_records),dtype=np.int16)
	#-- Corrections Status Flag
	L2_c_final_corrections['C_status'] = np.zeros((n_records),dtype=np.uint32)
	#-- Significant Wave Height (SWH) packed units (mm, 1e-3)
	L2_c_final_corrections['SWH'] = np.zeros((n_records),dtype=np.int16)
	#-- Wind Speed packed units (mm/s, 1e-3 m/s)
	L2_c_final_corrections['Wind_speed'] = np.zeros((n_records),dtype=np.uint16)
	L2_c_final_corrections['Spare3'] = np.zeros((n_records),dtype=np.int16)
	L2_c_final_corrections['Spare4'] = np.zeros((n_records),dtype=np.int16)
	L2_c_final_corrections['Spare5'] = np.zeros((n_records),dtype=np.int16)
	L2_c_final_corrections['Spare6'] = np.zeros((n_records),dtype=np.int16)

	#-- CryoSat-2 20 Hz data fields (Measurement Group)
	#-- Derived from instrument measurement parameters
	n_blocks = 20
	L2_c_final_measurements = {}
	#-- Delta between the timestamps for 20Hz record and the 1Hz record
	#-- D_time_mics packed units (microseconds)
	L2_c_final_measurements['D_time_mics'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Lat: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_c_final_measurements['Lat'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Lon: packed units (0.1 micro-degree, 1e-7 degrees)
	L2_c_final_measurements['Lon'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Measured elevation above ellipsoid from retracker 1: packed units (mm, 1e-3 m)
	L2_c_final_measurements['Elev_1'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Measured elevation above ellipsoid from retracker 2: packed units (mm, 1e-3 m)
	L2_c_final_measurements['Elev_2'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Measured elevation above ellipsoid from retracker 3: packed units (mm, 1e-3 m)
	L2_c_final_measurements['Elev_3'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Sigma Zero Backscatter for retracker 1: packed units (1e-2 dB)
	L2_c_final_measurements['Sig0_1'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Sigma Zero Backscatter for retracker 2: packed units (1e-2 dB)
	L2_c_final_measurements['Sig0_2'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Sigma Zero Backscatter for retracker 3: packed units (1e-2 dB)
	L2_c_final_measurements['Sig0_3'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Freeboard: packed units (mm, 1e-3 m)
	#-- -9999 default value indicates computation has not been performed
	L2_c_final_measurements['Freeboard'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Interpolated Sea Surface Height Anomaly: packed units (mm, 1e-3 m)
	L2_c_final_measurements['SSHA_interp'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Interpolated Sea Surface Height measurement count
	L2_c_final_measurements['SSHA_num'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Interpolation quality estimate RSS: packed units (mm, 1e-3 m)
	L2_c_final_measurements['SSHA_qual'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Peakiness: packed units (1e-2)
	L2_c_final_measurements['Peakiness'] = np.zeros((n_records,n_blocks),dtype=np.uint16)
	#-- Number of averaged echoes or beams
	L2_c_final_measurements['N_avg'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	L2_c_final_measurements['Spare1'] = np.zeros((n_records,n_blocks),dtype=np.int16)
	#-- Quality flags
	L2_c_final_measurements['Quality_Flg'] = np.zeros((n_records,n_blocks),dtype=np.uint32)
	#-- Corrections Application Flag
	L2_c_final_measurements['Corrections_Flg'] = np.zeros((n_records,n_blocks),dtype=np.uint32)
	#-- Quality metric for retracker 1
	L2_c_final_measurements['Quality_1'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Quality metric for retracker 2
	L2_c_final_measurements['Quality_2'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- Quality metric for retracker 3
	L2_c_final_measurements['Quality_3'] = np.zeros((n_records,n_blocks),dtype=np.int32)
	#-- for each record in the CryoSat file
	for r in range(n_records):
		#-- CryoSat-2 Location Group for record r
		L2_c_1Hz_parameters['Day'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Second'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Micsec'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Siral_mode'][r] = np.fromfile(fid,dtype='>u8',count=1)
		L2_c_1Hz_parameters['Lat_1Hz'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Lon_1Hz'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Alt_1Hz'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Roll'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Pitch'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Yaw'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_1Hz_parameters['Spare'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_1Hz_parameters['N_valid'][r] = np.fromfile(fid,dtype='>i2',count=1)
		#-- CryoSat-2 External Corrections Group for record r
		L2_c_final_corrections['dryTrop'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['wetTrop'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['InvBar'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['DynAtm'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Iono'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['SSB'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['ocTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['lpeTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['olTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['seTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['gpTideElv'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Spare1'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Surf_type'][r] = np.fromfile(fid,dtype='>u8',count=1)
		L2_c_final_corrections['MSS_Geoid'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_final_corrections['ODLE'][r] = np.fromfile(fid,dtype='>i4',count=1)
		L2_c_final_corrections['Ice_conc'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Snow_depth'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Snow_density'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Spare2'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['C_status'][r] = np.fromfile(fid,dtype='>u4',count=1)
		L2_c_final_corrections['SWH'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Wind_speed'][r] = np.fromfile(fid,dtype='>u2',count=1)
		L2_c_final_corrections['Spare3'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Spare4'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Spare5'][r] = np.fromfile(fid,dtype='>i2',count=1)
		L2_c_final_corrections['Spare6'][r] = np.fromfile(fid,dtype='>i2',count=1)
		#-- CryoSat-2 Measurements Group for record r and block b
		for b in range(n_blocks):
			L2_c_final_measurements['D_time_mics'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_c_final_measurements['Lat'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_c_final_measurements['Lon'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_c_final_measurements['Elev_1'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_c_final_measurements['Elev_2'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_c_final_measurements['Elev_3'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_c_final_measurements['Sig0_1'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['Sig0_2'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['Sig0_3'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['Freeboard'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['SSHA_interp'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['SSHA_num'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['SSHA_qual'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['Peakiness'][r,b] = np.fromfile(fid,dtype='>u2',count=1)
			L2_c_final_measurements['N_avg'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['Spare1'][r,b] = np.fromfile(fid,dtype='>i2',count=1)
			L2_c_final_measurements['Quality_Flg'][r,b] = np.fromfile(fid,dtype='>u4',count=1)
			L2_c_final_measurements['Corrections_Flg'][r,b] = np.fromfile(fid,dtype='>u4',count=1)
			L2_c_final_measurements['Quality_1'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_c_final_measurements['Quality_2'][r,b] = np.fromfile(fid,dtype='>i4',count=1)
			L2_c_final_measurements['Quality_3'][r,b] = np.fromfile(fid,dtype='>i4',count=1)

	#-- Bind all the bits of the l2_mds together into a single dictionary
	CS_l2_c_mds = {}
	CS_l2_c_mds['Data_1Hz'] = L2_c_1Hz_parameters
	CS_l2_c_mds['Corrections'] = L2_c_final_corrections
	CS_l2_c_mds['Data_20Hz'] = L2_c_final_measurements
	#-- return the output dictionary
	return CS_l2_c_mds

#-- PURPOSE: Read ASCII Main Product Header (MPH) block from an ESA PDS file
def read_MPH(full_filename):
	#-- read input data file
	with open(full_filename, 'rb') as fid:
		file_contents = fid.read().splitlines()

	#-- Define constant values associated with PDS file formats
	#-- number of text lines in standard MPH
	n_MPH_lines	= 41
	#-- check that first line of header matches PRODUCT
	if not bool(re.match(b'PRODUCT\=\"(.*)(?=\")',file_contents[0])):
		raise IOError('File does not start with a valid PDS MPH')
	#-- read MPH header text
	s_MPH_fields = {}
	for i in range(n_MPH_lines):
		#-- use regular expression operators to read headers
		if bool(re.match(b'(.*?)\=\"(.*)(?=\")',file_contents[i])):
			#-- data fields within quotes
			field,value=re.findall(b'(.*?)\=\"(.*)(?=\")',file_contents[i]).pop()
			s_MPH_fields[field.decode('utf-8')] = value.decode('utf-8').rstrip()
		elif bool(re.match(b'(.*?)\=(.*)',file_contents[i])):
			#-- data fields without quotes
			field,value=re.findall(b'(.*?)\=(.*)',file_contents[i]).pop()
			s_MPH_fields[field.decode('utf-8')] = value.decode('utf-8').rstrip()

	#-- Return block name array to calling function
	return s_MPH_fields

#-- PURPOSE: Read ASCII Specific Product Header (SPH) block from a PDS file
def read_SPH(full_filename,j_sph_size):
	#-- read input data file
	with open(full_filename, 'rb') as fid:
		file_contents = fid.read().splitlines()

	#-- Define constant values associated with PDS file formats
	#-- number of text lines in standard MPH
	n_MPH_lines	= 41
	#-- compile regular expression operator for reading headers
	rx = re.compile(b'(.*?)\=\"?(.*)',re.VERBOSE)
	#-- check first line of header matches SPH_DESCRIPTOR
	if not bool(re.match(b'SPH\_DESCRIPTOR\=',file_contents[n_MPH_lines+1])):
		raise IOError('File does not have a valid PDS DSD')
	#-- read SPH header text (no binary control characters)
	s_SPH_lines = [li for li in file_contents[n_MPH_lines+1:] if rx.match(li)
		and not re.search(b'[^\x20-\x7e]+',li)]

	#-- extract SPH header text
	s_SPH_fields = {}
	c = 0
	while (c < len(s_SPH_lines)):
		#-- check if line is within DS_NAME portion of SPH header
		if bool(re.match(b'DS_NAME',s_SPH_lines[c])):
			#-- add dictionary for DS_NAME
			field,value=re.findall(b'(.*?)\=\"(.*)(?=\")',s_SPH_lines[c]).pop()
			key = value.decode('utf-8').rstrip()
			s_SPH_fields[key] = {}
			for line in s_SPH_lines[c+1:c+7]:
				if bool(re.match(b'(.*?)\=\"(.*)(?=\")',line)):
					#-- data fields within quotes
					dsfield,dsvalue=re.findall(b'(.*?)\=\"(.*)(?=\")',line).pop()
					s_SPH_fields[key][dsfield.decode('utf-8')] = dsvalue.decode('utf-8').rstrip()
				elif bool(re.match(b'(.*?)\=(.*)',line)):
					#-- data fields without quotes
					dsfield,dsvalue=re.findall(b'(.*?)\=(.*)',line).pop()
					s_SPH_fields[key][dsfield.decode('utf-8')] = dsvalue.decode('utf-8').rstrip()
			#-- add 6 to counter to go to next entry
			c += 6
		#-- use regular expression operators to read headers
		elif bool(re.match(b'(.*?)\=\"(.*)(?=\")',s_SPH_lines[c])):
			#-- data fields within quotes
			field,value=re.findall(b'(.*?)\=\"(.*)(?=\")',s_SPH_lines[c]).pop()
			s_SPH_fields[field.decode('utf-8')] = value.decode('utf-8').rstrip()
		elif bool(re.match(b'(.*?)\=(.*)',s_SPH_lines[c])):
			#-- data fields without quotes
			field,value=re.findall(b'(.*?)\=(.*)',s_SPH_lines[c]).pop()
			s_SPH_fields[field.decode('utf-8')] = value.decode('utf-8').rstrip()
		#-- add 1 to counter to go to next line
		c += 1

	#-- Return block name array to calling function
	return s_SPH_fields

#-- PURPOSE: Read ASCII Data Set Descriptors (DSD) block from a PDS file
def read_DSD(full_filename):
	#-- read input data file
	with open(full_filename, 'rb') as fid:
		file_contents = fid.read().splitlines()

	#-- Define constant values associated with PDS file formats
	#-- number of text lines in standard MPH
	n_MPH_lines	= 41
	#-- number of text lines in a DSD header
	n_DSD_lines = 8

	#-- Level-2 CryoSat DS_NAMES within files
	regex_patterns = []
	regex_patterns.append(b'DS_NAME\="SIR_LRM_L2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SAR_L2B[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SAR_L2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_FDM_L2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SARIL2B[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SARIL2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SAR_L2B_I[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SAR_L2A[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SIN_L2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SID_L2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_LRMIL2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_LRM_L2_I[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SARIL2A[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SAR_L2A_I[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SAR_L2_I[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SINIL2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SIN_L2_I[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SIDIL2[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_SID_L2_I[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_GDR_2A[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_GDR_2B[\s+]*"')
	regex_patterns.append(b'DS_NAME\="SIR_GDR_2[\s+]*"')
	#-- find the DSD starting line within the SPH header
	c = 0
	Flag = False
	while ((Flag is False) and (c < len(regex_patterns))):
		#-- find indice within
		indice = [i for i,line in enumerate(file_contents[n_MPH_lines+1:]) if
			re.search(regex_patterns[c],line)]
		if indice:
			Flag = True
		else:
			c+=1
	#-- check that valid indice was found within header
	if not indice:
		raise IOError('Can not find correct DSD field')

	#-- extract s_DSD_fields info
	DSD_START = n_MPH_lines + indice[0] + 1
	s_DSD_fields = {}
	for i in range(DSD_START,DSD_START+n_DSD_lines):
		#-- use regular expression operators to read headers
		if bool(re.match(b'(.*?)\=\"(.*)(?=\")',file_contents[i])):
			#-- data fields within quotes
			field,value=re.findall(b'(.*?)\=\"(.*)(?=\")',file_contents[i]).pop()
			s_DSD_fields[field.decode('utf-8')] = value.decode('utf-8').rstrip()
		elif bool(re.match(b'(.*?)\=(.*)',file_contents[i])):
			#-- data fields without quotes
			field,value=re.findall(b'(.*?)\=(.*)',file_contents[i]).pop()
			s_DSD_fields[field.decode('utf-8')] = value.decode('utf-8').rstrip()

	#-- Return block name array to calling function
	return s_DSD_fields

#-- PURPOSE: read CryoSat Level-2 data
def read_cryosat_L2(full_filename, VERBOSE=False):
	#-- file basename and file extension of input file
	fileBasename,fileExtension=os.path.splitext(os.path.basename(full_filename))

	#-- CryoSat file class
	#-- OFFL (Off Line Processing/Systematic)
	#-- NRT_ (Near Real Time)
	#-- RPRO (ReProcessing)
	#-- TEST (Testing)
	#-- LTA_ (Long Term Archive)
	regex_class = 'OFFL|NRT_|RPRO|TEST|LTA_'
	#-- CryoSat mission products
	#-- SIR_LRM_2 L2 Product from Low Resolution Mode Processing
	#-- SIR_FDM_2 L2 Product from Fast Delivery Marine Mode Processing
	#-- SIR_SIN_2 L2 Product from SAR Interferometric Processing
	#-- SIR_SID_2 L2 Product from SIN Degraded Processing
	#-- SIR_SAR_2 L2 Product from SAR Processing
	#-- SIR_GDR_2 L2 Consolidated Product
	#-- SIR_LRMI2 In-depth L2 Product from LRM Processing
	#-- SIR_SINI2 In-depth L2 Product from SIN Processing
	#-- SIR_SIDI2 In-depth L2 Product from SIN Degraded Process.
	#-- SIR_SARI2 In-depth L2 Product from SAR Processing
	regex_products = ('SIR_LRM_2|SIR_FDM_2|SIR_SIN_2|SIR_SID_2|'
	'SIR_SAR_2|SIR_GDR_2|SIR_LRMI2|SIR_SINI2|SIR_SIDI2|SIR_SARI2')
	#-- CRYOSAT LEVEL-2 PRODUCTS NAMING RULES
	#-- Mission Identifier
	#-- File Class
	#-- File Product
	#-- Validity Start Date and Time
	#-- Validity Stop Date and Time
	#-- Baseline Identifier
	#-- Version Number
	regex_pattern = '(.*?)_({0})_({1})__(\d+T?\d+)_(\d+T?\d+)_(.*?)(\d+)'.format(
		regex_class, regex_products)
	rx = re.compile(regex_pattern, re.VERBOSE)
	#-- extract file information from filename
	MI,CLASS,PRODUCT,START,STOP,BASELINE,VERSION=rx.findall(fileBasename).pop()
	#-- Extract Date information
	start_yr,start_mon,start_day=np.array([START[:4],START[4:6],START[6:8]],dtype=np.uint16)
	start_hh,start_mm,start_ss=np.array([START[-6:-4],START[-4:-2],START[-2:]],dtype=np.uint8)
	stop_yr,stop_mon,stop_day=np.array([STOP[:4],STOP[4:6],STOP[6:8]],dtype=np.uint16)
	stop_hh,stop_mm,stop_ss=np.array([STOP[-6:-4],STOP[-4:-2],STOP[-2:]],dtype=np.uint8)

	#-- Record sizes
	CS_L2_MDS_REC_SIZE = 980
	CS_L2_C_MDS_REC_SIZE = 1392
	#-- check baseline from file to set i_record_size and allocation function
	if (BASELINE == 'C'):
		i_record_size = CS_L2_C_MDS_REC_SIZE
		read_cryosat_variables = cryosat_baseline_C
	else:
		i_record_size = CS_L2_MDS_REC_SIZE
		read_cryosat_variables = cryosat_baseline_AB

	#-- read the input file to get file information
	fid = os.open(os.path.expanduser(full_filename),os.O_RDONLY)
	file_info = os.fstat(fid)
	os.close(fid)

	#-- num DSRs from SPH
	j_num_DSR = np.int32(file_info.st_size//i_record_size)
	#-- print file information
	if VERBOSE:
		print(fileBasename)
		print('{0:d} {1:d} {2:d}'.format(j_num_DSR,file_info.st_size,i_record_size))
		#-- Check if MPH/SPH/DSD headers
		if (j_num_DSR*i_record_size == file_info.st_size):
			print('No Header on file')
			print('The number of DSRs is: {0:d}'.format(j_num_DSR))
		else:
			print('Header on file')

	#-- Check if MPH/SPH/DSD headers
	if (j_num_DSR*i_record_size != file_info.st_size):
		#-- If there are MPH/SPH/DSD headers
		s_MPH_fields = read_MPH(full_filename)
		j_sph_size = np.int32(re.findall('[-+]?\d+',s_MPH_fields['SPH_SIZE']).pop())
		s_SPH_fields = read_SPH(full_filename,j_sph_size)
		#-- extract information from DSD fields
		s_DSD_fields = read_DSD(full_filename)
		#-- extract DS_OFFSET
		j_DS_start = np.int32(re.findall('[-+]?\d+',s_DSD_fields['DS_OFFSET']).pop())
		#-- extract number of DSR in the file
		j_num_DSR = np.int32(re.findall('[-+]?\d+',s_DSD_fields['NUM_DSR']).pop())
		#-- check the record size
		j_DSR_size = np.int32(re.findall('[-+]?\d+',s_DSD_fields['DSR_SIZE']).pop())
		#--  minimum size is start of the read plus number of records to read
		j_check_size = j_DS_start +(j_DSR_size*j_num_DSR)
		if VERBOSE:
			print('The offset of the DSD is: {0:d} bytes'.format(j_DS_start))
			print('The number of DSRs is {0:d}'.format(j_num_DSR))
			print('The size of the DSR is {0:d}'.format(j_DSR_size))
		#-- check if invalid file size
		if (j_check_size > file_info.st_size):
			raise IOError('File size error')
		#-- extract binary data from input CryoSat data file (skip headers)
		fid = open(full_filename, 'rb')
		cryosat_header = fid.read(j_DS_start)
		#-- iterate through CryoSat file and fill output variables
		CS_L2_mds = read_cryosat_variables(fid,i_record_size,j_num_DSR)
		#-- add headers to output dictionary as METADATA
		CS_L2_mds['METADATA'] = {}
		CS_L2_mds['METADATA']['MPH'] = s_MPH_fields
		CS_L2_mds['METADATA']['SPH'] = s_SPH_fields
		CS_L2_mds['METADATA']['DSD'] = s_DSD_fields
		#-- add absolute orbit number to 1Hz data
		CS_L2_mds['Data_1Hz']['Abs_Orbit']=np.zeros((j_num_DSR),dtype=np.uint32)
		CS_L2_mds['Data_1Hz']['Abs_Orbit'][:]=np.uint32(s_MPH_fields['ABS_ORBIT'])
		#-- add ascending/descending flag to 1Hz data (A=ascending,D=descending)
		CS_L2_mds['Data_1Hz']['Ascending_Flg']=np.zeros((j_num_DSR),dtype=np.bool)
		if (s_SPH_fields['ASCENDING_FLAG'] == 'A'):
			CS_L2_mds['Data_1Hz']['Ascending_Flg'][:] = True
		#-- close the input CryoSat binary file
		fid.close()
	else:
		#-- If there are not MPH/SPH/DSD headers
		#-- extract binary data from input CryoSat data file
		fid = open(full_filename, 'rb')
		#-- iterate through CryoSat file and fill output variables
		CS_L2_mds = read_cryosat_variables(fid,i_record_size,j_num_DSR)
		#-- close the input CryoSat binary file
		fid.close()

	#-- return the data and headers
	return CS_L2_mds


# coding: utf-8

# ## Working with Matplotlib
# Introduction:
#
# Matplotlib is a Python 2D plotting library which produces publication quality figures in a variety of hardcopy formats and interactive environments across platforms. Matplotlib can be used in Python scripts, the Python and IPython shell, the jupyter notebook, web application servers, and four graphical user interface toolkits.
#
# ### Here are the main steps we will go through
# * How to save your plot using matplotlib?
# * How to create plot using matplotlib?
# * How to make your graph look pretty?
#
# This is Just a little illustration.
#
#

# #### How to save your plot using matplotlib?
# I am gonna show you how to save your graph, you can refer the docstring for complete information on the various ways it can be used.

# In[1]:

# import matplotlib, numpy
import matplotlib.pyplot as plt
import numpy as np

# In[2]:

#saves plot/figure to image
plt.savefig('name_your_graph.png')

# In[11]:

# we can create simple plot as follows, linear graph
# Prepare the data
linear = np.linspace(0, 10, 100)

# Plot the data, set label and color, b=blue, g=green
plt.plot(linear, linear, label='linear', color='b')

# Add a legend
plt.legend()

# Show the plot
plt.show()

# In[23]:

#by calling
plt.style.available

# In[26]:

x = np.linspace(0, 4 * np.pi)
y = np.sin(x)

plt.plot(x, y)

plt.show()

# In[27]:

plt.style.use('ggplot')

plt.plot(x, y)

plt.show()

# In[32]:

with plt.style.context(('seaborn-darkgrid')):
    plt.plot(x, y, 'r-o')
    plt.show()

# #### Working with text inside your graph

# However, you can use subplots to set up and place your Axes on a regular grid. So that means that in most cases, Axes and subplot are synonymous, they will designate the same thing. When you do call subplot to add Axes to your figure, do so with the add_subplots() function.

# In[14]:

# we can create scatter plot for this
fig = plt.figure()

# Set up Axes
ax = fig.add_subplot(111)

# Scatter the data
ax.scatter(np.linspace(0, 1, 5), np.linspace(0, 5, 5), label='scatter')

# Add a legend
plt.legend()

# Show the plot
plt.show()

# #### How To Change The Size of Figures
#
# Now that you have seen how to initialize a Figure and Axes from scratch, you will also want to know how you can change certain small details that the package sets up for you, such as the figure size.
# Let’s say you don’t to have your figure size to be default size and you want to change this. How do you set the size of your figures manually?

# In[20]:

# easy! you can just call
fig = plt.figure(figsize=(10,8))
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)

# In[21]:

# let's plot simple bar graph
# Plot the data
ax1.bar([1,2,3],[3,4,5])
ax2.barh([0.5,1,2.5],[0,1,2])

# Show the plot
plt.show()

# In[3]:

x = np.linspace(0, 2, 100)

plt.plot(x, x, label='linear')
plt.plot(x, x**2, label='quadratic')
plt.plot(x, x**3, label='cubic')

plt.xlabel('x label')
plt.ylabel('y label')

plt.title("Simple Plot")

plt.legend()

plt.show()

# In[4]:

#plot data connected by lines
lines = plt.plot(x,x)
plt.show()

# ## More to come soon!


import unittest

import numpy
import pytest

import chainer.testing
import chainerx
import chainerx.testing

from chainerx_tests import array_utils
from chainerx_tests import dtype_utils
from chainerx_tests import math_utils
from chainerx_tests import op_utils

@op_utils.op_test(['native:0', 'cuda:0'])
@chainer.testing.parameterize_pytest('shape,indices', [
    # empty indexing
    ((), ()),
    ((3,), ()),
    ((2, 2, 2), ()),
    # integer indexing - non-tuple indexing
    ((3,), 0),
    ((3,), 1),
    ((3,), 2),
    ((3,), -1),
    ((2, 3), 0),
    ((2, 3), 1),
    ((2, 3), numpy.int8(-1)),
    ((2, 3), numpy.int32(0)),
    ((2, 3), numpy.uint64(1)),
    # integer indexining - tuple indexing
    ((3,), (0,)),
    ((3,), (1,)),
    ((3,), (2,)),
    ((3,), (-1,)),
    ((2, 3), (0,)),
    ((2, 3), (1,)),
    ((2, 3), (0, 0)),
    ((2, 3), (1, 1)),
    ((2, 3, 4), (0, -2, 3)),
    ((2, 3, 4), (1, 0)),
    # slice indexing - non-tuple indexing
    ((3,), slice(None)),
    ((3,), slice(2)),
    ((3,), slice(0, 3)),
    ((3,), slice(0, 2)),
    ((3,), slice(1, 3)),
    ((3,), slice(0, 0)),
    ((3,), slice(0, 1)),
    ((3,), slice(2, 0, -1)),
    ((3,), slice(-2, -1)),
    ((3,), slice(2, None, -1)),
    ((3,), slice(None, 0, 1)),
    ((3,), slice(None, -1, -1)),
    ((3,), slice(None, -2, -1)),
    ((6,), slice(0, 6, 2)),
    ((6,), slice(1, 6, 2)),
    ((6,), slice(5, None, -2)),
    # slice indexing - tuple indexing
    ((3,), (slice(None),)),
    ((3,), (slice(2),)),
    ((3,), (slice(0, 3),)),
    ((3,), (slice(0, 2),)),
    ((3,), (slice(1, 3),)),
    ((3,), (slice(0, 0),)),
    ((3,), (slice(0, 1),)),
    ((3,), (slice(2, 0, -1),)),
    ((3,), (slice(-2, -1),)),
    ((3,), (slice(2, None, -1),)),
    ((3,), (slice(None, 0, 1),)),
    ((3,), (slice(None, -1, -1),)),
    ((3,), (slice(None, -2, -1),)),
    ((6,), (slice(0, 6, 2),)),
    ((6,), (slice(1, 6, 2),)),
    ((6,), (slice(5, None, -2),)),
    ((6,), (slice(50, 1, -1),)),
    ((6,), (slice(3, 3, 1),)),
    ((6,), (slice(3, 3, -2),)),
    ((6,), (slice(50, 50, 1),)),
    ((6,), (slice(50, 50, -2),)),
    ((6,), (slice(-50, -50, 1),)),
    ((6,), (slice(-50, -50, -2),)),
    ((2, 3), (slice(None), slice(None))),
    ((2, 3), (slice(1), slice(2))),
    ((2, 3), (slice(0, 2), slice(0, 3))),
    ((2, 3), (slice(0, 2), slice(0, -1))),
    ((2, 3), (slice(0, None, -1), slice(2, 3))),
    ((2, 3), (slice(0, None, None), slice(-2, 0, -1))),
    ((2, 3), (slice(1, 2), slice(0, 2))),
    ((2, 3), (slice(-2, None, -1), slice(0, 3))),
    ((2, 3), (slice(-2, None, -1), slice(-3, None, -1))),
    ((2, 3), (slice(-2, None, -1), slice(None, None, -2))),
    ((2, 3), (slice(1, 2), slice(None, None, 1))),
    ((2, 3), (slice(1, 2), slice(None, None, 2))),
    ((2, 3, 4), (slice(1), slice(-2, 3), slice(1, None, -1))),
    # newaxis indexing - non-tuple indexing
    ((), chainerx.newaxis),
    ((3,), chainerx.newaxis),
    # newaxis indexing - tuple indexing
    ((), (chainerx.newaxis,)),
    ((3,), (chainerx.newaxis,)),
    ((2, 3), (chainerx.newaxis, chainerx.newaxis)),
    # mixed indexing - tuple indexing
    ((2, 3), (0, slice(1, 3))),
    ((4, 3), (slice(1, 3), 1)),
    ((2, 3, 4), (1, slice(2,), slice(1, 3))),
    ((2, 3), (1, chainerx.newaxis, slice(1, 3))),
    ((2, 3, 4), (slice(0, 1), slice(1, 2), slice(1, 3), chainerx.newaxis)),
    ((2, 3, 4), (slice(0, 1), slice(1, 2), chainerx.newaxis, slice(1, 3))),
    ((2, 3, 4), (slice(0, 1), chainerx.newaxis, slice(1, 2), slice(1, 3))),
    ((2, 3, 4), (chainerx.newaxis, slice(0, 1), slice(1, 2), slice(1, 3))),
    ((2, 3, 4),
     (1, slice(2,), chainerx.newaxis, slice(1, 3), chainerx.newaxis)),
])
class TestGetitem(op_utils.NumpyOpTest):
    # TODO(niboshi): Remove this
    check_numpy_strides_compliance = False

    def generate_inputs(self):
        x = numpy.random.uniform(-1, 1, self.shape).astype('float32')
        return x,

    def forward_xp(self, inputs, xp):
        x, = inputs
        y = x[self.indices]
        return y,

@pytest.mark.parametrize_device(['native:0', 'cuda:0'])
def test_getitem_zero_sized_offsets(device):
    a = chainerx.arange(6)

    b = a[3:3]
    # Test pre-conditions.
    assert b.size == 0
    assert b.offset == 12

    # The offset of `c` should be the same as `b` since `b` is empty.
    c = b[2:]
    assert c.size == 0
    assert c.offset == b.offset

@op_utils.op_test(['native:0', 'cuda:0'])
# TODO(hvy): Add cases where axis=None, when supported.
@chainer.testing.parameterize_pytest('shape,indices,axis', [
    # Valid parameters
    ((3,), [0], 0),
    ((3,), [1], 0),
    ((2, 3), [0], 0),
    ((2, 3), [0], 1),
    ((2, 3), [0], -1),
    ((2, 3), [1], 0),
    ((2, 3), [0, -1], 0),
    ((2, 3), [1, 0], 0),
    ((2, 3), [1, 2], 1),
    ((2, 3), [2, 1], 1),
    ((2, 3), [[0], [1]], 0),
    # Invalid: Axis out of bounds
    ((2, 3), [0], 2),
    ((2, 3), [0], -3),
])
@chainer.testing.parameterize_pytest('is_module', [True, False])
@chainer.testing.parameterize_pytest(
    'indices_type', ['list', 'numpy', 'xp'])
# TODO(niboshi): indices_dtype is ignored if indices_type == 'list', which is
# wasteful.
@chainer.testing.parameterize_pytest(
    'indices_dtype', chainerx.testing.integral_dtypes)
class TestTake(op_utils.NumpyOpTest):

    check_numpy_strides_compliance = False
    forward_accept_errors = (chainerx.DimensionError, numpy.AxisError)

    def setup(self):
        if (numpy.dtype(self.indices_dtype).kind == 'u'
                and (numpy.array(self.indices, 'int64') < 0).any()):
            raise unittest.SkipTest(
                'Indices underflows and index out of bounds cannot be tested.')

    def generate_inputs(self):
        a = numpy.random.uniform(-1, 1, self.shape).astype('float32')
        return a,

    def forward_xp(self, inputs, xp):
        indices = self.indices
        axis = self.axis
        indices_type = self.indices_type
        a, = inputs

        assert isinstance(indices, list)
        if indices_type == 'list':
            pass
        elif indices_type == 'numpy':
            indices = numpy.array(indices).astype(self.indices_dtype)
        elif indices_type == 'xp':
            indices = xp.array(indices).astype(self.indices_dtype)
        else:
            assert False, indices_type

        if self.is_module:
            b = xp.take(a, indices, axis)
        else:
            b = a.take(indices, axis)
        return b,

def _random_condition(shape, dtype):
    size = int(numpy.prod(shape))
    mask = numpy.random.randint(0, 1, size).astype('bool_').reshape(shape)
    pos = array_utils.uniform(shape, dtype)
    pos[numpy.logical_not(pos)] = True  # All elements are True
    return pos * mask

@op_utils.op_test(['native:0', 'cuda:0'])
@chainer.testing.parameterize(*(
    # Special shapes
    chainer.testing.product({
        'cond_shape,in_shapes': [
            # Same Shapes
            ((2, 3), ((2, 3), (2, 3))),
            # Broadcast Shapes
            ((2, 3), ((1, 3), (1, 3))),
            ((2, 3), ((2, 1), (1, 3))),
            ((2, 3), ((2, 3), (1, 3))),
            ((4, 5), ((3, 4, 1), (1, 5))),
            ((1, 4, 5), ((3, 4, 1), (3, 1, 5))),
        ],
        'cond_dtype': ['bool_'],
        'in_dtypes,out_dtype': dtype_utils.result_dtypes_two_arrays,
    })
    # Dtype combinations
    + chainer.testing.product({
        'cond_shape,in_shapes': [((2, 3), ((2, 3), (2, 3)))],
        'cond_dtype': chainerx.testing.all_dtypes,
        'in_dtypes,out_dtype': dtype_utils.result_dtypes_two_arrays,
    })
))
class TestWhere(math_utils.BinaryMathTestBase, op_utils.NumpyOpTest):

    check_numpy_strides_compliance = False
    dodge_nondifferentiable = True
    input_lhs = 'random'
    input_rhs = 'random'

    def generate_inputs(self):
        self.condition = _random_condition(self.cond_shape, self.cond_dtype)
        return super().generate_inputs()

    def func(self, xp, x, y):
        condition = xp.array(self.condition)
        return xp.where(condition, x, y)

@chainerx.testing.numpy_chainerx_array_equal(
    accept_error=(
        chainerx.DimensionError, ValueError))
@pytest.mark.parametrize('cond_shape,x_shape,y_shape', [
    ((2, 3), (3, 4), (2, 3)),
    ((2, 3), (2, 3), (3, 4)),
    ((2, 3), (1, 3), (2, 4))
])
def test_where_invalid_shapes(xp, cond_shape, x_shape, y_shape):
    x = array_utils.create_dummy_ndarray(xp, x_shape, 'float32')
    y = array_utils.create_dummy_ndarray(xp, y_shape, 'float32')
    c = array_utils.create_dummy_ndarray(xp, cond_shape, 'float32')
    return xp.where(c, x, y)

@op_utils.op_test(['native:0', 'cuda:0'])
@chainer.testing.parameterize(*(
    # Special shapes
    chainer.testing.product({
        'cond_shape,shape': math_utils.shapes_combination_inplace_binary,
        'cond_dtype': ['bool_'],
        'in_dtypes,scalar_type,out_dtype': (
            dtype_utils.result_dtypes_array_scalar),
        'is_scalar_rhs': [True, False],
    })
    # Dtype combinations
    + chainer.testing.product({
        'cond_shape,shape': [((2, 3), (2, 3))],
        'cond_dtype': chainerx.testing.all_dtypes,
        'in_dtypes,scalar_type,out_dtype': (
            dtype_utils.result_dtypes_array_scalar),
        'is_scalar_rhs': [True, False],
    })
))
class TestWhereScalar(math_utils.MathScalarTestBase, op_utils.NumpyOpTest):

    check_numpy_strides_compliance = False
    input = 'random'
    scalar_value = 3

    def generate_inputs(self):
        self.condition = _random_condition(self.cond_shape, self.cond_dtype)
        return super().generate_inputs()

    def func_scalar(self, xp, a, scalar):
        condition = xp.array(self.condition)
        if self.is_scalar_rhs:
            return xp.where(condition, a, scalar)
        else:
            return xp.where(condition, scalar, a)

_in_out_dtypes_where_scalar = [
    ((bool, bool), 'bool_'),
    ((bool, int), 'int32'),
    ((bool, float), 'float32'),
    ((int, bool), 'int32'),
    ((int, int), 'int32'),
    ((int, float), 'float32'),
    ((float, bool), 'float32'),
    ((float, int), 'float32'),
    ((float, float), 'float32'),
]

@chainerx.testing.numpy_chainerx_array_equal()
@pytest.mark.parametrize('cond_shape', [(2, 3)])
@pytest.mark.parametrize('cond_dtype', chainerx.testing.all_dtypes)
@pytest.mark.parametrize('in_types,out_dtype', _in_out_dtypes_where_scalar)
def test_where_scalar_scalar(xp, cond_shape, cond_dtype, in_types, out_dtype):
    cond = xp.array(_random_condition(cond_shape, cond_dtype))
    x_type, y_type = in_types
    x = x_type(0)
    y = y_type(2)
    out = xp.where(cond, x, y)
    return dtype_utils.cast_if_numpy_array(xp, out, out_dtype)


###################################################################################################
# (c) 2022 Birgit Hillebrecht
#
# This code has been developed as part of
# 	  Certified machine learning: A posteriori error estimation for physics-informed neural networks
#     https://doi.org/10.48550/arXiv.2203.17055
# please kindly consider citing this publication when using this code.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
###################################################################################################

import json
import numpy as np

def load_training_params(filepath):
    """
    Loads parameters to configure the training from the json file provided via param filepath

    returns
    - epochs: number of epochs used for training
    - n_phys: number of collocation points for the PINN
    """

    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    opt = None
    try:
        opt = str(data['optimizer'])
    except KeyError as e:
        opt = None

    return int(data['epochs']), int(data['n_phys']), opt

def load_nn_params(filepath):
    """
    Loads parameters for the neural network from the json file provided via param filepath.

    returns
    - input_dim: number of input nodes
    - output_dim: number of output nodes
    - num_layers: number of layers
    - num_neurons: number of nodes per layer
    - lower_bound: lower bounds on input node values
    - upper_bound: upper bounds on input node values
    - af (None): if set in json, the value for the activation function
    """

    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    lb = np.array(data['lower_bound'])
    ub = np.array(data['upper_bound'])

    af = None
    try:
        af = str(data['activation_function'])
    except KeyError as e:
        af = None

    return int(data['input_dim']), int(data['output_dim']), int(data['num_layers']), int(data['num_neurons']), lb.astype(np.float), ub.astype(np.float), af

def load_ee_params(filepath):
    """
    Loads parameters for a posteriori error estimation for the PINN from the json file provided via param filepath.

    returns
    - K: as used for trapezoidal rule
    - mu: smoothing parameter for delta function
    - L_f: Lipschitz constant or spectral abscissa
    - delta_mean: average deviation of approximated ODE/PDE from target ODE/PDE
    """
    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    return float(data['K']), float(data['mu']), float(data['L_f']), float(data['delta_mean'])

def has_param(filepath, param_name):
    """
    Checks if keyword param_name exists in json file

    :param string filepath: path to json file
    :param string param_name: keyword used for parameter in json file.
    """
    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    try:
        data[param_name]
    except KeyError as e:
        return False

    return True

def get_param_as_float(filepath, param_name):
    """
    Extracts parameter as float from json file. Does not check existance, misuse may lead to exceptions.
    Call has_param first.

    :param string filepath: path to json file
    :param string param_name: keyword used for parameter in json file.
    """
    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    return float( data[param_name] )

def get_param_as_int(filepath, param_name):
    """
    Extracts parameter as integer from json file. Does not check existance, misuse may lead to exceptions.
    Call has_param first.

    :param string filepath: path to json file
    :param string param_name: keyword used for parameter in json file.
    """
    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    return int( data[param_name] )

def get_param_as_array(filepath, param_name):
    """
    Extracts parameter as array from json file. Does not check existance, misuse may lead to exceptions.
    Call has_param first.

    :param string filepath: path to json file
    :param string param_name: keyword used for parameter in json file.
    """
    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    return np.array( data[param_name] )

def get_param_as_string(filepath, param_name):
    """
    Extracts parameter as string from json file. Does not check existance, misuse may lead to exceptions.
    Call has_param first.

    :param string filepath: path to json file
    :param string param_name: keyword used for parameter in json file.
    """
    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    return str( data[param_name] )

def get_param_as_boolean(filepath, param_name):
    """
    Extracts parameter as boolean from json file. Does not check existance, misuse may lead to exceptions.
    Call has_param first.

    :param string filepath: path to json file
    :param string param_name: keyword used for parameter in json file.
    """
    with open(filepath, "r") as jsonfile:
        data = json.load(jsonfile)
        jsonfile.close()

    return bool( (data[param_name] == "True") or (data[param_name] == "true") )


import os
import nltk
import trax
from trax import layers as tl
from trax.supervised import training
from trax.fastmath import numpy as fastnp
import numpy as np
import pandas as pd
import random as rnd

# set random seeds
trax.supervised.trainer_lib.init_random_number_generators(34)
rnd.seed(34)

data = pd.read_csv("questions.csv")
N=len(data)
print('Number of question pairs: ', N)
data.head()

N_train = 300000
N_test  = 10*1024
data_train = data[:N_train]
data_test  = data[N_train:N_train+N_test]
print("Train set:", len(data_train), "Test set:", len(data_test))
del(data) # remove to free memory

td_index = (data_train['is_duplicate'] == 1).to_numpy()
print(td_index[:10,])
td_index = [i for i, x in enumerate(td_index) if x]
print('number of duplicate questions: ', len(td_index))
print('indexes of first ten duplicate questions:', td_index[:10])

print(data_train['question1'][5])  #  Example of question duplicates (first one in data)
print(data_train['question2'][5])
print('is_duplicate: ', data_train['is_duplicate'][5])

Q1_train_words = np.array(data_train['question1'][td_index])
Q2_train_words = np.array(data_train['question2'][td_index])

Q1_test_words = np.array(data_test['question1'])
Q2_test_words = np.array(data_test['question2'])
y_test  = np.array(data_test['is_duplicate'])

print('TRAINING QUESTIONS:\n')
print('Question 1: ', Q1_train_words[0])
print('Question 2: ', Q2_train_words[0], '\n')
print('Question 1: ', Q1_train_words[5])
print('Question 2: ', Q2_train_words[5], '\n')

print('TESTING QUESTIONS:\n')
print('Question 1: ', Q1_test_words[0])
print('Question 2: ', Q2_test_words[0], '\n')
print('is_duplicate =', y_test[0], '\n')

#create arrays
Q1_train = np.empty_like(Q1_train_words)
Q2_train = np.empty_like(Q2_train_words)

Q1_test = np.empty_like(Q1_test_words)
Q2_test = np.empty_like(Q2_test_words)

# Building the vocabulary with the train set         (this might take a minute)
from collections import defaultdict

vocab = defaultdict(lambda: 0)
vocab['<PAD>'] = 1

for idx in range(len(Q1_train_words)):
    Q1_train[idx] = nltk.word_tokenize(Q1_train_words[idx])
    Q2_train[idx] = nltk.word_tokenize(Q2_train_words[idx])
    q = Q1_train[idx] + Q2_train[idx]
    for word in q:
        if word not in vocab:
            vocab[word] = len(vocab) + 1
print('The length of the vocabulary is: ', len(vocab))

print(vocab['<PAD>'])
print(vocab['Astrology'])
print(vocab['Astronomy'])  #not in vocabulary, returns 0

for idx in range(len(Q1_test_words)):
    Q1_test[idx] = nltk.word_tokenize(Q1_test_words[idx])
    Q2_test[idx] = nltk.word_tokenize(Q2_test_words[idx])

print('Train set has reduced to: ', len(Q1_train) )
print('Test set length: ', len(Q1_test) )

# Converting questions to array of integers
for i in range(len(Q1_train)):
    Q1_train[i] = [vocab[word] for word in Q1_train[i]]
    Q2_train[i] = [vocab[word] for word in Q2_train[i]]

for i in range(len(Q1_test)):
    Q1_test[i] = [vocab[word] for word in Q1_test[i]]
    Q2_test[i] = [vocab[word] for word in Q2_test[i]]

print('first question in the train set:\n')
print(Q1_train_words[0], '\n')
print('encoded version:')
print(Q1_train[0],'\n')

print('first question in the test set:\n')
print(Q1_test_words[0], '\n')
print('encoded version:')
print(Q1_test[0])

# Splitting the data
cut_off = int(len(Q1_train)*.8)
train_Q1, train_Q2 = Q1_train[:cut_off], Q2_train[:cut_off]
val_Q1, val_Q2 = Q1_train[cut_off: ], Q2_train[cut_off:]
print('Number of duplicate questions: ', len(Q1_train))
print("The length of the training set is:  ", len(train_Q1))
print("The length of the validation set is: ", len(val_Q1))

# UNQ_C1
def data_generator(Q1, Q2, batch_size, pad=1, shuffle=True):
    """Generator function that yields batches of data

    Args:
        Q1 (list): List of transformed (to tensor) questions.
        Q2 (list): List of transformed (to tensor) questions.
        batch_size (int): Number of elements per batch.
        pad (int, optional): Pad character from the vocab. Defaults to 1.
        shuffle (bool, optional): If the batches should be randomnized or not. Defaults to True.
    Yields:
        tuple: Of the form (input1, input2) with types (numpy.ndarray, numpy.ndarray)
        NOTE: input1: inputs to your model [q1a, q2a, q3a, ...] i.e. (q1a,q1b) are duplicates
              input2: targets to your model [q1b, q2b,q3b, ...] i.e. (q1a,q2i) i!=a are not duplicates
    """

    input1 = []
    input2 = []
    idx = 0
    len_q = len(Q1)
    question_indexes = [*range(len_q)]

    if shuffle:
        rnd.shuffle(question_indexes)

    while True:
        if idx >= len_q:
            # if idx is greater than or equal to len_q, set idx accordingly
            # (Hint: look at the instructions above)
            idx = 0
            # shuffle to get random batches if shuffle is set to True
            if shuffle:
                rnd.shuffle(question_indexes)

        # get questions at the `question_indexes[idx]` position in Q1 and Q2
        q1 = Q1[question_indexes[idx]]
        q2 = Q2[question_indexes[idx]]

        # increment idx by 1
        idx += 1
        # append q1
        input1.append(q1)
        # append q2
        input2.append(q2)
        if len(input1) == batch_size:
            # determine max_len as the longest question in input1 & input 2
            # Hint: use the `max` function.
            # take max of input1 & input2 and then max out of the two of them.
            max_len = max(max([len(_) for _ in input1]),max([len(_) for _ in input2]))
            # pad to power-of-2 (Hint: look at the instructions above)
            max_len = 2**int(np.ceil(np.log2(max_len)))
            b1 = []
            b2 = []
            for q1, q2 in zip(input1, input2):
                # add [pad] to q1 until it reaches max_len
                q1 = q1 + [pad] * (max_len - len(q1))
                # add [pad] to q2 until it reaches max_len
                q2 = q2 + [pad] * (max_len - len(q2))
                # append q1
                b1.append(q1)
                # append q2
                b2.append(q2)
            # use b1 and b2
            yield np.array(b1), np.array(b2)

            # reset the batches
            input1, input2 = [], []  # reset the batches

batch_size = 2
res1, res2 = next(data_generator(train_Q1, train_Q2, batch_size))
print("First questions  : ",'\n', res1, '\n')
print("Second questions : ",'\n', res2)

# UNQ_C2
def Siamese(vocab_size=len(vocab), d_model=128, mode='train'):
    """Returns a Siamese model.

    Args:
        vocab_size (int, optional): Length of the vocabulary. Defaults to len(vocab).
        d_model (int, optional): Depth of the model. Defaults to 128.
        mode (str, optional): 'train', 'eval' or 'predict', predict mode is for fast inference. Defaults to 'train'.

    Returns:
        trax.layers.combinators.Parallel: A Siamese model.
    """

    def normalize(x):  # normalizes the vectors to have L2 norm 1
        return x / fastnp.sqrt(fastnp.sum(x * x, axis=-1, keepdims=True))

    q_processor = tl.Serial(  # Processor will run on Q1 and Q2.
        tl.Embedding(vocab_size, d_model), # Embedding layer
        tl.LSTM(d_model), # LSTM layer
        tl.Mean(axis=1), # Mean over columns
        tl.Fn('Normalize', lambda x: normalize(x))  # Apply normalize function
    )  # Returns one vector of shape [batch_size, d_model].

    # Run on Q1 and Q2 in parallel.
    model = tl.Parallel(q_processor, q_processor)
    return model

# check your model
model = Siamese()
print(model)

# UNQ_C3
def TripletLossFn(v1, v2, margin=0.25):
    """Custom Loss function.

    Args:
        v1 (numpy.ndarray): Array with dimension (batch_size, model_dimension) associated to Q1.
        v2 (numpy.ndarray): Array with dimension (batch_size, model_dimension) associated to Q2.
        margin (float, optional): Desired margin. Defaults to 0.25.

    Returns:
        jax.interpreters.xla.DeviceArray: Triplet Loss.
    """

    # use fastnp to take the dot product of the two batches (don't forget to transpose the second argument)
    scores = fastnp.dot(v1,v2.T)  # pairwise cosine sim
    # calculate new batch size
    batch_size = len(scores)
    # use fastnp to grab all postive `diagonal` entries in `scores`
    positive = fastnp.diagonal(scores)  # the positive ones (duplicates)
    # multiply `fastnp.eye(batch_size)` with 2.0 and subtract it out of `scores`
    negative_without_positive = scores - fastnp.eye(batch_size) * 2.0
    # take the row by row `max` of `negative_without_positive`.
    # Hint: negative_without_positive.max(axis = [?])
    closest_negative = negative_without_positive.max(axis = 1)
    # subtract `fastnp.eye(batch_size)` out of 1.0 and do element-wise multiplication with `scores`
    negative_zero_on_duplicate = (1.0 - fastnp.eye(batch_size)) * scores
    # use `fastnp.sum` on `negative_zero_on_duplicate` for `axis=1` and divide it by `(batch_size - 1)`
    mean_negative = fastnp.sum(negative_zero_on_duplicate, axis=1) / (batch_size - 1)
    # compute `fastnp.maximum` among 0.0 and `A`
    # A = subtract `positive` from `margin` and add `closest_negative`
    triplet_loss1 = fastnp.maximum(margin - positive + closest_negative, 0 )
    # compute `fastnp.maximum` among 0.0 and `B`
    # B = subtract `positive` from `margin` and add `mean_negative`
    triplet_loss2 = fastnp.maximum(margin - positive + mean_negative, 0 )
    # add the two losses together and take the `fastnp.mean` of it
    triplet_loss = fastnp.mean(triplet_loss1 + triplet_loss2)
    return triplet_loss

v1 = np.array([[0.26726124, 0.53452248, 0.80178373],[0.5178918 , 0.57543534, 0.63297887]])
v2 = np.array([[ 0.26726124,  0.53452248,  0.80178373],[-0.5178918 , -0.57543534, -0.63297887]])
TripletLossFn(v2,v1)
print("Triplet Loss:", TripletLossFn(v2,v1))

from functools import partial
def TripletLoss(margin=0.25):
    triplet_loss_fn = partial(TripletLossFn, margin=margin)
    return tl.Fn('TripletLoss', triplet_loss_fn)

batch_size = 256
train_generator = data_generator(train_Q1, train_Q2, batch_size, vocab['<PAD>'])
val_generator = data_generator(val_Q1, val_Q2, batch_size, vocab['<PAD>'])
print('train_Q1.shape ', train_Q1.shape)
print('val_Q1.shape   ', val_Q1.shape)

lr_schedule = trax.lr.warmup_and_rsqrt_decay(400, 0.01)

# UNQ_C4
def train_model(Siamese, TripletLoss, lr_schedule, train_generator=train_generator, val_generator=val_generator, output_dir='model/'):
    """Training the Siamese Model

    Args:
        Siamese (function): Function that returns the Siamese model.
        TripletLoss (function): Function that defines the TripletLoss loss function.
        lr_schedule (function): Trax multifactor schedule function.
        train_generator (generator, optional): Training generator. Defaults to train_generator.
        val_generator (generator, optional): Validation generator. Defaults to val_generator.
        output_dir (str, optional): Path to save model to. Defaults to 'model/'.

    Returns:
        trax.supervised.training.Loop: Training loop for the model.
    """
    output_dir = os.path.expanduser(output_dir)

    train_task = training.TrainTask(
        labeled_data=train_generator,       # Use generator (train)
        loss_layer=TripletLoss(),         # Use triplet loss. Don't forget to instantiate this object
        optimizer=trax.optimizers.Adam(learning_rate = 0.01),          # Don't forget to add the learning rate parameter
        lr_schedule=lr_schedule, # Use Trax multifactor schedule function
    )

    eval_task = training.EvalTask(
        labeled_data=val_generator,       # Use generator (val)
        metrics=[TripletLoss()],          # Use triplet loss. Don't forget to instantiate this object
    )

    training_loop = training.Loop(Siamese(),
                                  train_task,
                                  eval_task=eval_task,
                                  output_dir=output_dir)

    return training_loop

train_steps = 5
training_loop = train_model(Siamese, TripletLoss, lr_schedule)
training_loop.run(train_steps)

# Loading in the saved model
model = Siamese()
model.init_from_file('model.pkl.gz')

# UNQ_C5
def classify(test_Q1, test_Q2, y, threshold, model, vocab, data_generator=data_generator, batch_size=64):
    """Function to test the accuracy of the model.

    Args:
        test_Q1 (numpy.ndarray): Array of Q1 questions.
        test_Q2 (numpy.ndarray): Array of Q2 questions.
        y (numpy.ndarray): Array of actual target.
        threshold (float): Desired threshold.
        model (trax.layers.combinators.Parallel): The Siamese model.
        vocab (collections.defaultdict): The vocabulary used.
        data_generator (function): Data generator function. Defaults to data_generator.
        batch_size (int, optional): Size of the batches. Defaults to 64.

    Returns:
        float: Accuracy of the model.
    """
    accuracy = 0
    for i in range(0, len(test_Q1), batch_size):
        # Call the data generator (built in Ex 01) with shuffle=False using next()
        # use batch size chuncks of questions as Q1 & Q2 arguments of the data generator. e.g x[i:i + batch_size]
        # Hint: use `vocab['<PAD>']` for the `pad` argument of the data generator
        q1, q2 = next(data_generator(test_Q1[i: i + batch_size], test_Q2[i: i+batch_size], batch_size, pad=vocab['<PAD>'], shuffle=False))
        # use batch size chuncks of actual output targets (same syntax as example above)
        y_test = y[i: i + batch_size]
        # Call the model
        v1, v2 = model((q1,q2))

        for j in range(batch_size):
            # take dot product to compute cos similarity of each pair of entries, v1[j], v2[j]
            # don't forget to transpose the second argument
            d = fastnp.dot(v1[j],v2[j].T)
            # is d greater than the threshold?
            res = d > threshold
            # increment accurancy if y_test is equal `res`
            accuracy += (y_test[j] == res)
    # compute accuracy using accuracy and total length of test questions
    accuracy = accuracy / len(test_Q1)

    return accuracy

# this takes around 1 minute
accuracy = classify(Q1_test,Q2_test, y_test, 0.7, model, vocab, batch_size = 512)
print("Accuracy", accuracy)

# UNQ_C6
def predict(question1, question2, threshold, model, vocab, data_generator=data_generator, verbose=False):
    """Function for predicting if two questions are duplicates.

    Args:
        question1 (str): First question.
        question2 (str): Second question.
        threshold (float): Desired threshold.
        model (trax.layers.combinators.Parallel): The Siamese model.
        vocab (collections.defaultdict): The vocabulary used.
        data_generator (function): Data generator function. Defaults to data_generator.
        verbose (bool, optional): If the results should be printed out. Defaults to False.

    Returns:
        bool: True if the questions are duplicates, False otherwise.
    """
    # use `nltk` word tokenize function to tokenize
    q1 = nltk.word_tokenize(question1)  # tokenize
    q2 = nltk.word_tokenize(question2)  # tokenize
    Q1, Q2 = [], []
    for word in q1:  # encode q1
        # increment by checking the 'word' index in `vocab`
        Q1 += [vocab[word]]
    for word in q2:  # encode q2
        # increment by checking the 'word' index in `vocab`
        Q2 += [vocab[word]]

    # Call the data generator (built in Ex 01) using next()
    # pass [Q1] & [Q2] as Q1 & Q2 arguments of the data generator. Set batch size as 1
    # Hint: use `vocab['<PAD>']` for the `pad` argument of the data generator
    Q1, Q2 = next(data_generator([Q1], [Q2], 1, vocab['<PAD>']))
    # Call the model
    v1, v2 = model((Q1,Q2))
    # take dot product to compute cos similarity of each pair of entries, v1, v2
    # don't forget to transpose the second argument
    d = fastnp.dot(v1, v2.T)
    # is d greater than the threshold?
    res = d > threshold

    if(verbose):
        print("Q1  = ", Q1, "\nQ2  = ", Q2)
        print("d   = ", d)
        print("res = ", res)

    return res

# Feel free to try with your own questions
question1 = "When will I see you?"
question2 = "When can I see you again?"
# 1 means it is duplicated, 0 otherwise
predict(question1 , question2, 0.7, model, vocab, verbose = True)

# Feel free to try with your own questions
question1 = "Do they enjoy eating the dessert?"
question2 = "Do they like hiking in the desert?"
# 1 means it is duplicated, 0 otherwise
predict(question1 , question2, 0.7, model, vocab, verbose=True)


# -*- coding: utf-8 -*-

"""
This module implements:

1. particle history classes,  which store the full or partial history
   of a SMC algorithm.

2. off-line smoothing algorithms as methods of these classes.

For on-line smoothing, see instead the `collectors` module.

History classes
===============

A `SMC` object has a ``hist`` attribute, which is used to record at *certain*
times t:

* the N current particles :math:`X_t^n`;
* their weights;
* (optionally, see below), the ancestor variables :math:`A_t^n`.

The frequency at which history is recorded depends on option ``store_history``
of class `SMC`. Possible options are:

* ``True``: records full history (at every time t);
* ``False``: no history (attribute `hist` set to ``None``);
* callable ``f``: history is recorded at time t if ``f(t)`` returns True
* int k: records a rolling window history of length k (may be used
  to perform fixed-lag smoothing)

This module implements different classes that correspond to the different cases:

* `ParticleHistory`: full history (based on lists)
* `PartialParticleHistory`: partial history (based on dictionaries)
* `RollingParticleHistory`: rolling window history (based on `deques`_)

.. _deques: https://docs.python.org/3/library/collections.html#collections.deque

All these classes provide a similar interface. If ``smc`` is a `SMC` object,
then:

* ``smc.hist.X[t]`` returns the N particles at time t
* ``smc.hist.wgts[t]`` returns the N weights at time t (see `resampling.weights`)
* ``smc.hist.A[t]`` returns the N ancestor variables at time t

Partial History
===============

Here are some examples on one may record history only at certain times::

    # store every other 10 iterations
    smc = SMC(fk=fk, N=100, store_history=lambda t: (t % 10) == 0)

    # store at certain times given by a list
    times = [10, 30, 84]
    smc = SMC(fk=fk, N=100, store_history=lambda t: t in times)

Once the algorithm is run, ``smc.hist.X`` and ``smc.hist.wgts`` are
dictionaries, the keys of which are the times where history was recorded. The
ancestor variables are not recorded in that case::

    smc.run()
    smc.hist.X[10]  # the N particles at time 10
    smc.hist.A[10]  # raises an error

Full history, off-line smoothing algorithms
===========================================

For a given state-space model, off-line smoothing amounts to approximate the
distribution of the complete trajectory :math:`X_{0:T}`, given data
:math:`y_{0:T}`, at some fixed time horizon T. The corresponding algorithms
take as an input the complete history of a particle filter, run until time T
(forward pass). Say::

    # forward pass
    fk = ssm.Bootstrap(ssm=my_ssm, data=y)
    pf = particles.SMC(fk=fk, N=100, store_history=True)
    pf.run()

Then, ``pf.hist`` is an instance of class `ParticleHistory`, which has the
following methods:

    * `backward_sampling`: implements O(N) and O(N^2) FFBS algorithms,
      which generates smoothing trajectories from the history of the forward
      pass;
    * `backward_sampling_qmc`: same as above, but for when the forward pass
      was based on QMC (quasi-Monte Carlo).
    * `two_filter_smoothing`: to estimate expectations of marginal smoothing
      distributions (using the two-filter smoothing approach).

For more details, see the documentation of `ParticleHistory`, the ipython
notebook on smoothing, and Chapter 12 of the book.

.. warning:: the complete history of a particle filter may take a lot of
  memory.

Rolling history, Fixed-lag smoothing
====================================

To obtain a rolling window (fixed-length) history::

    smc = SMC(fk=fk, N=100, store_history=10)
    smc.run()

In that case, fields ``smc.hist.X``, ``smc.hist.wgts`` and ``smc.hist.A`` are
`deques`_  of max length 10.  Using negative indices::

    smc.hist.X[-1]  # the particles at final time T
    smc.hist.X[-2]  # the particles at time T - 1
    # ...
    smc.hist.X[-10] # the N particles at time T - 9
    smc.hist.X[-11] # raises an error

Note that this type of history makes it possible to perform fixed-lag smoothing
as follows::

    B = smc.hist.compute_trajectories()
    # B[t, n] is index of ancestor of X_T^n at time t
    phi = lambda x: x  # any test function
    est = np.average(phi(smc.hist.X[-10][B[-10, :]]), weights=smc.W)
    # est is an estimate of E[ phi(X_{T-9}) | Y_{0:T}]

.. note:: recall that it is possible to run `SMC` algorithms step by step,
   since they are iterators. Hence it is possible to do fixed-lag smoothing
   step-by-step as well.

"""

from __future__ import absolute_import, division, print_function

from collections import deque
from itertools import islice
import numpy as np
from numpy import random
from scipy import stats  # worker
import time

import particles # worker
from particles import hilbert
from particles import qmc
from particles import resampling as rs

def generate_hist_obj(option, smc):
    if option is True:
        return ParticleHistory(smc.fk, smc.qmc)
    elif option is False:
        return None
    elif callable(option):
        return PartialParticleHistory(option)
    elif isinstance(option, int) and option >= 0:
        return RollingParticleHistory(option)
    else:
        raise ValueError('store_history: invalid option')

class PartialParticleHistory(object):
    """Partial history.

    History that records the particle system only at certain times.
    See `smoothing` module doc for more details.
    """
    def __init__(self, func):
        self.is_save_time = func
        self.X, self.wgts = {}, {}

    def save(self, smc):
        t = smc.t
        if self.is_save_time(t):
            self.X[t] = smc.X
            self.wgts[t] = smc.wgts

class RollingParticleHistory(object):
    """Rolling window history.

    History that keeps only the k most recent particle systems. Based on
    deques. See `smoothing` module doc for more details.

    """
    def __init__(self, length):
        self.X = deque([], length)
        self.A = deque([], length)
        self.wgts = deque([], length)

    @property
    def N(self):
        """Number of particles at each time step.
        """
        return self.X[0].shape[0]

    @property
    def T(self):
        """Current length of history.
        """
        return len(self.X)

    def save(self, smc):
        self.X.append(smc.X)
        self.A.append(smc.A)
        self.wgts.append(smc.wgts)

    def compute_trajectories(self):
        """Compute the N trajectories that constitute the current genealogy.

        Returns a (T, N) int array, such that B[t, n] is the index of ancestor
        at time t of particle X_T^n, where T is the current length of history.
        """
        Bs = [np.arange(self.N)]
        for A in list(self.A)[-1:0:-1]:  # list in case self.A is a deque
            Bs.append(A[Bs[-1]])
        Bs.reverse()
        return np.array(Bs)

class ParticleHistory(RollingParticleHistory):
    """Particle history.

    A class to store the full history of a particle algorithm, i.e.
    at each time t=0,...T, the N particles, their weights, and their ancestors.
    Off-line smoothing algorithms are methods of this class.

    `SMC` creates an object of this class when invoked with
    ``store_history=True``, and then save at every time t the set of particles,
    their weights (and their logarithm), and the ancestor variables.

    Attributes
    ----------
    X: list
        X[t] is the object that represents the N particles at iteration t
    wgts: list
        wgts[t] is a `Weights` object (see module `resampling`) that represents
        the N weights at time t
    A: list
        A[t] is the vector of ancestor indices at time t

    """

    def __init__(self, fk, qmc):
        self.X, self.A, self.wgts = [], [], []
        if qmc:
            self.h_orders = []
        self.fk = fk

    def save(self, smc):
        RollingParticleHistory.save(self, smc)
        if hasattr(smc, 'h_order'):
            self.h_orders.append(smc.h_order)

    def extract_one_trajectory(self):
        """Extract a single trajectory from the particle history.

        The final state is chosen randomly, then the corresponding trajectory
        is constructed backwards, until time t=0.
        """
        traj = []
        for t in reversed(range(self.T)):
            if t == self.T - 1:
                n = rs.multinomial_once(self.wgts[-1].W)
            else:
                n = self.A[t + 1][n]
            traj.append(self.X[t][n])
        return traj[::-1]

    def _check_h_orders(self):
        if not hasattr(self, 'h_orders'):
            raise ValueError('QMC FFBS requires particles have been Hilbert\
                             ordered during the forward pass')

    def backward_sampling(self, M, linear_cost=False, return_ar=False):
        """Generate smoothing trajectories using FFBS.

        FFBS (forward filtering backward smoothing) is a class of off-line
        smoothing algorithms, which generate smoothing trajectories constructed
        from the history of a particle filter.

        Parameters
        ----------
        M: int
            number of trajectories we want to generate
        linear_cost: bool
            if set to True, the O(N) version is used, see below.

        return_ar: bool (default=False)
            if set to True, change the output, see below.

        Returns
        -------
        paths: a list of ndarrays
            paths[t][n] is component t of trajectory m.
        ar: float
            the overall acceptance rate of the rejection procedure

        Notes
        -----

        1. if ``linear_cost=False``, complexity is O(TMN); i.e. O(TN^2) for M=N;
           if ``linear_cost=True``, complexity is O(T(M+N)), i.e. O(TN) for M=N.
           This requires that model has method `upper_bound_trans`, which
           provides the log of a constant C_t such that
           :math:`p_t(x_t|x_{t-1}) \leq C_t`.

        2. main output is ``paths``, a list of T arrays such that
           ``paths[t][m]`` is component t of trajectory m.

        3. if ``linear_cost=True`` and ``return_ar=True``, output is tuple
           ``(paths, ar)``, where ``paths`` is as above, and ``ar`` is the overall
           acceptance rate (of the rejection steps that choose the ancestors);
           otherwise output is simply ``paths``.
        """
        idx = np.empty((self.T, M), dtype=int)
        idx[-1, :] = rs.multinomial(self.wgts[-1].W, M=M)
        if linear_cost:
            ar = self._backward_sampling_ON(M, idx)
        else:
            self._backward_sampling_ON2(M, idx)
        # When M=1, we want a list of states, not a list of arrays containing
        # one state
        if M == 1:
            idx = idx.squeeze(axis=1)
        paths = [self.X[t][idx[t]] for t in range(self.T)]
        if linear_cost and return_ar:
            return (paths, ar)
        else:
            return paths

    def _backward_sampling_ON(self, M, idx):
        """O(N) version of backward sampling.

        not meant to be called directly, see backward_sampling.
        """
        nattempts = 0
        for t in reversed(range(self.T - 1)):
            where_rejected = np.arange(M)
            who_rejected = self.X[t + 1][idx[t + 1, :]]
            nrejected = M
            gen = rs.MultinomialQueue(self.wgts[t].W, M=M)
            while nrejected > 0:
                nattempts += nrejected
                nprop = gen.dequeue(nrejected)
                lpr_acc = (self.fk.logpt(t + 1, self.X[t][nprop],
                                            who_rejected)
                           - self.fk.upper_bound_trans(t + 1))
                newly_accepted = np.log(random.rand(nrejected)) < lpr_acc
                still_rejected = np.logical_not(newly_accepted)
                idx[t, where_rejected[newly_accepted]] = nprop[newly_accepted]
                where_rejected = where_rejected[still_rejected]
                who_rejected = who_rejected[still_rejected]
                nrejected -= sum(newly_accepted)
        return (M * (self.T - 1)) / nattempts

    def _backward_sampling_ON2(self, M, idx):
        """O(N^2) version of backward sampling.

        not meant to be called directly, see backward_sampling.
        """
        for m in range(M):
            for t in reversed(range(self.T - 1)):
                lwm = (self.wgts[t].lw + self.fk.logpt(t + 1, self.X[t],
                                                     self.X[t + 1][idx[t + 1, m]]))
                idx[t, m] = rs.multinomial_once(rs.exp_and_normalise(lwm))

    def backward_sampling_qmc(self, M):
        """QMC version of backward sampling.

        Parameters
        ----------
        M : int
            number of trajectories

        Note
        ----
        Use this only on the history of a SQMC algorithm.
        """
        self._check_h_orders()
        u = qmc.sobol(M, self.T)
        # the final particles have not been sorted
        hT = hilbert.hilbert_sort(self.X[-1])
        # searchsorted to avoid having to sort in place u according to u[:,T-1]
        idx = np.searchsorted(np.cumsum(self.wgts[-1].W[hT]), u[:, -1])
        paths = [self.X[-1][hT][idx], ]
        for t in reversed(range(self.T - 1)):
            idx = np.empty(M, dtype=np.int64)
            for m, xn in enumerate(paths[-1]):
                lwm = self.wgts[t].lw + self.fk.logpt(t + 1, self.X[t], xn)
                # use ordered version here
                cw = np.cumsum(rs.exp_and_normalise(lwm[self.h_orders[t]]))
                idx[m] = np.searchsorted(cw, u[m, t])
            paths.append(self.X[t][self.h_orders[t]][idx])
        paths.reverse()
        return paths

#     def backward_sampling_lincost_pedagogical(self, M):
#         """ O(N) FFBS
#
#             Don't use this! This is the *pedagogical* version of O(N) FFBS:
#             code is simpler to understand, but quite slow, because it has loops
#         """
#         if not hasattr(self.fk, 'upper_bound_trans'):
#             raise ValueError('O(N) version of backward smoothing'
#                              +'requires to specify constant upper_bound_trans(t)'
#                              +' s.t. log p_t(x_t|x_{t-1})<upper_bound_trans(t)')
#         idx = np.empty((self.T, M), dtype=int)
#         idx[-1, :] = rs.multinomial(M, self.wgts[-1].W)
#         nattempts = 0
#         for t in xrange(self.T - 2, -1, -1):
#             gen = rs.MulinomialQueue(M, self.wgts[t].W)
#             for m in xrange(M):
#                 while True:
#                     nattempts += 1
#                     nprop = gen.dequeue(1)
#                     lpr_acc = (self.fk.logpt(t+1, self.X[t][nprop],
#                                                 self.X[t+1][idx[t+1, m]])
#                                -self.fk.upper_bound_trans(t+1))
#                     if np.log(random.rand()) < lpr_acc:
#                         break
#                 idx[t, m] = nprop
#         print('O(N) FFBS: acceptance rate is %1.2f' %
#               (M * (self.T - 1) / nattempts))
#         return [self.X[t][idx[t, :]] for t in range(self.T)]

    def two_filter_smoothing(self, t, info, phi, loggamma, linear_cost=False,
                             return_ess=False, modif_forward=None,
                             modif_info=None):
        """Two-filter smoothing.

        Parameters
        ----------
        t: time, in range 0 <= t < T-1
        info: SMC object
            the information filter
        phi: function
            test function, a function of (X_t,X_{t+1})
        loggamma: function
            a function of (X_{t+1})
        linear_cost: bool
            if True, use the O(N) variant (basic version is O(N^2))

        Returns
        -------
        Two-filter estimate of the smoothing expectation of phi(X_t,x_{t+1})
        """
        ti = self.T - 2 - t  # t+1 in reverse
        if t < 0 or t >= self.T - 1:
            raise ValueError(
                'two-filter smoothing: t must be in range 0,...,T-2')
        lwinfo = info.hist.wgts[ti].lw - loggamma(info.hist.X[ti])
        if linear_cost:
            return self._two_filter_smoothing_ON(t, ti, info, phi, lwinfo,
                                               return_ess,
                                               modif_forward, modif_info)
        else:
            return self._two_filter_smoothing_ON2(t, ti, info, phi, lwinfo)

    def _two_filter_smoothing_ON2(self, t, ti, info, phi, lwinfo):
        """O(N^2) version of two-filter smoothing.

        This method should not be called directly, see two_filter_smoothing.
        """
        sp, sw = 0., 0.
        upb = lwinfo.max() + self.wgts[t].lw.max()
        if hasattr(self.fk, 'upper_bound_trans'):
            upb += self.fk.upper_bound_trans(t + 1)
        # Loop over n, to avoid having in memory a NxN matrix
        for n in range(self.N):
            omegan = np.exp(lwinfo + self.wgts[t].lw[n] - upb
                            + self.fk.logpt(t + 1, self.X[t][n],
                                               info.hist.X[ti]))
            sp += np.sum(omegan * phi(self.X[t][n], info.hist.X[ti]))
            sw += np.sum(omegan)
        return sp / sw

    def _two_filter_smoothing_ON(self, t, ti, info, phi, lwinfo, return_ess,
                               modif_forward, modif_info):
        """O(N) version of two-filter smoothing.

        This method should not be called directly, see two_filter_smoothing.
        """
        if modif_info is not None:
            lwinfo += modif_info
        Winfo = rs.exp_and_normalise(lwinfo)
        I = rs.multinomial(Winfo)
        if modif_forward is not None:
            lw = self.wgts[t].lw + modif_forward
            W = rs.exp_and_normalise(lw)
        else:
            W = self.wgts[t].W
        J = rs.multinomial(W)
        log_omega = self.fk.logpt(t + 1, self.X[t][J], info.hist.X[ti][I])
        if modif_forward is not None:
            log_omega -= modif_forward[J]
        if modif_info is not None:
            log_omega -= modif_info[I]
        Om = rs.exp_and_normalise(log_omega)
        est = np.average(phi(self.X[t][J], info.hist.X[ti][I]), axis=0,
                         weights=Om)
        if return_ess:
            return (est, 1. / np.sum(Om**2))
        else:
            return est

def smoothing_worker(method=None, N=100, fk=None, fk_info=None,
                     add_func=None, log_gamma=None):
    """Generic worker for off-line smoothing algorithms.

    This worker may be used in conjunction with utils.multiplexer in order to
    run in parallel off-line smoothing algorithms.

    Parameters
    ----------
    method: string
         ['FFBS_ON', 'FFBS_ON2', 'FFBS_QMC',
           'two-filter_ON', 'two-filter_ON_prop', 'two-filter_ON2']
    N: int
        number of particles
    fk: Feynman-Kac object
        The Feynman-Kac model for the forward filter
    fk_info: Feynman-Kac object (default=None)
        the Feynman-Kac model for the information filter; if None,
        set to the same Feynman-Kac model as fk, with data in reverse
    add_func: function, with signature (t, x, xf)
        additive function, at time t, for particles x=x_t and xf=x_{t+1}
    log_gamma: function
        log of function gamma (see book)

    Returns
    -------
    a dict with fields:
        est: a ndarray of length T
        cpu_time
    """
    T = fk.T
    if fk_info is None:
        fk_info = fk.__class__(ssm=fk.ssm, data=fk.data[::-1])
    est = np.zeros(T - 1)
    if method=='FFBS_QMC':
        pf = particles.SQMC(fk=fk, N=N, store_history=True)
    else:
        pf = particles.SMC(fk=fk, N=N, store_history=True)
    tic = time.perf_counter()
    pf.run()
    if method in ['FFBS_ON', 'FFBS_ON2', 'FFBS_QMC']:
        if method.startswith('FFBS_ON'):
            z = pf.hist.backward_sampling(N, linear_cost=(method == 'FFBS_ON'))
        else:
            z = pf.hist.backward_sampling_qmc(N)
        for t in range(T - 1):
            est[t] = np.mean(add_func(t, z[t], z[t + 1]))
    elif method in ['two-filter_ON2', 'two-filter_ON', 'two-filter_ON_prop']:
        infopf = particles.SMC(fk=fk_info, N=N, store_history=True)
        infopf.run()
        for t in range(T - 1):
            psi = lambda x, xf: add_func(t, x, xf)
            if method == 'two-filter_ON2':
                est[t] = pf.hist.two_filter_smoothing(t, infopf, psi, log_gamma)
            else:
                ti = T - 2 - t  # t+1 for info filter
                if method == 'two-filter_ON_prop':
                    modif_fwd = stats.norm.logpdf(pf.hist.X[t],
                                          loc=np.mean(infopf.hist.X[ti + 1]),
                                          scale=np.std(infopf.hist.X[ti + 1]))
                    modif_info = stats.norm.logpdf(infopf.hist.X[ti],
                                           loc=np.mean(pf.hist.X[t + 1]),
                                           scale=np.std(pf.hist.X[t + 1]))
                else:
                    modif_fwd, modif_info = None, None
                est[t] = pf.hist.two_filter_smoothing(t, infopf, psi, log_gamma,
                                                     linear_cost=True,
                                                     modif_forward=modif_fwd,
                                                     modif_info=modif_info)
    else:
        print('no such method?')
    cpu_time = time.perf_counter() - tic
    print(method + ' took %.2f s for N=%i' % (cpu_time, N))
    return {'est': est, 'cpu': cpu_time}


from __future__ import print_function
import sys
import re
from operator import add
import numpy as np
from pyspark import SparkContext

if __name__ == "__main__":

    sc = SparkContext(appName="LogisticRegression")

    # Read the dataset
    d_corpus = sc.textFile(sys.argv[1])

    # Each entry in validLines will be a line from the text file
    validDocLines = d_corpus.filter(lambda x : 'id' in x and 'url=' in x)

    # Now, we transform it into a set of (docID, text) pairs
    keyAndText = validDocLines.map(lambda x : (x[x.index('id="') + 4 : x.index('" url=')], x[x.index('">') + 2:][:-6]))

    # leveraged the code from assignment 2
    # remove all non letter characters
    regex = re.compile('[^a-zA-Z]')
    keyAndWordsList = keyAndText.map(lambda x : (str(x[0]), regex.sub(' ', x[1]).lower().split()))

    # Now get the top 20,000 words... first change (docID, ["word1", "word2", "word3", ...])
    # to ("word1", 1) ("word2", 1)...
    conslidatedWords = keyAndWordsList.flatMap(lambda x: x[1]).map(lambda x: (x,1))

    # Now, count all of the words, giving us ("word1", 1433), ("word2", 3423423), etc.
    allCounts = conslidatedWords.reduceByKey(add)

    # Get the top 20,000 words in a local array in a sorted format based on frequency
    topWordsinDict = allCounts.top(20000, key = lambda x : x[1])

    # We'll create a RDD that has a set of (word, dictNum) pairs
    # start by creating an RDD that has the number 0 through 20000
    # 20000 is the number of words that will be in our dictionary
    top20000Words = sc.parallelize(range(20000))

    # Now, we transform (0), (1), (2), ... to ("MostCommonWord", 1)
    # ("NextMostCommon", 2), ...
    # the number will be the spot in the dictionary used to tell us
    # where the word is located
    dictionary = top20000Words.map (lambda x : (topWordsinDict[x][0], x))

    # Filter out the required five words
    theFiveWords = dictionary.filter(lambda x: x[0] in {'applicant', 'and', 'attack', 'protein', 'car'})
    listofFInalFiveWordsWithCount = theFiveWords.collect()

    # Set of required words
    requiredWords  = {'applicant', 'and', 'attack', 'protein', 'car'}

    print ('#'*5,'Word Postions in our Dictionary:','#'*5,'\n')

    # List to store the position of each of the five words
    IntermediateResult = []

    # Function to find and print the relevant positions of words from our dictionary
    def findPos(requiredWords, listofFInalFiveWordsWithCount):
      for i in requiredWords:
        # Initialize a variable to check whether the required word is found or not
        currItemFound = 0

        # Initialize the position as -1
        positionFound = -1

        for j in range(len(listofFInalFiveWordsWithCount)):
          # If the required word is found, change the value of the flag and assign the position found
          if i in listofFInalFiveWordsWithCount[j]:
            positionFound = listofFInalFiveWordsWithCount[j][1]
            currItemFound = 1

        IntermediateResult.append((i, '->', positionFound))

        # Print the positions and the required word
        print(i, '->', positionFound)
      return IntermediateResult

    # Call the function to print answers for task 1
    ansForTask1 = findPos(requiredWords, listofFInalFiveWordsWithCount)

    # Save the results in a file
    sc.parallelize(ansForTask1).coalesce(1, shuffle = False).saveAsTextFile(sys.argv[2])

    sc.stop()