%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")

# ctem_switch = ipywidgets.Checkbox(
#     value=True,
#     description="CTEM ray diagram",
#     style=style,
#     layout = layout,
# )

# stem_switch = ipywidgets.Checkbox(
#     value=False,
#     description="STEM ray diagram",
#     style=style,
#     layout = layout,
# )

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

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

invert_source_switch = ipywidgets.Checkbox(
    value=True,
    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
    # if stem_switch.value:
    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')

# # toggle STEM
# def toggle_stem(change):
#     show_stem = change['new']
#     for stem_el in ax_stem.get_children():
#         stem_el.set_visible(show_stem)

#     if show_stem:
#         ax_stem.spines["top"].set_visible(False)
#         ax_stem.spines["right"].set_visible(False)
#         ax_stem.spines["bottom"].set_visible(False)
#         ax_stem.spines["left"].set_visible(False)
        
#         if not add_traces_switch.value:
#             for stem_ray in ax_stem.lines[:33] + ax_stem.lines[37:39] + ax_stem.lines[40:41]:
#                 stem_ray.set_visible(False)
                
#     fig.canvas.draw_idle()
#     return None
    
# stem_switch.observe(toggle_stem,names='value')

# # toggle CTEM
# def toggle_ctem(change):
#     show_tem = change['new']
#     for tem_el in ax.get_children():
#         tem_el.set_visible(show_tem)

#     if show_tem:
#         ax.spines["top"].set_visible(False)
#         ax.spines["right"].set_visible(False)
#         ax.spines["bottom"].set_visible(False)
#         ax.spines["left"].set_visible(False)
#     fig.canvas.draw_idle()
#     return None
    
# ctem_switch.observe(toggle_ctem,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("CTEM", 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])
ax_stem.invert_yaxis()

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})
# toggle_stem({'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")
        )
    ])
)