!===============================================================================
! pdsi_main.f90              University of Nebraska - Lincoln            Mar 2026
!===============================================================================
!
! Calculates Palmer Drought Severity Index (PDSI) with options for:
!   - Original or Self-calibrated
!   - Monthly or Weekly data
!   - User-defined calibration period
!
! Based on C++ version by Rob Reed and Nate Wells
! Converted to Fortran 90/95
!
!===============================================================================

module pdsi_parameters
!===============================================================================
! Module chứa các tham số và hằng số chung
!===============================================================================
    implicit none
    
    ! Định nghĩa kiểu số
    integer, parameter :: dp = kind(1.0d0)
    real(dp), parameter :: MISSING = -99.0_dp
    real(dp), parameter :: TOLERANCE = 0.00001_dp
    real(dp), parameter :: PI = 3.141592653589793_dp
    
    ! Các options cho tính toán
    integer, parameter :: OPT_ORIGINAL = 1
    integer, parameter :: OPT_SELF_CAL = 2
    integer, parameter :: OPT_MONTHLY = 1
    integer, parameter :: OPT_WEEKLY = 2
    
    ! Output modes
    integer, parameter :: OUT_TABLE = 0
    integer, parameter :: OUT_COLUMN = 1
    integer, parameter :: OUT_BOTH = 2
    
    ! Extra output options
    integer, parameter :: EXTRA_NONE = 0
    integer, parameter :: EXTRA_WB = 1
    integer, parameter :: EXTRA_XTABLE = 2
    integer, parameter :: EXTRA_POTENTIALS = 3
    integer, parameter :: EXTRA_ALL = 9
    
    ! Maximum array sizes
    integer, parameter :: MAX_PERIODS = 52
    integer, parameter :: MAX_YEARS = 200
    integer, parameter :: MAX_DATA = MAX_PERIODS * MAX_YEARS
    
end module pdsi_parameters

!===============================================================================
module pdsi_types
!===============================================================================
! Module chứa các cấu trúc dữ liệu
!===============================================================================
    use pdsi_parameters
    implicit none
    
    ! Kiểu dữ liệu cho danh sách liên kết đơn giản (dùng mảng động)
    type :: list_type
        real(dp), allocatable :: data(:)
        integer :: size = 0
        integer :: capacity = 0
    end type list_type
    
    ! Kiểu dữ liệu chính cho PDSI
    type :: pdsi_type
        ! Flags
        logical :: weekly = .false.
        logical :: monthly = .false.
        logical :: scmonthly = .false.
        logical :: bug = .false.
        logical :: verbose = .false.
        logical :: metric = .false.
        logical :: south = .false.
        logical :: nadss = .false.
        
        ! Thời gian
        integer :: startyear
        integer :: endyear
        integer :: totalyears
        integer :: period_length = 1
        integer :: num_of_periods = 52
        
        ! Calibration
        logical :: setCalibrationStartYear = .false.
        logical :: setCalibrationEndYear = .false.
        integer :: calibrationStartYear
        integer :: calibrationEndYear
        integer :: nStartYearsToSkip
        integer :: nEndYearsToSkip
        integer :: nCalibrationYears
        integer :: nStartPeriodsToSkip
        integer :: nEndPeriodsToSkip
        integer :: nCalibrationPeriods
        
        ! Các hằng số
        real(dp) :: TLA
        real(dp) :: AWC
        real(dp) :: Ss = 1.0_dp
        real(dp) :: Su
        real(dp) :: I
        real(dp) :: A
        
        ! Các mảng dữ liệu
        real(dp) :: TNorm(MAX_PERIODS)
        real(dp) :: T(MAX_PERIODS)
        real(dp) :: P(MAX_PERIODS)
        
        ! Water balance variables
        real(dp) :: ET, R, L, RO, PE, PR, PL, PRO
        
        ! Sum arrays
        real(dp) :: ETSum(MAX_PERIODS)
        real(dp) :: RSum(MAX_PERIODS)
        real(dp) :: LSum(MAX_PERIODS)
        real(dp) :: ROSum(MAX_PERIODS)
        real(dp) :: PESum(MAX_PERIODS)
        real(dp) :: PRSum(MAX_PERIODS)
        real(dp) :: PLSum(MAX_PERIODS)
        real(dp) :: PROSum(MAX_PERIODS)
        real(dp) :: PSum(MAX_PERIODS)
        
        ! Coefficients
        real(dp) :: Alpha(MAX_PERIODS)
        real(dp) :: Beta(MAX_PERIODS)
        real(dp) :: Gamma(MAX_PERIODS)
        real(dp) :: Delta(MAX_PERIODS)
        
        ! PDSI variables
        real(dp) :: Phat
        real(dp) :: d
        real(dp) :: D(MAX_PERIODS)
        real(dp) :: k(MAX_PERIODS)
        real(dp) :: K
        real(dp) :: Z
        real(dp) :: drym, dryb, wetm, wetb
        real(dp) :: dry_ratio, wet_ratio
        
        ! Book keeping
        real(dp) :: X1, X2, X3, X
        real(dp) :: Prob, V, Q
        
        ! Statistical variables
        real(dp) :: DSSqr(MAX_PERIODS)
        real(dp) :: DEPSum(MAX_PERIODS)
        real(dp) :: DKSum
        real(dp) :: SD, SD2
        
        ! Linked lists (implemented as arrays)
        type(list_type) :: Xlist
        type(list_type) :: altX1
        type(list_type) :: altX2
        type(list_type) :: XL1
        type(list_type) :: XL2
        type(list_type) :: XL3
        type(list_type) :: ProbL
        type(list_type) :: ZIND
        type(list_type) :: PeriodList
        type(list_type) :: YearList
        type(list_type) :: CMIList
        
        ! Directories
        character(len=256) :: input_dir
        character(len=256) :: output_dir
        
    end type pdsi_type
    
end module pdsi_types

!===============================================================================
module pdsi_io
!===============================================================================
! Module xử lý input/output
!===============================================================================
    use pdsi_parameters
    use pdsi_types
    implicit none
    
contains

!-------------------------------------------------------------------------------
! Subroutine đọc Namelist để cấu hình options
!-------------------------------------------------------------------------------
subroutine read_namelist(config_file, calc_type, data_type, cal_start, cal_end, &
                         input_dir, output_dir, verbose)
    implicit none
    character(len=*), intent(in) :: config_file
    integer, intent(out) :: calc_type      ! OPT_ORIGINAL or OPT_SELF_CAL
    integer, intent(out) :: data_type      ! OPT_MONTHLY or OPT_WEEKLY
    integer, intent(out) :: cal_start      ! Calibration start year
    integer, intent(out) :: cal_end        ! Calibration end year
    character(len=*), intent(out) :: input_dir
    character(len=*), intent(out) :: output_dir
    logical, intent(out) :: verbose
    
    ! Namelist definition
    namelist /pdsi_options/ calc_type, data_type, cal_start, cal_end, &
                            input_dir, output_dir, verbose
    
    integer :: unit
    
    ! Default values
    calc_type = OPT_SELF_CAL
    data_type = OPT_MONTHLY
    cal_start = 0
    cal_end = 0
    input_dir = './'
    output_dir = './results/'
    verbose = .false.
    
    ! Read namelist file
    open(newunit=unit, file=config_file, status='old', action='read', iostat=cal_start)
    if (cal_start == 0) then
        read(unit, nml=pdsi_options, iostat=cal_end)
        if (cal_end /= 0) then
            print *, "Warning: Error reading namelist, using defaults"
        end if
        close(unit)
    else
        print *, "Warning: Config file not found, using defaults"
    end if
    
end subroutine read_namelist

