"""
Read and Visualize Horizontal/Vertical Slices in 2/3D
"""
import numpy as np
import os
from Utility import timer
from scipy.interpolate import griddata
from PlottingTool import Plot2D, Plot2D_InsetZoom, PlotSurfaceSlices3D, PlotContourSlices3D, pathpatch_translate, pathpatch_2d_to_3d
try:
    import PostProcess_AnisotropyTensor as PPAT
except ImportError:
    raise ImportError('\nNo module named PostProcess_AnisotropyTensor. Check setup.py and run'
                      '\npython setup.py build_ext --inplace')

try:
    import cpickle as pickle
except ModuleNotFoundError:
    import pickle

import matplotlib.pyplot as plt
from matplotlib.path import Path
from matplotlib.patches import Circle, PathPatch
import SliceData as PPSD
from DataBase import *
from copy import copy

"""
User Inputs
"""
time = 'latestTime'  #'23243.2156219'
casedir = '/media/yluan'
# casename = 'RANS/N_H_OneTurb_LowZ_Rwall2'  #'RANS/N_H_OneTurb_Simple_ABL'  #'URANS/N_H_OneTurb'  # 'ALM_N_H_ParTurb'
casename = 'ALM_N_H_OneTurb'
# properties = ('kResolved', 'kSGSmean')
properties = ('UAvg',)
# slicenames = ('oneDupstreamTurbine', 'rotorPlane', 'oneDdownstreamTurbine')
# slicenames = ('threeDdownstreamTurbine', 'fiveDdownstreamTurbine', 'sevenDdownstreamTurbine')
slicenames = ('hubHeight', 'quarterDaboveHub', 'turbineApexHeight')
# slicenames = ('groundHeight', 'hubHeight', 'turbineApex')
# Subscript for the slice names
slicenames_sub = 'Slice'
# Height of the horizontal slices, only used for 3D horizontal slices plot
horslice_offsets = (90., 121.5, 153.)
horslice_offsets2 = ((90., 90.), (121.5, 121.5), (153., 153.))
result_folder = 'Result'
# Orientation of x-axis in x-y plane, in case of angled flow direction
# Only used for values decomposition and confinebox
# Angle in rad and counter-clockwise
rot_z = np.pi/6.
# Turbine radius, only used for confinebox
r = 63
# For calculating total <epsilon> only
nu = 1e-5

"""
Plot Settings
"""
# Which type(s) of plot to make
plot_type = '3D'  # '2D', '3D', 'all'
# Total number cells intended to plot via interpolation
target_meshsize = 1e5
interp_method = 'linear'
# Number of contours, only for 2D plots or 3D horizontal slice plots
contour_lvl = 200
# Label of the property, could be overridden below
val_label = 'Data'
ext = 'png'
show, save = False, True
dpi = 1000

"""
Process User Inputs
"""
# Ensure slicenames is a tuple
slicenames = (slicenames,) if isinstance(slicenames, str) else slicenames
# Confined region auto definition
if 'OneTurb' in casename:
    # For rotor plane vertical slices
    if 'oneDupstream' in slicenames[0] or 'threeDdownstream' in slicenames[0]:
        turb_borders, turb_centers_frontview, confinebox, confinebox2 = OneTurb('vert')
        if 'three' in slicenames[0]:
            confinebox = confinebox2
            turb_centers_frontview = turb_centers_frontview[3:]
        else:
            turb_centers_frontview = turb_centers_frontview[:3]

    # For horizontal slices
    elif 'hubHeight' in slicenames[0] or 'groundHeight' in slicenames[0]:
        # Confinement for z doesn't matter since the slices are horizontal
        confinebox = ((800, 2400, 800, 2400, 0, 216),)*len(slicenames)
        turb_borders, turb_centers_frontview, confinebox, _ = OneTurb('hor')

elif 'ParTurb' in casename:
    if 'oneDupstream' in slicenames[0] or 'threeDdownstream' in slicenames[0]:
        # Read coor info from database
        turb_borders, turb_centers_frontview, confinebox, confinebox2 = ParTurb('vert')
        if 'threeDdownstream' in slicenames[0]:
            confinebox = confinebox2
            turb_centers_frontview = turb_centers_frontview[6:]
        else:
            turb_centers_frontview = turb_centers_frontview[:6]

    elif 'hubHeight' in slicenames[0] or 'groundHeight' in slicenames[0]:
        turb_borders, turb_centers_frontview, confinebox, _ = ParTurb('hor')

else:
    turb_borders = ((99999,)*4,)
    turb_centers_frontview = ((99999,)*3,)*6
    confinebox = confinebox2 = [[5., 2995., 5., 2995., 5., 995.]]*10

# If you don't want confinement
# confinebox = confinebox2 = [[5., 2995., 5., 2995., 5., 995.]]*10

# Automatic view_angle and figure settings, only for 3D plots
if 'oneDupstream' in slicenames[0] or 'threeDdownstream' in slicenames[0]:
    view_angle = (20, -80) if 'one' in slicenames[0] else (20, -95)
    equalaxis, figwidth = True,  'half'
elif 'groundHeight' in slicenames[0] or 'hubHeight' in slicenames[0]:
    view_angle = (25, -115)
    equalaxis, figwidth = False, 'half'
else:
    view_angle = (20, -100)
    equalaxis, figwidth = True, 'full'

# Unify plot_type user inputs
if plot_type in ('2D', '2d'):
    plot_type = '2D'
elif plot_type in ('3D', '3d'):
    plot_type = '3D'
elif plot_type in ('all', 'All', '*'):
    plot_type = 'all'

if 'U' in properties[0]:
    val_lim = (0, 12)
    val_lim_z = (-2, 2)
    val_label = [r'$\langle U_\mathrm{hor} \rangle$ [m/s]', r'$\langle w \rangle$ [m/s]']
elif 'k' in properties[0]:
    val_lim = (0, 2.5)
    val_lim_z = None
    if 'Resolved' in properties[0]:
        val_label = (r'$\langle k_\mathrm{resolved} \rangle$ [m$^2$/s$^2$]',) if len(properties) == 1 else (r'$\langle k \rangle$ [m$^2$/s$^2$]',)
    elif 'SGS' in properties[0]:
        val_label = (r'$\langle k_\mathrm{SGS} \rangle$ [m$^2$/s$^2$]',)
    else:
        val_label = (r'$\langle k \rangle$ [m$^2$/s$^2$]',)

elif 'uuPrime2' in properties[0] or 'R' in properties[0]:
    val_lim = (-0.5, 2/3.)
    val_lim_z = None
    val_label = (r"$\langle u'u' \rangle$ [-]", r"$\langle u'v' \rangle$ [-]", r"$\langle u'w' \rangle$ [-]",
                 r"$\langle v'v' \rangle$ [-]", r"$\langle v'w' \rangle$ [-]",
                                                r"$\langle w'w' \rangle$ [-]")
elif "epsilon" in properties[0]:
    val_lim = None
    val_lim_z = None
    if 'SGS' in properties[0]:
        val_label = (r'$\langle \epsilon_{\mathrm{SGS}} \rangle$ [m$^2$/s$^3$]',) if 'mean' in properties[0] else (r'$\epsilon_{\mathrm{SGS}}$ [m$^2$/s$^3$]',)
    elif 'Resolved' in properties[0]:
        val_label = (r'$\langle \epsilon_{\mathrm{resolved}} \rangle$ [m$^2$/s$^3$]',)
    else:
        val_label = (r'$\langle \epsilon \rangle$ [m$^2$/s$^3$]',)

elif 'G' in properties[0]:
    val_lim = (-0.05, 0.17)
    val_lim_z = None
    val_label = (r'$\langle P_k \rangle$ [m$^2$/s$^3$]',)

else:
    val_lim = None
    val_lim_z = None
    val_label = ('data',)

"""
Read Slice Data
"""
# Initialize case
case = PPSD.SliceProperties(time=time, casedir=casedir, casename=casename, rot_z=rot_z, result_folder=result_folder)
# Read slices
case.readSlices(properties=properties, slicenames=slicenames, slicenames_sub=slicenames_sub)

list_x2d, list_y2d, list_z2d, list_val3d, list_val3d_z = [], [], [], [], []
# Go through specified slices and flow properties
for i in range(len(case.slicenames)):
    # for i, slicename in enumerate(case.slicenames):
    slicename = case.slicenames[i]
    vals2d = case.slices_val[slicename]
    # If kResolved and kSGSmean in properties, get total kMean
    if 'kResolved' in properties and 'kSGSmean' in properties:
        print(' Calculating total <k> for {}...'.format(slicenames[i]))
        slicename2 = case.slicenames[i + len(slicenames)]
        vals2d += case.slices_val[slicename2]
    # # Else if epsilonSGSmean and nuSGSmean in properties then get total epsilonMean
    # # By assuming isotropic homogeneous turbulence and
    # # <epsilon> = <epsilonSGS>/(1 - 1/(1 + <nuSGS>/nu))
    # elif 'epsilonSGSmean' in properties and 'nuSGSmean' in properties:
    #     print(' Calculating total <epsilon> for {}...'.format(slicenames[i]))
    #     slicename2 = case.slicenames[i + len(slicenames)]
    #     vals2d_2 = case.slices_val[slicename2]
    #     # Determine which vals2d is epsilonSGSmean or nuSGSmean
    #     nusgs_mean, epsilonSGSmean = (vals2d, vals2d_2) if 'epsilonSGSmean' in slicename2 else (vals2d_2, vals2d)
    #     # Calculate epsilonMean
    #     vals2d = case.calcSliceMeanDissipationRate(epsilonSGSmean = epsilonSGSmean, nusgs_mean = nusgs_mean, nu = nu)

    # Interpolation
    x2d, y2d, z2d, vals3d = case.interpolateDecomposedSliceData_Fast(case.slices_coor[slicename][:, 0], case.slices_coor[slicename][:, 1], case.slices_coor[slicename][:, 2], vals2d,
                                                                     slice_orient=case.slices_orient[slicename], rot_z=case.rot_z,
                                                                     target_meshsize=target_meshsize,
                                                                     interp_method=interp_method,
                                                                     confinebox=confinebox[i])

    # Flatten if vals3d only have one component like a scalar field
    if vals3d.shape[2] == 1:
        vals3d = vals3d.reshape((vals3d.shape[0], vals3d.shape[1]))

    # Calculate magnitude if U
    if 'U' in properties or 'UAvg' in properties:
        vals3d_hor = np.sqrt(vals3d[:, :, 0]**2 + vals3d[:, :, 1]**2)
        vals3d_z =  vals3d[:, :, 2]
    else:
        vals3d_hor = vals3d
        vals3d_z = None

    # Append 2D mesh to a list for 3D plots
    if plot_type in ('3D', 'all'):
        list_x2d.append(x2d)
        list_y2d.append(y2d)
        list_z2d.append(z2d)
        list_val3d.append(vals3d_hor)
        list_val3d_z.append(vals3d_z)

    # Determine the unit along the vertical slice since it's angled, only for 2D plots of vertical slices
    if case.slices_orient[slicename] == 'vertical':
        # # If angle from x-axis is 45 deg or less
        # if lx >= ly:
        #     rot_z = np.arctan(lx/ly)
        if confinebox is None:
            lx = np.max(x2d) - np.min(x2d)
            ly = np.max(y2d) - np.min(y2d)
        else:
            lx = confinebox[i][1] - confinebox[i][0]
            ly = confinebox[i][3] - confinebox[i][2]

        r2d = np.linspace(0, np.sqrt(lx**2 + ly**2), x2d.shape[0])

    # Break if i finishes all kResolved or kSGSmean
    if 'kResolved' in properties and 'kSGSmean' in properties and i == (len(slicenames) - 1):
        break
    elif 'epsilonSGSmean' in properties and 'nuSGSmean' in properties and i == (len(slicenames) - 1):
        break

    """
    Plotting
    """
    xlabel, ylabel = (r'$x$ [m]', r'$y$ [m]') \
        if case.slices_orient[slicename] == 'vertical' else \
        (r'$x$ [m]', r'$y$ [m]')
    # Figure name
    if 'kResolved' in properties and 'kSGSmean' in properties:
        figname = 'kMean_' + slicenames[i] + slicenames_sub
    elif 'epsilonSGSmean' in properties and 'nuSGSmean' in properties:
        figname = 'epsilonMean_' + slicenames[i] + slicenames_sub
    else:
        figname = slicename

    if plot_type in ('2D', 'all'):
        slicePlot = Plot2D(x2d, y2d, vals3d, name=figname, xlabel=xlabel, ylabel=ylabel, val_label= val_label, save=save, show=show, figdir=case.result_path)
        slicePlot.initializeFigure()
        slicePlot.plotFigure(contour_lvl=contour_lvl)
        slicePlot.finalizeFigure()

