!===========================================================
! STL Decomposition with LOESS-based seasonal smoothing
! Fortran prototype (simplified): Y(t) = T(t) + S(t) + R(t)
! Uses LOESS smoothing for seasonal component (nonparametric)
! Note: Real LOESS requires robust fitting and interpolation,
!       which are simplified here.
!===========================================================
program stl_loess
  implicit none
  integer, parameter :: n = 240, p = 12, span = 7, max_iter = 10
  real :: Y(n), T(n), S(n), R(n), Y_detrend(n), Y_deseason(n)
  real :: seasonal_cycle(p), smoothed_cycle(n)
  integer :: i, j, k, iter

  !--- Sample synthetic data (nonlinear trend + seasonality + noise)
  do i = 1, n
     Y(i) = 0.005*i + 0.00005*i*i + 0.8*sin(2.0*3.1416*i/12.0) + 0.5*sin(0.5*i)
  end do

  T(:) = 0.0
  S(:) = 0.0
  R(:) = 0.0

  do iter = 1, max_iter

     !------------------------
     ! Step 1: Detrend
     !------------------------
     do i = 1, n
        Y_detrend(i) = Y(i) - T(i)
     end do

     !------------------------
     ! Step 2: Seasonal Subseries and LOESS smoothing
     !------------------------
     seasonal_cycle(:) = 0.0
     do k = 1, p
        ! Collect subseries for seasonal month k
        j = 0
        do i = k, n, p
           j = j + 1
           smoothed_cycle(j) = Y_detrend(i)
        end do

        ! Apply simple local LOESS smoothing (tri-cube kernel, linear fit)
        call loess_smooth(smoothed_cycle, j, span)

        ! Assign back to seasonal_cycle(k): average of smoothed points
        seasonal_cycle(k) = sum(smoothed_cycle(1:j)) / j
     end do

     ! Fill seasonal series
     do i = 1, n
        S(i) = seasonal_cycle(mod(i-1, p) + 1)
     end do

     !------------------------
     ! Step 3: Deseasonalize
     !------------------------
     do i = 1, n
        Y_deseason(i) = Y(i) - S(i)
     end do

     !------------------------
     ! Step 4: Trend via 13-point moving average
     !------------------------
     do i = 1, n
        T(i) = 0.0
        j = 0
        do k = max(1, i-6), min(n, i+6)
           T(i) = T(i) + Y_deseason(k)
           j = j + 1
        end do
        if (j > 0) T(i) = T(i) / j
     end do

     !------------------------
     ! Step 5: Residual
     !------------------------
     do i = 1, n
        R(i) = Y(i) - T(i) - S(i)
     end do
  end do

  !--- Output (print to screen)
  print *, "Time", "Y", "Trend", "Seasonal", "Residual"
  do i = 1, n
     print *, i, Y(i), T(i), S(i), R(i)
  end do

contains

  !---------------------------------------------------
  ! Simple LOESS smoothing: moving linear regression
  ! with tricube kernel (span: window size)
  !---------------------------------------------------
  subroutine loess_smooth(X, m, span)
    real, intent(inout) :: X(*)
    integer, intent(in) :: m, span
    integer :: i, j, left, right
    real :: xw(span), yw(span), w(span), dist, u, xbar, ybar, beta

    do i = 1, m
       left = max(1, i - span/2)
       right = min(m, i + span/2)
       j = 0
       do k = left, right
          j = j + 1
          xw(j) = k
          yw(j) = X(k)
          dist = abs(real(k - i)) / (span/2)
          if (dist < 1.0) then
             w(j) = (1.0 - dist**3)**3
          else
             w(j) = 0.0
          end if
       end do

       ! Local linear fit: slope (beta) and intercept
       call weighted_linear_fit(xw, yw, w, j, xbar, ybar, beta)
       X(i) = ybar + beta * (i - xbar)
    end do
  end subroutine loess_smooth

  !---------------------------------------------------
  ! Weighted linear regression
  !---------------------------------------------------
  subroutine weighted_linear_fit(x, y, w, n, xbar, ybar, beta)
    real, intent(in) :: x(*), y(*), w(*)
    integer, intent(in) :: n
    real, intent(out) :: xbar, ybar, beta
    real :: sumw, sumwx, sumwy, sumwx2, sumwxy
    integer :: i

    sumw = 0.0; sumwx = 0.0; sumwy = 0.0
    sumwx2 = 0.0; sumwxy = 0.0

    do i = 1, n
       sumw = sumw + w(i)
       sumwx = sumwx + w(i) * x(i)
       sumwy = sumwy + w(i) * y(i)
       sumwx2 = sumwx2 + w(i) * x(i)**2
       sumwxy = sumwxy + w(i) * x(i) * y(i)
    end do

    if (sumw == 0.0) then
       xbar = 0.0; ybar = 0.0; beta = 0.0
       return
    end if

    xbar = sumwx / sumw
    ybar = sumwy / sumw
    beta = (sumwxy - sumwx * ybar) / (sumwx2 - sumwx * xbar + 1.0e-6)
  end subroutine weighted_linear_fit

end program stl_loess