!-------------------------------------------------------------------------------
! Subroutine đọc số liệu dựa trên options
!-------------------------------------------------------------------------------
subroutine read_data(pdsi, data_type, iostat)
    type(pdsi_type), intent(inout) :: pdsi
    integer, intent(in) :: data_type
    integer, intent(out) :: iostat
    
    character(len=256) :: temp_file, precip_file, norm_file, param_file
    integer :: unit_temp, unit_precip, unit_norm, unit_param
    
    iostat = 0
    
    ! Xác định tên file dựa trên data_type
    if (data_type == OPT_MONTHLY) then
        temp_file = trim(pdsi%input_dir) // 'monthly_T.txt'
        precip_file = trim(pdsi%input_dir) // 'monthly_P.txt'
        norm_file = trim(pdsi%input_dir) // 'mon_T_normal.txt'
    else  ! Weekly
        temp_file = trim(pdsi%input_dir) // 'weekly_T.txt'
        precip_file = trim(pdsi%input_dir) // 'weekly_P.txt'
        norm_file = trim(pdsi%input_dir) // 'wk_T_normal.txt'
    end if
    param_file = trim(pdsi%input_dir) // 'parameter.txt'
    
    ! Đọc normal temperatures
    open(newunit=unit_norm, file=norm_file, status='old', action='read', iostat=iostat)
    if (iostat /= 0) then
        print *, "Error: Cannot open normal temperature file: ", trim(norm_file)
        return
    end if
    call read_normal_temp(unit_norm, pdsi%TNorm, pdsi%num_of_periods, iostat)
    close(unit_norm)
    if (iostat /= 0) return
    
    ! Đọc parameter file (AWC và latitude)
    open(newunit=unit_param, file=param_file, status='old', action='read', iostat=iostat)
    if (iostat /= 0) then
        print *, "Error: Cannot open parameter file: ", trim(param_file)
        return
    end if
    call read_parameter(unit_param, pdsi%AWC, pdsi%TLA, pdsi%metric, pdsi%nadss, &
                        pdsi%south, iostat)
    close(unit_param)
    if (iostat /= 0) return
    
    ! Đọc temperature và precipitation data
    open(newunit=unit_temp, file=temp_file, status='old', action='read', iostat=iostat)
    if (iostat /= 0) then
        print *, "Error: Cannot open temperature file: ", trim(temp_file)
        return
    end if
    
    open(newunit=unit_precip, file=precip_file, status='old', action='read', iostat=iostat)
    if (iostat /= 0) then
        print *, "Error: Cannot open precipitation file: ", trim(precip_file)
        close(unit_temp)
        return
    end if
    
    ! Đọc dữ liệu theo năm
    call read_yearly_data(unit_temp, unit_precip, pdsi, data_type, iostat)
    
    close(unit_temp)
    close(unit_precip)
    
end subroutine read_data

!-------------------------------------------------------------------------------
! Subroutine đọc normal temperatures
!-------------------------------------------------------------------------------
subroutine read_normal_temp(unit, TNorm, nperiods, iostat)
    integer, intent(in) :: unit
    real(dp), intent(out) :: TNorm(:)
    integer, intent(in) :: nperiods
    integer, intent(out) :: iostat
    
    real :: temp
    integer :: i
    
    do i = 1, nperiods
        read(unit, *, iostat=iostat) temp
        if (iostat /= 0) exit
        TNorm(i) = real(temp, dp)
    end do
    
end subroutine read_normal_temp

!-------------------------------------------------------------------------------
! Subroutine đọc parameter file
!-------------------------------------------------------------------------------
subroutine read_parameter(unit, AWC, TLA, metric, nadss, south, iostat)
    integer, intent(in) :: unit
    real(dp), intent(out) :: AWC, TLA
    logical, intent(in) :: metric, nadss
    logical, intent(out) :: south
    integer, intent(out) :: iostat
    
    real :: scn1, scn2
    real(dp) :: lat
    
    read(unit, *, iostat=iostat) scn1, scn2
    if (iostat /= 0) return
    
    AWC = real(scn1, dp)
    if (metric) AWC = AWC / 25.4_dp
    
    if (nadss) then
        TLA = real(scn2, dp)
        if (TLA > 0) then
            south = .true.
            TLA = -TLA
        else
            south = .false.
        end if
    else
        lat = real(scn2, dp)
        south = (lat >= 0)
        TLA = -tan(PI * lat / 180.0_dp)
        if (.not. south) TLA = -TLA
    end if
    
end subroutine read_parameter

!-------------------------------------------------------------------------------
! Subroutine đọc dữ liệu theo năm
!-------------------------------------------------------------------------------
subroutine read_yearly_data(unit_temp, unit_precip, pdsi, data_type, iostat)
    integer, intent(in) :: unit_temp, unit_precip
    type(pdsi_type), intent(inout) :: pdsi
    integer, intent(in) :: data_type
    integer, intent(out) :: iostat
    
    integer :: year, i, j, nweeks
    real :: t(MAX_PERIODS), p(MAX_PERIODS)
    real(dp) :: temp_sum
    integer :: valid_count
    
    iostat = 0
    nweeks = 52  ! Weekly data has 52 values
    
    do
        ! Đọc năm và dữ liệu temperature
        if (data_type == OPT_MONTHLY) then
            read(unit_temp, *, iostat=iostat) year, (t(i), i=1,12)
        else
            read(unit_temp, *, iostat=iostat) year, (t(i), i=1,nweeks)
        end if
        if (iostat /= 0) exit
        
        ! Đọc dữ liệu precipitation
        if (data_type == OPT_MONTHLY) then
            read(unit_precip, *, iostat=iostat) year, (p(i), i=1,12)
        else
            read(unit_precip, *, iostat=iostat) year, (p(i), i=1,nweeks)
        end if
        if (iostat /= 0) exit
        
        ! Xử lý dữ liệu theo period_length
        do i = 1, pdsi%num_of_periods
            temp_sum = 0.0_dp
            valid_count = 0
            
            ! Sum/Average data for this period
            do j = 1, pdsi%period_length
                if (t((i-1)*pdsi%period_length + j) > -90.0) then  ! Not missing
                    temp_sum = temp_sum + real(t((i-1)*pdsi%period_length + j), dp)
                    valid_count = valid_count + 1
                end if
            end do
            
            if (valid_count > 0) then
                pdsi%T(i) = temp_sum / real(valid_count, dp)
            else
                pdsi%T(i) = MISSING
            end if
            
            ! Tương tự cho precipitation (sum, không average)
            temp_sum = 0.0_dp
            valid_count = 0
            do j = 1, pdsi%period_length
                if (p((i-1)*pdsi%period_length + j) >= 0.0) then
                    temp_sum = temp_sum + real(p((i-1)*pdsi%period_length + j), dp)
                    valid_count = valid_count + 1
                end if
            end do
            
            if (valid_count > 0) then
                pdsi%P(i) = temp_sum
            else
                pdsi%P(i) = MISSING
            end if
        end do
        
        ! Lưu năm và period
        call list_append(pdsi%YearList, real(year, dp))
        do i = 1, pdsi%num_of_periods
            call list_append(pdsi%PeriodList, real(i, dp))
        end do
        
    end do
    
end subroutine read_yearly_data

!-------------------------------------------------------------------------------
! Subroutine lưu kết quả ra file text
!-------------------------------------------------------------------------------
subroutine write_results(pdsi, output_dir, calc_type, data_type)
    type(pdsi_type), intent(in) :: pdsi
    character(len=*), intent(in) :: output_dir
    integer, intent(in) :: calc_type, data_type
    
    character(len=256) :: filename
    integer :: unit, i, n
    real(dp), allocatable :: years(:), periods(:), values(:)
    
    ! Tạo directory nếu chưa tồn tại
    call create_directory(output_dir)
    
    ! Ghi PDSI
    filename = trim(output_dir) // 'PDSI.txt'
    open(newunit=unit, file=filename, status='replace', action='write')
    write(unit, '(A)') '# Year  Period  PDSI'
    n = pdsi%Xlist%size
    if (n > 0) then
        allocate(years(n), periods(n), values(n))
        call list_to_array(pdsi%YearList, years)
        call list_to_array(pdsi%PeriodList, periods)
        call list_to_array(pdsi%Xlist, values)
        do i = 1, n
            write(unit, '(I6, I6, F12.2)') int(years(i)), int(periods(i)), values(i)
        end do
        deallocate(years, periods, values)
    end if
    close(unit)
    
    ! Ghi PHDI
    filename = trim(output_dir) // 'PHDI.txt'
    open(newunit=unit, file=filename, status='replace', action='write')
    write(unit, '(A)') '# Year  Period  PHDI'
    n = pdsi%Xlist%size
    if (n > 0) then
        allocate(years(n), periods(n), values(n))
        call list_to_array(pdsi%YearList, years)
        call list_to_array(pdsi%PeriodList, periods)
        do i = 1, n
            values(i) = calculate_phdi(pdsi, i)
        end do
        do i = 1, n
            write(unit, '(I6, I6, F12.2)') int(years(i)), int(periods(i)), values(i)
        end do
        deallocate(years, periods, values)
    end if
    close(unit)
    
    ! Ghi WPLM
    filename = trim(output_dir) // 'WPLM.txt'
    open(newunit=unit, file=filename, status='replace', action='write')
    write(unit, '(A)') '# Year  Period  WPLM'
    n = pdsi%Xlist%size
    if (n > 0) then
        allocate(years(n), periods(n), values(n))
        call list_to_array(pdsi%YearList, years)
        call list_to_array(pdsi%PeriodList, periods)
        do i = 1, n
            values(i) = calculate_wplm(pdsi, i)
        end do
        do i = 1, n
            write(unit, '(I6, I6, F12.2)') int(years(i)), int(periods(i)), values(i)
        end do
        deallocate(years, periods, values)
    end if
    close(unit)
    
    ! Ghi ZIND
    filename = trim(output_dir) // 'ZIND.txt'
    open(newunit=unit, file=filename, status='replace', action='write')
    write(unit, '(A)') '# Year  Period  ZIND'
    n = pdsi%ZIND%size
    if (n > 0) then
        allocate(years(n), periods(n), values(n))
        call list_to_array(pdsi%YearList, years)
        call list_to_array(pdsi%PeriodList, periods)
        call list_to_array(pdsi%ZIND, values)
        do i = 1, n
            write(unit, '(I6, I6, F12.2)') int(years(i)), int(periods(i)), values(i)
        end do
        deallocate(years, periods, values)
    end if
    close(unit)
    
    print *, "Results written to ", trim(output_dir)
    