if plot_type in ('3D', 'all'):
    zlabel = r'$z$ [m]'
    # Figure name for 3D plots
    if 'kResolved' in properties and 'kSGSmean' in properties:
        figname_3d = 'kMean_' + str(slicenames)
    elif 'epsilonSGSmean' in properties and 'nuSGSmean' in properties:
        figname_3d = 'epsilonMean_' + str(slicenames)
    else:
        figname_3d = str(case.slicenames)

    if case.slices_orient[slicename] == 'horizontal':
        show_xylabel = (False, False)
        show_zlabel = True
        show_ticks = (False, False, True)
        # Initialize plot object for horizontal contour slices
        plot3d = PlotContourSlices3D(list_x2d, list_y2d, list_val3d, horslice_offsets, gradient_bg=False, name=figname_3d, xlabel=xlabel, ylabel=ylabel, zlabel=zlabel, val_label=val_label[0], save=save, show=show, figdir=case.result_path, viewangle=view_angle, figwidth=figwidth, equalaxis=equalaxis, cbar_orient='vertical',
                                          figheight_multiplier=1.75,
                                          val_lim=val_lim,
                                     zlim=None)
        # If there's a z component e.g. Uz, initialize it separately
        if list_val3d_z[0] is not None:
            plot3d_z = PlotContourSlices3D(list_x2d, list_y2d, list_val3d_z, horslice_offsets, gradient_bg=False,
                                              name=figname_3d + '_z', xlabel=xlabel, ylabel=ylabel, zlabel=zlabel,
                                              val_label=val_label[1], save=save, show=show, figdir=case.result_path,
                                              viewangle=view_angle, figwidth=figwidth, equalaxis=equalaxis,
                                              cbar_orient='vertical',
                                              figheight_multiplier=1.75,
                                              val_lim=val_lim_z,
                                           zlim=None)
    elif case.slices_orient[slicename] == 'vertical':
        show_xylabel = (True, True)
        show_zlabel = False
        show_ticks = (True, True, False)
        patch = Circle((0., 0.), 63., alpha=0.5, fill=False, edgecolor=(0.25, 0.25, 0.25), zorder=100)
        patches = []
        for i in range(10):
            patches.append(copy(patch))

        patches = iter(patches)
        # Initialize vertical surface plot instance
        plot3d = PlotSurfaceSlices3D(list_x2d, list_y2d, list_z2d, list_val3d, xlabel=xlabel, ylabel=ylabel, zlabel=zlabel, val_label=val_label[0], name=figname_3d, save=save, show=show, figdir=case.result_path, viewangle=view_angle, figwidth=figwidth, equalaxis=equalaxis, cbar_orient='horizontal',
                                          val_lim=val_lim)
        # Again separate instance for z component
        if list_val3d_z[0] is not None:
            plot3d_z = PlotSurfaceSlices3D(list_x2d, list_y2d, list_z2d, list_val3d_z, xlabel=xlabel, ylabel=ylabel,
                                              zlabel=zlabel, val_label=val_label[1], name=figname_3d + '_z', save=save,
                                              show=show, figdir=case.result_path, viewangle=view_angle,
                                              figwidth=figwidth, equalaxis=equalaxis, cbar_orient='horizontal',
                                              val_lim=val_lim_z)

    plot3d.initializeFigure(constrained_layout=True)
    plot3d.plotFigure(contour_lvl=contour_lvl)
    if casename not in ('ABL_N_H', 'ABL_N_L'):
        if case.slices_orient[slicename] == 'horizontal':
            for i in range(len(horslice_offsets)):
                plot3d.axes.plot([turb_borders[0][0], turb_borders[0][2]], [turb_borders[0][1], turb_borders[0][3]], zs=horslice_offsets2[i], alpha=0.5, color=(0.25, 0.25, 0.25),
                                 # Very important to set a super larger value
                                 zorder=500 + i*500)
        else:
            for i in range(len(list_x2d)):
                p = next(patches)
                plot3d.axes.add_patch(p)
                pathpatch_2d_to_3d(p, z=0, normal=(0.8660254037844, 0.5, 0.))
                pathpatch_translate(p, turb_centers_frontview[i])

    plot3d.finalizeFigure(tight_layout=False, show_ticks=show_ticks, show_xylabel=show_xylabel, show_zlabel=show_zlabel)
    # For Uz or any other z component
    if list_val3d_z[0] is not None:
        plot3d_z.initializeFigure()
        plot3d_z.plotFigure(contour_lvl=contour_lvl)
        if casename not in ('ABL_N_H', 'ABL_N_L'):
            if case.slices_orient[slicename] == 'horizontal':
                for i in range(len(horslice_offsets)):
                    plot3d_z.axes.plot([turb_borders[0][0], turb_borders[0][2]], [turb_borders[0][1], turb_borders[0][3]],
                                     zs=horslice_offsets2[i], alpha=0.5, color=(0.25, 0.25, 0.25), zorder=500 + i*500)
            else:
                for i in range(len(list_x2d)):
                    p = next(patches)
                    plot3d_z.axes.add_patch(p)
                    pathpatch_2d_to_3d(p, z=0, normal=(0.8660254037844, 0.5, 0.))
                    pathpatch_translate(p, turb_centers_frontview[i])

        plot3d_z.finalizeFigure(show_ticks=show_ticks, show_xylabel=show_xylabel, show_zlabel=show_zlabel)

