import os
import sys
import glob
import numpy as np

try:
    import tklib.tkimport as imp
except Exception as e:
    print()
    print("######################################################################")
    print("###########  ERROR ERROR ERROR ERROR ERROR ERROR #####################")
    print("######################################################################")
    print(f"# Failed to import [tklib.tkimport] module ({e}).")
    print(f"#  Add [tkProg]{os.sep}tklib{os.sep}python to PYTHONPATH variable")
    print(f"#  Current PYTHONPATH:", sys.path)
    print("######################################################################")
    input("Press ENTER to terminate>>")
    exit()

np  = imp.import_lib("numpy", stop_by_error = False)
mpl = imp.import_lib("matplotlib", stop_by_error = False)
tk  = imp.import_lib("tkinter", stop_by_error = False)
from numpy import sin, cos, tan, pi, exp, log, sqrt
from matplotlib import pyplot as plt
import matplotlib.patches as patches
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
from tkinter import ttk
from tkinter import simpledialog
imp.messages(stop_by_error = True)


from tklib.tkutils import print_data, pint, pfloat, index2val, pconv_by_type
from tklib.tkapplication import tkApplication
from tklib.tkparams import tkParams
from tklib.tkvariousdata import tkVariousData
from tklib.tkgraphic.tkplotevent import tkPlotEvent, RangeSelector
from tklib.tkgraphic.tkplot_pyplot import select_plt, tkPlot_pyplot
from tklib.tkgraphic.tkplot_tkinter import tkPlot_tkinter
from tklib.tkgui.tksimple_gui import get_window_from_plt, CustomDialog_with_config
from tklib.tkwrapper.tktqdm import tqdm



"""
TFT simulation
"""


#===================================
# physical constants
#===================================
pi   = 3.14159265358979323846
pi2  = 2.0 * pi
h    = 6.6260755e-34    # Js";
hbar = 1.05459e-34      # "Js";
c    = 2.99792458e8     # m/s";
e    = 1.60218e-19      # C";
e0   = 8.854418782e-12; # C<sup>2</sup>N<sup>-1</sup>m<sup>-2</sup>";
kB   = 1.380658e-23     # JK<sup>-1</sup>";
me   = 9.1093897e-31    # kg";
R    = 8.314462618      # J/K/mol
a0   = 5.29177e-11      # m";


def initialize():
    app = tkApplication()
    app.config.plugin_dir = app.replace_path(None, template = ["{dirname}", "plugin"])
    appconfig_path, config = app.read_app_config(print_level = 1)

    cfg = tkParams()
    app.cparams = cfg

    app.add_argument(opt = "--mode", type = "str", opt_str = "--mode=[model|model3d|IV|Vch|plotz|plotx]", defval = "IV")
    app.add_argument(opt = "--use_tkplt", type = "int", defval = 0)

    app.add_argument(opt = "--mu", type = "float", defval = 10.0e-4) # m^2/Vs

    app.add_argument(opt = "--EV", type = "float", defval = 0.0)
    app.add_argument(opt = "--EC", type = "float", defval = 3.0)
    app.add_argument(opt = "--ND", type = "float", defval = 1.0e15)
    app.add_argument(opt = "--ED", type = "float", defval = 2.95)
    app.add_argument(opt = "--NA", type = "float", defval = 1.0e10)
    app.add_argument(opt = "--EA", type = "float", defval = 0.05)

    app.add_argument(opt = "--dg",  type = "float", defval = 100.0e-9)  #m
    app.add_argument(opt = "--erg", type = "float", defval = 11.9)

    app.add_argument(opt = "--W", type = "float", defval = 300.0e-6) #m
    app.add_argument(opt = "--L", type = "float", defval =  50.0e-6) #m
    app.add_argument(opt = "--Vth", type = "float", defval =  0.0)

    app.add_argument(opt = "--Vg_min", type = "float", defval =  0.0) 
    app.add_argument(opt = "--Vg_max", type = "float", defval = 20.0)
    app.add_argument(opt = "--nVg",    type = "int",   defval = 51)

    app.add_argument(opt = "--Vd_min", type = "float", defval =  0.0) 
    app.add_argument(opt = "--Vd_max", type = "float", defval = 20.0)
    app.add_argument(opt = "--nVd",    type = "int",   defval = 51)

#===================================
# figure configuration
#===================================
    cfg.figsize = (8, 6)
    cfg.fontsize        = 12
    cfg.legend_fontsize = 6

    return app, cfg
    