end subroutine write_results

!-------------------------------------------------------------------------------
! Function tính PHDI
!-------------------------------------------------------------------------------
function calculate_phdi(pdsi, idx) result(phdi)
    type(pdsi_type), intent(in) :: pdsi
    integer, intent(in) :: idx
    real(dp) :: phdi
    
    real(dp) :: x, x3
    
    if (idx <= pdsi%Xlist%size .and. idx <= pdsi%XL3%size) then
        x = pdsi%Xlist%data(idx)
        x3 = pdsi%XL3%data(idx)
        if (abs(x3) < TOLERANCE) then
            phdi = x
        else
            phdi = x3
        end if
    else
        phdi = MISSING
    end if
    
end function calculate_phdi

!-------------------------------------------------------------------------------
! Function tính WPLM
!-------------------------------------------------------------------------------
function calculate_wplm(pdsi, idx) result(wplm)
    type(pdsi_type), intent(in) :: pdsi
    integer, intent(in) :: idx
    real(dp) :: wplm
    
    real(dp) :: x1, x2, x3, p
    
    if (idx <= pdsi%XL1%size .and. idx <= pdsi%XL2%size .and. &
        idx <= pdsi%XL3%size .and. idx <= pdsi%ProbL%size) then
        
        x1 = pdsi%XL1%data(idx)
        x2 = pdsi%XL2%data(idx)
        x3 = pdsi%XL3%data(idx)
        p = pdsi%ProbL%data(idx) / 100.0_dp
        
        if (abs(x3) < TOLERANCE) then
            ! No established spell
            wplm = x1
            if ((-x2) > (x1 + TOLERANCE)) wplm = x2
        else if (p > TOLERANCE/100.0_dp .and. p < (1.0_dp - TOLERANCE/100.0_dp)) then
            ! Spell with possibility of ending
            if (x3 < 0) then
                wplm = (1.0_dp - p) * x3 + p * x1
            else
                wplm = (1.0_dp - p) * x3 + p * x2
            end if
        else
            ! Established spell
            wplm = x3
        end if
    else
        wplm = MISSING
    end if
    
end function calculate_wplm

end module pdsi_io

!===============================================================================
module pdsi_core
!===============================================================================
! Module chứa các tính toán chính của PDSI
!===============================================================================
    use pdsi_parameters
    use pdsi_types
    implicit none
    
contains

!-------------------------------------------------------------------------------
! Subroutine khởi tạo PDSI
!-------------------------------------------------------------------------------
subroutine pdsi_init(pdsi, calc_type, data_type, cal_start, cal_end)
    type(pdsi_type), intent(inout) :: pdsi
    integer, intent(in) :: calc_type
    integer, intent(in) :: data_type
    integer, intent(in) :: cal_start, cal_end
    
    ! Set calculation type
    if (calc_type == OPT_ORIGINAL) then
        pdsi%monthly = (data_type == OPT_MONTHLY)
        pdsi%weekly = (data_type == OPT_WEEKLY)
        pdsi%scmonthly = .false.
    else  ! Self-calibrated
        pdsi%monthly = .false.
        pdsi%weekly = (data_type == OPT_WEEKLY)
        pdsi%scmonthly = (data_type == OPT_MONTHLY)
    end if
    
    ! Set data type specific parameters
    if (data_type == OPT_MONTHLY) then
        pdsi%period_length = 1
        pdsi%num_of_periods = 12
    else  ! Weekly
        pdsi%period_length = 1
        pdsi%num_of_periods = 52
    end if
    
    ! Set calibration years
    if (cal_start > 0) then
        pdsi%setCalibrationStartYear = .true.
        pdsi%calibrationStartYear = cal_start
    else
        pdsi%calibrationStartYear = pdsi%startyear
    end if
    
    if (cal_end > 0) then
        pdsi%setCalibrationEndYear = .true.
        pdsi%calibrationEndYear = cal_end
    else
        pdsi%calibrationEndYear = pdsi%endyear
    end if
    
    ! Initialize soil moisture
    pdsi%Ss = 1.0_dp
    pdsi%Su = max(pdsi%AWC - pdsi%Ss, 0.0_dp)
    
    ! Initialize arrays
    call init_arrays(pdsi)
    
end subroutine pdsi_init