# """
# User Inputs
# """
# casedir = 'J:'  # '/media/yluan/Toshiba External Drive/'
# casedir = '/media/yluan/Toshiba External Drive/'
# casename = 'ALM_N_H_OneTurb'  # 'ALM_N_H_ParTurb'
# time = 23275.1388025  # 22000.0918025 20000.9038025
# # slicenames = ['alongWind', 'groundHeight', 'hubHeight', 'oneDaboveHubHeight', 'oneDdownstreamTurbineOne', 'oneDdownstreamTurbineTwo', 'rotorPlaneOne', 'rotorPlaneTwo', 'sixDdownstreamTurbineTwo', 'threeDdownstreamTurbineOne', 'threeDdownstreamTurbineTwo', 'twoDupstreamTurbineOne']
# # For Upwind and Downwind turbines
# # slicenames = ['oneDdownstreamTurbineOne', 'oneDdownstreamTurbineTwo', 'rotorPlaneOne', 'rotorPlaneTwo', 'sixDdownstreamTurbineTwo', 'threeDdownstreamTurbineOne', 'threeDdownstreamTurbineTwo', 'twoDupstreamTurbineOne']
# # # For Parallel Turbines
# # slicenames = ['alongWindRotorOne', 'alongWindRotorTwo', 'twoDupstreamTurbines', 'rotorPlane', 'oneDdownstreamTurbines', 'threeDdownstreamTurbines', 'sixDdownstreamTurbines']
# # slicenames = ['groundHeight', 'hubHeight', 'oneDaboveHubHeight']
# # slicenames = ['rotorPlane','sixDdownstreamTurbines']
# slicenames = ['alongWind']
# # Only for PlotContourSlices3D
# sliceOffsets = (5, 90, 153)
# propertyName = 'uuPrime2'
# fileExt = '.raw'
# precisionX, precisionY, precisionZ = 1000j, 1000j, 333j
# interp_method = 'nearest'
#
#
# """
# Plot Settings
# """
# figwidth = 'full'
# # View angle best (15, -40) for vertical slices in rotor plane
# view_angle, equalaxis = (15, -45), True
# xLim, yLim, zLim = (0, 3000), (0, 3000), (0, 1000)
# show, save = False, True
# xlabel, ylabel, zlabel = r'$x$ [m]', r'$y$ [m]', r'$z$ [m]'
# # valLabels = (r'$b_{11}$ [-]', r'$b_{12}$ [-]', r'$b_{13}$ [-]', r'$b_{22}$ [-]', r'$b_{23}$ [-]', r'$b_{33}$ [-]')
# # valLabels = (r'$\langle u\rangle$ [m/s]', r'$\langle v\rangle$ [m/s]', r'$\langle w\rangle$ [m/s]')
# if propertyName == 'U':
#     valLabels = (r'$U$ [m/s]', r'$U$ [m/s]', r'$U$ [m/s]')
# elif propertyName == 'uuPrime2':
#     valLabels = (r'$b_{11}$ [-]', r'$b_{12}$ [-]', r'$b_{13}$ [-]', r'$b_{22}$ [-]', r'$b_{23}$ [-]', r'$b_{33}$ [-]', r'$k_{\rm{resolved}}$ [m$^2$/s$^2$]')
#
#
# """
# Process User Inputs
# """
# # Combine propertyName with slicenames and Subscript to form the full file names
# # Don't know why I had to copy it...
# fileNames = slicenames.copy()
# for i, name in enumerate(slicenames):
#     slicenames[i] = propertyName + '_' + name + '_Slice'
#     fileNames[i] = slicenames[i] + fileExt
#
# figDir = casedir + casename + '/Slices/Result/' + str(time)
# try:
#     os.makedirs(figDir)
# except FileExistsError:
#     pass
#
#
# """
# Functions
# """
# @timer
# @jit
# def readSlices(time, casedir = '/media/yluan/Toshiba External Drive', casename = 'ALM_N_H', fileNames = ('*',), skipCol = 3, skipRow = 0):
#     caseFullPath = casedir + '/' + casename + '/Slices/' + str(time) + '/'
#     fileNames = os.listdir(caseFullPath) if fileNames[0] in ('*', 'all') else fileNames
#     slices_val, slicesDir, slices_coor = {}, {}, {}
#     for fileName in fileNames:
#         vals = np.genfromtxt(caseFullPath + fileName)
#         # partition('.') removes anything after '.'
#         slices_coor[fileName.partition('.')[0]] = vals[skipRow:, :skipCol]
#         # If max(z) - min(z) < 1 then it's assumed horizontal
#         slicesDir[fileName.partition('.')[0]] = 'vertical' if (vals[skipRow:, skipCol - 1]).max() - (vals[skipRow:, skipCol - 1]).min() > 1. else 'horizontal'
#         slices_val[fileName.partition('.')[0]] = vals[skipRow:, skipCol:]
#
#     print('\n' + str(fileNames) + ' read')
#     return slices_coor, slicesDir, slices_val
#
#
# @timer
# @jit
# def interpolateSlices(x, y, z, vals, sliceDir = 'vertical', precisionX = 1500j, precisionY = 1500j, precisionZ = 500j, interp_method = 'cubic'):
#     # Bound the coordinates to be interpolated in case data wasn't available in those borders
#     bnd = (1.00001, 0.99999)
#     if sliceDir is 'vertical':
#         # Known x and z coordinates, to be interpolated later
#         knownPoints = np.vstack((x, z)).T
#         # Interpolate x and z according to precisions
#         x2d, z2d = np.mgrid[x.min()*bnd[0]:x.max()*bnd[1]:precisionX, z.min()*bnd[0]:z.max()*bnd[1]:precisionZ]
#         # Then interpolate y in the same fashion of x
#         y2d, _ = np.mgrid[y.min()*bnd[0]:y.max()*bnd[1]:precisionY, z.min()*bnd[0]:z.max()*bnd[1]:precisionZ]
#         # In case the vertical slice is at a negative angle,
#         # i.e. when x goes from low to high, y goes from high to low,
#         # flip y2d from low to high to high to low
#         y2d = np.flipud(y2d) if x[0] > x[1] else y2d
#     else:
#         knownPoints = np.vstack((x, y)).T
#         x2d, y2d = np.mgrid[x.min()*bnd[0]:x.max()*bnd[1]:precisionX, y.min()*bnd[0]:y.max()*bnd[1]:precisionY]
#         _, z2d = np.mgrid[x.min()*bnd[0]:x.max()*bnd[1]:precisionX, z.min()*bnd[0]:z.max()*bnd[1]:precisionZ]
#
#     # Decompose the vector/tensor of slice values
#     # If vector, order is x, y, z
#     # If symmetric tensor, order is xx, xy, xz, yy, yz, zz
#     valsDecomp = {}
#     for i in range(vals.shape[1]):
#         if sliceDir is 'vertical':
#             # Each component is interpolated from the known locations pointsXZ to refined fields (x2d, z2d)
#             valsDecomp[str(i)] = griddata(knownPoints, vals[:, i].ravel(), (x2d, z2d), method = interp_method)
#         else:
#             valsDecomp[str(i)] = griddata(knownPoints, vals[:, i].ravel(), (x2d, y2d), method = interp_method)
#
#     return x2d, y2d, z2d, valsDecomp
#
#
# @timer
# @jit
# def calculateAnisotropicTensor(valsDecomp):
#     # k in the interpolated mesh
#     # xx is '0', xy is '1', xz is '2', yy is '3', yz is '4', zz is '5'
#     k = 0.5*(valsDecomp['0'] + valsDecomp['3'] + valsDecomp['5'])
#     # Convert Rij to bij
#     for key, val in valsDecomp.items():
#         valsDecomp[key] = val/(2.*k) - 1/3. if key in ('0', '3', '5') else val/(2.*k)
#
#     return valsDecomp, k
#
#
# @timer
# @jit
# def mergeHorizontalComponent(valsDecomp):
#     valsDecomp['hor'] = np.sqrt(valsDecomp['0']**2 + valsDecomp['1']**2)
#     return valsDecomp
#
#
# """
# Read, Decompose and Plot 2/3D Slices
# """
# slices_coor, slicesDir, slices_val = readSlices(time = time, casedir = casedir, casename = casename, fileNames = fileNames)
#
# # Initialize slice lists for multple slice plots in one 3D figure
# horSliceLst, zSliceLst, list_x2d, list_y2d, list_z2d = [], [], [], [], []
# # Go through slices
# for slicename in slicenames:
#     x2d, y2d, z2d, valsDecomp = interpolateSlices(slices_coor[slicename][:, 0], slices_coor[slicename][:, 1], slices_coor[slicename][:, 2], slices_val[slicename], sliceDir = slicesDir[slicename], precisionX = precisionX, precisionY = precisionY, precisionZ = precisionZ, interp_method = interp_method)
#
#     # For anisotropic stress tensor bij
#     # bij = Rij/(2k) - 1/3*deltaij
#     # where Rij is uuPrime2, k = 1/2trace(Rij), deltaij is Kronecker delta
#     if propertyName == 'uuPrime2':
#         valsDecomp, k = calculateAnisotropicTensor(valsDecomp)
#         valsDecomp['kResolved'] = k
#     elif propertyName == 'U':
#         valsDecomp = mergeHorizontalComponent(valsDecomp)
#
#
#     """
#     2D Contourf Plots
#     """
#     xLim, yLim, zLim = (x2d.min(), x2d.max()), (y2d.min(), y2d.max()), (z2d.min(), z2d.max())
#     plotsLabel = iter(valLabels)
#     for key, val in valsDecomp.items():
#         # if slicesDir[slicename] is 'vertical':
#         #     slicePlot = Plot2D(x2d, z2d, z2d = val, equalaxis = True,
#         #                                  name = slicename + '_' + key, figDir = figDir, xLim = xLim, yLim = zLim,
#         #                                  show = show, xlabel = xlabel, ylabel = zlabel, save = save,
#         #                                  zlabel = next(plotsLabel))
#         #
#         # else:
#         #     slicePlot = Plot2D(x2d, y2d, z2d = val, equalaxis = True,
#         #                        name = slicename + '_' + key, figDir = figDir, xLim = xLim, yLim = yLim,
#         #                        show = show, xlabel = xlabel, ylabel = ylabel, save = save,
#         #                        zlabel = next(plotsLabel))
#         # slicePlot = Plot2D_InsetZoom(x2d, z2d, zoomBox = (1000, 2500, 0, 500), z2d = val, equalaxis = True, name = slicename + '_' + key, figDir = figDir, xLim = xLim, yLim = zLim, show = show, xlabel = xlabel, ylabel = zlabel, save = save, zlabel = next(plotsLabel))
#         # plot_type = 'contour2D'
#
#         slicePlot = PlotSurfaceSlices3D(x2d, y2d, z2d, val, name = slicename + '_' + key + '_3d', figDir = figDir, xLim = xLim, yLim = yLim, zLim = zLim, show = show, xlabel = xlabel, ylabel = ylabel, zlabel = zlabel, save = save, cmapLabel = next(plotsLabel), viewAngles = view_angle, figwidth = figwidth)
#         plot_type = 'surface3D'
#
#         slicePlot.initializeFigure()
#         if plot_type == 'contour2D':
#             slicePlot.plotFigure(contour_lvl = 100)
#         else:
#             slicePlot.plotFigure()
#
#         slicePlot.finalizeFigure()
#
#     if propertyName == 'U':
#         horSliceLst.append(valsDecomp['hor'])
#         zSliceLst.append(valsDecomp['2'])
#
#     list_x2d.append(x2d)
#     list_y2d.append(y2d)
#     list_z2d.append(z2d)

"""
Multiple Slices of Horizontal Component 3D Plot
"""
# if slicesDir[slicename] is 'horizontal':
#     slicePlot = PlotContourSlices3D(x2d, y2d, horSliceLst, sliceOffsets = sliceOffsets, contour_lvl = 100, zLim = (0, 216), gradientBg = False, name = str(slicenames) + '_hor', figDir = figDir, show = show, xlabel = xlabel, ylabel = ylabel, zlabel = zlabel, cmapLabel = r'$U_{\rm{hor}}$ [m/s]', save = save, cbarOrientate = 'vertical')
# else:
#     slicePlot = PlotSurfaceSlices3D(list_x2d, list_y2d, list_z2d, horSliceLst, name = str(slicenames) + '_hor', figDir = figDir, show = show, xlabel = xlabel,
#                                     ylabel = ylabel, zlabel = zlabel, save = save, cmapLabel = r'$U_{\rm{hor}}$ [m/s]', viewAngles = view_angle, figwidth = figwidth, xLim = xLim, yLim = yLim, zLim = zLim, equalaxis = equalaxis)
#
# slicePlot.initializeFigure()
# slicePlot.plotFigure()
# slicePlot.finalizeFigure()

"""
Multiple Slices of Z Component 3D Plot
"""
# if slicesDir[slicename] is 'horizontal':
#     slicePlot = PlotContourSlices3D(x2d, y2d, zSliceLst, sliceOffsets = sliceOffsets, contour_lvl = 100,
#                                     xLim = (0, 3000), yLim = (0, 3000), zLim = (0, 216), gradientBg = False,
#                                     name = str(slicenames) + '_z', figDir = figDir, show = show,
#                                     xlabel = xlabel, ylabel = ylabel, zlabel = zlabel,
#                                     cmapLabel = r'$U_{z}$ [m/s]', save = save, cbarOrientate = 'vertical')
# else:
#     slicePlot = PlotSurfaceSlices3D(list_x2d, list_y2d, list_z2d, zSliceLst,
#                                     name = str(slicenames) + '_z', figDir = figDir, show = show, xlabel = xlabel,
#                                     ylabel = ylabel, zlabel = zlabel, save = save, cmapLabel = r'$U_{z}$ [m/s]', viewAngles = view_angle, figwidth = figwidth, xLim = xLim, yLim = yLim, zLim = zLim, equalaxis = equalaxis)
#
# slicePlot.initializeFigure()
# slicePlot.plotFigure()
# slicePlot.finalizeFigure()


import nmrglue.fileio.simpson as simpson
import numpy as np
from numpy.testing import assert_allclose, assert_raises
import os.path

from setup import DATA_DIR

DD_1D = os.path.join(DATA_DIR, 'simpson_1d')
DD_2D = os.path.join(DATA_DIR, 'simpson_2d')

def test_1d_time():
    """ reading 1D time domain files """
    # read the text, binary, xreim, and rawbin data
    text_dic, text_data = simpson.read(os.path.join(DD_1D, '1d_text.fid'))
    bin_dic, bin_data = simpson.read(os.path.join(DD_1D, '1d_bin.fid'))
    xreim_units, xreim_data = simpson.read(os.path.join(DD_1D, '1d_ftext.fid'))
    rd, rawbin_data = simpson.read(os.path.join(DD_1D, '1d_rawbin.fid'), spe=False, ndim=1)

    # check data in text file
    assert text_data.shape == (4096, )
    assert text_data.dtype == 'complex64'
    assert round(text_data[0].real, 2) == 2.0
    assert round(text_data[0].imag, 2) == 0.0
    assert round(text_data[1].real, 2) == 1.78
    assert round(text_data[1].imag, 2) == -0.01

    # data in all files should be close
    assert np.allclose(rawbin_data, text_data)
    assert np.allclose(rawbin_data, bin_data)
    assert np.allclose(rawbin_data, xreim_data)

def test_1d_freq():
    """ reading 1D freq domain files """
    # read the text, binary, xreim, and rawbin data
    text_dic, text_data = simpson.read(os.path.join(DD_1D, '1d_text.spe'))
    bin_dic, bin_data = simpson.read(os.path.join(DD_1D, '1d_bin.spe'))
    xreim_units, xreim_data = simpson.read(os.path.join(DD_1D, '1d_ftext.spe'))
    rd, rawbin_data = simpson.read(os.path.join(DD_1D, '1d_rawbin.spe'), spe=True, ndim=1)

    # check data in text file
    assert text_data.shape == (4096, )
    assert text_data.dtype == 'complex64'
    assert round(text_data[2048].real, 2) == 40.34
    assert round(text_data[2048].imag, 2) == -1.51
    assert round(text_data[2049].real, 2) == 39.58
    assert round(text_data[2049].imag, 2) == -3.97

    # data in all file should be close
    assert np.allclose(rawbin_data, text_data)
    assert np.allclose(rawbin_data, bin_data)
    assert np.allclose(rawbin_data, xreim_data)

