module type_mod
    implicit none
    integer, parameter :: dp = selected_real_kind(15, 307)
    real(dp), parameter :: pi = 3.14159265358979323846_dp
end module type_mod

module legendre_mod
    use type_mod
    implicit none
contains

    ! レジャンドル多項式 P_n(x) とその微分 P'_n(x)
    subroutine leg_poly(n, x, p, dp_val)
        integer, intent(in) :: n
        real(dp), intent(in) :: x
        real(dp), intent(out) :: p, dp_val
        real(dp) :: p_prev, p_curr, p_next
        integer :: k

        if (n == 0) then
            p = 1.0_dp; dp_val = 0.0_dp; return
        end if
        p_prev = 1.0_dp
        p_curr = x
        do k = 1, n - 1
            p_next = ((2*k + 1)*x*p_curr - k*p_prev) / (k + 1)
            p_prev = p_curr
            p_curr = p_next
        end do
        p = p_curr
        dp_val = n * (x * p_curr - p_prev) / (x**2 - 1.0_dp)
    end subroutine leg_poly

    ! Gauss-Lobatto-Legendre (GLL) 点と重みの計算
    subroutine get_gll_points(n, x, w)
        integer, intent(in) :: n
        real(dp), intent(out) :: x(0:n), w(0:n)
        integer :: i, iter
        real(dp) :: z, z_new, p, dp_val, d2p_val

        x(0) = -1.0_dp
        x(n) = 1.0_dp
        do i = 1, n - 1
            z = -cos(pi * real(i, dp) / real(n, dp))
            do iter = 1, 100
                call leg_poly(n, z, p, dp_val)
                d2p_val = (2.0_dp * z * dp_val - real(n*(n+1), dp) * p) / (1.0_dp - z**2)
                z_new = z - dp_val / d2p_val
                if (abs(z_new - z) < 1.0e-14_dp) exit
                z = z_new
            end do
            x(i) = z
        end do
        do i = 0, n
            call leg_poly(n, x(i), p, dp_val)
            w(i) = 2.0_dp / (real(n*(n+1), dp) * p**2)
        end do
    end subroutine get_gll_points

    ! 微分行列 (Local)
    subroutine calc_d_matrix(n, x, d_mat)
        integer, intent(in) :: n
        real(dp), intent(in) :: x(0:n)
        real(dp), intent(out) :: d_mat(0:n, 0:n)
        integer :: i, j
        real(dp) :: p_i, p_j, dp_i, dp_j

        do i = 0, n
            do j = 0, n
                if (i /= j) then
                    call leg_poly(n, x(i), p_i, dp_i)
                    call leg_poly(n, x(j), p_j, dp_j)
                    d_mat(i, j) = (p_i / p_j) / (x(i) - x(j))
                else
                    if (i == 0) then
                        d_mat(i, j) = -real(n*(n+1), dp) / 4.0_dp
                    else if (i == n) then
                        d_mat(i, j) = real(n*(n+1), dp) / 4.0_dp
                    else
                        d_mat(i, j) = 0.0_dp
                    end if
                end if
            end do
        end do
    end subroutine calc_d_matrix
end module legendre_mod