!-------------------------------------------------------------------------------
! Subroutine khởi tạo các mảng
!-------------------------------------------------------------------------------
subroutine init_arrays(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    integer :: i
    
    do i = 1, MAX_PERIODS
        pdsi%ETSum(i) = 0.0_dp
        pdsi%RSum(i) = 0.0_dp
        pdsi%LSum(i) = 0.0_dp
        pdsi%ROSum(i) = 0.0_dp
        pdsi%PESum(i) = 0.0_dp
        pdsi%PRSum(i) = 0.0_dp
        pdsi%PLSum(i) = 0.0_dp
        pdsi%PROSum(i) = 0.0_dp
        pdsi%PSum(i) = 0.0_dp
        pdsi%DSSqr(i) = 0.0_dp
        pdsi%DEPSum(i) = 0.0_dp
        pdsi%D(i) = 0.0_dp
        pdsi%k(i) = 0.0_dp
    end do
    
    pdsi%SD = 0.0_dp
    pdsi%SD2 = 0.0_dp
    pdsi%DKSum = 0.0_dp
    
end subroutine init_arrays

!-------------------------------------------------------------------------------
! Subroutine tính toán các chỉ số PDSI
!-------------------------------------------------------------------------------
subroutine calculate_pdsi(pdsi, calc_type)
    type(pdsi_type), intent(inout) :: pdsi
    integer, intent(in) :: calc_type
    
    ! Tính Thornthwaite heat index và exponent
    if (pdsi%weekly) then
        pdsi%I = calculate_weekly_thorn_i(pdsi)
    else
        pdsi%I = calculate_monthly_thorn_i(pdsi)
    end if
    pdsi%A = calculate_thorn_a(pdsi%I)
    
    ! Sum all water balance variables
    call sum_all(pdsi)
    
    ! Calculate water balance coefficients
    call calc_wb_coef(pdsi)
    
    ! Calculate d and D
    call calc_d(pdsi)
    
    ! Calculate K
    call calc_k(pdsi)
    
    ! Calculate Z
    call calc_z(pdsi)
    
    ! Calculate PDSI based on type
    if (calc_type == OPT_ORIGINAL) then
        call calc_original_pdsi(pdsi)
    else  ! Self-calibrated
        call calc_self_cal_pdsi(pdsi)
    end if
    
end subroutine calculate_pdsi

!-------------------------------------------------------------------------------
! Function tính Thornthwaite heat index I (weekly)
!-------------------------------------------------------------------------------
function calculate_weekly_thorn_i(pdsi) result(I)
    type(pdsi_type), intent(in) :: pdsi
    real(dp) :: I
    
    integer :: i
    real(dp) :: temp
    
    I = 0.0_dp
    do i = 1, 52
        if (pdsi%TNorm(i) > 32.0_dp) then
            temp = (pdsi%TNorm(i) - 32.0_dp) / 9.0_dp
            I = I + temp**1.514_dp
        end if
    end do
    
    ! Adjust for weekly to monthly
    I = I / 52.0_dp * 12.0_dp
    
end function calculate_weekly_thorn_i

!-------------------------------------------------------------------------------
! Function tính Thornthwaite heat index I (monthly)
!-------------------------------------------------------------------------------
function calculate_monthly_thorn_i(pdsi) result(I)
    type(pdsi_type), intent(in) :: pdsi
    real(dp) :: I
    
    integer :: i
    real(dp) :: temp
    
    I = 0.0_dp
    do i = 1, 12
        if (pdsi%TNorm(i) > 32.0_dp) then
            temp = (pdsi%TNorm(i) - 32.0_dp) / 9.0_dp
            I = I + temp**1.514_dp
        end if
    end do
    
end function calculate_monthly_thorn_i

!-------------------------------------------------------------------------------
! Function tính Thornthwaite exponent A
!-------------------------------------------------------------------------------
function calculate_thorn_a(I) result(A)
    real(dp), intent(in) :: I
    real(dp) :: A
    
    A = 6.75e-7_dp * I**3 - 7.71e-5_dp * I**2 + 0.0179_dp * I + 0.49_dp
    
end function calculate_thorn_a

!-------------------------------------------------------------------------------
! Subroutine tính tổng các water balance variables
!-------------------------------------------------------------------------------
subroutine sum_all(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    integer :: year, per, i
    real(dp) :: DEP
    integer :: nCalibrationPeriodsLeft
    
    nCalibrationPeriodsLeft = pdsi%nCalibrationPeriods
    DEP = 0.0_dp
    
    ! Loop through all years
    do year = 1, pdsi%totalyears
        ! Loop through all periods
        do per = 1, pdsi%num_of_periods
            if (pdsi%P(per) >= 0.0_dp .and. pdsi%T(per) /= MISSING) then
                ! Calculate potentials
                if (pdsi%weekly) then
                    call calc_weekly_pe(pdsi, per, year)
                else
                    call calc_monthly_pe(pdsi, per, year)
                end if
                
                call calc_pr(pdsi)
                call calc_pro(pdsi)
                call calc_pl(pdsi)
                call calc_actual(pdsi, per)
                
                ! Statistical calculations for output
                if (pdsi%weekly) then
                    if (per > 17/pdsi%period_length .and. per < 35/pdsi%period_length) then
                        DEP = DEP + pdsi%P(per) + pdsi%L - pdsi%PE
                        if (per > 30/pdsi%period_length .and. per < 35/pdsi%period_length) then
                            pdsi%SD = pdsi%SD + DEP
                            pdsi%SD2 = pdsi%SD2 + DEP * DEP
                            DEP = 0.0_dp
                        end if
                    end if
                else  ! Monthly
                    if (per > 4 .and. per < 8) then
                        DEP = DEP + pdsi%P(per) + pdsi%L - pdsi%PE
                        if (per == 7) then
                            pdsi%SD = pdsi%SD + DEP
                            pdsi%SD2 = pdsi%SD2 + DEP * DEP
                            DEP = 0.0_dp
                        end if
                    end if
                end if
                
                ! Update sums if within calibration period
                if (year > pdsi%nStartYearsToSkip .and. nCalibrationPeriodsLeft > 0) then
                    nCalibrationPeriodsLeft = nCalibrationPeriodsLeft - 1
                    
                    pdsi%ETSum(per) = pdsi%ETSum(per) + pdsi%ET
                    pdsi%RSum(per) = pdsi%RSum(per) + pdsi%R
                    pdsi%ROSum(per) = pdsi%ROSum(per) + pdsi%RO
                    pdsi%LSum(per) = pdsi%LSum(per) + pdsi%L
                    pdsi%PSum(per) = pdsi%PSum(per) + pdsi%P(per)
                    pdsi%PESum(per) = pdsi%PESum(per) + pdsi%PE
                    pdsi%PRSum(per) = pdsi%PRSum(per) + pdsi%PR
                    pdsi%PROSum(per) = pdsi%PROSum(per) + pdsi%PRO
                    pdsi%PLSum(per) = pdsi%PLSum(per) + pdsi%PL
                end if
            end if
        end do
    end do
    
    ! Calculate DEPSum
    do per = 1, pdsi%num_of_periods
        pdsi%DEPSum(per) = pdsi%ETSum(per) + pdsi%RSum(per) - pdsi%PESum(per) + pdsi%ROSum(per)
    end do
    
end subroutine sum_all

!-------------------------------------------------------------------------------
! Subroutine tính Potential Evapotranspiration (weekly)
!-------------------------------------------------------------------------------
subroutine calc_weekly_pe(pdsi, period, year)
    type(pdsi_type), intent(inout) :: pdsi
    integer, intent(in) :: period, year
    
    real(dp) :: Phi(52)
    real(dp) :: adjusted_Phi(52), Dum, Dk
    integer :: i, j, offset
    
    ! Phi constants (daylight factor)
    Phi = [ -0.39796983_dp, -0.38279708_dp, -0.36154002_dp, -0.33468535_dp, &
            -0.30281503_dp, -0.26657407_dp, -0.22664177_dp, -0.18370885_dp, &
            -0.13846106_dp, -0.09156812_dp, -0.04368365_dp, 0.00456585_dp, &
            0.05256951_dp, 0.09973421_dp, 0.14548036_dp, 0.18923931_dp, &
            0.23045239_dp, 0.26857222_dp, 0.30306721_dp, 0.33342965_dp, &
            0.35918741_dp, 0.37991898_dp, 0.39527047_dp, 0.40497267_dp, &
            0.40885593_dp, 0.40686015_dp, 0.39903840_dp, 0.38555316_dp, &
            0.36666599_dp, 0.34272229_dp, 0.31413363_dp, 0.28136006_dp, &
            0.24489444_dp, 0.20524994_dp, 0.16295153_dp, 0.11853109_dp, &
            0.07252590_dp, 0.02547986_dp, -0.02205355_dp, -0.06950714_dp, &
            -0.11629575_dp, -0.16181268_dp, -0.20542804_dp, -0.24649005_dp, &
            -0.28433068_dp, -0.31827697_dp, -0.34766880_dp, -0.37188342_dp, &
            -0.39036560_dp, -0.40266061_dp, -0.40844587_dp, -0.40755601_dp ]
    
    if (pdsi%south) then
        offset = 26
    else
        offset = 0
    end if
    
    ! Average Phi for the period
    do i = 1, pdsi%num_of_periods
        adjusted_Phi(i) = 0.0_dp
        do j = 1, pdsi%period_length
            adjusted_Phi(i) = adjusted_Phi(i) + Phi(mod((i-1)*pdsi%period_length + j + offset - 1, 52) + 1)
        end do
        adjusted_Phi(i) = adjusted_Phi(i) / pdsi%period_length
    end do
    
    if (pdsi%T(period) <= 32.0_dp) then
        pdsi%PE = 0.0_dp
    else
        Dum = adjusted_Phi(period) * pdsi%TLA
        Dk = atan2(sqrt(1.0_dp - Dum*Dum), Dum)
        if (Dk < 0.0_dp) Dk = Dk + PI
        Dk = (Dk + 0.0157_dp) / 1.57_dp
        
        if (pdsi%T(period) >= 80.0_dp) then
            pdsi%PE = (sin(pdsi%T(period)/57.3_dp - 0.166_dp) - 0.76_dp) * Dk
        else
            Dum = log(pdsi%T(period) - 32.0_dp)
            pdsi%PE = exp(-3.863233_dp + pdsi%A*1.715598_dp - pdsi%A*log(pdsi%I) + pdsi%A*Dum) * Dk
        end if
    end if
    
    pdsi%PE = pdsi%PE * pdsi%period_length * 7.0_dp
    
end subroutine calc_weekly_pe

!-------------------------------------------------------------------------------
! Subroutine tính Potential Evapotranspiration (monthly)
!-------------------------------------------------------------------------------
subroutine calc_monthly_pe(pdsi, month, year)
    type(pdsi_type), intent(inout) :: pdsi
    integer, intent(in) :: month, year
    
    real(dp) :: Phi(12)
    real(dp) :: Dum, Dk
    integer :: offset
    integer :: Days(12)
    
    ! Phi constants for monthly
    Phi = [ -0.3865982_dp, -0.2316132_dp, -0.0378180_dp, 0.1715539_dp, &
            0.3458803_dp, 0.4308320_dp, 0.3916645_dp, 0.2452467_dp, &
            0.0535511_dp, -0.15583436_dp, -0.3340551_dp, -0.4310691_dp ]
    
    ! Days in month
    Days = [ 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31 ]
    
    if (pdsi%south) then
        offset = 6
    else
        offset = 0
    end if
    
    if (pdsi%T(month) <= 32.0_dp) then
        pdsi%PE = 0.0_dp
    else
        Dum = Phi(mod(month + offset - 1, 12) + 1) * pdsi%TLA
        Dk = atan2(sqrt(1.0_dp - Dum*Dum), Dum)
        if (Dk < 0.0_dp) Dk = Dk + PI
        Dk = (Dk + 0.0157_dp) / 1.57_dp
        
        if (pdsi%T(month) >= 80.0_dp) then
            pdsi%PE = (sin(pdsi%T(month)/57.3_dp - 0.166_dp) - 0.76_dp) * Dk
        else
            Dum = log(pdsi%T(month) - 32.0_dp)
            pdsi%PE = exp(-3.863233_dp + pdsi%A*1.715598_dp - pdsi%A*log(pdsi%I) + pdsi%A*Dum) * Dk
        end if
    end if
    
    ! Adjust for month length and leap year
    if (month == 2) then
        if (mod(year, 400) == 0 .or. (mod(year, 4) == 0 .and. mod(year, 100) /= 0)) then
            pdsi%PE = pdsi%PE * 29.0_dp
        else
            pdsi%PE = pdsi%PE * 28.0_dp
        end if
    else
        pdsi%PE = pdsi%PE * Days(month)
    end if
    
end subroutine calc_monthly_pe

!-------------------------------------------------------------------------------
! Subroutine tính Potential Recharge
!-------------------------------------------------------------------------------
subroutine calc_pr(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    pdsi%PR = pdsi%AWC - (pdsi%Su + pdsi%Ss)
    
end subroutine calc_pr

!-------------------------------------------------------------------------------
! Subroutine tính Potential Runoff
!-------------------------------------------------------------------------------
subroutine calc_pro(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    pdsi%PRO = pdsi%Ss + pdsi%Su
    
end subroutine calc_pro

!-------------------------------------------------------------------------------
! Subroutine tính Potential Loss
!-------------------------------------------------------------------------------
subroutine calc_pl(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    if (pdsi%Ss >= pdsi%PE) then
        pdsi%PL = pdsi%PE
    else
        pdsi%PL = ((pdsi%PE - pdsi%Ss) * pdsi%Su) / pdsi%AWC + pdsi%Ss
        if (pdsi%PL > pdsi%PRO) pdsi%PL = pdsi%PRO
    end if
    
end subroutine calc_pl

!-------------------------------------------------------------------------------
! Subroutine tính Actual values
!-------------------------------------------------------------------------------
subroutine calc_actual(pdsi, per)
    type(pdsi_type), intent(inout) :: pdsi
    integer, intent(in) :: per
    
    real(dp) :: R_surface, R_under, surface_L, under_L
    real(dp) :: new_Su, new_Ss
    
    if (pdsi%P(per) >= pdsi%PE) then
        ! Excess moisture
        pdsi%ET = pdsi%PE
        pdsi%L = 0.0_dp
        
        if ((pdsi%P(per) - pdsi%PE) > (1.0_dp - pdsi%Ss)) then
            ! Recharge both layers
            R_surface = 1.0_dp - pdsi%Ss
            new_Ss = 1.0_dp
            
            if ((pdsi%P(per) - pdsi%PE - R_surface) < ((pdsi%AWC - 1.0_dp) - pdsi%Su)) then
                R_under = pdsi%P(per) - pdsi%PE - R_surface
                pdsi%RO = 0.0_dp
            else
                R_under = (pdsi%AWC - 1.0_dp) - pdsi%Su
                pdsi%RO = pdsi%P(per) - pdsi%PE - (R_surface + R_under)
            end if
            new_Su = pdsi%Su + R_under
            pdsi%R = R_surface + R_under
        else
            ! Recharge only surface layer
            pdsi%R = pdsi%P(per) - pdsi%PE
            new_Ss = pdsi%Ss + pdsi%R
            new_Su = pdsi%Su
            pdsi%RO = 0.0_dp
        end if
    else
        ! Moisture deficit
        if (pdsi%Ss > (pdsi%PE - pdsi%P(per))) then
            ! Loss only from surface
            surface_L = pdsi%PE - pdsi%P(per)
            under_L = 0.0_dp
            new_Ss = pdsi%Ss - surface_L
            new_Su = pdsi%Su
        else
            ! Loss from both layers
            surface_L = pdsi%Ss
            under_L = (pdsi%PE - pdsi%P(per) - surface_L) * pdsi%Su / pdsi%AWC
            if (pdsi%Su < under_L) under_L = pdsi%Su
            new_Ss = 0.0_dp
            new_Su = pdsi%Su - under_L
        end if
        pdsi%R = 0.0_dp
        pdsi%L = under_L + surface_L
        pdsi%RO = 0.0_dp
        pdsi%ET = pdsi%P(per) + pdsi%L
    end if
    
    ! Update soil moisture
    pdsi%Ss = new_Ss
    pdsi%Su = new_Su
    
end subroutine calc_actual

!-------------------------------------------------------------------------------
! Subroutine tính water balance coefficients
!-------------------------------------------------------------------------------
subroutine calc_wb_coef(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    integer :: per
    
    do per = 1, pdsi%num_of_periods
        ! Alpha (ET/PE)
        if (abs(pdsi%PESum(per)) > TOLERANCE) then
            pdsi%Alpha(per) = pdsi%ETSum(per) / pdsi%PESum(per)
        else if (abs(pdsi%ETSum(per)) < TOLERANCE) then
            pdsi%Alpha(per) = 1.0_dp
        else
            pdsi%Alpha(per) = 0.0_dp
        end if
        
        ! Beta (R/PR)
        if (abs(pdsi%PRSum(per)) > TOLERANCE) then
            pdsi%Beta(per) = pdsi%RSum(per) / pdsi%PRSum(per)
        else if (abs(pdsi%RSum(per)) < TOLERANCE) then
            pdsi%Beta(per) = 1.0_dp
        else
            pdsi%Beta(per) = 0.0_dp
        end if
        
        ! Gamma (RO/PRO)
        if (abs(pdsi%PROSum(per)) > TOLERANCE) then
            pdsi%Gamma(per) = pdsi%ROSum(per) / pdsi%PROSum(per)
        else if (abs(pdsi%ROSum(per)) < TOLERANCE) then
            pdsi%Gamma(per) = 1.0_dp
        else
            pdsi%Gamma(per) = 0.0_dp
        end if
        
        ! Delta (L/PL)
        if (abs(pdsi%PLSum(per)) > TOLERANCE) then
            pdsi%Delta(per) = pdsi%LSum(per) / pdsi%PLSum(per)
        else
            pdsi%Delta(per) = 0.0_dp
        end if
    end do
    
end subroutine calc_wb_coef

!-------------------------------------------------------------------------------
! Subroutine tính d (moisture departure)
!-------------------------------------------------------------------------------
subroutine calc_d(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    integer :: yr, per, i
    real(dp) :: d_abs_sum(MAX_PERIODS)
    
    d_abs_sum = 0.0_dp
    pdsi%DSSqr = 0.0_dp
    
    ! Simulate reading from potentials file - in Fortran we'll use arrays
    do yr = 1, pdsi%totalyears
        do per = 1, pdsi%num_of_periods
            if (pdsi%P(per) /= MISSING .and. pdsi%PE /= MISSING) then
                ! Calculate Phat
                pdsi%Phat = pdsi%Alpha(per) * pdsi%PE + &
                           pdsi%Beta(per) * pdsi%PR + &
                           pdsi%Gamma(per) * pdsi%PRO - &
                           pdsi%Delta(per) * pdsi%PL
                
                pdsi%d = pdsi%P(per) - pdsi%Phat
                
                ! Update statistics if in calibration period
                if (yr >= pdsi%calibrationStartYear .and. yr <= pdsi%calibrationEndYear) then
                    d_abs_sum(per) = d_abs_sum(per) + abs(pdsi%d)
                    pdsi%DSSqr(per) = pdsi%DSSqr(per) + pdsi%d * pdsi%d
                end if
            else
                pdsi%d = MISSING
            end if
        end do
    end do
    
    ! Calculate D (mean of absolute d)
    do per = 1, pdsi%num_of_periods
        if (pdsi%nCalibrationYears > 0) then
            pdsi%D(per) = d_abs_sum(per) / pdsi%nCalibrationYears
        else
            pdsi%D(per) = 0.0_dp
        end if
    end do
    
end subroutine calc_d

!-------------------------------------------------------------------------------
! Subroutine tính K
!-------------------------------------------------------------------------------
subroutine calc_k(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    integer :: per
    real(dp) :: sums
    
    pdsi%DKSum = 0.0_dp
    
    do per = 1, pdsi%num_of_periods
        if (abs(pdsi%PSum(per) + pdsi%LSum(per)) > TOLERANCE) then
            sums = (pdsi%PESum(per) + pdsi%RSum(per) + pdsi%ROSum(per)) / &
                   (pdsi%PSum(per) + pdsi%LSum(per))
        else
            sums = 0.0_dp
        end if
        
        if (abs(pdsi%D(per)) > TOLERANCE) then
            pdsi%k(per) = 1.5_dp * log10((sums + 2.8_dp) / pdsi%D(per)) + 0.5_dp
        else
            pdsi%k(per) = 0.5_dp
        end if
        
        pdsi%DKSum = pdsi%DKSum + pdsi%D(per) * pdsi%k(per)
    end do
    
end subroutine calc_k

!-------------------------------------------------------------------------------
! Subroutine tính Z index
!-------------------------------------------------------------------------------
subroutine calc_z(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    integer :: yr, per
    
    ! Clear ZIND list
    call list_clear(pdsi%ZIND)
    
    do yr = 1, pdsi%totalyears
        do per = 1, pdsi%num_of_periods
            if (pdsi%d /= MISSING) then
                pdsi%Z = pdsi%d * pdsi%k(per)
            else
                pdsi%Z = MISSING
            end if
            call list_append(pdsi%ZIND, pdsi%Z)
        end do
    end do
    
end subroutine calc_z

!-------------------------------------------------------------------------------
! Subroutine tính original PDSI
!-------------------------------------------------------------------------------
subroutine calc_original_pdsi(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    ! Set duration factors
    if (pdsi%weekly) then
        pdsi%drym = 2.925_dp
        pdsi%dryb = 0.075_dp
    else
        pdsi%drym = 0.309_dp
        pdsi%dryb = 2.691_dp
    end if
    pdsi%wetm = pdsi%drym
    pdsi%wetb = pdsi%dryb
    
    ! Initialize bookkeeping
    pdsi%Prob = 0.0_dp
    pdsi%X1 = 0.0_dp
    pdsi%X2 = 0.0_dp
    pdsi%X3 = 0.0_dp
    pdsi%X = 0.0_dp
    pdsi%V = 0.0_dp
    pdsi%Q = 0.0_dp
    
    ! Calculate X values
    call calc_x_values(pdsi)
    
end subroutine calc_original_pdsi

!-------------------------------------------------------------------------------
! Subroutine tính self-calibrated PDSI
!-------------------------------------------------------------------------------
subroutine calc_self_cal_pdsi(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    ! Calculate duration factors
    call calc_duration_factors(pdsi, pdsi%wetm, pdsi%wetb, 1)
    call calc_duration_factors(pdsi, pdsi%drym, pdsi%dryb, -1)
    
    ! Initialize bookkeeping
    pdsi%Prob = 0.0_dp
    pdsi%X1 = 0.0_dp
    pdsi%X2 = 0.0_dp
    pdsi%X3 = 0.0_dp
    pdsi%X = 0.0_dp
    pdsi%V = 0.0_dp
    pdsi%Q = 0.0_dp
    
    ! Calculate X values
    call calc_x_values(pdsi)
    
    ! Calibrate
    call calibrate_index(pdsi)
    
end subroutine calc_self_cal_pdsi

!-------------------------------------------------------------------------------
! Subroutine tính duration factors
!-------------------------------------------------------------------------------
subroutine calc_duration_factors(pdsi, slope, intercept, sign)
    type(pdsi_type), intent(inout) :: pdsi
    real(dp), intent(out) :: slope, intercept
    integer, intent(in) :: sign
    
    integer :: num_list, i
    integer :: length(10)
    real(dp) :: sum_z(10)
    
    num_list = 10
    
    ! Set lengths based on period
    if (pdsi%weekly) then
        if (pdsi%period_length == 1) then
            length = [ 13, 26, 39, 52, 78, 104, 130, 156, 182, 208 ]
        else if (pdsi%period_length == 2) then
            length = [ 6, 13, 19, 26, 39, 52, 65, 78, 91, 104 ]
        else if (pdsi%period_length == 4) then
            length = [ 3, 6, 10, 13, 20, 26, 33, 39, 46, 52 ]
        else
            length = [ 2, 3, 4, 5, 6, 8, 10, 12, 14, 16 ]
        end if
    else
        length = [ 3, 6, 9, 12, 18, 24, 30, 36, 42, 48 ]
    end if
    
    ! Calculate Z sums
    do i = 1, num_list
        sum_z(i) = get_z_sum(pdsi, length(i), sign)
    end do
    
    ! Perform linear regression
    call least_squares(length, sum_z, num_list, sign, slope, intercept)
    
    ! Adjust by 4.0
    slope = slope / (sign * 4.0_dp)
    intercept = intercept / (sign * 4.0_dp)
    
end subroutine calc_duration_factors

!-------------------------------------------------------------------------------
! Function tính tổng Z cho một độ dài nhất định
!-------------------------------------------------------------------------------
function get_z_sum(pdsi, length, sign) result(z_sum)
    type(pdsi_type), intent(in) :: pdsi
    integer, intent(in) :: length, sign
    real(dp) :: z_sum
    
    real(dp), allocatable :: z_values(:)
    real(dp), allocatable :: sums(:)
    real(dp) :: current_sum
    integer :: i, n, n_cal
    
    n = pdsi%ZIND%size
    if (n < length) then
        z_sum = 0.0_dp
        return
    end if
    
    allocate(z_values(n))
    call list_to_array(pdsi%ZIND, z_values)
    
    ! Skip pre-calibration period
    n_cal = pdsi%nCalibrationPeriods
    if (n_cal < length) then
        z_sum = 0.0_dp
        deallocate(z_values)
        return
    end if
    
    ! Calculate initial sum
    current_sum = 0.0_dp
    do i = 1, length
        if (z_values(i) /= MISSING) then
            current_sum = current_sum + z_values(i)
        end if
    end do
    
    ! Find maximum sum
    z_sum = current_sum
    do i = length + 1, n_cal
        if (z_values(i - length) /= MISSING) then
            current_sum = current_sum - z_values(i - length)
        end if
        if (z_values(i) /= MISSING) then
            current_sum = current_sum + z_values(i)
        end if
        if (sign * current_sum > sign * z_sum) then
            z_sum = current_sum
        end if
    end do
    
    deallocate(z_values)
    
end function get_z_sum

!-------------------------------------------------------------------------------
! Subroutine linear regression
!-------------------------------------------------------------------------------
subroutine least_squares(x, y, n, sign, slope, intercept)
    integer, intent(in) :: x(:), n, sign
    real(dp), intent(in) :: y(:)
    real(dp), intent(out) :: slope, intercept
    
    real(dp) :: sumX, sumY, sumX2, sumXY
    real(dp) :: SSX, SSXY
    integer :: i, nn
    real(dp) :: corr, corr_tol
    real(dp) :: max_diff, max_y
    integer :: max_i
    
    corr_tol = 0.85_dp
    nn = n
    
    do
        sumX = 0.0_dp; sumY = 0.0_dp; sumX2 = 0.0_dp; sumXY = 0.0_dp
        
        do i = 1, nn
            sumX = sumX + real(x(i), dp)
            sumY = sumY + y(i)
            sumX2 = sumX2 + real(x(i), dp)**2
            sumXY = sumXY + real(x(i), dp) * y(i)
        end do
        
        SSX = sumX2 - sumX * sumX / real(nn, dp)
        SSXY = sumXY - sumX * sumY / real(nn, dp)
        
        if (SSX > TOLERANCE .and. nn > 3) then
            corr = SSXY / sqrt(SSX * (sumY**2/real(nn, dp)))
            if (sign * corr < corr_tol) then
                nn = nn - 1
                cycle
            end if
        end if
        exit
    end do
    
    slope = SSXY / SSX
    
    ! Find best intercept
    max_diff = -1.0e30_dp
    max_y = 0.0_dp
    max_i = 1
    
    do i = 1, nn
        if (sign * (y(i) - slope * real(x(i), dp)) > sign * max_diff) then
            max_diff = y(i) - slope * real(x(i), dp)
            max_i = i
            max_y = y(i)
        end if
    end do
    
    intercept = max_y - slope * real(x(max_i), dp)
    
end subroutine least_squares

!-------------------------------------------------------------------------------
! Subroutine tính các giá trị X
!-------------------------------------------------------------------------------
subroutine calc_x_values(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    integer :: i, n
    real(dp) :: z_val
    
    ! Clear lists
    call list_clear(pdsi%Xlist)
    call list_clear(pdsi%XL1)
    call list_clear(pdsi%XL2)
    call list_clear(pdsi%XL3)
    call list_clear(pdsi%ProbL)
    call list_clear(pdsi%altX1)
    call list_clear(pdsi%altX2)
    
    ! Reset bookkeeping
    pdsi%Prob = 0.0_dp
    pdsi%X1 = 0.0_dp
    pdsi%X2 = 0.0_dp
    pdsi%X3 = 0.0_dp
    pdsi%X = 0.0_dp
    pdsi%V = 0.0_dp
    pdsi%Q = 0.0_dp
    
    n = pdsi%ZIND%size
    do i = 1, n
        z_val = pdsi%ZIND%data(i)
        call calc_one_x(pdsi, z_val, i)
    end do
    
end subroutine calc_x_values

!-------------------------------------------------------------------------------
! Subroutine tính một giá trị X
!-------------------------------------------------------------------------------
subroutine calc_one_x(pdsi, z_val, idx)
    type(pdsi_type), intent(inout) :: pdsi
    real(dp), intent(in) :: z_val
    integer, intent(in) :: idx
    
    real(dp) :: newV, newProb, newX, newX1, newX2, newX3
    real(dp) :: ZE, m, b, c
    integer :: wd
    
    if (z_val /= MISSING) then
        ! Set wd flag based on X3
        if (pdsi%X3 >= 0.0_dp) then
            wd = 1
            m = pdsi%wetm
            b = pdsi%wetb
        else
            wd = -1
            m = pdsi%drym
            b = pdsi%dryb
        end if
        c = 1.0_dp - (m / (m + b))
        
        if (abs(pdsi%X3) < TOLERANCE) then
            newX3 = 0.0_dp
            newV = 0.0_dp
            newProb = 0.0_dp
            call choose_x(pdsi, newX, newX1, newX2, newX3)
        else
            newX3 = c * pdsi%X3 + z_val / (m + b)
            ZE = (m + b) * (wd * 0.5_dp - c * pdsi%X3)
            pdsi%Q = ZE + pdsi%V
            newV = z_val - wd * (m * 0.5_dp) + wd * min(wd * pdsi%V + TOLERANCE, 0.0_dp)
            
            if (wd * newV > 0.0_dp) then
                newV = 0.0_dp
                newProb = 0.0_dp
                newX1 = 0.0_dp
                newX2 = 0.0_dp
                newX = newX3
                call list_clear(pdsi%altX1)
                call list_clear(pdsi%altX2)
            else
                newProb = (newV / pdsi%Q) * 100.0_dp
                if (newProb >= 100.0_dp - TOLERANCE) then
                    newX3 = 0.0_dp
                    newV = 0.0_dp
                    newProb = 100.0_dp
                end if
                call choose_x(pdsi, newX, newX1, newX2, newX3)
            end if
        end if
        
        ! Update variables
        pdsi%V = newV
        pdsi%Prob = newProb
        pdsi%X1 = newX1
        pdsi%X2 = newX2
        pdsi%X3 = newX3
        
        ! Append to lists
        call list_append(pdsi%Xlist, newX)
        call list_append(pdsi%XL1, pdsi%X1)
        call list_append(pdsi%XL2, pdsi%X2)
        call list_append(pdsi%XL3, pdsi%X3)
        call list_append(pdsi%ProbL, pdsi%Prob)
    else
        ! Missing data
        call list_append(pdsi%Xlist, MISSING)
        call list_append(pdsi%XL1, MISSING)
        call list_append(pdsi%XL2, MISSING)
        call list_append(pdsi%XL3, MISSING)
        call list_append(pdsi%ProbL, MISSING)
    end if
    
end subroutine calc_one_x

!-------------------------------------------------------------------------------
! Subroutine chọn giá trị X
!-------------------------------------------------------------------------------
subroutine choose_x(pdsi, newX, newX1, newX2, newX3)
    type(pdsi_type), intent(inout) :: pdsi
    real(dp), intent(out) :: newX, newX1, newX2, newX3
    real(dp), intent(inout) :: newX3
    
    real(dp) :: wetc, dryc, m, b
    
    if (pdsi%X3 >= 0.0_dp) then
        m = pdsi%wetm
        b = pdsi%wetb
    else
        m = pdsi%drym
        b = pdsi%dryb
    end if
    
    wetc = 1.0_dp - (pdsi%wetm / (pdsi%wetm + pdsi%wetb))
    dryc = 1.0_dp - (pdsi%drym / (pdsi%drym + pdsi%dryb))
    
    newX1 = wetc * pdsi%X1 + pdsi%Z / (pdsi%wetm + pdsi%wetb)
    if (newX1 < 0.0_dp) newX1 = 0.0_dp
    
    if (pdsi%bug) then
        newX2 = pdsi%X2
    else
        newX2 = dryc * pdsi%X2 + pdsi%Z / (pdsi%drym + pdsi%dryb)
        if (newX2 > 0.0_dp) newX2 = 0.0_dp
    end if
    
    if (newX1 >= 0.5_dp .and. abs(newX3) < TOLERANCE) then
        call backtrack(pdsi, newX1, newX2)
        newX = newX1
        newX3 = newX1
        newX1 = 0.0_dp
    else
        newX2 = dryc * pdsi%X2 + pdsi%Z / (pdsi%drym + pdsi%dryb)
        if (newX2 > 0.0_dp) newX2 = 0.0_dp
        
        if (newX2 <= -0.5_dp .and. abs(newX3) < TOLERANCE) then
            call backtrack(pdsi, newX2, newX1)
            newX = newX2
            newX3 = newX2
            newX2 = 0.0_dp
        else if (abs(newX3) < TOLERANCE) then
            if (abs(newX1) < TOLERANCE) then
                call backtrack(pdsi, newX2, newX1)
                newX = newX2
            else if (abs(newX2) < TOLERANCE) then
                call backtrack(pdsi, newX1, newX2)
                newX = newX1
            else
                call list_append(pdsi%altX1, newX1)
                call list_append(pdsi%altX2, newX2)
                newX = newX3
            end if
        else
            call list_append(pdsi%altX1, newX1)
            call list_append(pdsi%altX2, newX2)
            newX = newX3
        end if
    end if
    
end subroutine choose_x

!-------------------------------------------------------------------------------
! Subroutine backtrack
!-------------------------------------------------------------------------------
subroutine backtrack(pdsi, X1, X2)
    type(pdsi_type), intent(inout) :: pdsi
    real(dp), intent(in) :: X1, X2
    
    real(dp) :: num1, num2
    integer :: i
    
    num1 = X1
    do i = 1, min(pdsi%altX1%size, pdsi%altX2%size)
        if (num1 > 0.0_dp) then
            num1 = pdsi%altX1%data(i)
            num2 = pdsi%altX2%data(i)
        else
            num1 = pdsi%altX2%data(i)
            num2 = pdsi%altX1%data(i)
        end if
        
        if (abs(num1) <= TOLERANCE) num1 = num2
        pdsi%Xlist%data(pdsi%Xlist%size - i) = num1
    end do
    
    call list_clear(pdsi%altX1)
    call list_clear(pdsi%altX2)
    
end subroutine backtrack

!-------------------------------------------------------------------------------
! Subroutine calibrate index
!-------------------------------------------------------------------------------
subroutine calibrate_index(pdsi)
    type(pdsi_type), intent(inout) :: pdsi
    
    real(dp) :: cal_range, p02, p98
    real(dp), allocatable :: tmp_data(:)
    integer :: i, n
    
    cal_range = 4.0_dp
    
    ! Get percentiles
    if (pdsi%setCalibrationStartYear .or. pdsi%setCalibrationEndYear) then
        ! Use calibration period
        n = pdsi%nCalibrationPeriods
        if (n > 0) then
            allocate(tmp_data(n))
            do i = 1, n
                tmp_data(i) = pdsi%Xlist%data(pdsi%nStartPeriodsToSkip + i)
            end do
            call percentile(tmp_data, n, 0.02_dp, p02)
            call percentile(tmp_data, n, 0.98_dp, p98)
            deallocate(tmp_data)
        else
            p02 = -4.0_dp
            p98 = 4.0_dp
        end if
    else
        ! Use full period
        n = pdsi%Xlist%size
        allocate(tmp_data(n))
        call list_to_array(pdsi%Xlist, tmp_data)
        call percentile(tmp_data, n, 0.02_dp, p02)
        call percentile(tmp_data, n, 0.98_dp, p98)
        deallocate(tmp_data)
    end if
    
    pdsi%dry_ratio = -cal_range / p02
    pdsi%wet_ratio = cal_range / p98
    
    ! Adjust Z-index
    do i = 1, pdsi%ZIND%size
        if (pdsi%ZIND%data(i) /= MISSING) then
            if (pdsi%ZIND%data(i) >= 0.0_dp) then
                pdsi%ZIND%data(i) = pdsi%ZIND%data(i) * pdsi%wet_ratio
            else
                pdsi%ZIND%data(i) = pdsi%ZIND%data(i) * pdsi%dry_ratio
            end if
        end if
    end do
    
    ! Recalculate X values
    call calc_x_values(pdsi)
    
end subroutine calibrate_index

!-------------------------------------------------------------------------------
! Subroutine tính percentile
!-------------------------------------------------------------------------------
subroutine percentile(data, n, p, result)
    real(dp), intent(in) :: data(:)
    integer, intent(in) :: n
    real(dp), intent(in) :: p
    real(dp), intent(out) :: result
    
    real(dp), allocatable :: sorted(:)
    real(dp) :: pos
    integer :: k
    
    allocate(sorted(n))
    sorted = data(1:n)
    call sort_array(sorted, n)
    
    pos = p * real(n + 1, dp)
    k = int(pos)
    
    if (k < 1) then
        result = sorted(1)
    else if (k >= n) then
        result = sorted(n)
    else
        result = sorted(k) + (pos - real(k, dp)) * (sorted(k+1) - sorted(k))
    end if
    
    deallocate(sorted)
    
end subroutine percentile

!-------------------------------------------------------------------------------
! Subroutine sắp xếp mảng
!-------------------------------------------------------------------------------
subroutine sort_array(arr, n)
    real(dp), intent(inout) :: arr(:)
    integer, intent(in) :: n
    
    real(dp) :: temp
    integer :: i, j
    
    do i = 1, n-1
        do j = i+1, n
            if (arr(i) > arr(j)) then
                temp = arr(i)
                arr(i) = arr(j)
                arr(j) = temp
            end if
        end do
    end do
    
end subroutine sort_array

end module pdsi_core

!===============================================================================
module list_operations
!===============================================================================
! Module chứa các operations cho list_type
!===============================================================================
    use pdsi_parameters
    use pdsi_types
    implicit none
    
contains

!-------------------------------------------------------------------------------
! Subroutine khởi tạo list
!-------------------------------------------------------------------------------
subroutine list_init(list, capacity)
    type(list_type), intent(inout) :: list
    integer, optional, intent(in) :: capacity
    
    if (allocated(list%data)) deallocate(list%data)
    
    if (present(capacity)) then
        list%capacity = capacity
    else
        list%capacity = 1000  ! Default capacity
    end if
    
    allocate(list%data(list%capacity))
    list%size = 0
    
end subroutine list_init

!-------------------------------------------------------------------------------
! Subroutine thêm phần tử vào list
!-------------------------------------------------------------------------------
subroutine list_append(list, value)
    type(list_type), intent(inout) :: list
    real(dp), intent(in) :: value
    
    real(dp), allocatable :: temp(:)
    
    if (.not. allocated(list%data)) then
        call list_init(list)
    end if
    
    if (list%size >= list%capacity) then
        ! Expand array
        list%capacity = list%capacity * 2
        allocate(temp(list%capacity))
        temp(1:list%size) = list%data(1:list%size)
        call move_alloc(temp, list%data)
    end if
    
    list%size = list%size + 1
    list%data(list%size) = value
    
end subroutine list_append

!-------------------------------------------------------------------------------
! Subroutine xóa list
!-------------------------------------------------------------------------------
subroutine list_clear(list)
    type(list_type), intent(inout) :: list
    
    if (allocated(list%data)) deallocate(list%data)
    list%size = 0
    list%capacity = 0
    
end subroutine list_clear

!-------------------------------------------------------------------------------
! Subroutine copy list
!-------------------------------------------------------------------------------
subroutine list_copy(dest, src)
    type(list_type), intent(inout) :: dest
    type(list_type), intent(in) :: src
    
    call list_clear(dest)
    if (src%size > 0) then
        dest%capacity = src%capacity
        allocate(dest%data(dest%capacity))
        dest%data(1:src%size) = src%data(1:src%size)
        dest%size = src%size
    end if
    
end subroutine list_copy

!-------------------------------------------------------------------------------
! Subroutine chuyển list sang array
!-------------------------------------------------------------------------------
subroutine list_to_array(list, arr)
    type(list_type), intent(in) :: list
    real(dp), intent(out) :: arr(:)
    
    integer :: i
    
    do i = 1, list%size
        arr(i) = list%data(i)
    end do
    
end subroutine list_to_array

end module list_operations

!===============================================================================
module file_operations
!===============================================================================
! Module xử lý file system operations
!===============================================================================
    implicit none
    
contains

!-------------------------------------------------------------------------------
! Subroutine tạo directory
!-------------------------------------------------------------------------------
subroutine create_directory(path)
    character(len=*), intent(in) :: path
    
    character(len=256) :: cmd
    
    ! Use system command to create directory
    cmd = 'mkdir -p ' // trim(path)
    call system(cmd)
    
end subroutine create_directory

end module file_operations

!===============================================================================
! Chương trình chính
!===============================================================================
program pdsi_main
    use pdsi_parameters
    use pdsi_types
    use pdsi_io
    use pdsi_core
    use list_operations
    use file_operations
    implicit none
    
    type(pdsi_type) :: pdsi
    integer :: calc_type, data_type
    integer :: cal_start, cal_end
    character(len=256) :: input_dir, output_dir, config_file
    logical :: verbose
    integer :: iostat
    
    ! Mặc định config file
    config_file = 'pdsi.nml'
    
    ! Đọc command line arguments nếu có
    call parse_command_line(config_file)
    
    ! Đọc namelist
    call read_namelist(config_file, calc_type, data_type, cal_start, cal_end, &
                       input_dir, output_dir, verbose)
    
    print *, "========================================"
    print *, "PDSI Calculation Program"
    print *, "========================================"
    print *, "Calculation type: ", merge("Self-calibrated", "Original       ", calc_type == OPT_SELF_CAL)
    print *, "Data type:       ", merge("Weekly", "Monthly", data_type == OPT_WEEKLY)
    print *, "Input directory: ", trim(input_dir)
    print *, "Output directory:", trim(output_dir)
    if (cal_start > 0) print *, "Calibration start:", cal_start
    if (cal_end > 0)   print *, "Calibration end:  ", cal_end
    print *, "========================================"
    
    ! Thiết lập PDSI
    pdsi%input_dir = input_dir
    pdsi%output_dir = output_dir
    pdsi%verbose = verbose
    
    ! Đọc số liệu
    print *, "Reading data..."
    call read_data(pdsi, data_type, iostat)
    if (iostat /= 0) then
        print *, "Error reading data. Exiting."
        stop
    end if
    
    ! Xác định năm bắt đầu và kết thúc từ dữ liệu
    pdsi%startyear = int(pdsi%YearList%data(1))
    pdsi%endyear = int(pdsi%YearList%data(pdsi%YearList%size))
    pdsi%totalyears = pdsi%endyear - pdsi%startyear + 1
    
    ! Khởi tạo PDSI
    call pdsi_init(pdsi, calc_type, data_type, cal_start, cal_end)
    
    ! Tính toán PDSI
    print *, "Calculating PDSI..."
    call calculate_pdsi(pdsi, calc_type)
    
    ! Ghi kết quả
    print *, "Writing results..."
    call write_results(pdsi, output_dir, calc_type, data_type)
    
    print *, "Calculation complete!"
    
contains

!-------------------------------------------------------------------------------
! Subroutine xử lý command line arguments
!-------------------------------------------------------------------------------
subroutine parse_command_line(config_file)
    character(len=*), intent(out) :: config_file
    
    integer :: i, nargs
    character(len=256) :: arg
    
    config_file = 'pdsi.nml'
    nargs = command_argument_count()
    
    do i = 1, nargs
        call get_command_argument(i, arg)
        if (arg == '-c' .or. arg == '--config') then
            if (i < nargs) then
                call get_command_argument(i+1, config_file)
            end if
        else if (arg == '-h' .or. arg == '--help') then
            print *, "Usage: pdsi [options]"
            print *, "Options:"
            print *, "  -c, --config FILE  Use specified configuration file"
            print *, "  -h, --help         Show this help message"
            stop
        end if
    end do
    
end subroutine parse_command_line

end program pdsi_main