def test_2d_time():
    """ reading 2D time domain files """
    # read the text, binary, xreim, and rawbin data
    text_dic, text_data = simpson.read(os.path.join(DD_2D, '2d_text.fid'))
    bin_dic, bin_data = simpson.read(os.path.join(DD_2D, '2d.fid'))
    xyreim_units, xyreim_data = simpson.read(os.path.join(DD_2D, '2d_ftext.fid'))
    rd, rawbin_data = simpson.read(os.path.join(DD_2D, '2d_raw.fid'), NP=128, NI=48, ndim=2,
                                    spe=False)

    # check data in text file
    assert text_data.shape == (48, 128)
    assert text_data.dtype == 'complex64'
    assert round(text_data[0,0].real, 2) == 1.00
    assert round(text_data[0,0].imag, 2) == 0.03
    assert round(text_data[0,1].real, 2) == 0.75
    assert round(text_data[0,1].imag, 2) == 0.59
    assert round(text_data[1,0].real, 2) == 0.89
    assert round(text_data[1,0].imag, 2) == 0.03

    # data in all files should be close
    assert np.allclose(rawbin_data, text_data)
    assert np.allclose(rawbin_data, bin_data)
    assert np.allclose(rawbin_data, xyreim_data)

def test_2d_freq():
    """ reading 2D freq domain files """
    # read the text, binary, xreim, and rawbin data
    text_dic, text_data = simpson.read(os.path.join(DD_2D, '2d_text.spe'))
    bin_dic, bin_data = simpson.read(os.path.join(DD_2D, '2d.spe'))
    xyreim_units, xyreim_data = simpson.read(os.path.join(DD_2D, '2d_ftext.spe'))
    rd, rawbin_data = simpson.read(os.path.join(DD_2D, '2d_raw.spe'), ndim=2, NP=256,
                                    NI=512, spe=True)

    # check data in text file
    assert text_data.shape == (512, 256)
    assert text_data.dtype == 'complex64'
    assert round(text_data[4, 150].real, 2) == 0.29
    assert round(text_data[4, 150].imag, 2) == 0.34
    assert round(text_data[4, 151].real, 2) == 0.13
    assert round(text_data[4, 151].imag, 2) == 0.16
    assert round(text_data[5, 150].real, 2) == 0.41
    assert round(text_data[5, 150].imag, 2) == 0.14

    # data in text, bin and xyreim files should all be close
    assert np.allclose(text_data, bin_data)
    assert np.allclose(text_data, xyreim_data)

    # rawbin should be close except for first point along each vector
    assert np.allclose(rawbin_data[:, 1:], text_data[:, 1:])

def test_exceptions_read():
    """ raising exceptions due to missing read parameters """

    # missing spe parameter
    assert_raises(ValueError, simpson.read, os.path.join(DD_1D, '1d_rawbin.fid'))

    # missing ndim parameter
    assert_raises(ValueError, simpson.read, os.path.join(DD_1D, '1d_rawbin.fid'), spe=False)

    # missing NP/NI parameter
    assert_raises(ValueError, simpson.read, os.path.join(DD_2D, '2d_raw.fid'), spe=False,
            ndim=2)

    # bad ftype
    assert_raises(ValueError, simpson.read, os.path.join(DD_1D, '1d_rawbin.fid'), ftype='a')


#!/usr/bin/python3
# -*- coding: utf-8 -*-

import numpy as np

def test_save_load():# {{{
    """
    存储多个数组到文件中
    """
    a1 = np.arange(8)
    a2 = np.add.accumulate(a1)
    a3 = a1 + a2
    with open("/tmp/save_load.npy", "wb") as fw:
        np.save(fw, a1)
        np.save(fw, a2)
        np.save(fw, a3)

    with open("/tmp/save_load.npy", "rb") as fr:
        b1 = np.load(fr)
        b2 = np.load(fr)
        b3 = np.load(fr)
        # 再次np.load(r)会有异常

    print(a1, a2, a3)
    print(b1, b2, b3)
# }}}

def test_save_load2():# {{{
    """
    低级保存, 不保存数组的结构(shape,dtype)
    """
    arr = np.arange(12)
    arr.shape = (3, 4)
    print(arr, arr.shape, arr.dtype)

    arr.tofile("/tmp/arr.bin")
    arr2 = np.fromfile("/tmp/arr.bin")
    # 错误, 默认是float型
    print(arr2)
    arr3 = np.fromfile("/tmp/arr.bin", dtype=np.int64)
    # 错误, 一维的
    print(arr3)
    arr3.shape = (3, 4)
    # 正确
    print(arr3)
# }}}

def test_read_csv():# {{{
    """
    读取csv文件, 列类型不统一的处理
    """
    # loadtxt 对编码有个bug, 临时如下解决
    with open("./test.csv", encoding='latin1') as fp:
        arr = np.loadtxt(fp, dtype='S', delimiter=',')
        print("arr.shape = " , arr.shape)
        for i in range(arr.shape[0]):
            for j in range(arr.shape[1]):
                print(arr[i,j].decode('utf-8'))
        data = arr[1:,1:].astype(np.float)
        print(data)

    with open("./test.csv", encoding='latin1') as fp:
        dt = np.dtype({
            'names':['name', 'age', 'weight', 'height'],
            'formats':['S32', 'i', 'f', 'f']
            })
        # 如果列中数组类型不统一, 则返回所有列作为一维数据
        arr = np.loadtxt(fp, dtype=dt, delimiter=',', skiprows=1)
        print("arr.shape = " , arr.shape)

        # 方法1 转置
        data = np.transpose(np.array([arr['age'], arr['weight'], arr['height']]))
        print("methd1:", data)

        # 方法2
        data = np.zeros((3, 3))
        data[:,0] = arr['age']
        data[:,1] = arr['weight']
        data[:,2] = arr['height']
        print("method2:", data)

    # 方法3
    with open("./test.csv", encoding='latin1') as fp:
        arr = np.loadtxt(fp, delimiter=',', skiprows=1, usecols=(1,2,3))
        print(arr, arr.shape)

# }}}

def test_data_table():# {{{
    """
    unpack, NAN, MISSING, genfromtxt
    """
    arr = np.loadtxt("./data_table.txt", skiprows=1)
    print(arr, arr.shape)
    # unpack: 返回一个列独立数组
    c1, c2, c3 = np.loadtxt("./data_table.txt", skiprows=1, unpack=True)
    print(c1, c2, c3)

    # 如果文件中有NAN, 自动处理, 不会增加dtype类型
    arr = np.loadtxt("./data_table2.txt", skiprows=1)
    print(arr, arr.shape)

    # 如果skiprows=1,忽略
    arr = np.genfromtxt("./data_table3.txt", skip_header=1)
    print(arr, arr.shape)

    # 对特殊值进行填充 (每一列可以补充不同的值)
    arr = np.genfromtxt("./data_table3.txt", skip_header=1,
            missing_values=('MISSING', 'MISSING', 'MISSING'),
            filling_values=(-999, -888, -777))
    print(arr, arr.shape)
# }}}

def main():
    test_save_load()
    test_save_load2()
    test_read_csv()
    test_data_table()

if __name__ == "__main__":
    main()


from __future__ import absolute_import, print_function
import os
import logging
import pickle
import tempfile

from IPython.display import Image, SVG
import numpy as np
import pytest

from oct2py import Oct2Py, Oct2PyError, Struct, Cell
from oct2py.io import MatlabFunction
from oct2py.compat import StringIO