def simulate_IV(app, cfg, pause = True):
    print()
    print("#==============================================")
    print("# TFT simulator")
    print("#==============================================")
    cfg.print_parameters()

    Cox = cfg.erg * e0 / cfg.dg  #(F/m^2)
    print("")
    print(f"Cox = {Cox:12.6g} [F/m^2]")

    eps = 1.0e-6
    K = Cox * cfg.mu * cfg.W / cfg.L
    Vd_list = np.linspace(cfg.Vd_min, cfg.Vd_max + eps, cfg.nVd)
    Vg_list = np.linspace(cfg.Vg_min, cfg.Vg_max + eps, cfg.nVg)
    Id_VdVg_list = np.zeros([cfg.nVd, cfg.nVg])
    Id_VgVd_list = np.zeros([cfg.nVg, cfg.nVd])
    Vp_list = []
    Isat_list = []
    Vd_max = max(Vd_list)
    print("calculate Id-Vg-Vd list:")
    for ig, Vg in tqdm(enumerate(Vg_list), total = len(Vg_list)):
#    for ig in tqdm(range(len(Vg_list))):
#        Vg = Vg_list[ig]
        Vp = Vg - cfg.Vth
        Isat = K * 0.5 * Vp * Vp
        if Vp <= Vd_max:
            Vp_list.append(Vp)
            Isat_list.append(Isat)
        for id, Vd in enumerate(Vd_list):
            if Vd >= Vp:
                Vd = Vp
                I = Isat
            else:
                I = K * (Vp - 0.5 * Vd) * Vd
            
            Id_VdVg_list[id, ig] = I
            Id_VgVd_list[ig, id] = I

    outfile = 'TFT_IV.xlsx'
    print()
    print(f"Save IV characteristics to {outfile}")
    labels = ["Vd [V]", "Vg [V]"] + [f"Id@Vg={Vg} [A]" for Vg in Vg_list] 
    data_list = [Vd_list, Vg_list, *Id_VgVd_list]

    tkVariousData().to_excel(outfile, labels = labels, data_list = data_list)

    print()
    print("Plot")
    tkplt, root, tkpyplot = select_plt(use_tkinter = cfg.use_tkplt, plt = plt, parent = None, title = "Matplotlib with Tkinter")
    fig, axes = tkplt.subplots(1, 2, figsize = cfg.figsize, dpi = 100, tight_layout = True)
    if cfg.use_tkplt:
        tkplt.add_toolbar()
        canvas = tkplt.add_canvas(fig)

    plot_event = tkPlotEvent(plt)
    plot_event.prepare_annotation()
    plot_event.prepare_move_points()
    plot_event.prepare_popup_menu(fig, parent = root)
    plot_event.prepare_move_text(fig)

    def inf_format(plotevent, sel, iline, idata, label, inf, x = 0, y = 0, opt = None):
        if inf is None:
            return

        args = {}
        for key, val in inf.items():
            args[key] = val[idata]
        text = "Vg={Vg:8.4f} Vd={Vd:8.4f} Id={Id:10.4g}".format(**args)
        print(text)

    def annotation_format(plotevent, sel, iline, idata, label, inf, x, y, opt = None):
        text = f"{label} Vg={x:g} Id={y:g}"
        sel.annotation.set(text = text)
#        print(text)

    def inf_format_Vp(plotevent, sel, iline, idata, label, inf, x = 0, y = 0, opt = None):
        if inf is None:
            return

        args = {}
        for key, val in inf.items():
            args[key] = val[idata]
        text = "Vp={Vp:8.4f} V  Isat={Isat:10.4g} A".format(**args)
        print(text)

    def annotation_format_Vp(plotevent, sel, iline, idata, label, inf, x = 0, y = 0, opt = None):
        text = f"Vp={x:g} Isat={y:g}"
        sel.annotation.set(text = text)
