# -*- coding: utf-8 -*-
"""
Created on Sat Nov 25 12:30:47 2023

@author: user
"""
import matplotlib.pyplot as plt
import numpy as np

from .Params import Params, pwem_params
from .Device import device

plt.rc('text', usetex=True)
plt.rc('font', family='serif')

"""
create_figs_for_an_app
"""
class Create_Fig():
    
    def __init__(self, app, Device, params, Pwem_params):
        
        app.get_params()
        self.app = app
        
        self.params = params
        self.Device = Device
        self.Pwem_params = Pwem_params
    
    def plot_device(self):
        
        fig1, ax1 = plt.subplots()
        
        if self.app.background == 'Dark':
            ax1.set_facecolor('black')
            
        # set ER = ERxx for anisotropic case
        if self.app.sel_anisotropy == 'Yes':
            self.Device.ER = self.Device.ERxx
        
        if self.app.sel_import == 'Yes':
            if np.shape(self.Device.X0) == np.shape(self.Device.ER):
                img1 = ax1.pcolor(self.Device.X0, self.Device.Y0, self.Device.ER, cmap = 'jet')
            else:
                img1 = ax1.imshow(np.real(self.Device.ER), extent = [-self.params.Lx/2, self.params.Lx/2, -self.params.Ly/2, self.params.Ly/2],
                           aspect = 'auto',  cmap = 'jet', vmin=1, vmax=np.max([self.params.er1, self.params.er2]))
        else:
            if self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
                img1 = ax1.imshow(np.real(self.Device.ER), extent = [-self.params.Lx/2, self.params.Lx/2, -self.params.Ly/2, self.params.Ly/2],
                                  aspect = 'auto',  cmap = 'jet', vmin=1, vmax=np.max([self.params.er1, self.params.er2]))
            else:
                img1 = ax1.pcolor(self.Device.X0, self.Device.Y0, self.Device.ER, cmap = 'jet')
        
        img1.set_rasterized(True)
                
    ######################################################################################################
        if self.app.background == 'Dark':
            fg_color = 'white'
            bg_color = 'black'
        else:
            fg_color = 'black'
            bg_color = 'white'
            
        ax1.set_xlabel(r'$x/a$, a.u.', fontsize=self.app.FontSize, color = fg_color)
        ax1.set_ylabel(r'$y/a$, a.u.', fontsize=self.app.FontSize, color = fg_color)
        
        ax1.set_title(r'Device unit cell', fontsize=self.app.FontSize, color = fg_color)
           
        cb = plt.colorbar(img1, fraction=0.046, pad=0.04) # equalize colorbar size
        ax1.set_aspect('equal')
        
        # set tick and ticklabel color
        img1.axes.tick_params(color=fg_color, labelcolor=fg_color)
        """Maybe the code below can be fixed"""
        ax1.tick_params(axis='both', which='major', labelsize=self.app.FontSize)
        ax1.tick_params(axis='both', which='minor', labelsize=self.app.FontSize)
        
        # set imshow outline
        for spine in img1.axes.spines.values():
            spine.set_edgecolor(fg_color)    
        
        # COLORBAR
        # set colorbar label plus label color
        cb.set_label(r'$\varepsilon_r$', color=fg_color)
        
        # set colorbar tick color
        cb.ax.yaxis.set_tick_params(color=fg_color, labelsize = self.app.FontSize)
        
        # set colorbar edgecolor 
        cb.outline.set_edgecolor(fg_color)
        
        # set colorbar ticklabels
        plt.setp(plt.getp(cb.ax.axes, 'yticklabels'), color=fg_color)
        
        fig1.patch.set_facecolor(bg_color)
        
        plt.close()
    
        return fig1


    def plot_band_diagram(self, WE, WH):
            
        if self.app.background == 'Light':
            fg_color = 'white'
            bg_color = 'black'
                    
            E_mode_color_bck = 'red'
            H_mode_color_bck = 'blue'
    
        else:
            fg_color = 'black'
            bg_color = 'white'
            
            E_mode_color_bck = '#8E5B68'
            H_mode_color_bck = '#B7D0E1'
            
    ########################################## HANDLE X AXIS #############################################            
                
        beta = self.Pwem_params.KP
        Ticks = self.Pwem_params.KT
    
        fig2, ax2 = plt.subplots(facecolor=fg_color)
            
        for i in range(0,self.app.NBANDS):
            ax2.plot(WH[i,:], H_mode_color_bck, linestyle = '-')
            ax2.plot(WE[i,:], E_mode_color_bck, linestyle = '--')
            
        for i in range(1,len(Ticks) - 1):
            ax2.axvline(x =  Ticks[i],   color = 'gray',ls='--')
          
    ######################################################################################################
        
        plt.xticks(Ticks, beta, fontsize=self.app.FontSize, color = fg_color)
            
    ######################################################################################################
            
        ax2.text(Ticks[-1] * 14/40, self.app.omega_Lo+(self.app.omega_Hi-self.app.omega_Lo)/20, "H mode", 
                 bbox=dict(facecolor=H_mode_color_bck, alpha=0.25), fontsize=self.app.FontSize, color = bg_color)
        ax2.text(Ticks[-1] * 21/40, self.app.omega_Lo+(self.app.omega_Hi-self.app.omega_Lo)/20, "E mode", 
                 bbox=dict(facecolor=E_mode_color_bck, alpha=0.25), fontsize=self.app.FontSize, color = bg_color)
        
        
        ax2.set_xlim(0, Ticks[-1])
        ax2.set_ylim(self.app.omega_Lo, self.app.omega_Hi)
    
        ax2.set_xlabel(r'$\vec{\beta}$', color = bg_color)
        ax2.set_ylabel(r'$a/\lambda_0$', color = bg_color)
        
        ax2.set_facecolor(fg_color)
        fig2.patch.set_facecolor(fg_color)
    
        
        """
        AXES VALUE COLOR
        """
        [t.set_color(bg_color) for t in ax2.xaxis.get_ticklines()]
        [t.set_color(bg_color) for t in ax2.xaxis.get_ticklabels()]
        [t.set_color(bg_color) for t in ax2.yaxis.get_ticklines()]
        [t.set_color(bg_color) for t in ax2.yaxis.get_ticklabels()]
        
        ax2.xaxis.set_tick_params(color=bg_color)
        ax2.yaxis.set_tick_params(color=bg_color)
    
        ax2.set_title(r'Band diagram', color = bg_color, fontsize=self.app.FontSize) 
        
        fig2.patch.set_facecolor(fg_color)   
        
        """Maybe the code below can be fixed"""
        ax2.tick_params(axis='both', which='major', labelsize=self.app.FontSize)
        ax2.tick_params(axis='both', which='minor', labelsize=self.app.FontSize)
        
        if self.app.background:
            ax2.set_facecolor(fg_color)
            fig2.patch.set_facecolor(fg_color)
        
        # set imshow outline
        for spine in ax2.axes.spines.values():
            spine.set_edgecolor(bg_color)  
            
        plt.close()
    
        return fig2
    
    """
    create_frame for gifs below
    """


    def create_frame(self, WE, WH, t):
    
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(17, 6), gridspec_kw={'width_ratios': [4, 5]})
        plt.subplots_adjust(wspace=0.2)
        
        # plot the unit cell
        
        if self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
            img1 = ax1.imshow(np.real(self.Device.ER), extent = [-self.params.Lx/2, self.params.Lx/2, -self.params.Ly/2, self.params.Ly/2],
                           aspect = 'auto',  cmap = 'jet', vmin=1, vmax=np.max([self.params.er1, self.params.er2]))
        else:
            img1 = ax1.pcolor(self.Device.X0, self.Device.Y0, self.Device.ER, vmin=1, vmax=np.max([self.params.er1, self.params.er2]), 
                              cmap = 'jet', rasterized=True)
            
    ######################################################################################################
            
        ax1.set_xlabel('$x$, a.u.')
        ax1.set_ylabel('$y$, a.u.')
        ax1.set_title(r'$r = {}a$'.format(f'{self.params.r:.3f}'))
        ax1.set_aspect('equal')
        plt.colorbar(img1, ax=ax1)
       
    ###################################################################################################### 
    
        '''
        Plot band diagram
        '''
            
    ########################################## HANDLE X AXIS #############################################            
        
        if self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
            beta  = ["$\Gamma$", "$X$", "$M$", "$\Gamma$"]
        else:
            beta  = ["$\Gamma$", "$M$", "$K$", "$\Gamma$"]
                
    ######################################################################################################
        Ticks = self.Pwem_params.KT
        
        for i in range(0,self.app.no_of_bands):
            ax2.plot(WH[i,:], 'blue')
            ax2.plot(WE[i,:], 'r--')
            
        for i in range(1,len(Ticks) - 1):
            ax2.axvline(x =  Ticks[i],   color = 'gray',ls='--')
          
    ######################################################################################################
            
        plt.xticks(Ticks,beta)
            
    ######################################################################################################
            
        ax2.text(Ticks[-1] * 15/40, self.app.lower_freq_lim+(self.app.upper_freq_lim-self.app.lower_freq_lim)/20, "H mode",
                 bbox=dict(facecolor='blue', alpha=0.25), fontsize=self.app.FontSize)
        ax2.text(Ticks[-1] * 21/40, self.app.lower_freq_lim+(self.app.upper_freq_lim-self.app.lower_freq_lim)/20, "E mode", 
                 bbox=dict(facecolor='red', alpha=0.25), fontsize=self.app.FontSize)
        
        ax2.set_xlim(0, Ticks[-1])
        ax2.set_ylim(self.app.lower_freq_lim, self.app.upper_freq_lim)
    
        ax2.set_xlabel(r'$\vec{\beta}$')
        ax2.set_ylabel(r'$a/\lambda_0$')
           
        ax2.set_title(r'Band Diagram')
     
        plt.savefig(f'./img_er{self.params.er1}/img_{t}.png', bbox_inches='tight', dpi = 200)
        
        plt.close()
    

    def Bands(self, WH, WE, BG_min, BG_max, omega_min_E, omega_min_H, omega_max_E, omega_max_H, gap_E, gap_H):
        
        if self.app.background == 'Light':
            fg_color = 'white'
            bg_color = 'black'
                    
            E_mode_color_bck = 'red'
            H_mode_color_bck = 'blue'
            
        else:
            fg_color = 'black'
            bg_color = 'white'
            
            E_mode_color_bck = '#8E5B68'
            H_mode_color_bck = '#B7D0E1'
            
        fig = plt.figure(facecolor = fg_color, figsize = (7, 5))
        ax = fig.add_subplot()
        
        colors_e = ['#440a02', '#881405', '#cd1f08', '#f53f27', '#f87c6c', '#fa9a8e']
        colors_h = ['#08203f', '#10407e', '#1860bd', '#3883e5', '#77abed', '#97bef1']
            
    ########################################## HANDLE X AXIS #############################################            
        ticks = self.Pwem_params.KT
        beta  = self.Pwem_params.KP
                
        ######################################################################################################
        
        for i in range(0,self.app.NBANDS):
            ax.plot(WH[i,:], H_mode_color_bck, linestyle = '-')
            ax.plot(WE[i,:], E_mode_color_bck, linestyle = '--')
            
        for i in range(1,len(ticks) - 1):
            ax.axvline(x =  ticks[i],   color = 'gray', ls='--')
    
        for n in range(len(omega_min_E)):
            if gap_E[n] > 0.01:
                ax.fill_between([0, ticks[-1]], omega_min_E[n], omega_max_E[n], color = colors_e[n % 6], alpha = 0.5)
        for n in range(len(omega_min_H)):
            if gap_H[n] > 0.01:
                ax.fill_between([0, ticks[-1]], omega_min_H[n], omega_max_H[n], color = colors_h[n % 6], alpha = 0.5)
                
        for i in range(len(BG_min)):
            ax.fill_between([0, ticks[-1]], BG_min[i], BG_max[i], color = 'white')
            ax.fill_between([0, ticks[-1]], BG_min[i], BG_max[i], color = 'yellow', alpha = 0.5)
          
        ######################################################################################################
            
        plt.xticks(ticks,beta, fontsize = self.app.FontSize, color = fg_color)
        plt.yticks(fontsize = self.app.FontSize)
            
        ######################################################################################################
            
        ax.text(ticks[-1] * 14/40, self.app.omega_Lo+(self.app.omega_Hi-self.app.omega_Lo)/20, "H mode", 
                 bbox=dict(facecolor=H_mode_color_bck, alpha=0.25), fontsize=self.app.FontSize, color = bg_color)
        ax.text(ticks[-1] * 21/40, self.app.omega_Lo+(self.app.omega_Hi-self.app.omega_Lo)/20, "E mode", 
                 bbox=dict(facecolor=E_mode_color_bck, alpha=0.25), fontsize=self.app.FontSize, color = bg_color)
        
        ax.set_xlim(0, ticks[-1])
        ax.set_ylim(self.app.omega_Lo, self.app.omega_Hi)
        
        ax.set_xlabel(r'$\vec{\beta}$', fontsize = self.app.FontSize, color = bg_color)
        ax.set_ylabel(r'$a/\lambda_0$', fontsize = self.app.FontSize, color = bg_color)
        
        """
        AXES VALUE COLOR
        """
        [t.set_color(bg_color) for t in ax.xaxis.get_ticklines()]
        [t.set_color(bg_color) for t in ax.xaxis.get_ticklabels()]
        [t.set_color(bg_color) for t in ax.yaxis.get_ticklines()]
        [t.set_color(bg_color) for t in ax.yaxis.get_ticklabels()]
        
        ax.xaxis.set_tick_params(color=bg_color)
        ax.yaxis.set_tick_params(color=bg_color)
           
        ax.set_title(r'Band gap diagram', fontsize = self.app.FontSize, color = bg_color)   
        
        fig.patch.set_facecolor(fg_color)   
        
        """Maybe the code below can be fixed"""
        ax.tick_params(axis='both', which='major', labelsize=self.app.FontSize)
        ax.tick_params(axis='both', which='minor', labelsize=self.app.FontSize)
        
        if self.app.background:
            ax.set_facecolor(fg_color)
            fig.patch.set_facecolor(fg_color)
        
        # set imshow outline
        for spine in ax.axes.spines.values():
            spine.set_edgecolor(bg_color)
            
        plt.close()
        
        return fig
    
    
    def Plot_Contours(self, W, band):
        
        if self.app.background == 'Light':
            fg_color = 'white'
            bg_color = 'black'
        else:
            fg_color = 'black'
            bg_color = 'white'
         
    ######################################################################################################
            
        maks = max(W[band,:])
        minimum = min(W[band,:])
    
        Harmonic = int(self.params.Harmonics[0] * self.params.Harmonics[1])
        sheet    = np.zeros((self.Pwem_params.N_Points, self.Pwem_params.N_Points, Harmonic))
        for i in range(0, Harmonic):
            sheet[:,:,i] = np.reshape(W[i,:], (self.Pwem_params.N_Points, self.Pwem_params.N_Points))
            
    ######################################################################################################
        if self.app.sel_import == 'No':
            x    = np.linspace(-np.pi/self.params.Lx,  np.pi/self.params.Lx, self.Pwem_params.N_Points)
            y    = np.linspace( np.pi/self.params.Ly, -np.pi/self.params.Ly, self.Pwem_params.N_Points)
        else:
            x    = np.linspace(min(self.Pwem_params.beta[0,:]), max(self.Pwem_params.beta[0,:]), self.Pwem_params.N_Points)
            y    = np.linspace(max(self.Pwem_params.beta[1,:]), min(self.Pwem_params.beta[1,:]), self.Pwem_params.N_Points)
    
        x, y = np.meshgrid(x,y)
        
        fig, ax = plt.subplots()
        
        if self.app.sel_import == 'Yes':
            grad = ax.imshow(sheet[:,:,band], extent=[np.min(x), np.max(x), np.min(y), np.max(y)],
                             origin='lower', cmap='jet', alpha=0.3, aspect='equal')
            
            contours = ax.contour(x, y, sheet[:,:,band], self.app.Line_num, colors='black')
            ax.clabel(contours, inline=True, fontsize=8)
    
        elif self.app.device_selection == 'Hex' or self.app.device_selection == 'Honeycomb':
            grad = ax.imshow(sheet[:,:,band], extent=[np.min(self.Pwem_params.X0), np.max(self.Pwem_params.X0), 
                             np.min(self.Pwem_params.Y0), np.max(self.Pwem_params.Y0)], origin='lower',
                             cmap='jet', alpha=0.3, aspect='equal')
    
            contours = ax.contour(self.Pwem_params.X0, self.Pwem_params.Y0, sheet[:,:,band], self.app.Line_num, colors='black')
            ax.clabel(contours, inline=True, fontsize=8)
            
            # add lines to distinguish hexagon
            if self.params.Lx != 1 or self.params.Ly != 1:
                pass
            elif self.app.device_selection == 'Hex':
                x_line = np.array([-4*np.pi/3, 4*np.pi/3])
                y1_line =  x_line*np.sqrt(3) + self.params.Lx/2 * 8 * np.sqrt(3) * np.pi/3
                y2_line =  x_line*np.sqrt(3) - self.params.Lx/2 * 8 * np.sqrt(3) * np.pi/3
                y3_line = -x_line*np.sqrt(3) + self.params.Lx/2 * 8 * np.sqrt(3) * np.pi/3
                y4_line = -x_line*np.sqrt(3) - self.params.Lx/2 * 8 * np.sqrt(3) * np.pi/3
                ax.plot(x_line, y1_line, color='black', linewidth=2.5)
                ax.plot(x_line, y2_line, color='black', linewidth=2.5)
                ax.plot(x_line, y3_line, color='black', linewidth=2.5)
                ax.plot(x_line, y4_line, color='black', linewidth=2.5)
            elif self.app.device_selection == 'Honeycomb':
                x_line = np.array([-4/3/np.sqrt(3)*np.pi, 4/3/np.sqrt(3)*np.pi])
                y1_line =  x_line*np.sqrt(3) + self.params.Lx/2 * 2 * 4 * np.pi/3
                y2_line =  x_line*np.sqrt(3) - self.params.Lx/2 * 2 * 4 * np.pi/3
                y3_line = -x_line*np.sqrt(3) + self.params.Lx/2 * 2 * 4 * np.pi/3
                y4_line = -x_line*np.sqrt(3) - self.params.Lx/2 * 2 * 4 * np.pi/3
                ax.plot(x_line, y1_line, color='black', linewidth=2.5)
                ax.plot(x_line, y2_line, color='black', linewidth=2.5)
                ax.plot(x_line, y3_line, color='black', linewidth=2.5)
                ax.plot(x_line, y4_line, color='black', linewidth=2.5)
        else:
            grad = ax.imshow(sheet[:,:,band], extent=[np.min(x), np.max(x), np.min(y), np.max(y)], origin='lower',
                       cmap='jet', alpha=0.3) # cmap RdGy or jet
            
            contours = ax.contour(x, y, sheet[:,:,band], self.app.Line_num, colors='black')
            ax.clabel(contours, inline=True, fontsize=8)
            
        cbar = fig.colorbar(grad, shrink = 0.7);
        ticks = []
        for i in np.linspace(minimum, maks, 5):
            ticks.append(float(('{:.3f}'.format(i))))
        cbar.set_ticks(ticks)
            
    ######################################################################################################
        
        if self.app.sel_import == 'Yes':
            pass
        elif self.params.Lx != 1 or self.params.Ly != 1:
            pass
        elif self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
            ax.set_xticks([-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
                       [r'$-\pi$',r'$-\frac{\pi}{2}$','0',r'$\frac{\pi}{2}$', r'$\pi$'])
        
            ax.set_yticks([-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
                       [r'$-\pi$',r'$-\frac{\pi}{2}$','0',r'$\frac{\pi}{2}$', r'$\pi$'])
            # pass
        elif self.app.device_selection == 'Hex':
            ax.set_xticks([-4/3*np.pi, -2*np.pi/3, 0, 2/3*np.pi, 4/3*np.pi],
                       [r'$-\frac{4\pi}{3}$',r'$-\frac{2\pi}{3}$','0',r'$\frac{2\pi}{3}$', r'$\frac{4\pi}{3}$'])
        
            ax.set_yticks([-np.pi*2/np.sqrt(3), -np.pi/np.sqrt(3), 0, np.pi/np.sqrt(3), np.pi*2/np.sqrt(3)],
                       [r'$-\frac{2\pi}{\sqrt{3}}$',r'$-\frac{\pi}{\sqrt{3}}$','0',r'$\frac{\pi}{\sqrt{3}}$', r'$\frac{2\pi}{\sqrt{3}}$'])
            ax.set_ylim([-np.pi*2/np.sqrt(3), np.pi*2/np.sqrt(3)])
        elif self.app.device_selection == 'Honeycomb':
            ax.set_xticks([-4/3/np.sqrt(3)*np.pi, -2*np.pi/np.sqrt(3)/3, 0, 2/np.sqrt(3)/3*np.pi, 4/3/np.sqrt(3)*np.pi],
                       [r'$-\frac{4\pi}{3\sqrt{3}}$',r'$-\frac{2\pi}{3\sqrt{3}}$','0',r'$\frac{2\pi}{3\sqrt{3}}$', r'$\frac{4\pi}{3\sqrt{3}}$'])
        
            ax.set_yticks([-np.pi*2/3, -np.pi/3, 0, np.pi/3, np.pi*2/3],
                       [r'$-\frac{2\pi}{3}$',r'$-\frac{\pi}{3}$','0',r'$\frac{\pi}{3}$', r'$\frac{2\pi}{3}$'])
            ax.set_ylim([-np.pi*2/3, np.pi*2/3])
        
        ax.set_xlabel(r'$\beta_x\Lambda_x$, rad', color=bg_color)
        ax.set_ylabel(r'$\beta_y\Lambda_y$, rad', color=bg_color)
        
        if   band == 1:
            ax.set_title('${}$nd Bloch mode, ${}$ $mode$'.format(band+1, self.params.Mode), color = bg_color)
        elif band == 0:
            ax.set_title('${}$st Bloch mode, ${}$ $mode$'.format(band+1, self.params.Mode), color = bg_color)
        elif band == 2:
            ax.set_title('${}$rd Bloch mode, ${}$ $mode$'.format(band+1, self.params.Mode), color = bg_color)
        else:
            ax.set_title('${}$th Bloch mode, ${}$ $mode$'.format(band+1, self.params.Mode), color = bg_color)
           
        ax.set_aspect('equal')
                
        """
        AXES VALUE COLOR
        """
        [t.set_color(bg_color) for t in ax.xaxis.get_ticklines()]
        [t.set_color(bg_color) for t in ax.xaxis.get_ticklabels()]
        [t.set_color(bg_color) for t in ax.yaxis.get_ticklines()]
        [t.set_color(bg_color) for t in ax.yaxis.get_ticklabels()]
        
        ax.xaxis.set_tick_params(color=bg_color)
        ax.yaxis.set_tick_params(color=bg_color)  
        
        fig.patch.set_facecolor(fg_color)  
            
        if self.app.background:
            fig.patch.set_facecolor(fg_color)
        
        # set imshow outline
        for spine in ax.axes.spines.values():
            spine.set_edgecolor(bg_color) 
            
        # COLORBAR
        # set colorbar label plus label color
        cbar.set_label(r'$a/\lambda_0$', color=bg_color)
        
        # set colorbar tick color
        cbar.ax.yaxis.set_tick_params(color=bg_color, labelsize = self.app.FontSize)
        
        # set colorbar edgecolor 
        cbar.outline.set_edgecolor(bg_color)
        
        # set colorbar ticklabels
        plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color=bg_color)
        
        plt.close()
        
        return fig
   

    def Plot_Contours3D(self, W, band):
        
        if self.app.background == 'Light':
            fg_color = 'white'
            bg_color = 'black'
        else:
            fg_color = 'black'
            bg_color = 'white'
         
        #####################################################################################################
        
        maks = max(W[band,:])
                
        Harmonic = int(self.params.Harmonics[0] * self.params.Harmonics[1])
        sheet    = np.zeros((self.Pwem_params.N_Points, self.Pwem_params.N_Points, Harmonic))
        for i in range(0, Harmonic):
            sheet[:,:,i] = np.reshape(W[i,:], (self.Pwem_params.N_Points, self.Pwem_params.N_Points))
            
        ######################################################################################################
    
        if self.app.sel_import == 'Yes':
            x    = np.linspace(min(self.Pwem_params.beta[0,:]), max(self.Pwem_params.beta[0,:]), self.Pwem_params.N_Points)
            y    = np.linspace(max(self.Pwem_params.beta[1,:]), min(self.Pwem_params.beta[1,:]), self.Pwem_params.N_Points)
                
        elif self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
            x    = np.linspace(-np.pi/self.params.Lx,  np.pi/self.params.Lx, self.Pwem_params.N_Points)
            y    = np.linspace( np.pi/self.params.Ly, -np.pi/self.params.Ly, self.Pwem_params.N_Points)
        elif self.app.device_selection == 'Hex':
            x    = np.linspace(-4/3*np.pi,  4/3*np.pi, self.Pwem_params.N_Points)
            y    = np.linspace( np.pi*2/np.sqrt(3), -np.pi*2/np.sqrt(3), self.Pwem_params.N_Points)
        elif self.app.device_selection == 'Honeycomb':
            x    = np.linspace(-4/3/np.sqrt(3)*np.pi,  4/3/np.sqrt(3)*np.pi, self.Pwem_params.N_Points)
            y    = np.linspace( np.pi*2/3, -np.pi*2/3, self.Pwem_params.N_Points)
                
        x, y = np.meshgrid(x,y)
    
        ######################################################################################################
        
        fig = plt.Figure()
        ax = fig.add_subplot(111, projection='3d')
        
        ax.view_init(20, -30)
        
        ######################################################################################################
        
        surf = ax.plot_surface(x, y, sheet[:,:,band], cmap="jet", alpha  = 0.5)
        cp   = ax.contour(x, y, sheet[:,:,band], self.app.Line_num, cmap="jet", offset = 0)
        
        ######################################################################################################
        
        ax.clabel(cp, inline=True, 
                  fontsize=self.app.FontSize, offset = 0)
        ax.set_zlim(0, maks)
        ax.set_xlim(min(x[0,:]), max(x[0,:]))
        ax.set_ylim(min(y[:,0]), max(y[:,0]))
        
        #####################################################################################################
        
        ax.set_xlabel(r'$\beta_x\Lambda_x$', color=bg_color, fontsize=self.app.FontSize)
        ax.set_ylabel(r'$\beta_y\Lambda_y$', color=bg_color, fontsize=self.app.FontSize)
        ax.set_zlabel(r'$a/\lambda_0$', color=bg_color, fontsize=self.app.FontSize)
        
        if self.app.sel_import == 'Yes':
            pass
        elif self.params.Lx != 1 or self.params.Ly != 1:
            pass    
        elif self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
            ax.set_xticks([-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
                       [r'$-\pi$',r'$-\frac{\pi}{2}$','0',r'$\frac{\pi}{2}$', r'$\pi$'])
        
            ax.set_yticks([-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
                       [r'$-\pi$',r'$-\frac{\pi}{2}$','0',r'$\frac{\pi}{2}$', r'$\pi$'])
            
        elif self.app.device_selection == 'Hex':
            ax.set_xticks([-4/3*np.pi, -2*np.pi/3, 0, 2/3*np.pi, 4/3*np.pi],
                       [r'$-\frac{4\pi}{3}$',r'$-\frac{2\pi}{3}$','0',r'$\frac{2\pi}{3}$', r'$\frac{4\pi}{3}$'])
        
            ax.set_yticks([-np.pi*2/np.sqrt(3), -np.pi/np.sqrt(3), 0, np.pi/np.sqrt(3), np.pi*2/np.sqrt(3)],
                       [r'$-\frac{2\pi}{\sqrt{3}}$',r'$-\frac{\pi}{\sqrt{3}}$','0',r'$\frac{\pi}{\sqrt{3}}$', r'$\frac{2\pi}{\sqrt{3}}$'])
            
        elif self.app.device_selection == 'Honeycomb':
            ax.set_xticks([-4/3/np.sqrt(3)*np.pi, -2*np.pi/np.sqrt(3)/3, 0, 2/np.sqrt(3)/3*np.pi, 4/3/np.sqrt(3)*np.pi],
                       [r'$-\frac{4\pi}{3\sqrt{3}}$',r'$-\frac{2\pi}{3\sqrt{3}}$','0',r'$\frac{2\pi}{3\sqrt{3}}$', r'$\frac{4\pi}{3\sqrt{3}}$'])
            ax.set_xlim([-4/3/np.sqrt(3)*np.pi, 4/3/np.sqrt(3)*np.pi])
        
            ax.set_yticks([-np.pi*2/3, -np.pi/3, 0, np.pi/3, np.pi*2/3],
                       [r'$-\frac{2\pi}{3}$',r'$-\frac{\pi}{3}$','0',r'$\frac{\pi}{3}$', r'$\frac{2\pi}{3}$'])
            ax.set_ylim([-np.pi*2/3, np.pi*2/3])
            
        ######################################################################################################
        
        ax.set_aspect("auto")
        cbar = fig.colorbar(surf, pad = 0.05, shrink=0.5, aspect=20, orientation = 'horizontal')
        
        ######################################################################################################
        
        if   band == 1:
            ax.set_title('${}$nd Bloch mode, ${}$ $mode$'.format(band+1, self.params.Mode), color = bg_color, fontsize=self.app.FontSize)
        elif band == 0:
            ax.set_title('${}$st Bloch mode, ${}$ $mode$'.format(band+1, self.params.Mode), color = bg_color, fontsize=self.app.FontSize)
        elif band == 2:
            ax.set_title('${}$rd Bloch mode, ${}$ $mode$'.format(band+1, self.params.Mode), color = bg_color, fontsize=self.app.FontSize)
        else:
            ax.set_title('${}$th Bloch mode, ${}$ $mode$'.format(band+1, self.params.Mode), color = bg_color, fontsize=self.app.FontSize)
        
    ######################################################################################################
                
        """
        AXES VALUE COLOR
        """
        [t.set_color(bg_color) for t in ax.xaxis.get_ticklines()]
        [t.set_color(bg_color) for t in ax.xaxis.get_ticklabels()]
        [t.set_color(bg_color) for t in ax.yaxis.get_ticklines()]
        [t.set_color(bg_color) for t in ax.yaxis.get_ticklabels()]
        
        ax.xaxis.set_tick_params(color=bg_color, labelsize=self.app.FontSize)
        ax.yaxis.set_tick_params(color=bg_color, labelsize=self.app.FontSize)  
        ax.zaxis.set_tick_params(color=bg_color, labelsize=self.app.FontSize)
        
        ax.tick_params(axis='z', colors=bg_color)
        
        fig.patch.set_facecolor(fg_color)  
        ax.set_facecolor(fg_color)
            
        if self.app.background:
            fig.patch.set_facecolor(fg_color)
        
        # set imshow outline
        for spine in ax.axes.spines.values():
            spine.set_edgecolor(bg_color)
            
        # COLORBAR
        # set colorbar label plus label color
        cbar.set_label(r'$a/\lambda_0$', color=bg_color, fontsize=self.app.FontSize)
        
        # set colorbar tick color
        cbar.ax.xaxis.set_tick_params(color=bg_color, labelsize = self.app.FontSize)
        
        # set colorbar edgecolor 
        cbar.outline.set_edgecolor(bg_color)
        
        # set colorbar ticklabels
        plt.setp(plt.getp(cbar.ax.axes, 'xticklabels'), color=bg_color)
    
        plt.close()
        
        return fig


    def gif_field(self, t, WE, field_1):
    
        if self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11,4))
        else:
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9,4), gridspec_kw={'width_ratios': [2, 1]})
            
        plt.subplots_adjust(wspace=0.05)
        
    ###################################################################################################### 
        '''
        Plot band diagram
        '''
    ########################################## HANDLE X AXIS #############################################            
        
        if self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
            beta  = ["$\Gamma$", "$X$", "$M$", "$\Gamma$"]
        elif self.app.device_selection == 'Hex' or self.app.device_selection == 'Honeycomb':
            beta  = ["$\Gamma$", "$M$", "$K$", "$\Gamma$"]
           
        ticks = self.Pwem_params.KT
    ######################################################################################################
    
        for i in range(0, t+1):
            # ax1.plot(WH[i,:], 'blue')
            ax1.plot(np.ones(len(WE[:,i]))*i, WE[:,i], 'r*')
            
        for i in range(1,len(ticks) - 1):
            ax1.axvline(x =  ticks[i],   color = 'gray',ls='--')
          
    ######################################################################################################
            
        ax1.set_xticks(ticks,beta)
            
    ######################################################################################################
            
        # ax1.text(ticks[-1] * 15/40, ymin+0.05, "H mode", bbox=dict(facecolor='blue', alpha=0.25), fontsize=font_size)
        # ax1.text(ticks[-1] * 21/40, ymin+0.05, "E mode", bbox=dict(facecolor='red', alpha=0.25), fontsize=font_size)
        
        ax1.set_xlim(0, ticks[-1])
        ax1.set_ylim(self.app.lower_freq_lim, self.app.upper_freq_lim)
    
        ax1.set_xlabel(r'$\vec{\beta}$')
        ax1.set_ylabel(r'$a/\lambda_0$')
           
        ax1.set_title(r'$\varepsilon_{} = {}$'.format('r', max([self.params.er1, self.params.er2])))
        
    ###############################################################################
        '''
        Plot fields
        '''
    ###############################################################################
        
        field_11 = np.append(field_1, field_1, axis=0)
        field_11 = np.append(field_1, field_11, axis=0)
        field_12 = np.append(field_11, field_11, axis=1)
        field_12 = np.append(field_12, field_11, axis=1)
        field_12 = field_12[::3,::3]
        
        ER1 = np.append(self.Device.ER, self.Device.ER, axis=0)
        ER1 = np.append(self.Device.ER, ER1, axis=0)
        ER2 = np.append(ER1, ER1, axis=1)
        ER2 = np.append(ER2, ER1, axis=1)
        ER2 = ER2[::3,::3]
        
        if self.app.sel_import == 'No':
            if self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
                x = np.linspace(self.params.Lx*3/2, -self.params.Lx*3/2, self.params.dim[0])
                y = np.linspace(self.params.Ly*3/2, -self.params.Ly*3/2, self.params.dim[1])
                
                X, Y = np.meshgrid(x, y)
                
                # normalize the field
                f_max = np.array([max(abs(sublist)) for sublist in field_1])
                field_12 = field_12/max(f_max)
                
                a = ax2.imshow(abs(field_12)**2, extent=[self.params.x[0]*3, self.params.x[-1]*3, 
                               self.params.y[0]*3, self.params.y[-1]*3], cmap='jet', vmin=0, vmax=1)
            else:
                X = self.Device.X0
                Y = self.Device.Y0
                
                # normalize the field
                f_max = np.array([max(abs(sublist)) for sublist in field_1])
                field_12 = field_12/max(f_max)
                
                a = ax2.pcolor(X, Y, abs(field_12)**2, vmin=-1, vmax=1, cmap='jet')
        else:
    
            X = self.Device.X0
            Y = self.Device.Y0
            
            # normalize the field
            f_max = np.array([max(abs(sublist)) for sublist in field_1])
            field_12 = field_12/max(f_max)
            
            a = ax2.pcolor(X, Y, abs(field_12)**2, vmin=0, vmax=1, cmap='jet')
            
        ax2.set_aspect('equal')
        cb = fig.colorbar(a)
        cb.set_label(r'I, a.u.')
        ax2.contour(X, Y, ER2, colors = 'black')
        ax2.set_xticks([])
        ax2.set_yticks([])
        
        if self.app.Bloch_mode_no == 1:
            ax2.set_title(r'1st Bloch mode')
        elif self.app.Bloch_mode_no == 2:
            ax2.set_title(r'2nd Bloch mode')
        elif self.app.Bloch_mode_no == 3:
            ax2.set_title(r'3rd Bloch mode')
        else:
            ax2.set_title(r'{}th Bloch mode'.format(self.app.Bloch_mode_no))
        
        plt.savefig(f'./field_img/img_{t}.png', bbox_inches='tight', dpi = 200)
        
        plt.close()
    
    
    def single_field(self, field_1, beta_x, beta_y):
        
        if self.app.background == 'Dark':
            fg_color = 'white'
            bg_color = 'black'
        else:
            fg_color = 'black'
            bg_color = 'white'        
        
        fig, ax = plt.subplots(1, 1)
        
        if self.app.background == 'Dark':
            ax.set_facecolor('black')
            
        field_11 = np.append(field_1, field_1, axis=0)
        field_11 = np.append(field_1, field_11, axis=0)
        field_12 = np.append(field_11, field_11, axis=1)
        field_12 = np.append(field_12, field_11, axis=1)
        field_12 = field_12[::3,::3]
        
        ER1 = np.append(self.Device.ER, self.Device.ER, axis=0)
        ER1 = np.append(self.Device.ER, ER1, axis=0)
        ER2 = np.append(ER1, ER1, axis=1)
        ER2 = np.append(ER2, ER1, axis=1)
        ER2 = ER2[::3,::3]
        
        if self.app.sel_import == 'No':
            if self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
                x = np.linspace(self.params.Lx*3/2, -self.params.Lx*3/2, self.params.dim[0])
                y = np.linspace(self.params.Ly*3/2, -self.params.Ly*3/2, self.params.dim[1])
                X, Y = np.meshgrid(x, y)
                    
                if self.app.sel_field == 'Re(f)':
                    # add phase
                    phi = np.exp(-1j * (beta_x*Y + beta_y*X + self.app.phase))
                    
                    # normalize the field
                    field_12 = np.real(field_12 * phi)
                    f_max = np.array([max(np.real(sublist)) for sublist in field_12])
                    field_12 = field_12/max(f_max)
                
                    a = ax.imshow(field_12, extent=[self.params.x[0]*3, self.params.x[-1]*3, self.params.y[0]*3, self.params.y[-1]*3], cmap='jet')
                else:
                    # normalize the field
                    f_max = np.array([max(abs(sublist)) for sublist in field_1])
                    field_12 = field_12/max(f_max)
                    
                    a = ax.imshow(abs(field_12)**2, extent=[self.params.x[0]*3, self.params.x[-1]*3, self.params.y[0]*3, self.params.y[-1]*3], cmap='jet')
                    
                a.set_rasterized(True) # helps save memory
                
            else:
                X = 3*self.Device.X0
                Y = 3*self.Device.Y0
                if self.app.sel_field == 'Re(f)':
                    # add phase
                    phi = np.exp(1j * (beta_x*X + beta_y*Y + self.app.phase))
                    
                    # normalize the field
                    field_12 = np.real(field_12 * phi)
                    f_max = np.array([max(np.real(sublist)) for sublist in field_12])
                    field_12 = field_12/max(f_max)
                
                    a = ax.pcolor(X, Y, field_12, cmap='jet')
                else:
                    # normalize the field
                    f_max = np.array([max(abs(sublist)) for sublist in field_1])
                    field_12 = abs(field_12)/max(f_max)
                    
                    a = ax.pcolor(X, Y, abs(field_12)**2, cmap='jet')
    
                a.set_rasterized(True)
            
        else:
    
            X = self.Device.X0
            Y = self.Device.Y0
            if self.app.sel_field == 'Re(f)':
                # add phase
                phi = np.exp(-1j * (beta_x*X + beta_y*Y + self.app.phase))
                
                # normalize the field
                field_12 = np.real(field_12 * phi)
                f_max = np.array([max(np.real(sublist)) for sublist in field_12])
                field_12 = field_12/max(f_max)
                
                a = ax.pcolor(X, Y, field_12, cmap='jet')
            else:
                # normalize the field
                f_max = np.array([max(abs(sublist)) for sublist in field_1])
                field_12 = field_12/max(f_max)
                
                a = ax.pcolor(X, Y, abs(field_12)**2, cmap='jet')
    
            a.set_rasterized(True)
                    
        cb = fig.colorbar(a)
        ax.contour(X, Y, ER2, colors = 'black')
        # ax.contour(X, Y, ER, colors = 'black')
        ax.set_aspect('equal')
        ax.set_xticks([])
        ax.set_yticks([])
        
        if self.app.Bloch_mode_no == 1:
            ax.set_title(r'1st Bloch mode', color=fg_color)
        elif self.app.Bloch_mode_no == 2:
            ax.set_title(r'2nd Bloch mode', color=fg_color)
        elif self.app.Bloch_mode_no == 3:
            ax.set_title(r'3rd Bloch mode', color=fg_color)
        else:
            ax.set_title(r'{}th Bloch mode'.format(self.app.Bloch_mode_no), color=fg_color)
            
        a.axes.tick_params(color=fg_color, labelcolor=fg_color)
        ax.tick_params(axis='both', which='major', labelsize=self.app.FontSize)
        ax.tick_params(axis='both', which='minor', labelsize=self.app.FontSize)
        
        
        # set imshow outline
        for spine in a.axes.spines.values():
            spine.set_edgecolor(fg_color)    
        
        # COLORBAR
        # set colorbar label plus label color
        if self.app.sel_field == 'Re(f)':
            cb.set_label(r'Re(f), a.u.', color=fg_color)
        else:
            cb.set_label(r'I, a.u.', color=fg_color)
        
        # set colorbar tick color
        cb.ax.yaxis.set_tick_params(color=fg_color, labelsize = self.app.FontSize)
        
        # set colorbar edgecolor 
        cb.outline.set_edgecolor(fg_color)
        
        # set colorbar ticklabels
        plt.setp(plt.getp(cb.ax.axes, 'yticklabels'), color=fg_color)
        
        fig.patch.set_facecolor(bg_color)
            
        plt.close()
            
        return fig


    def Plot_Berry_Curvature(self, F, bx, by):
    
        if self.app.background == 'Dark':
            fg_color = 'white'
            bg_color = 'black'
        else:
            fg_color = 'black'
            bg_color = 'white' 
    
        fig, ax = plt.subplots()
        
        if self.app.background == 'Dark':
            ax.set_facecolor('black')
            
        if self.app.device_selection == 'Square' or self.app.device_selection == 'Frame' or self.app.device_selection == 'Ring':
            img = ax.pcolor(bx, by, F, vmin=np.min(F), vmax=np.max(F))
        else:
            img = ax.pcolor(bx, by, F, vmin=np.min(F), vmax=np.max(F))
        
        ax.set_aspect('equal')
        cbar = fig.colorbar(img, ax=ax, fraction=0.035, ticks=[np.min(F), np.max(F)])
        cbar.ax.set_yticklabels(['min', 'max'])
        plt.title(r'Berry Curvature', color=fg_color)
        
        img.axes.tick_params(color=fg_color, labelcolor=fg_color)
        ax.tick_params(axis='both', which='major', labelsize=self.app.FontSize)
        ax.tick_params(axis='both', which='minor', labelsize=self.app.FontSize)
        
        # set imshow outline
        for spine in img.axes.spines.values():
            spine.set_edgecolor(fg_color)    
        
        # set colorbar tick color
        cbar.ax.yaxis.set_tick_params(color=fg_color, labelsize = self.app.FontSize)
        
        # set colorbar edgecolor 
        cbar.outline.set_edgecolor(fg_color)
        
        # set colorbar ticklabels
        plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color=fg_color)
        
        fig.patch.set_facecolor(bg_color)
            
        plt.close()
        
        return fig