class TestUsage:
    """Excercise the basic interface of the package
    """
    @classmethod
    def setup_class(cls):
        cls.oc = Oct2Py()
        cls.oc.addpath(os.path.realpath(os.path.dirname(__file__)))

    @classmethod
    def teardown_class(cls):
        cls.oc.exit()

    def test_run(self):
        """Test the run command
        """
        out = self.oc.eval('ones(3,3)')
        desired = np.ones((3, 3))
        assert np.allclose(out, desired)
        out = self.oc.eval('ans = mean([[1, 2], [3, 4]])', verbose=True)
        assert out == 2.5
        with pytest.raises(Oct2PyError):
            self.oc.eval('_spam')

    def test_dynamic_functions(self):
        """Test some dynamic functions
        """
        out = self.oc.ones(1, 2)
        assert np.allclose(out, np.ones((1, 2)))

        U, S, V = self.oc.svd([[1, 2], [1, 3]], nout=3)
        assert np.allclose(U, ([[-0.57604844, -0.81741556],
                           [-0.81741556, 0.57604844]]))
        assert np.allclose(S, ([[3.86432845, 0.],
                           [0., 0.25877718]]))
        assert np.allclose(V, ([[-0.36059668, -0.93272184],
                           [-0.93272184, 0.36059668]]))
        out = self.oc.roundtrip(1)
        assert out == 1
        with pytest.raises(Oct2PyError):
            self.oc.eval('_spam')

    def test_push_pull(self):
        self.oc.push('spam', [1, 2])
        out = self.oc.pull('spam')
        assert np.allclose(out, np.array([1, 2]))
        self.oc.push(['spam', 'eggs'], ['foo', [1, 2, 3, 4]])
        spam, eggs = self.oc.pull(['spam', 'eggs'])
        assert spam == 'foo'
        assert np.allclose(eggs, np.array([[1, 2, 3, 4]]))

    def test_help(self):
        """Testing help command
        """
        doc = self.oc.cos.__doc__
        assert 'Compute the cosine for each element of X in radians.' in doc

    def test_dynamic(self):
        """Test the creation of a dynamic function
        """
        tests = [self.oc.zeros, self.oc.ones, self.oc.plot]
        for item in tests:
            assert "class 'oct2py.dynamic" in repr(type(item))
        with pytest.raises(Oct2PyError):
            self.oc.__getattr__('aaldkfasd')
        with pytest.raises(Oct2PyError):
            self.oc.__getattr__('_foo')
        with pytest.raises(Oct2PyError):
            self.oc.__getattr__('foo\W')

    def test_open_close(self):
        """Test opening and closing the Octave session
        """
        self.oc.exit()
        with pytest.raises(Oct2PyError):
            self.oc.push(name=['a'], var=[1.0])
        self.oc.restart()
        self.oc.push('a', 5)
        a = self.oc.pull('a')
        assert a == 5

    def test_struct(self):
        """Test Struct construct
        """
        test = Struct()
        test.spam = 'eggs'
        test.eggs.spam = 'eggs'
        assert test['spam'] == 'eggs'
        assert test['eggs']['spam'] == 'eggs'
        test["foo"]["bar"] = 10
        assert test.foo.bar == 10
        p = pickle.dumps(test)
        test2 = pickle.loads(p)
        assert test2['spam'] == 'eggs'
        assert test2['eggs']['spam'] == 'eggs'
        assert test2.foo.bar == 10
        assert 'spam' in test.__dict__

    def test_syntax_error(self):
        """Make sure a syntax error in Octave throws an Oct2PyError
        """
        with pytest.raises(Oct2PyError):
            self.oc.eval("a='1")

        if os.name == 'nt':
            self.oc.restart()

        with pytest.raises(Oct2PyError):
            self.oc.eval("a=1++3")

        if os.name == 'nt':
            self.oc.restart()

        self.oc.push('a', 1)
        a = self.oc.pull('a')
        assert a == 1

    def test_extract_figures(self):
        plot_dir = tempfile.mkdtemp().replace('\\', '/')
        code = """
        figure 1
        plot([1,2,3])
        figure 2
        temp=rand(100,100);
        imshow(temp)
        """
        self.oc.eval(code, plot_dir=plot_dir, plot_format='svg')
        imgs = self.oc.extract_figures(plot_dir)
        assert len(imgs) == 2
        assert isinstance(imgs[0], SVG) or isinstance(imgs[1], SVG)

    def test_quit(self):
        with pytest.raises(Oct2PyError):
            self.oc.eval("quit")
        self.oc.eval('a=1')

    def test_octave_error(self):
        with pytest.raises(Oct2PyError):
            self.oc.eval("a = ones2(1)")

    def test_keyword_arguments(self):
        self.oc.set(0, DefaultFigureColor='b', nout=0)
        plot_dir = tempfile.mkdtemp().replace('\\', '/')
        self.oc.plot([1, 2, 3], linewidth=3, plot_dir=plot_dir)
        assert self.oc.extract_figures(plot_dir)

    def test_octave_function(self):
        func = MatlabFunction([1])
        with pytest.raises(Oct2PyError):
            self.oc.push('x', func)

    def test_bad_getattr(self):
        self.oc.eval('foo = 1')
        with pytest.raises(Oct2PyError):
            self.oc.__getattr__('foo')

    def test_octave_class(self):
        self.oc.addpath(os.path.realpath(os.path.dirname(__file__)))
        polynomial = self.oc.polynomial
        p0 = polynomial([1, 2, 3])
        assert np.allclose(p0.poly, [[1, 2, 3]])

        p1 = polynomial([0, 1, 2])
        sobj = StringIO()
        hdlr = logging.StreamHandler(sobj)
        hdlr.setLevel(logging.DEBUG)
        self.oc.logger.addHandler(hdlr)
        self.oc.logger.setLevel(logging.DEBUG)
        p1.display(verbose=True, nout=0)
        text = hdlr.stream.getvalue().strip()
        self.oc.logger.removeHandler(hdlr)
        assert 'in poly display' in text

        self.oc.push('y', p0)
        p2 = self.oc.pull('y')
        assert np.allclose(p2.poly, [1, 2, 3])

        p2.poly = [2, 3, 4]
        assert np.allclose(p2.poly, [2, 3, 4])

        assert 'Display a polynomial object' in p2.display.__doc__

        self.oc.eval('p3 = polynomial([1,2,3])')
        p3 = self.oc.pull('p3')
        assert np.allclose(p3.poly, [1, 2, 3])

    def test_get_pointer(self):
        self.oc.addpath(os.path.realpath(os.path.dirname(__file__)))
        self.oc.push('y', 1)
        yptr = self.oc.get_pointer('y')
        assert yptr.name == 'y'
        assert yptr.value == 1
        assert yptr.address == 'y'
        yptr.value = 2
        assert yptr.value == 2
        assert self.oc.pull('y') == 2
        assert 'is a variable' in yptr.__doc__
        ones = self.oc.ones(yptr)
        assert ones.shape == (2, 2)

        onesptr = self.oc.get_pointer('ones')
        assert onesptr.name == 'ones'
        assert onesptr.address == '@ones'
        assert 'ones' in onesptr.__doc__

        sin = self.oc.get_pointer('sin')
        x = self.oc.quad(sin, 0, self.oc.pi())
        assert x == 2

        self.oc.eval('p = polynomial([1,2,3])')
        ppter = self.oc.get_pointer('p')
        assert ppter.name == 'p'
        assert ppter.address == 'p'
        p = ppter.value
        assert np.allclose(p.poly, [1, 2, 3])

        clsptr = self.oc.get_pointer('polynomial')
        value = clsptr([1, 2, 3])
        assert np.allclose(value.poly, [1, 2, 3])

        with pytest.raises(Oct2PyError):
            self.oc.get_pointer('foo123')

    def test_feval(self):
        self.oc.addpath(os.path.realpath(os.path.dirname(__file__)))
        a = self.oc.feval('ones', 3)
        assert np.allclose(a, np.ones((3, 3)))

        self.oc.feval('ones', 3, store_as='foo')
        b = self.oc.pull('foo')
        assert np.allclose(b, np.ones((3, 3)))

        self.oc.push('x', 3)
        ptr = self.oc.get_pointer('x')
        c = self.oc.feval('ones', ptr)
        assert np.allclose(c, np.ones((3, 3)))

        p = self.oc.polynomial([1, 2, 3])
        poly = self.oc.feval('get', p, 'poly')
        assert np.allclose(poly, [1, 2, 3])

        val = self.oc.feval('disp', self.oc.zeros)
        assert val.strip() == '@zeros'

        lines = []
        self.oc.feval('evalin', 'base', 'disp(1);disp(2);disp(3)',
                      nout=0,
                      stream_handler=lines.append)
        assert lines == [' 1', ' 2', ' 3'], lines

        val = self.oc.feval('svd', [[1, 2], [1, 3]])
        u, v, d = self.oc.feval('svd', [[1, 2], [1, 3]], nout=3)
        assert isinstance(val, np.ndarray)
        assert isinstance(u, np.ndarray)

        self.oc.feval('test_nodocstring.m', 1)
        with pytest.raises(TypeError):
            self.oc.feval('test_usage.py')

    def test_eval(self):
        a = self.oc.eval('ones(3);')
        assert np.allclose(a, np.ones((3, 3)))

        lines = []
        self.oc.eval('disp(1);disp(2);disp(3)',
                     nout=0,
                     stream_handler=lines.append)
        assert lines == [' 1', ' 2', ' 3'], lines

        a = self.oc.eval(['zeros(3);', 'ones(3);'])
        assert np.allclose(a, np.ones((3, 3)))

        U, S, V = self.oc.eval('svd(hilb(3))', nout=3)
        assert isinstance(U, np.ndarray)

    def test_no_args_returned(self):
        # Test a function that only works when nargout=0
        here = os.path.dirname(__file__)
        self.oc.source(os.path.join(here, 'roundtrip.m'))

    def test_script_error(self):
        here = os.path.dirname(__file__)
        with pytest.raises(Oct2PyError) as exec_info:
            self.oc.source(os.path.join(here, 'script_error.m'))
        msg = str(exec_info.value)
        assert msg == (
            "Octave evaluation error:\nerror: "
            "'b' undefined near line 2 column 3\nerror: called from:\n    script_error at line 2, column 2"
        )

    def test_pkg_load(self):
        self.oc.eval('pkg load signal')
        t = np.linspace(0, 1, num=100)
        x = np.cos(2*np.pi*t*3)
        # on Travis CI this is giving a dimension mismatch error
        try:
            y = self.oc.sgolayfilt(x, 3, 5)
        except Oct2PyError as e:
            if 'dimensions mismatch' in str(e):
                return
        assert y.shape == (1, 100)

    def test_passing_integer_args(self):
        self.oc.eval("""
function [res, a, b] = foo(a, b)
res = a * b;
end
""")
        res, a, b = self.oc.foo(np.nan, 2, nout=3)
        assert np.isnan(res)
        assert np.isnan(a)
        assert b == 2

    def test_carriage_return(self):
        self.oc.eval(r"disp('hi\rthere')")


"""
Structured array exercise
==========================

Creating a structured array.
"""

import numpy as np
x = np.zeros((10, 10, 4), dtype=np.int8)
x[:,:,0] = 1
x[:,:,1] = 2
x[:,:,2] = 3
x[:,:,3] = 4

# How to make a (10, 10) structured array with fields 'r', 'g', 'b', 'a',
# without copying?

# y = ...

assert (y['r'] == 1).all()
assert (y['g'] == 2).all()
assert (y['b'] == 3).all()
assert (y['a'] == 4).all()


#twoDimensionalListAndNumpyArray.py
import numpy as np
two_dim_list = [[1,2,3,4],[2,3,4,5]]
print(two_dim_list)
two_dim_array = np.array(two_dim_list)
print(two_dim_array)

for i in range(len(two_dim_list)):
    print(two_dim_list[i])
    print(two_dim_array[i])
    for j in range(len(two_dim_list[i])):
        print(two_dim_list[i][j])
        print(two_dim_array[i][j])


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [raw] raw_mimetype="text/restructuredtext"
#
# .. _polar: https://matplotlib.org/3.1.0/gallery/pie_and_polar_charts/polar_demo.html

# .. _cartopy: https://scitools.org.uk/cartopy/docs/latest/

# .. _basemap: https://matplotlib.org/basemap/index.html
#
# .. _ug_proj:
#
# Geographic and polar plots
# ==========================
#
# ProPlot includes features for working with `polar axes <polar_>`_
# and the `cartopy`_ and `basemap`_ map projection packages. These features
# are optional -- installation of cartopy and basemap are not required.
#
# To change the axes projection, pass ``proj='name'`` to
# `~proplot.ui.subplots`. To use different projections for different
# subplots, pass a dictionary of projection names with the subplot number as
# the key -- for example, ``proj={1: 'name'}``. The default "projection" is
# always `~proplot.axes.CartesianAxes`.

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_polar:
#
# Polar axes
# ----------
#
# To draw `polar axes <polar_>`_, pass ``proj='polar'`` or e.g. ``proj={1:'polar'}``
# to `~proplot.ui.subplots`. This generates a `~proplot.axes.PolarAxes`
# instance with its own `proplot.axes.PolarAxes.format` command. This
# command permits polar-specific modifications like changing the central radius `r0`,
# the zero azimuth location `theta0`, and the positive azimuthal direction `thetadir`.
# It also supports changing the radial and azimuthal limits `rlim` and `thetalim`,
# which can be used to make sector plots and annular plots.
#
# For details, see `proplot.axes.PolarAxes.format`.

# %%
import proplot as plot
import numpy as np
N = 200
state = np.random.RandomState(51423)
x = np.linspace(0, 2 * np.pi, N)
y = 100 * (state.rand(N, 5) - 0.3).cumsum(axis=0) / N
fig, axs = plot.subplots([[1, 1, 2, 2], [0, 3, 3, 0]], proj='polar')
axs.format(
    suptitle='Polar axes demo', linewidth=1, titlepad='1em',
    ticklabelsize=9, rlines=0.5, rlim=(0, 19),
)
for i in range(5):
    xi = x + i * 2 * np.pi / 5
    axs.plot(xi, y[:, i], cycle='FlatUI', zorder=0, lw=3)

# Standard polar plot
axs[0].format(
    title='Normal plot', thetaformatter='tau',
    rlabelpos=225, rlines=plot.arange(5, 30, 5),
    color='red8', tickpad='1em',
)

# Sector plot
axs[1].format(
    title='Sector plot', thetadir=-1, thetalines=90, thetalim=(0, 270), theta0='N',
    rlim=(0, 22), rlines=plot.arange(5, 30, 5),
)

