%matplotlib widget

import numpy as np
import matplotlib.pyplot as plt
import temgymlite

from IPython.display import display
import ipywidgets

# controls
style = {
    'description_width': 'initial',
}

layout = ipywidgets.Layout(width="300px",height="30px")

add_traces_switch = ipywidgets.Checkbox(
    value=False,
    description="show ray traces",
    style=style,
    layout = layout,
)

tilted_beams_switch = ipywidgets.Checkbox(
    value=False,
    description="show tilted beams",
    style=style,
    layout = layout,
)

invert_source_switch = ipywidgets.Checkbox(
    value=False,
    description="invert STEM source",
    style=style,
    layout = layout,
)


# toggle rays
def toggle_rays(change):
    show_traces = change['new']
    # TEM
    for ray in ax.lines[:39] + ax.lines[42:45] + ax.lines[48:54]:
        ray.set_visible(show_traces)
    # STEM
    for stem_ray in ax_stem.lines[:33] + ax_stem.lines[37:39] + ax_stem.lines[40:41]:
        stem_ray.set_visible(show_traces)
    fig.canvas.draw_idle()
    return None

add_traces_switch.observe(toggle_rays,names='value')

# toggle beams
def toggle_tilted_beams(change):
    tilted_beams = change['new']
    for beam in ax.collections[:12]:
        beam.set_visible(tilted_beams)
    fig.canvas.draw_idle()
    return None

tilted_beams_switch.observe(toggle_tilted_beams,names='value')

# invert STEM source
def invert_stem_source(change):
    ax_stem.invert_yaxis()
    fig.canvas.draw_idle()
    return None
    
invert_source_switch.observe(invert_stem_source,names='value')

# tem
components = [
    temgymlite.Lens(name="Condenser Lens", z=1.25, f=-0.25),
    temgymlite.Sample(name="Sample", z=0.75),
    temgymlite.Lens(name="Objective Lens", z=0.5, f=-0.205),
    temgymlite.Aperture(
        name="Objective Aperture", z=0.2875, aperture_radius_inner=0.0875
    ),
]

components_scattered_left = [
    temgymlite.Lens(name="Condenser Lens", z=1.25, f=-0.25),
    temgymlite.DoubleDeflector(
        name="Deflector",
        z_up=1.05,
        z_low=0.95,
        updefx=0.3,
        lowdefx=-0.6,
    ),
    temgymlite.Sample(name="Sample", z=0.75),
    temgymlite.Lens(name="Objective Lens", z=0.5, f=-0.205),
    temgymlite.Aperture(
        name="Objective Aperture", z=0.2875, aperture_radius_inner=0.0875
    ),
]

components_scattered_right = [
    temgymlite.Lens(z=1.25, f=-0.25),
    temgymlite.DoubleDeflector(
        z_up=1.05,
        z_low=0.95,
        updefx=-0.3,
        lowdefx=0.6,
    ),
    temgymlite.Lens(z=0.5, f=-0.205),
]

unscattered_model = temgymlite.Model(
    components,
    beam_z=1.5,
    beam_type="x_axial",
    num_rays=3,
    gun_beam_semi_angle=0.35,
)

scattered_model_left = temgymlite.Model(
    components_scattered_left,
    beam_z=1.5,
    beam_type="x_axial",
    num_rays=3,
    gun_beam_semi_angle=0.35,
)

scattered_model_right = temgymlite.Model(
    components_scattered_right,
    beam_z=1.5,
    beam_type="x_axial",
    num_rays=3,
    gun_beam_semi_angle=0.35,
)

# stem
components_stem = [
    temgymlite.Aperture(
        name="Objective Aperture", z=1.5 - 0.2875, aperture_radius_inner=0.0875
    ),
    temgymlite.Lens(name="Objective Lens", z=1, f=-0.205),
    temgymlite.Sample(name="Sample", z=0.75),
]

