# -*- coding: utf-8 -*-
"""
Created on Sun Jul 13 16:10:11 2025

@author: 2704879w
"""


import legume
from scipy.io import savemat
import numpy as np
import matplotlib.pyplot as plt
import copy
j = 0.25
P = 2800*j
D =500/P# slab thickness in units of a

w = j*1200/P  # hole radius in units of a
g1 = j*0/P
g2 = j*00/P
g3 =0/P
w1 = w + g1/2;
w2 = w - g1/2;
l = j*2100/P
k = 1
# D = D/j
w = w*k
l = l*k
P = P*k
g1 = g1*k
g2= g2*k
g3 = g3*k

# a1 = np.array([0.5, 0.0])   # 横向周期为 1
# a2 = np.array([0.0, 1.0])   # 纵向周期为 2
y_half = l / 2
eps_c = 3.37**2      # permittivity of circular holes
eps_b = 1    # background permittivity of slab
eps_lower, eps_upper =  1.4**2,1 # permittivities of lower and upper claddings

x0 = -0.25
g3 = g3/2
rectA = legume.Poly(
    eps=eps_c,
    x_edges=np.array([x0+g2 - w1/2, x0+g2+ w1/2, x0+g2 + w1/2, x0+g2 - w1/2]),
    y_edges=np.array([-y_half+g3/2, -y_half+g3/2, y_half-g3/2, y_half-g3/2])
)

# 矩形 B：位置 x = +0.25
rectB = legume.Poly(
    eps=eps_c,
    x_edges=np.array([ 0.25-g2 - w2/2,  0.25-g2 + w2/2,  0.25-g2 + w2/2,  0.25-g2 - w2/2]),
    y_edges=np.array([-y_half-g3/2, -y_half-g3/2, y_half+g3/2, y_half+g3/2])
)




lattice = legume.Lattice('square')
lattice = legume.Lattice(
    
        [1, 0],  # x方向周期减半
        [0, 1]   # y方向周期不变
    )


phc = legume.PhotCryst(lattice, eps_l=eps_lower, eps_u=eps_upper)
phc.add_layer(d=D, eps_b=eps_b)
# phc.layers[-1].add_shape([rectA])
phc.layers[-1].add_shape([rectA,rectB])

gme = legume.GuidedModeExp(phc, gmax=5.5, truncate_g='abs')
npw = np.shape(gme.gvec)[1]  # number of plane waves in the expansion
print('Number of reciprocal lattice vectors in the expansion: npw = ', npw)

legume.viz.structure(phc, xz=True, yz=True, figsize=3.)

# Run the guided-mode expansion
numeig, verbose = 10, True
nk =200
X1 = [np.pi, 0]
Y1_e = [0,-1*np.pi]
Y1 = [0, 1*np.pi]
M1 = [2*np.pi, 1*np.pi]
# path = lattice.bz_path([Y1,'G',X1,M1], [nk, 2*nk, nk])
path = lattice.bz_path([X1,'G',Y1], [nk,nk])
gme.run(kpoints=path['kpoints'], gmode_inds=[0],numeig=numeig, verbose=True)
freqs = gme.freqs
freqs_im = gme.freqs_im
Qs = freqs / (2 * freqs_im)
# angles=path['angles'], kz_symmetry = 'odd',
nkappa, nfreq = freqs.shape[0], freqs.shape[1]
print(f'Number of wavevectros = {nkappa}, number of frequencies = {nfreq}')

def plot_bands(gme):
    fig, ax = plt.subplots(1, figsize = (6, 6))
    legume.viz.bands(gme, Q=True, ax=ax, cone=True, k_units=True, conecolor='lightgrey',
                      markersize=3, markeredgecolor='black', markeredgewidth=0)
    ax.set_xticks(path['k_indexes'])
    ax.set_xticklabels(path['labels'])
    ax.xaxis.grid('True')
    ax.set_ylim([0, 0.6])

plot_bands(gme)

savemat("bands_folding.mat", {
    "freqs": gme.freqs,
    "kpoints": gme.kpoints,
    'Qs':Qs,
    "labels": path['labels'],
    "k_indexes": path['k_indexes']
})


# savemat("bands_with_Q_FB_FB.mat", {
#     "freqs": gme.freqs,
#     "kpoints": gme.kpoints,
#     'Qs':Qs,
#     "labels": path['labels'],
#     "k_indexes": path['k_indexes']
# })
# # # # Run the guided-mode expansion
# # numeig, verbose = 20, True
# # nk = 40

# # X1 = [0.2*np.pi, 0]
# # M1 = [0.2*np.pi, 0.1*np.pi]
# # path = lattice.bz_path([X1, 'G', X1], [nk])

# # gme = legume.GuidedModeExp(phc, gmax=4.5, truncate_g='abs')   
# # gme.run(kpoints=path['kpoints'], gmode_inds=[0, 1, 2, 3], numeig=numeig, verbose=True)
# # freqs = gme.freqs
# # freqs_im = gme.freqs_im
# # nkappa, nfreq = freqs.shape[0], freqs.shape[1]
# # print(f'Number of wavevectors = {nkappa}, number of frequencies = {nfreq}')