# -*- coding: utf-8 -*-
"""
Created on Wed Feb  4 12:06:39 2026

@author: andri
"""

import numpy as np
from tqdm import tqdm

import customtkinter as tk
from tkinter import Toplevel
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk

from .Params import Params, pwem_params
from .Device import device
from .PWEM_model_fields import calc_E_mode_field, calc_H_mode_field, \
                              calc_E_mode_field_anisotropic, calc_H_mode_field_anisotropic
from .create_fig import Create_Fig

class Chern():
    
    def __init__(self, app):
        
        app.get_params()
        self.app = app
        
        self.params = Params(app);
        self.Device = device();        

    def vectors(self):
        # define lattice and beta (Bloch wave) vectors 
        Pwem_params    = pwem_params(self.app, self.params)
        Pwem_params.Symmetry(self.params);
        
        self.Create_Fig = Create_Fig(self.app, self.Device, self.params, Pwem_params)
        
        # calculate device's unit cell grid
        if self.app.sel_anisotropy == 'No':
            if self.app.device_selection == 'Square':
                self.Device.Ellipse(self.params)
            elif self.app.device_selection == 'Frame':
                self.Device.Frame(self.params)
            elif self.app.device_selection == 'Ring':
                self.Device.Ring(self.params)
            elif self.app.device_selection == 'Hex':
                self.Device.oblique(self.params, Pwem_params.t1, Pwem_params.t2)
            elif self.app.device_selection == 'Honeycomb':
                self.Device.honeycomb(self.params, Pwem_params.t1, Pwem_params.t2)
        else:
            if self.app.device_selection == 'Square':
                self.Device.Ellipse_anisotropic(self.params)
            elif self.app.device_selection == 'Frame':
                self.Device.Frame_anisotropic(self.params)
            elif self.app.device_selection == 'Ring':
                self.Device.Ring_anisotropic(self.params)
            elif self.app.device_selection == 'Hex':
                self.Device.oblique_anisotropic(self.params, Pwem_params.t1, Pwem_params.t2)
            elif self.app.device_selection == 'Honeycomb':
                self.Device.honeycomb_anisotropic(self.params, Pwem_params.t1, Pwem_params.t2)
                
        # use imported device    
        if self.app.sel_import == 'Yes':
            self.app.sel_anisotropy = 'No'
            self.Device.imported_uc(self.params, self.app.imported_data)
        
        # create an array for eigenvectors
        eigvecs = np.zeros((self.app.BZ_Points, self.app.BZ_Points, self.params.Harmonics[0]*self.params.Harmonics[1]), dtype=complex)
        
        # define the Brillouin zone
        Q, P  = np.meshgrid(np.linspace(-1/2, 1/2, self.app.BZ_Points), np.linspace(-1/2, 1/2, self.app.BZ_Points))
        self.bx    = Q * Pwem_params.T1[0] + P * Pwem_params.T2[0]; bx_flat = self.bx.flatten()
        self.by    = Q * Pwem_params.T1[1] + P * Pwem_params.T2[1]; by_flat = self.by.flatten()

        for nx in tqdm(range(0, self.app.BZ_Points)):
            for ny in range(0, self.app.BZ_Points):
                
                if self.app.sel_anisotropy == 'No':
                    
                    if self.app.sel_mode == 'E':
                        V, m, WE = calc_E_mode_field(self.params.Harmonics[0], self.params.Harmonics[1], Pwem_params.T1, Pwem_params.T2, 
                                                     bx_flat[nx], by_flat[ny], self.Device.ERC, self.Device.URC, self.params.norm, self.app.Berry_Bloch_mode)
                    else:
                        V, m, WE = calc_H_mode_field(self.params.Harmonics[0], self.params.Harmonics[1], Pwem_params.T1, Pwem_params.T2, 
                                                     bx_flat[nx], by_flat[ny], self.Device.ERC, self.Device.URC, self.params.norm, self.app.Berry_Bloch_mode)
                else:
                    
                    if self.app.sel_mode == 'E':
                        V, m, WE = calc_E_mode_field_anisotropic(self.params.Harmonics[0], self.params.Harmonics[1], Pwem_params.T1, Pwem_params.T2, 
                                                     bx_flat[nx], by_flat[ny], self.Device.ERCzz, self.Device.URC, self.Device.URC, self.params.norm,
                                                     self.app.Berry_Bloch_mode)
                    else:
                        V, m, WE = calc_H_mode_field_anisotropic(self.params.Harmonics[0], self.params.Harmonics[1], Pwem_params.T1, Pwem_params.T2, 
                                                     bx_flat[nx], by_flat[ny], self.Device.ERCxx, self.Device.ERCyy, self.Device.URC, self.params.norm,
                                                     self.app.Berry_Bloch_mode)
                
                s = V[:,m]
                s = np.reshape(s, (self.params.Harmonics[0]*self.params.Harmonics[1]))
                    
                eigvecs[nx, ny, :] = s

        self.eigvecs = eigvecs
        
                
    def link_variable(self, u_k, u_kp):
        overlap = np.vdot(u_k, u_kp)
        return overlap / np.abs(overlap)
                
    def berry_curvature(self):
        
        Nkx, Nky, NG = self.eigvecs.shape
        F = np.zeros((Nkx-1, Nky-1))
    
        for i in range(Nkx-1):
            for j in range(Nky-1):
    
                u00 = self.eigvecs[i,   j  ]
                u10 = self.eigvecs[i+1, j  ]
                u01 = self.eigvecs[i,   j+1]
                u11 = self.eigvecs[i+1, j+1]
    
                Ux_00 = self.link_variable(u00, u10)
                Uy_00 = self.link_variable(u00, u01)
                Ux_01 = self.link_variable(u01, u11)
                Uy_10 = self.link_variable(u10, u11)
    
                F[i, j] = np.angle(
                    Ux_00 * Uy_10 / (Ux_01 * Uy_00)
                )
    
        
        # display figure on TopLevel
        if self.app.plot_Berry == True:
            
            fig = self.Create_Fig.Plot_Berry_Curvature(F, self.bx, self.by)
            
            plot_window = Toplevel()
            
            canvas = FigureCanvasTkAgg(fig, master=plot_window)
            canvas.draw()
            canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
    
            # Add navigation toolbar
            color = "#505050"
     
            toolbar = NavigationToolbar2Tk(canvas, plot_window)
            toolbar.config(background=color)
            toolbar.update()
            toolbar.pack(side=tk.BOTTOM)
        
        self.F = F
        return self
    
    def Valley_Chern(self):
        # Chern number should equal 0 around the BZ
        # chern = np.sum(self.F) / (2* np.pi)
        
        if self.app.valley == "K'":
            curve    = np.argmin(self.F)
            curve_id = np.unravel_index(curve, self.F.shape)
        else:
            curve    = np.argmax(self.F)
            curve_id = np.unravel_index(curve, self.F.shape)
        
        # make a mask to integrate the valley only
        kc = 0.2
        dkx = self.bx - self.bx[curve_id[0], curve_id[1]]
        dky = self.by - self.by[curve_id[0], curve_id[1]]
        mask = dkx**2 + dky**2 < kc**2
        mask = mask[1:, 1:]
        
        # get valley Chern number
        self.Cv = np.sum(self.F[mask]) / (2*np.pi)
        
        if self.app.valley == "K'":
            self.app.textboxKp.delete('0.0', 'end')
            self.app.textboxKp.insert('0.0', str(self.Cv))
        else:
            self.app.textboxK.delete('0.0', 'end')
            self.app.textboxK.insert('0.0', str(self.Cv))
            
        return self