#        print(text)

    ax = axes[0]
    ax.tick_params(labelsize = cfg.fontsize)
    ax.set_title("Transfer characteristics")
    frac = 0.9
    for id, Vd in enumerate(Vd_list):
        label = f"Vd = {Vd:g}"
        _Vd_list = [Vd] * cfg.nVd
        color = 'blue'  if id <= 4 else 'black'
        linewidth = 1.0 if id <= 4 else 0.5

        line, = ax.plot(Vg_list, Id_VdVg_list[id], label = label, linestyle = '-', color = color, linewidth = linewidth)
        plot_event.annotation.add_line(label, ax, ax, Vg_list, Id_VdVg_list[id], line, 
                    inf_list = {"Vg": Vg_list, "Vd": _Vd_list, "Id": Id_VdVg_list[id]},
#                    annotation_format = "{label} Vg={x:g} Id={y:g}",
#                    inf_format = "Vg={Vg:8.4f} Vd={Vd:8.4f} Id={Id:10.4g}")
                    annotation_format = annotation_format, inf_format = inf_format)

        if id <= 4:
            plot_event.move_text.add_annotation(ax, ax, x_list = Vg_list, y_list = Id_VdVg_list[id], frac = frac, 
                     text = label, fontsize = 10, ha = 'right', va = 'center', color = color, fc = "w", ec = "none")

    ax.set_xlabel("Vg [V]", fontsize = cfg.fontsize)
    ax.set_ylabel("Id [A]", fontsize = cfg.fontsize)
    ax.set_yscale("log")

    ax = axes[1]
    ax.tick_params(labelsize = cfg.fontsize)
    ax.set_title("Output characteristics")
    ntext = 5
    nskip = int(len(Vg_list) / ntext + 1.0e-4)
    for ig, Vg in enumerate(Vg_list):
        label = f"Vg = {Vg:g}"
        _Vg_list = [Vg] * cfg.nVg
        color = 'blue' if ig % nskip == 0 else 'black'
        linewidth = 1.0 if ig % nskip == 0 else 0.5

        line, = ax.plot(Vd_list, Id_VgVd_list[ig], label = label, linestyle = '-', color = color, linewidth = linewidth)
        plot_event.annotation.add_line(label, ax, ax, Vd_list, Id_VgVd_list[ig], line, 
                    inf_list = {"Vg": _Vg_list, "Vd": Vd_list, "Id": Id_VgVd_list[ig]},
#                    annotation_format = "{label} Vd={x:g} Id={y:g}",
#                    inf_format = "Vg={Vg:8.4f} Vd={Vd:8.4f} Id={Id:10.4g}")
                    annotation_format = annotation_format, inf_format = inf_format)

        if ig % nskip == 0:
            plot_event.move_text.add_annotation(ax, ax, x_list = Vd_list, y_list = Id_VgVd_list[ig], frac = frac, 
                     text = label, fontsize = 10, ha = 'right', va = 'center', color = color, fc = "w", ec = "none")

    line, = ax.plot(Vp_list, Isat_list, label = f"Vp", linestyle = 'dashed', color = 'red', linewidth = 0.5)
    plot_event.annotation.add_line("Vp", ax, ax, Vp_list, Isat_list, line, 
                    inf_list = {"Vp": Vp_list, "Isat": Isat_list},
#                    annotation_format = "Vp={x:g} Isat={y:g}",
#                    inf_format = "Vp={Vp:8.4f} V  Isat={Isat:10.4g} A")
                    annotation_format = annotation_format_Vp, inf_format = inf_format_Vp)

    ax.set_xlabel("Vd [V]", fontsize = cfg.fontsize)
    ax.set_ylabel("Id [A]", fontsize = cfg.fontsize)
#    ax.legend()

    plot_event.register_annotation_event(fig, activate = False, print_level = 0)
    plot_event.register_move_text_event(fig, activate = False)
    plot_event.register_follow_mouse_event(fig, activate = False, print_level = 1)
    for ax in axes:
        plot_event.follow_mouse.add(ax, marker = '*', size = 10.0, color = 'green')

#    plot_event.move_points.add(axes[0], 0, 1e-4, "o", picker=15, color = 'blue')
#    plot_event.move_points.add(axes[1], 5, 1e-4, "o", picker=15, color = 'red')
#    plot_event.register_move_points_event(fig)

    selector0 = RangeSelector('', axes[0], color = 'green', print_level = 0)
    selector1 = RangeSelector('', axes[1], color = 'blue', print_level = 0)

# popup_menuをpluginから読み込む
    root = get_window_from_plt(plt)
    popup_menu = plot_event.popup_menu.menu

    vars = cfg.copy()
    vars.plot_event = plot_event
    vars.selector0 = selector0
    vars.selector1 = selector1

    if len(app.config.plugin_dir) >= 1 and app.config.plugin_dir[0] not in "/\\":
        plugin_dir = app.replace_path(None, template = ["{dirname}", app.config.plugin_dir])
    print()
    print(f"Read plugins from [{plugin_dir}]")
    module_names, modules = app.load_modules(plugin_dir, "*.py", target = "popup_menu", sort = True, is_print = True)
    for m in modules:
        if hasattr(m, "add_popup_menu"):
            m.add_popup_menu(popup_menu, app = app, vars = vars, parent = root)


    plot_event.register_popup_menu_event()

    if cfg.use_tkplt:
        def redraw():