# Annular plot
axs[2].format(
    title='Annular plot', thetadir=-1, thetalines=20, gridcolor='red',
    r0=-20, rlim=(0, 22), rformatter='null', rlocator=2
)

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_geo:
#
# Geographic axes
# ---------------
#
# ProPlot can turn any subplot into a geographic projection using
# the `cartopy`_ or `basemap`_ packages as "backends". To turn a subplot into
# a geographic projection, pass ``proj='name'`` or e.g. ``proj={2: 'name'}``
# (:ref:`see above <ug_proj>`) to  `~proplot.ui.subplots`
# where ``name`` is any valid :ref:`PROJ projection name <proj_included>`.
# You can also generate a `cartopy.crs.Projection` or `mpl_toolkits.basemap.Basemap`
# instance directly using the `~proplot.constructor.Proj` constructor function and
# pass the class instance with ``proj=<object>``.
#
# When you request a geographic projection,
# `~proplot.ui.subplots` returns instances of `proplot.axes.CartopyAxes`
# or `proplot.axes.BasemapAxes`, depending on whether ``basemap=True`` was used.
# Both of these derive from `proplot.axes.GeoAxes`, which includes a
# `~proplot.axes.GeoAxes.format` method. This method allows you to
# :ref:`modify geographic features <ug_geoformat>` with the same syntax whether
# cartopy or basemap is the "backend".
#
# * `proplot.axes.CartopyAxes` unifies the `cartopy.mpl.geoaxes.GeoAxes`
#   class with the `proplot.axes.Axes` class. The `~proplot.axes.GeoAxes.format`
#   method changes map bounds with
#   `~cartopy.mpl.geoaxes.GeoAxes.set_extent`, adds major and minor gridlines with
#   `~cartopy.mpl.geoaxes.GeoAxes.gridlines`, and adds geographic features with
#   `~cartopy.mpl.geoaxes.GeoAxes.add_feature`.
#
# * `~proplot.axes.BasemapAxes` redirects the plot, scatter, contour, contourf,
#   pcolor, pcolormesh, quiver, streamplot, and barb *axes methods* to
#   identically named methods on the `~mpl_toolkits.basemap.Basemap` instance.
#   This means you can work with *axes* plotting methods rather than the
#   basemap methods, just like cartopy. The `~proplot.axes.GeoAxes.format`
#   method adds major and minor gridlines with
#   `~mpl_toolkits.basemap.Basemap.drawmeridians` and
#   `~mpl_toolkits.basemap.Basemap.drawparallels` and adds geographic features
#   with commands like `~mpl_toolkits.basemap.Basemap.fillcontinents`
#   and `~mpl_toolkits.basemap.Basemap.drawcoastlines`. In case you need to
#   use it, the corresponding `~mpl_toolkits.basemap.Basemap` instance is
#   available via the `proplot.axes.BasemapAxes.projection` attribute.
#
# These features let you work with geographic data without having to invoke verbose
# cartopy classes like `~cartopy.crs.LambertAzimuthalEqualArea` and
# `~cartopy.feature.NaturalEarthFeature` or keep track of a separate
# `~mpl_toolkits.basemap.Basemap` object. They considerably reduce the amount of
# code needed to make geographic plots. In the below examples, we create a variety
# of geographic plots with both cartopy and basemap as the backends.
#
# .. note::
#
#    * ProPlot ensures that polar cartopy projections like
#      `~cartopy.crs.NorthPolarStereo` have circular boundaries (see `this example\
#      <https://scitools.org.uk/cartopy/docs/latest/gallery/always_circular_stereo>`__
#      from the cartopy website).
#    * By default, non-polar cartopy projections are forced to have global extent
#      with `~cartopy.mpl.geoaxes.GeoAxes.set_global` and polar cartopy projections
#      are bounded at the equator. This stands in contrast to the default cartopy
#      behavior, where map boundaries are determined automatically based on the
#      coordinates of the plotted content. To revert to cartopy's default behavior,
#      set :rcraw:`cartopy.autoextent` to ``True`` or pass ``autoextent=True``
#      to `~proplot.axes.CartopyAxes`.
#    * To make things more consistent between cartopy and basemap, the
#      `~proplot.constructor.Proj` constructor function lets you supply native
#      `PROJ <https://proj.org>`__ keyword names for the cartopy
#      `~cartopy.crs.Projection` classes (e.g. `lon_0` instead of `central_longitude`)
#      and instantiates `~mpl_toolkits.basemap.Basemap` projections with sensible
#      default PROJ parameters rather than raising an error when they are
#      omitted (e.g. ``lon_0=0`` as the default for most projections).
#
# .. warning::
#
#    Basemap is `no longer a maintained package\
#    <https://matplotlib.org/basemap/users/intro.html#cartopy-new-management-and-eol-announcement>`__.
#    However as shown below, gridline labels tend to look much nicer in basemap
#    than in cartopy -- especially when "inline" cartopy labels are disabled.
#    This is the main reason ProPlot continues to support both basemap and cartopy.
#    When cartopy catches up, basemap support may be deprecated.

# %%
# Simple figure with just one projection

# Option 1: Create a projection manually with plot.Proj()
# immport proplot as plot
# proj = plot.Proj('robin', lon_0=180)
# fig, axs = plot.subplots(nrows=2, axwidth=3, proj=proj)

# Option 2: Pass the name to 'proj' and keyword arguments to 'proj_kw'
import proplot as plot
fig, axs = plot.subplots(nrows=2, axwidth=3, proj='robin', proj_kw={'lon_0': 180})
axs.format(
    suptitle='Figure with single projection',
    coast=True, latlines=30, lonlines=60,
)

# %%
# Complex figure with different projections
import proplot as plot
fig, axs = plot.subplots(
    ncols=2, nrows=3,
    hratios=(1, 1, 1.4),
    basemap=(False, True, False, True, False, True),  # cartopy column 1
    proj=('cyl', 'cyl', 'hammer', 'hammer', 'npstere', 'npstere'),
)
axs.format(
    suptitle='Figure with several projections',
    collabels=['Cartopy projections', 'Basemap projections'],
    coast=True, latlines=20, lonlines=30,
    lonlabels='b', latlabels='r',  # or lonlabels=True, labels=True, etc.
)
axs[0, :].format(latlines=30, lonlines=60, labels=True)
plot.rc.reset()

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_geoplot:
#
# Plotting geographic data
# ------------------------
#
# In ProPlot, plotting in `~proplot.axes.GeoAxes` is not much different from
# plotting in `~proplot.axes.CartesianAxes`. ProPlot makes longitude-latitude
# (i.e. Plate Carrée) coordinates the *default* coordinate system for your
# datasets by passing ``transform=ccrs.PlateCarree()`` to cartopy plotting
# commands and ``latlon=True`` to basemap plotting commands. And again, basemap
# plotting commands are invoked from the `proplot.axes.GeoAxes` rather
# than from the `~mpl_toolkits.basemap.Basemap` instance.
#
# To ensure 2D plots like `~matplotlib.axes.Axes.contour` cover the entire globe,
# pass ``globe=True`` to the plotting command. This interpolates your data
# to the poles and across the longitude seams before plotting, having the same
# effect as cartopy's `~cartopy.util.add_cyclic_point` function and basemap's
# `~mpl_toolkits.basemap.addcyclic` function.
#
# Geographic feature can be drawn underneath data or on top of data by changing the
# corresponding `zorder <https://matplotlib.org/3.1.1/gallery/misc/zorder_demo.html>`__
# rc setting. For example, to draw land patches on top of all plotted content as
# a "land mask," use ``ax.format(land=True, landzorder=4)``.
# See the :ref:`next section <ug_geoformat>` for details.

# %%
import proplot as plot
import numpy as np

# Fake data with unusual longitude seam location and without coverage over poles
offset = -40
lon = plot.arange(offset, 360 + offset - 1, 60)
lat = plot.arange(-60, 60 + 1, 30)
state = np.random.RandomState(51423)
data = state.rand(len(lat), len(lon))

# Plot data both without and with globe=True
for globe in (False, True,):
    string = 'with' if globe else 'without'
    fig, axs = plot.subplots(
        ncols=2, nrows=2, axwidth=2.5,
        proj='kav7', basemap={(1, 3): False, (2, 4): True}
    )
    axs.format(
        suptitle=f'Geophysical data {string} global coverage',
        collabels=['Cartopy example', 'Basemap example'],
        rowlabels=['Contourf', 'Pcolormesh'],
        abc=True, abcstyle='a)', abcloc='ul', abcborder=False,
        coast=True, lonlines=90,
    )
    for i, ax in enumerate(axs):
        cmap = ('sunset', 'sunrise')[i % 2]
        if i < 2:
            m = ax.contourf(lon, lat, data, cmap=cmap, globe=globe, extend='both')
            fig.colorbar(m, loc='b', span=i + 1, label='values', extendsize='1.7em')
        else:
            ax.pcolor(lon, lat, data, cmap=cmap, globe=globe, extend='both')

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_geoformat:
#
# Formatting projections
# ----------------------
#
# `~proplot.axes.CartopyAxes` and `~proplot.axes.BasemapAxes` both derive from
# `proplot.axes.GeoAxes`, which provides a `~proplot.axes.GeoAxes.format` method.
# This can be used to draw "major" gridlines "minor" gridlines. Gridline locations
# and label formats can be configured with the `lonlocator`, `latlocator`,
# `lonformatter`, `latformatter`, `lonminorlocator`, and `latminorlocator` keywords.
# Major gridline labels and their positions can be configured with the `labels`,
# `lonlabels`, and `latlabels` keywords. Cartopy map bounds can be set with the
# `lonlim`, `latlim`, and `boundinglat` keywords. Geographic features like land masses,
# coastlines, and administrative borders can be toggled on and off and stylized with
# a variety of :ref:`rc settings <rc_proplot>`. Finally, `proplot.axes.GeoAxes.format`
# also calls `proplot.axes.Axes.format`, and so can be used to for subplot titles,
# a-b-c labels, and figure titles as before.
#
# For details, see the `proplot.axes.GeoAxes.format` documentation.

# %%
import proplot as plot
fig, axs = plot.subplots(
    [[1, 1, 2], [3, 3, 3]],
    axwidth=4, proj={1: 'eqearth', 2: 'ortho', 3: 'wintri'},
    wratios=(1, 1, 1.2), hratios=(1, 1.2),
)
axs.format(
    suptitle='Projection axes formatting demo',
    collabels=['Column 1', 'Column 2'],
    abc=True, abcstyle='A.', abcloc='ul', abcborder=False, linewidth=1.5
)

# Styling projections in different ways
ax = axs[0]
ax.format(
    title='Equal earth', land=True, landcolor='navy', facecolor='pale blue',
    coastcolor='gray5', borderscolor='gray5', innerborderscolor='gray5',
    gridlinewidth=1.5, gridcolor='gray5', gridalpha=0.5,
    gridminor=True, gridminorlinewidth=0.5,
    coast=True, borders=True, borderslinewidth=0.8,
)
ax = axs[1]
ax.format(
    title='Orthographic', reso='med', land=True, coast=True, latlines=10, lonlines=15,
    landcolor='mushroom', suptitle='Projection axes formatting demo',
    facecolor='petrol', coastcolor='charcoal', coastlinewidth=0.8, gridlinewidth=1
)
ax = axs[2]
ax.format(
    land=True, facecolor='ocean blue', landcolor='bisque', title='Winkel tripel',
    lonlines=60, latlines=15,
    gridlinewidth=0.8, gridminor=True, gridminorlinestyle=':',
    lonlabels=True, latlabels='r', loninline=True,
    gridlabelcolor='gray8', gridlabelsize='med-large',
)

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_zoom:
#
# Zooming into projections
# ------------------------
#
# To zoom into cartopy projections, use
# `~cartopy.mpl.geoaxes.GeoAxes.set_extent` or pass `lonlim`,
# `latlim`, or `boundinglat` to `~proplot.axes.GeoAxes.format`. The `boundinglat`
# keyword controls the circular latitude boundary for North Polar and
# South Polar Stereographic, Azimuthal Equidistant, Lambert Azimuthal
# Equal-Area, and Gnomonic projections. By default, ProPlot tries to use the
# degree-minute-second cartopy locators and formatters made available in cartopy
# 0.18. You can switch from minute-second subintervals to traditional decimal
# subintervals by passing ``dms=False`` to `~proplot.axes.GeoAxes.format`.
#
# To zoom into basemap projections, pass any of the `boundinglat`,
# `llcrnrlon`, `llcrnrlat`, `urcrnrlon`, `urcrnrlat`, `llcrnrx`, `llcrnry`,
# `urcrnrx`, `urcrnry`, `width`, or `height` keyword arguments to
# the `~proplot.constructor.Proj` constructor function either directly or via
# the `proj_kw` `~proplot.ui.subplots` keyword argument. You can also pass
# `lonlim` and `latlim` to `~proplot.constructor.Proj` and these arguments
# will be used for `llcrnrlon`, `llcrnrlat`, etc. You can not zoom into basemap
# projections with `format` after they have already been created.