bf_stem_center = temgymlite.Model(
    components_stem,
    beam_z=1.5,
    beam_type="x_axial",
    num_rays=3,
    gun_beam_semi_angle=0.35,
)

bf_stem_left = temgymlite.Model(
    components_stem,
    beam_z=1.5,
    beam_type="x_axial",
    num_rays=2,
    gun_beam_semi_angle=0.2,
)

bf_stem_right = temgymlite.Model(
    components_stem,
    beam_z=1.5,
    beam_type="x_axial",
    num_rays=1,
    gun_beam_semi_angle=0.2,
)

bf_stem = temgymlite.Model(
    components_stem,
    beam_z=1.5,
    beam_type="x_axial",
    num_rays=9,
    gun_beam_semi_angle=0.35,
)

# Static Figure

with plt.ioff():
    dpi = 72
    fig, (ax, ax_stem) = plt.subplots(1,2, figsize=(675/dpi, 450/dpi), dpi=dpi)

fig, ax = temgymlite.show_matplotlib(
    scattered_model_left,
    figax=(fig, ax),
    label_fontsize=12,
    plot_rays=True,
    fill_color="purple",
    ray_color="purple",
    fill_between=True,
    highlight_edges=False,
    show_labels=True,
    fill_alpha=0.5,
)

fig, ax = temgymlite.show_matplotlib(
    scattered_model_right,
    figax=(fig, ax),
    label_fontsize=12,
    plot_rays=True,
    fill_color="orange",
    ray_color="orange",
    fill_between=True,
    highlight_edges=False,
    show_labels=False,
    fill_alpha=0.5,
)

fig, ax = temgymlite.show_matplotlib(
    unscattered_model,
    figax=(fig, ax),
    label_fontsize=12,
    plot_rays=True,
    fill_color="green",
    ray_color="green",
    fill_between=True,
    highlight_edges=False,
    fill_alpha=0.5,
    show_labels=False,
)

ax.set_title("HRTEM", fontsize=16)
ax.set_xlim([-0.3, 0.3])

fig, ax_stem = temgymlite.show_matplotlib(
    bf_stem,
    figax=(fig, ax_stem),
    label_fontsize=12,
    plot_rays=False,
    fill_color="gray",
    fill_between=True,
    fill_alpha=0.5,
    highlight_edges=False,
    show_labels=True,
)

fig, ax_stem = temgymlite.show_matplotlib(
    bf_stem_center,
    figax=(fig, ax_stem),
    label_fontsize=12,
    plot_rays=True,
    ray_lw=1,
    ray_color="green",
    fill_between=False,
    highlight_edges=False,
    show_labels=False,
)

fig, ax_stem = temgymlite.show_matplotlib(
    bf_stem_left,
    figax=(fig, ax_stem),
    label_fontsize=12,
    plot_rays=True,
    ray_lw=1,
    ray_color="orange",
    fill_between=False,
    highlight_edges=False,
    show_labels=False,
)

fig, ax_stem = temgymlite.show_matplotlib(
    bf_stem_right,
    figax=(fig, ax_stem),
    label_fontsize=12,
    plot_rays=True,
    ray_lw=1,
    ray_color="purple",
    fill_between=False,
    highlight_edges=False,
    show_labels=False,
)

ax_stem.set_title("BF-STEM", fontsize=16)
ax_stem.set_xlim([-0.3, 0.3])

fig.canvas.resizable = False
fig.canvas.header_visible = False
fig.canvas.footer_visible = False
fig.canvas.toolbar_visible = True
fig.canvas.layout.width = '675px'
fig.canvas.layout.height = '480px'
fig.canvas.toolbar_position = 'bottom'
fig.tight_layout()

toggle_rays({'new':False})
toggle_tilted_beams({'new':False})

display(
    ipywidgets.VBox([
        fig.canvas,
        ipywidgets.HBox(
            [
                add_traces_switch,
                tilted_beams_switch,
                invert_source_switch
            ],
            layout=ipywidgets.Layout(justify_content="center",width="680px")
        )
    ])
)