#            self.ax.clear()  # 既存のグラフをクリア
            tkplt.draw()
            canvas.draw()

        def update_label():
            label.config(text = entry.get())

        def button_start_annotate(button, label1, label2):
            current_text = button.cget("text")
            new_text = label2 if current_text == label1 else label2
            button.config(text = new_text)
            menu_start_annotate(plot_event.popup_menu.menu, label1, label2)

        def button_start_follow_mouse(button, label1, label2):
            current_text = button.cget("text")
            new_text = label2 if current_text == label1 else label2
            button.config(text = new_text)
            menu_start_follow_mouse(plot_event.popup_menu.menu, label1, label2)

        def save_figure():
            outfile = "image.png"
            print(f"Save figure to {outfile}")
            fig.savefig(outfile)
            
        plot_event.button_annotate = tkplt.add_toolbar_button(text = "Start annotate")
        plot_event.button_annotate.config(
                command = lambda: button_start_annotate(plot_event.button_annotate, "Start annotate", "Stop annotate"))

        plot_event.button_follow_mouse = tkplt.add_toolbar_button(text = "Start follow mouse")
        plot_event.button_follow_mouse.config(
                command = lambda: button_start_follow_mouse(plot_event.button_follow_mouse, "Start follow mouse", "Stop follow mouse"))

        tkplt.add_toolbar_button(text = "Save image.png", command = save_figure)

        tkplt.add_toolbar_button(text = "Redraw", command = redraw)

        entry = ttk.Entry(root)
        entry.pack(side=tk.TOP, fill = tk.X, expand = 0)
        label = ttk.Label(root, text="エントリーの内容がここに表示されます")
        label.pack(side=tk.TOP, fill = tk.X, expand = 0)

        button = ttk.Button(root, text="更新", command = update_label)
        button.pack(side=tk.TOP, fill=tk.X, expand = 0)

        button = ttk.Button(root, text="Redraw", command = redraw)
        button.pack(side=tk.TOP, fill=tk.X, expand = 0)

    tkplt.pause(1.0e-6)
#    tk.mainloop()

    if pause:
        input("\nPress Enter to terminate>>\n")

def illustrate_TFT_model3d(app, cfg):
    label_sub  = 'substrate'
    label_gate = 'gate'
    label_ins  = 'insultor'
    label_semi = 'semiconductor'
    label_s    = 'source'
    label_d    = 'drain'

    color_sub  = 'cyan'
    color_gate = 'silver'
    color_ins  = 'lightblue'
    color_semi = 'orange'
    color_sd   = 'silver'

    W      = cfg.W * 1.0e9

    d_sub  = 300.0
    d_gate = 100.0
    d_ins  = cfg.dg * 1.0e9  # nm
    d_semi = 100.0
    L_semi = cfg.L * 1.0e9
    L_sd   = 50.0e3 # nm
    d_sd   = 100.0
    
    x_tot = 2.0 * L_sd + L_semi
    z_tot = d_sd + d_semi + d_ins + d_gate + d_sub

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.view_init(elev=30, azim=60)
    ax.dist = 10  # 視点距離を調整

# 各レイヤーのサイズ設定
    width = depth = 8
    heights = [0.5, 0.2, 0.7, 0.2, 0.2]

    z_top = 0.0
    gate = ax.bar3d(x=0, y=0, z=z_top, dx=x_tot, dy=W, dz=d_sub, color=color_sub, zorder=1, label=label_sub)
    z_top += d_sub
    dielectric = ax.bar3d(x=0, y=0, z=z_top, dx=x_tot, dy=W, dz=d_ins, color=color_gate, zorder=2, label=label_gate)
    z_top += d_ins
    semiconductor = ax.bar3d(x=0, y=0, z=z_top, dx=x_tot, dy=W, dz=d_semi, color=color_semi, zorder=3, label=label_semi)
    z_top += d_semi
    source = ax.bar3d(x=0, y=0, z=z_top, dx=L_sd, dy=W, dz=d_sd, color=color_sd, zorder=4, alpha=0.8, label=label_s)
    drain = ax.bar3d(x=x_tot - L_sd, y=0, z=z_top, dx=L_sd, dy=W, dz=d_sd, color=color_sd, zorder=5, alpha=0.8, label=label_d)

    ax.set_xlabel('Width')
    ax.set_ylabel('Depth')
    ax.set_zlabel('Height')