# %%
import proplot as plot

# Plate Carrée map projection
plot.rc.reso = 'med'  # use higher res for zoomed in geographic features
proj = plot.Proj('cyl', lonlim=(-20, 180), latlim=(-10, 50), basemap=True)
fig, axs = plot.subplots(nrows=2, axwidth=5, proj=('cyl', proj))
axs.format(
    land=True, labels=True, lonlines=20, latlines=20,
    gridminor=True, suptitle='Zooming into projections'
)
axs[0].format(
    lonlim=(-140, 60), latlim=(-10, 50),
    labels=True, title='Cartopy example'
)
axs[1].format(title='Basemap example')

# %%
import proplot as plot

# Pole-centered map projections
proj = plot.Proj('npaeqd', boundinglat=60, basemap=True)
fig, axs = plot.subplots(ncols=2, axwidth=2.7, proj=('splaea', proj))
axs.format(
    land=True, latmax=80,  # no gridlines poleward of 80 degrees
    suptitle='Zooming into polar projections'
)
axs[0].format(boundinglat=-60, title='Cartopy example')
axs[1].format(title='Basemap example')

# %%
import proplot as plot

# Zooming in on continents
proj1 = plot.Proj('lcc', lon_0=0)  # cartopy projection
proj2 = plot.Proj('lcc', lon_0=-100, lat_0=45, width=8e6, height=8e6, basemap=True)
fig, axs = plot.subplots(ncols=2, axwidth=3, proj=(proj1, proj2))
axs.format(suptitle='Zooming into specific regions', land=True)
axs[0].format(lonlim=(-20, 50), latlim=(30, 70), title='Cartopy example')
axs[1].format(lonlines=20, title='Basemap example')

# Zooming to very small scale with degree-minute-second labels
plot.rc.reso = 'hi'
fig, axs = plot.subplots(ncols=2, axwidth=2.5, proj='cyl')
axs.format(
    land=True, labels=True,
    borders=True, borderscolor='white',
    suptitle='Degree-minute-second labels',
)
axs[0].format(lonlim=(-7.5, 2), latlim=(49.5, 59))
axs[1].format(lonlim=(-6, -2), latlim=(54.5, 58.5))
plot.rc.reset()

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _proj_included:
#
# Included projections
# --------------------
#
# The available `cartopy <https://scitools.org.uk/cartopy/docs/latest/>`__
# and `basemap <https://matplotlib.org/basemap/index.html>`__ projections are
# plotted below. See `~proplot.constructor.Proj` for a table of projection
# names with links to the relevant `PROJ <https://proj.org>`__ documentation.
#
# ProPlot uses the cartopy API to add the Aitoff, Hammer, Winkel Tripel, and
# Kavrisky VII projections (i.e. ``'aitoff'``, ``'hammer'``, ``'wintri'``,
# and ``'kav7'``), as well as North and South polar versions of the Azimuthal
# Equidistant, Lambert Azimuthal Equal-Area, and Gnomic projections (i.e.
# ``'npaeqd'``, ``'spaeqd'``, ``'nplaea'``, ``'splaea'``, ``'npgnom'``, and
# ``'spgnom'``), modeled after the existing `~cartopy.crs.NorthPolarStereo`
# and `~cartopy.crs.SouthPolarStereo` projections.

# %%
import proplot as plot

# Table of cartopy projections
projs = [
    'cyl', 'merc', 'mill', 'lcyl', 'tmerc',
    'robin', 'hammer', 'moll', 'kav7', 'aitoff', 'wintri', 'sinu',
    'geos', 'ortho', 'nsper', 'aea', 'eqdc', 'lcc', 'gnom',
    'npstere', 'nplaea', 'npaeqd', 'npgnom', 'igh',
    'eck1', 'eck2', 'eck3', 'eck4', 'eck5', 'eck6'
]
fig, axs = plot.subplots(ncols=3, nrows=10, width=7, proj=projs)
axs.format(
    land=True, reso='lo', labels=False,
    suptitle='Table of cartopy projections'
)
for proj, ax in zip(projs, axs):
    ax.format(title=proj, titleweight='bold', labels=False)

# %%
import proplot as plot

# Table of basemap projections
projs = [
    'cyl', 'merc', 'mill', 'cea', 'gall', 'sinu',
    'eck4', 'robin', 'moll', 'kav7', 'hammer', 'mbtfpq',
    'geos', 'ortho', 'nsper',
    'vandg', 'aea', 'eqdc', 'gnom', 'cass', 'lcc',
    'npstere', 'npaeqd', 'nplaea'
]
fig, axs = plot.subplots(ncols=3, nrows=8, basemap=True, width=7, proj=projs)
axs.format(
    land=True, labels=False,
    suptitle='Table of basemap projections'
)
for proj, ax in zip(projs, axs):
    ax.format(title=proj, titleweight='bold', labels=False)


