import gradio as gr
import matplotlib.pyplot as plt
#from fraunhofer_test import generate_diff
from zern_generator import Zernike
from fraunhofer_test import Fraunhofer
css = """
.container {
    # height: 50vh;
    # width: 100%;
    # overflow-x: auto !important;
    # overflow-y: auto !important;
    # scrollbar-width: thin !important;
}
"""

zerns = {}


def update_zerns(key, value):
    global zerns
    zerns[key] = value


def get_zern_from_file(file):
    with open(file) as f:
        coeffs = [float(l) for l in f]
    return tuple(coeffs)


zernike = Zernike()
fraunhofer = Fraunhofer(zernike)

theme = gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.pink)
with gr.Blocks(theme=theme, css=css) as demo:
    with gr.Row():
        zern_sliders = []

        with gr.Column(scale=1):
            # file_inp = gr.File(label="Загрузка коэффициентов из файла", height="2em")
            with gr.Tab(label="Генератор"):
                
                defocus_amt = gr.Number(label="Величина дефокуса", value=2, interactive=True)
                with gr.Row():
                    num_zernikes = gr.Number(label="Количество коэффициентов Цернике", value=6, interactive=True)
                    show_zernikes_btn = gr.Button("Показать")
                sliders_container = gr.Column(elem_classes=["container"])
            
                @gr.render(inputs=[num_zernikes], triggers=[show_zernikes_btn.click])
                def render_count(count):
                    global zerns
                    zerns = {}
                    sliders_container.children.clear()
                    count = min(count, 100)
                    for i in range(count):
                        term_name = f" - {zernike.term[i]}" if i < len(zernike.term) else ""
                        sld = gr.Slider(key=i, label=f"Zernike {i+1}{term_name}", minimum=-100, interactive=True,
                                        step=0.01, value=0)
    
                        sld.change(update_zerns, inputs=[gr.Number(value=i, visible=False), sld])
                        sliders_container.add_child(sld)
    
                    return gr.update()
            
            with gr.Tab(label="Настройки"):
                resolution_num = gr.Number(value=256, label="Разрешение")
           
            # file_inp.upload(get_zern_from_file, inputs=[file_inp], outputs=[*sliders_container.children])
            
        with gr.Column(scale=5):
            with gr.Tab(label="Экран"):
                with gr.Row():
                    plot_2d = gr.Plot()
                    plot_3d = gr.Plot()
                btn = gr.Button("Рассчитать")
            with gr.Tab(label="Тестовые экраны"):
                plot_batch = gr.Plot(label="Тестовые экраны (ANSI scheme)")
                btn_batch = gr.Button("Рассчитать тестовые экраны")
            with gr.Tab(label="Разница у фокуса"):
                plot_diff = gr.Plot(label="Разница в фокусе")
                btn_diff = gr.Button("Рассчитать разницу в фокусе")

            def on_button_click(resolution):
                return zernike.generate_zern_wavefront_fig(*[zerns[key] for key in sorted(zerns)])
            
            def on_settings_click(resolution):
                print(resolution)
                zernike.set_image_params(npix=resolution)
            
            def on_diff_button_click(defocus_amount):
                pl, mn, diff = fraunhofer.generate_diff(defocus_amount, *[zerns[key] for key in sorted(zerns)])
                
                fig, axs = plt.subplots(1, 3, figsize = (20, 10))
                axs[0].imshow(pl, cmap='grey')
                axs[1].imshow(mn, cmap='grey')
                axs[2].imshow(diff, cmap='grey')
                axs[0].axis('off')
                axs[0].set_title(f"+{defocus_amount}")
                axs[1].axis('off')
                axs[1].set_title(f"-{defocus_amount}")
                axs[2].axis('off')
                axs[2].set_title(f"Difference")
                return fig


            btn.click(on_button_click, inputs=[resolution_num], outputs=[plot_2d, plot_3d])
            btn_batch.click(zernike.generate_batch, inputs=[resolution_num], outputs=[plot_batch])
            resolution_num.change(on_settings_click, inputs=[resolution_num])
            btn_diff.click(on_diff_button_click, inputs=[defocus_amt], outputs=[plot_diff])

demo.launch()