# メッシュ（グリッド）を消す
    ax.grid(False)

# 目盛りを消す
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])
#    ax.legend()
    
    plt.pause(1.0e-4)
    input("\nPress ENTER to terminate>>")

def illustrate_TFT_model(app, cfg):
    fontsize = 10.0

    label_sub  = 'substrate'
    label_gate = 'gate'
    label_ins  = 'insultor'
    label_semi = 'semiconductor'
    label_s    = 'source'
    label_d    = 'drain'

    color_sub  = 'cyan'
    color_gate = 'silver'
    color_ins  = 'lightblue'
    color_semi = 'orange'
    color_sd   = 'silver'

    d_sub  = 300.0
    d_gate = 100.0
    d_ins  = cfg.dg * 1.0e9  # nm
    d_semi = 100.0
    L_semi = cfg.L * 1.0e9
    L_sd   = 50.0e3 # nm
    d_sd   = 100.0

    x_tot = 2.0 * L_sd + L_semi
    z_tot = d_sd + d_semi + d_ins + d_gate + d_sub

    fig, ax = plt.subplots(figsize=(6, 4))
    plt.title('Thin film transistor structure')
    ax.set_xlim(0, x_tot)
    ax.set_ylim(0, z_tot)
    ax.axis('off')

    z_top = 0.0
    sub = patches.Rectangle((0.0, 0.0), x_tot, d_sub, edgecolor = color_sub, facecolor = color_sub, linewidth = 1.0)
    ax.add_patch(sub)
    ax.text(x_tot / 2.0, z_top + d_sub / 2.0, label_sub, fontsize = fontsize, va = 'center', ha = 'center')
    z_top = d_sub

    gate = patches.Rectangle((0.0, z_top), x_tot, d_gate, edgecolor = color_gate, facecolor = color_gate, linewidth = 1.0)
    ax.add_patch(gate)
    ax.text(x_tot / 2.0, z_top + d_gate / 2.0, label_gate, fontsize = fontsize, va = 'center', ha = 'center')
    z_top += d_gate

    dielectric = patches.Rectangle((0.0, z_top), x_tot, d_ins, edgecolor = color_ins, facecolor = color_ins, linewidth = 1.0)
    ax.add_patch(dielectric)
    ax.text(x_tot / 2.0, z_top + d_ins / 2.0, label_ins, fontsize = fontsize, va = 'center', ha = 'center')
    z_top += d_ins

    semiconductor = patches.Rectangle((0.0, z_top), x_tot, d_semi, edgecolor = color_semi, facecolor = color_semi, linewidth = 1.0)
    ax.add_patch(semiconductor)
    ax.text(x_tot / 2.0, z_top + d_semi / 2.0, label_semi, fontsize = fontsize, va = 'center', ha = 'center')
    z_top += d_semi

    source = patches.Rectangle((0.0, z_top),          L_sd, z_top, edgecolor = color_sd, facecolor = color_sd, linewidth = 1.0)
    drain  = patches.Rectangle((x_tot - L_sd, z_top), L_sd, z_top, edgecolor = color_sd, facecolor = color_sd, linewidth = 1.0)
    ax.add_patch(source)
    ax.add_patch(drain)
    ax.text(L_sd / 2.0,         z_top + d_sd / 2.0, label_s, fontsize = fontsize, va = 'center', ha = 'center')
    ax.text(x_tot - L_sd / 2.0, z_top + d_sd / 2.0, label_d,  fontsize = fontsize, va = 'center', ha = 'center')

    plt.pause(1.0e-4)
    input("\nPress ENTER to terminate>>")

def main():
    app, cfg = initialize()
    app.update_vars()
    cfg.Eg = cfg.EC - cfg.EV

    if cfg.mode == 'model':
        illustrate_TFT_model(app, cfg)
    elif cfg.mode == 'model3d':
        illustrate_TFT_model3d(app, cfg)
    elif cfg.mode == 'IV':
        simulate_IV(app, cfg)
    elif cfg.mode == 'plotx':
        plot_x(app, cfg)
    elif cfg.mode == 'Vch':
        simulate_Vch(app, cfg)
    else:
        print(f"Error: Unknown mode [{cfg.mode}]")
        sys.exit()

if __name__ == '__main__':
    main()