![Callysto.ca Banner](https://github.com/callysto/curriculum-notebooks/blob/master/callysto-notebook-banner-top.jpg?raw=true)

import matplotlib.pyplot as plt
import plotly as py
import plotly.graph_objs as go
import numpy as np
import math
import ipywidgets as widgets
from IPython.display import display, Math, Latex, HTML, IFrame
from astropy.table import Table, Column
from ipywidgets import interact, interactive

py.offline.init_notebook_mode(connected=True)
%matplotlib inline

font = {'family' : 'sans-serif',
        'weight' : 'normal',
        'size'   : 14}

plt.rc('font', **font)

'''Above, we are importing all the necessary modules in order to run the notebook.
Numpy allows us to define arrays of values for our variables to plot them
matplotlib is what we use to create the figures
the display and widgets are to make the notebook look neat
'''

HTML('''
<form action="javascript:code_toggle()"><input type="submit" id="toggleButton" value="Show Code"></form>''')

***
# Light and Optics
***

Gif taken from https://giphy.com/gifs/fandor-sun-eclipse-3o7OsM9vKFH2ESl0KA/links, August 1st, 2018.
 Figure 1: For hundreds of years, scientists have tried to understand the nature of light. With advances in technology, and inventions like telescopes, we have been able to see farther than ever before.

***

## Introduction

Throughout most of history, humans did not understand light as we do today. As science and technology have progressed over time, so too has our knowledge of the nature of light.

In this lesson, when we say the word "light", we will be referring to visible light (light that comes from the sun, lightbulbs etc.). We will go over how a few key experiments kickstarted a new way of thinking, and a few of the ways that we are able to manipulate light. We will also talk about how our eyes enable us to see.

## Background

If you had to describe to someone what light is, you may have a hard time. Some people think of light as the absence of darkness, but even that doesn't say much about light itself.

Our understanding of light truly began around the 17th century, when a few individuals started to realize that light was not a mystical substance. Scientists (or "natural philosophers", as they were called during that time) recognized that certain properties of light were measurable, and that some properties could be manipulated. Sir Isaac Newton and Ole Rømer were among the first scientists to take a step in this direction.

> ### Isaac Newton's Prism Experiment

Sir Isaac Newton has made contributions to many fields of science and mathematics. In 1666, while spending time at his childhood home in Lincolnshire, England, Newton began experimenting with light.

Using a small slit in his window shutters, Newton passed a narrow beam of sunlight through a glass prism. The light travelled through the prism, and projected a rainbow of color on the other side!

Picture taken from http://media.web.britannica.com/eb-media/10/7710-050-36C066AC.jpg.
 Figure 2: This picture shows how a prism can create a spectrum of color. This is what Newton would have seen in 1666.

Later on, scientists determined that the prism was actually splitting light into its component parts. This phenomenon is called **dispersion**.

Through this experiment, Newton demonstrated that white light was actually made up of all the individual colors of the rainbow!

> ### Ole Rømer and the Speed of Light

For many years, people thought that if somebody lit a match, the light from that match would be instantly visible to everyone, no matter how far away they were. However, in 1676 Ole Rømer proved that this is not the case.

Rømer spent a long time studying the orbit of Io, one of Jupiter's moons. As part of his study, he began predicting the times when Io should be hidden behind Jupiter's shadow (these periods are called eclipses). However, Rømer saw that his predictions for when these eclipses should occur were not always accurate.

Gif taken from https://giphy.com/gifs/timelapse-DXIa1beDspYRy, August 1st, 2018.
 Figure 3: Here we can see Jupiter as it looks through a telescope. You might be able to see a black spot move from the left to the right across Jupiter's surface. This is actually one of Jupiter's many moons!

Rømer then realized that these errors may be because the distance between Io and the Earth was always changing. Rømer thought that when the distance between Io and the Earth increased, it might take a longer time for light coming from Io to reach Earth. If this were the case, then the light must be travelling at a finite speed!

After taking many measurements and using some clever mathematics, Rømer calculated the speed of light to be roughly 220,000,000 m/s, or 792,000,000 km/h.

Today, we have measured the speed of light to be 299,792,458 m/s. Although he was not exactly right, Rømer provided one of the first mathematical calculations for the speed of light.

***
Since the time of Rømer and Newton, scientists have made many new discoveries about the nature of light. While not all of these discoveries agree with one another, here are two things we know for sure:
- Light is made up of a spectrum of color
- Light travels at a speed of 299,792,458 m/s

Now let's talk about some of the ways we can manipulate light.
***

## Reflection

We are all familiar with reflection; chances are, you look at your reflection more than once a day. But have you ever stopped to wonder what is really going on?

Reflection is the term used to describe how light can change direction when it comes into contact with certain surfaces.

When incoming light rays encounter a reflective surface, they bounce off the surface and continue moving in a new direction. The new direction in which it moves is determined by the **law of reflection**.

\begin{equation}
\rm Law\: of\: Reflection: Angle\: of\: Incidence = Angle\: of\: Reflection
\end{equation}

On the animation below, click on the flashlight to turn it on, and move your mouse to change the angle of incidence.

#IFrame('Animations/reflect.html',width=500,height=320)

%%html

As seen above, the **normal** is what we call the line that forms a 90$^{\circ}$ angle with  the surface. The **angle of incidence** is what we call the angle between the flash lights beam and the normal. Similarly, the **angle of reflection** is the angle that the newly reflected light beam makes with the normal. The law of reflection states that these two angles will always be equal.

## Refraction

Have you ever tried to reach down and grab an object sitting at the bottom of a pool of water? If you have, you may have noticed that the object isn't actually in the location that you thought it was.

 Image taken from http://legacy.sciencelearn.org.nz/Contexts/Light-and-Sight/Sci-Media/Video/Refraction/(quality)/hi on August 3rd, 2018.
 Figure 4: When you are looking into a body of water from above, the objects you see beneath the surface are not actually where they appear to be.

This phenomenon occurs because the light travelling to your eyes from the bottom of the pool **refracts**, or changes its direction of travel, when it transitions from water to air.

The **index of refraction** is a value that we use to show how much light will bend when travelling through a substance. For example, the index of refraction for air is approximately 1.00, and the index of refraction for water is about 1.33. Because these indexes are different, light will bend when passing from water to air, or vice versa.

Use the animation below to see how light refracts when passing from air to water. Click on the flashlight to turn it on.

#IFrame('Animations/refract.html',width=520,height=320)

%%html

Mathematically, reflection can be described using the following equation, known as Snell's Law:

\begin{equation}
\textrm{Snells Law:}\: n_1\sin(\theta_1) = n_2\sin(\theta_2)
\end{equation}

where $n_1$ is the index of refraction for the first medium, $\theta_1$ is the incident angle, $n_2$ is the index of refraction for the second medium, and $\theta_2$ is the angle of refraction.

Light will bend *towards* the normal when travelling from a medium with a *lower* index of refraction to one with a *higher* index of refraction, and vice versa.

***
Some of the most beautiful sights in nature are caused by reflection and refraction. Here are a couple of examples:

### Rainbows

Rainbows are a result of both reflection and refraction. As its raining, each water droplet acts like a tiny prism, just like the one we saw in Figure 2. The water droplets split visible light into colors, and these colors are then reflected back towards our eyes.

 Image taken from https://waterstories.nestle-waters.com/environment/how-does-a-rainbow-form/ on August 3rd, 2018.
 Figure 5: Water droplets use reflection and refraction to create the beautiful rainbows that we see while it is raining.

### Mirages

Have you ever been driving on a sunny day, and up ahead it looks as though a stream of water is running across the road? You are really seeing a mirage.
Mirages also occur because of refraction, but they do not result in a display of color like a rainbow. This type of refraction occurs due to a difference in temperature between separate layers of air.

As we were describing before, refraction occurs when light travels from one substance to another. Well, it turns out that hot air and cold air are actually different enough to act as different substances. Therefore, light will refract when passing through one to the other.

 Image taken from https://edexcellence.net/articles/what-the-mirage-gets-wrong-on-teacher-development on August 3rd, 2018.
 Figure 6: Although it may look like water running across the road, it is actually a mirage. These commonly occur in desert areas, where the road can become very hot.

When you are looking at a mirage, it can look as though the air is wavy and fluid, which is why it is common to think that you are looking at water. This appearance occurs when layers of hot and cold air are mixing together, and light passing through these layers is constantly being refracted in different directions.

You may see a mirage appear on top of a hot roadway, behind the exhaust pipe of a plane or car, or around any other source of heat.

## Applications of Reflection and Refraction

### Lenses

If you have glasses, or contact lenses, then you are constantly using refraction in order to help you see! Lenses use refraction to point light in specific directions.

Generally speaking, there are two types of lenses: **convex** and **concave**.

To see how each type of lens affects light, use the following animation.

#IFrame('Animations/convex.html',width=520,height=420)

%%html

As seen above, a convex lens focuses light towards a specific point, while a concave lens will spread light away from a point. These lenses can be combined in many ways in order to produce different effects. For example, a camera lens uses a series of both convex and concave lenses in order to direct incoming light towards the back of the camera.

 Image taken from https://www.reddit.com/r/pic/comments/3o3b7w/camera_lens_cut_in_half/ on August 3rd, 2018.
 Figure 5: This is what the inside of a camera lens looks like. The photographer can adjust how they want the picture to look by changing the distance between the individual lenses.

## Vision

Our eyes are very complex organs, but the process that enables us to see is actually pretty simple. The basic steps are as follows:

1. Light enters the eye through the **pupil**
2. The convex **lens** behind the pupil directs incoming light towards the **retina**, which is like a screen at the back of our eye.
3. The retina then sends this image to the brain.
4. The brain then interprets the image.

 Image taken from https://openclipart.org/detail/261647/eye-diagram on August 3rd, 2018.
 Figure 6: This diagram shows some of the key components of the eye that enable us to see.

However, the image that is projected onto the retina is actually upside down!

 Image taken from https://www.eetimes.com/author.asp?section_id=14&doc_id=1282795 on August 3rd, 2018.
 Figure 7: The convex lens at the front of our eye actually flips images upside down.

So the retina actually sends an upside down image to the brain, and the brain automatically flips the image rightside up.

Use the following link to see an animation showing how a convex lens flips images upside down:https://phet.colorado.edu/sims/geometric-optics/geometric-optics_en.html.

## Technology & Inventions

### The Telescope

The first telescope was made by Hans Lippershey in 1608, but it was Galileo Galilee who became famous by using it for astronomy. There are many different types of telescopes, but they all use reflection and refraction to make far away objects appear closer.

A telescope uses a large opening to collect incoming light, and then directs this light towards your eye by using mirrors and lenses.

 Image taken from https://www.skyandtelescope.com/press-releases/tips-for-first-time-telescope-buyers/ on August 3rd, 2018.
 Figure 8: Telescopes come in many different shapes and sizes.

The reason why things look bigger when looking through a telescope is because of the lenses.

### The Microscope

- talk about invention of microscope
- why it works

## Conclusion

Our understanding of light is the result of hundreds of years of research and innovation. Along the way, we have created incredible new technologies that have allowed us to look further than ever before.

 Image taken from https://xenlife.com.au/hubble-space-telescope-important/ on August 3rd, 2018.
 Figure 9: The Hubble Space Telescope has shown us pictures of galaxies that are billions of light years away.

[![Callysto.ca License](https://github.com/callysto/curriculum-notebooks/blob/master/callysto-notebook-banner-bottom.jpg?raw=true)](https://github.com/callysto/curriculum-notebooks/blob/master/LICENSE.md)


import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def importarDados(filepath):
    path = os.getcwd() + filepath
    data = pd.read_csv(path, header=None)

    X = data.iloc[:, 0:-1].values
    y = data.iloc[:, -1:].values

    return X, y


"""

    Script for comparing two classifications

    Input:  - The STAR files with the classification to compare

    Output: - Table with the similitude (number of shared particle / number of particles per class) among classes

"""

__author__ = 'Antonio Martinez-Sanchez'

###### Global variables

import pyseg as ps

CLASS_COL, PART_COL = '_rlnClassNumber', '_rlnImageName'

########################################################################################
# PARAMETERS
########################################################################################

ROOT_PATH = '/fs/pool/pool-lucic2/antonio/workspace/psd_an/ex/syn/sub/relion/fils/pst_t/ch'

# Input star files to compares
in_star_1 = ROOT_PATH + '/pst_cont5/class_batch/run1_c8_5_it030_data.star'
in_star_2 = ROOT_PATH + '/pst_cont5/class_batch/run1_c8_15_it030_data.star'

########################################################################################
# MAIN ROUTINE
########################################################################################

################# Package import

import sys
import copy
import time
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.pyplot import cm

########## Print initial message

print('Similitude comparison between two classifications.')
print('\tAuthor: ' + __author__)
print('\tDate: ' + time.strftime("%c") + '\n')
print('Options:')
print('\tInput file 1: ' + in_star_1)
print('\tInput file 2: ' + in_star_2)
print('')

######### Process

print('Main Routine: ')

print('\tLoading input STAR file 1: ' + in_star_1)
star_1 = ps.sub.Star()
try:
    star_1.load(in_star_1)
except ps.pexceptions.PySegInputError as e:
    print('ERROR: input STAR file could not be read because of "' + str(e.msg, e.expr) + '"')
    sys.exit(-1)
classes_1 = list(set(star_1.get_column_data(CLASS_COL)))
print('\t\t-' + str(len(classes_1)) + ' classes found: ' + str(classes_1))

print('\tLoading input STAR file 2: ' + in_star_2)
star_2 = ps.sub.Star()
try:
    star_2.load(in_star_2)
except ps.pexceptions.PySegInputError as e:
    print('ERROR: input STAR file could not be read because of "' + str(e.msg, e.expr) + '"')
    sys.exit(-1)
classes_2 = list(set(star_2.get_column_data(CLASS_COL)))
print('\t\t-' + str(len(classes_1)) + ' classes found: ' + str(classes_1))

print('\tComputing similitude matrix (1 to 2)...')
nparts_1 = np.zeros(shape=max(classes_1)+1, dtype=np.int)
sim_mat_1 = np.zeros(shape=(max(classes_1)+1,max(classes_2)+1), dtype=np.float)
parts_1, pclasses_1 = star_1.get_column_data(PART_COL), star_1.get_column_data(CLASS_COL)
parts_2, pclasses_2 = star_2.get_column_data(PART_COL), star_2.get_column_data(CLASS_COL)
for part_1, pclass_1 in zip(parts_1, pclasses_1):
    nparts_1[pclass_1] += 1
    try:
        idx = parts_2.index(part_1)
        pclass_2 = pclasses_2[idx]
        sim_mat_1[pclass_1, pclass_2] += 1
    except ValueError:
        pass
for i in range(1, sim_mat_1.shape[0]):
    sim_mat_1[i, :] /= float(nparts_1[i])

print('\tComputing similitude matrix (2 to 1)...')
nparts_2 = np.zeros(shape=max(classes_2)+1, dtype=np.int)
sim_mat_2 = np.zeros(shape=(max(classes_2)+1,max(classes_1)+1), dtype=np.float)
for part_2, pclass_2 in zip(parts_2, pclasses_2):
    nparts_2[pclass_2] += 1
    try:
        idx = parts_1.index(part_2)
        pclass_1 = pclasses_1[idx]
        sim_mat_2[pclass_2, pclass_1] += 1
    except ValueError:
        pass
for i in range(1, sim_mat_2.shape[0]):
    sim_mat_2[i, :] /= float(nparts_2[i])

print('\tSIMILITUDE FOR 1 TO 2 CASE:')
for i in range(1, sim_mat_1.shape[0]):
    print('\t\t-CLASS ' + str(classes_1[i-1]) + ': ' + str(sim_mat_1[i, 1:]))
    print('\t\t\t+Max: ' + str(sim_mat_1[i, 1:].max()))
    print('\t\t\t+Min: ' + str(sim_mat_1[i, 1:].min()))
    print('\t\t\t+Mean: ' + str(sim_mat_1[i, 1:].mean()))
    print('\t\t\t+Std: ' + str(sim_mat_1[i, 1:].std()))
    print('\t\t\t+Sum: ' + str(int(sim_mat_1[i, 1:].sum()) * nparts_1[i]))

print('\tSIMILITUDE FOR 2 TO 1 CASE:')
for i in range(1, sim_mat_2.shape[0]):
    print('\t\t-CLASS ' + str(classes_2[i-1]) + ': ' + str(sim_mat_2[i, 1:]))
    print('\t\t\t+Max: ' + str(sim_mat_2[i, 1:].max()))
    print('\t\t\t+Min: ' + str(sim_mat_2[i, 1:].min()))
    print('\t\t\t+Mean: ' + str(sim_mat_2[i, 1:].mean()))
    print('\t\t\t+Std: ' + str(sim_mat_2[i, 1:].std()))
    print('\t\t\t+Sum: ' + str(int(sim_mat_2[i, 1:].sum()) * nparts_2[i]))

print('Terminated. (' + time.strftime("%c") + ')')