program fedvr_pbc
    use type_mod
    use legendre_mod
    implicit none

    ! ==========================================
    ! ユーザーパラメータ設定
    ! ==========================================
    integer, parameter :: ng = 10          ! 要素内多項式の次数
    real(dp), parameter :: soft_a = 1.0_dp ! ソフトコアパラメータ

    ! 領域設定: [-L_box, L_box]
    ! Core領域: [-R_core, R_core]
    real(dp), parameter :: L_box  = 50.0_dp
    real(dp), parameter :: R_core = 1.0_dp

    ! 要素数設定
    integer, parameter :: Ne_core = 1     ! 中心領域の要素数
    integer, parameter :: Ne_side = 1     ! 片側の外側領域の要素数

    ! ==========================================

    integer, parameter :: ne = Ne_side + Ne_core + Ne_side ! 全要素数
    integer :: n_dof

    ! 配列
    real(dp), allocatable :: x_gll(:), w_gll(:)
    real(dp), allocatable :: d_local(:,:), ke_base(:,:)
    real(dp), allocatable :: h_mat(:,:), s_mat(:,:)
    real(dp), allocatable :: x_global(:), w_global(:)
    real(dp), allocatable :: boundaries(:)
    integer, allocatable :: iglob(:,:)
    real(dp), allocatable :: evals(:), work(:)

    ! 作業変数
    real(dp) :: h_len, jac_elem, current_x
    integer :: i, j, ie, idx, row, col
    integer :: info, lwork

    ! ★追加: 時間計測用変数
    real(dp) :: t_start, t_end

    ! --- 1. メモリ確保と初期化 ---
    n_dof = ne * ng ! PBCのため最後の点は最初の点と重複

    allocate(boundaries(0:ne))
    allocate(x_gll(0:ng), w_gll(0:ng))
    allocate(d_local(0:ng, 0:ng), ke_base(0:ng, 0:ng))
    allocate(h_mat(n_dof, n_dof), s_mat(n_dof, n_dof))
    allocate(x_global(n_dof), w_global(n_dof))
    allocate(iglob(ne, 0:ng))
    allocate(evals(n_dof))

    h_mat = 0.0_dp; s_mat = 0.0_dp; w_global = 0.0_dp

    ! --- 2. グリッド生成 (Manual) ---
    current_x = -L_box
    boundaries(0) = current_x

    ! 左側 (Side)
    h_len = (-R_core - (-L_box)) / real(Ne_side, dp)
    do ie = 1, Ne_side
        current_x = current_x + h_len
        boundaries(ie) = current_x
    end do

    ! 中心 (Core)
    h_len = (R_core - (-R_core)) / real(Ne_core, dp)
    do ie = Ne_side + 1, Ne_side + Ne_core
        current_x = current_x + h_len
        boundaries(ie) = current_x
    end do

    ! 右側 (Side)
    h_len = (L_box - R_core) / real(Ne_side, dp)
    do ie = Ne_side + Ne_core + 1, ne
        current_x = current_x + h_len
        boundaries(ie) = current_x
    end do

    print *, "Grid Generated: Total Elements =", ne

    ! --- 3. ローカル基底とKEベース行列 ---
    call get_gll_points(ng, x_gll, w_gll)
    call calc_d_matrix(ng, x_gll, d_local)

    ! 基底運動エネルギー行列 (ヤコビアン無し)
    do i = 0, ng
        do j = 0, ng
            ke_base(i, j) = 0.0_dp
            do idx = 0, ng
                ke_base(i, j) = ke_base(i, j) + d_local(idx, i) * w_gll(idx) * d_local(idx, j)
            end do
            ke_base(i, j) = ke_base(i, j) * 0.5_dp
        end do
    end do

    ! --- 4. グローバルマッピング (PBC) ---
    idx = 0
    do ie = 1, ne
        do i = 0, ng - 1
            idx = idx + 1
            iglob(ie, i) = idx
        end do
        if (ie < ne) then
            iglob(ie, ng) = idx + 1
        else
            iglob(ie, ng) = 1 ! PBC: Loop back to start
        end if
    end do

    ! --- 5. 行列アセンブリ ---
    do ie = 1, ne
        h_len = boundaries(ie) - boundaries(ie-1)
        jac_elem = h_len / 2.0_dp

        do i = 0, ng
            row = iglob(ie, i)
            x_global(row) = (boundaries(ie-1) + boundaries(ie)) / 2.0_dp + x_gll(i) * jac_elem

            do j = 0, ng
                col = iglob(ie, j)
                ! H += (1/jac) * KE_base
                h_mat(row, col) = h_mat(row, col) + ke_base(i, j) / jac_elem

                ! S += jac * weight
                if (i == j) then
                    s_mat(row, col) = s_mat(row, col) + w_gll(i) * jac_elem
                    w_global(row) = w_global(row) + w_gll(i) * jac_elem
                end if
            end do
        end do
    end do

    ! --- 6. ポテンシャル追加 ---
    !do i = 1, n_dof
    !   h_mat(i, i) = h_mat(i, i) + (-1.0_dp / sqrt(x_global(i)**2 + soft_a**2)) * w_global(i)
    !end do

    ! --- 7. 固有値計算 (LAPACK DSYGV) ---
    lwork = 3 * n_dof + 1000
    allocate(work(lwork))

    print *, "Solving Generalized Eigenvalue Problem (Size:", n_dof, ")..."

    ! ★追加: 計測開始
    call cpu_time(t_start)

    call dsygv(1, 'V', 'U', n_dof, h_mat, n_dof, s_mat, n_dof, evals, work, lwork, info)

    ! ★追加: 計測終了
    call cpu_time(t_end)

    if (info /= 0) then
        print *, "Error: LAPACK info =", info
        stop
    end if

    ! ★追加: 時間表示
    print *, "--------------------------------------------------"
    print *, " Diagonalization Time: ", t_end - t_start, " seconds"
    print *, "--------------------------------------------------"

    ! --- 8. 結果出力 ---

    ! (A) 固有値
    open(10, file='eigenvalues.dat', status='replace')
    write(10, '(A)') "# Index   Energy(a.u.)"
    do i = 1, min(20, n_dof)
        write(10, '(I5, 2X, ES20.12)') i-1, evals(i)
    end do
    close(10)
    print *, "Output: eigenvalues.dat"

    ! (B) 波動関数 (GS, 1st, 2nd, 3rd)
    open(20, file='wavefunctions.dat', status='replace')
    write(20, '(A)') "# x              Psi_0               Psi_1               Psi_2               Psi_3"

    ! 本体
    do i = 1, n_dof
        write(20, '(5ES20.12)') x_global(i), h_mat(i, 1), h_mat(i, 2), h_mat(i, 3), h_mat(i, 4)
    end do

    ! 右端の接続点 (PBCの描画用: x=+L の値を x=-L からコピー)
    write(20, '(5ES20.12)') boundaries(ne), h_mat(1, 1), h_mat(1, 2), h_mat(1, 3), h_mat(1, 4)

    close(20)
    print *, "Output: wavefunctions.dat"

end program fedvr_pbc