program main

implicit none

character(len=999) :: inputdir="input"
character(len=999) :: forcingdir="/Volumes/PASSPORT5T/data/L4_SM_runoff_daily/OL7000_Pfaf"
character(len=999) :: outputdir="output"

integer,parameter :: step_start=9221 !start day, 1 is for Jan-1-1990
integer,parameter :: step_end=11596 !end day

logical :: is_coldstart=.True.

real*8,parameter :: P1=0.010611D0
real*8,parameter :: P2=0.188556D0
real*8,parameter :: P3=0.096864D0
real*8,parameter :: P4=0.691310D0
real*8,parameter :: P5=0.1D0
real*8,parameter :: P6=0.009831D0

real*8,parameter :: small=1.D-48
integer,parameter :: ny=33

real*8,parameter :: M=0.10D0
real*8,parameter :: mm=0.40D0

real*8,parameter :: dt=86400.D0 !s

integer,parameter :: nmax=373
integer,parameter :: upmax=34
integer,parameter :: nc=291284
real*8,parameter :: rho=1.D3 !kg/m3 

integer,parameter :: nlon=964
integer,parameter :: nlat=406

integer :: i,j,n,dnID,iter,ntime
character(len=50) :: M_string,mm_string,iter_s

integer,allocatable,dimension(:) :: nts
real*8,allocatable,dimension(:) :: qstr_clmt,qri_clmt,qin_clmt,&
                      llc_ori,llc,lstr,&
                      Qrunf,&
                      nume,deno,&
                      alp_s,alp_r,&
                      K,&
                      catarea,wid,alp_s_low,alp_s_high 
real*8,allocatable,dimension(:) :: Ws,Wr
real*8,allocatable,dimension(:) :: Qs0,ks,Ws_last,Qs,&
                      Qr0,kr,Cl,Al
real*8,allocatable,dimension(:) :: C1,C2,Qout,Qin,A1,P,B1,B2
integer,allocatable,dimension(:)  :: nup
integer,allocatable,dimension(:,:) :: upID
real*8 :: co1,co2,co3 
integer :: ui

real*8,allocatable,dimension(:) :: lon,lat
real*8,allocatable,dimension(:,:,:) :: runoff,runoffr,baseflow

integer,dimension(ny) :: days_in_year=(/365,365,366,365,&
                                        365,365,366,365,&
                                        365,365,366,365,&
                                        365,365,366,365,&
                                        365,365,366,365,&
                                        365,365,366,365,&
                                        365,365,366,365,&
                                        365,365,366,365,365/) !from 1990 to 2020
integer,dimension(12) :: days_in_mon_noleap=(/31,28,31,30,31,30,31,31,30,31,30,31/)
integer,dimension(12) :: days_in_mon_leap=(/31,29,31,30,31,30,31,31,30,31,30,31/)
integer :: days_acc_year(ny),days_acc_noleap(12),days_acc_leap(12),days_acc_mon(12)
integer :: yr_cur,mon_cur,day_cur,d_res,step_prev,nstep
character(len=50) :: yr_s,mon_s,day_s

nstep=step_end-step_start+1
days_acc_year(1)=0
do i=2,ny
  days_acc_year(i)=days_acc_year(i-1)+days_in_year(i-1)
end do
days_acc_noleap(1)=0
days_acc_leap(1)=0
do i=2,12
  days_acc_noleap(i)=days_acc_noleap(i-1)+days_in_mon_noleap(i-1)
  days_acc_leap(i)=days_acc_leap(i-1)+days_in_mon_leap(i-1)
end do


allocate(nts(nc))
allocate(qstr_clmt(nc),qri_clmt(nc),qin_clmt(nc))
allocate(llc_ori(nc),llc(nc),lstr(nc))
allocate(Qrunf(nc),nume(nc),deno(nc),alp_s(nc),alp_r(nc))
allocate(Ws(nc),Wr(nc))
allocate(Qs0(nc),ks(nc),Ws_last(nc),Qs(nc))
allocate(Qr0(nc),kr(nc),Cl(nc),Al(nc))
allocate(C1(nc),C2(nc),Qout(nc),Qin(nc),A1(nc),P(nc),B1(nc),B2(nc))
allocate(nup(nc))
allocate(upID(upmax,nc))
allocate(K(nc))
allocate(catarea(nc),wid(nc),alp_s_low(nc),alp_s_high(nc))

!for output
write(M_string,'(f4.2)')M
write(mm_string,'(f4.2)')mm


!read input
open(77,file=trim(inputdir)//"/Pfaf_qstr.txt")
read(77,*)qstr_clmt !m3/s
qstr_clmt=qstr_clmt*rho !m3/s->kg/s
open(77,file=trim(inputdir)//"/Pfaf_qri.txt")
read(77,*)qri_clmt !m3/s
qri_clmt=qri_clmt*rho !m3/s->kg/s
open(77,file=trim(inputdir)//"/Pfaf_tosink.txt")
read(77,*)nts
open(77,file=trim(inputdir)//"/upstream_1D.txt")
read(77,*)upID
open(77,file=trim(inputdir)//"/Pfaf_upnum.txt")
read(77,*)nup
open(77,file=trim(inputdir)//"/Pfaf_llcM.txt")
read(77,*)llc_ori !km
llc_ori=llc_ori*1.D3 !km->m
open(77,file=trim(inputdir)//"/Pfaf_lstrM.txt")
read(77,*)lstr !km
lstr=lstr*1.D3 !km->m
open(77,file=trim(inputdir)//"/Pfaf_qin.txt")
read(77,*)qin_clmt !m3/s
qin_clmt=qin_clmt*rho !m3/s->kg/s
open(77,file=trim(inputdir)//"/Pfaf_Kv.txt")
read(77,*)K
open(77,file=trim(inputdir)//"/Pfaf_area.txt")
read(77,*)catarea
catarea=catarea*1.D6 !km2->m2
wid=catarea/llc_ori !m

!calculate llc
nume=qri_clmt**(2.D0-M)-qin_clmt**(2.D0-M)
deno=(2.D0-M)*(qri_clmt-qin_clmt)*(qri_clmt**(1.D0-M))
where(abs(deno)>small) llc=llc_ori*(nume/deno)
where(abs(deno)<=small) llc=llc_ori*0.5D0

alp_s_low=P1 * ( (1.D0/(wid/1.D3))**P2 ) /3600.D0
alp_s_high=P3 * ( (1.D0/(wid/1.D3))**P4 ) /3600.D0

alp_r=P5 * ( (1.D0/(llc_ori/1.D3))**P6 ) /3600.D0

allocate(runoff(nlon,nlat,1),runoffr(nlon,nlat,1),baseflow(nlon,nlat,1))

!temporal loop
DO iter=step_start,step_end

write(iter_s,'(I5.5)')iter
print *,trim(iter_s)

if(iter==1.or.is_coldstart)then
  open(77,file=trim(inputdir)//"/Pfaf_Ws_Sari_M0.10_mm0.40_20160330_OL7000.txt")
  read(77,*)Ws
  open(77,file=trim(inputdir)//"/Pfaf_Wr_Sari_M0.10_mm0.40_20160330_OL7000.txt")  
  read(77,*)Wr 
  is_coldstart=.False.
else
  step_prev=iter-1
  do i=ny,1,-1
    if(step_prev>days_acc_year(i))then
      yr_cur=1989+i
      d_res=step_prev-days_acc_year(i)
      exit
    endif
  enddo
  if(mod(yr_cur,4)==0)then
    days_acc_mon=days_acc_leap
  else
    days_acc_mon=days_acc_noleap
  endif
  do i=12,1,-1
    if(d_res>days_acc_mon(i))then
      mon_cur=i
      day_cur=d_res-days_acc_mon(i)
      exit
    endif
  enddo
  write(yr_s,'(I4)')yr_cur
  write(mon_s,'(I2.2)')mon_cur
  write(day_s,'(I2.2)')day_cur
  open(77,file=trim(outputdir)//"/Pfaf_Ws_Sari_M"//trim(M_string)//"_mm"//trim(mm_string)//"_"//trim(yr_s)//trim(mon_s)//trim(day_s)//"_OL7000.txt")
  read(77,*)Ws
  open(77,file=trim(outputdir)//"/Pfaf_Wr_Sari_M"//trim(M_string)//"_mm"//trim(mm_string)//"_"//trim(yr_s)//trim(mon_s)//trim(day_s)//"_OL7000.txt")
  read(77,*)Wr  
endif


do i=ny,1,-1
  if(iter>days_acc_year(i))then
    yr_cur=1989+i
    d_res=iter-days_acc_year(i)
    exit
  endif
enddo
if(mod(yr_cur,4)==0)then
  days_acc_mon=days_acc_leap
else
  days_acc_mon=days_acc_noleap
endif
do i=12,1,-1
  if(d_res>days_acc_mon(i))then
    mon_cur=i
    day_cur=d_res-days_acc_mon(i)
    exit
  endif
enddo

write(yr_s,'(I4)')yr_cur
write(mon_s,'(I2.2)')mon_cur
write(day_s,'(I2.2)')day_cur
print *,trim(yr_s)," ",trim(mon_s)," ",trim(day_s)

open(77,file=trim(forcingdir)//"/runoff_"//trim(yr_s)//trim(mon_s)//trim(day_s)//".txt")
read(77,*)Qrunf
Qrunf=Qrunf*rho !m3/s -> kg/s

alp_s=alp_s_low+((Qrunf-2000.D0)/8000.D0)*(alp_s_high-alp_s_low)
where(Qrunf<2000.D0) alp_s=alp_s_low
where(Qrunf>10000.D0) alp_s=alp_s_high

!$omp parallel default(shared)
!$omp workshare
!update ks Ws Qs
where(Qrunf<=small)Qrunf=0.D0
Qs0=max(0.D0,alp_s * Ws) !kg/s
ks=max(0.D0,alp_s) !s-1
Ws_last=Ws
where(ks>small) Ws=Ws + (Qrunf-Qs0)/ks*(1.D0-exp(-ks*dt)) !kg
where(ks<=small) Ws=Ws + (Qrunf-Qs0)*dt
Ws=max(0.D0,Ws)
Qs=max(0.D0,Qrunf-(Ws-Ws_last)/dt)

!cal Qr0 kr Cl Al
Qr0=max(0.D0,alp_r * Wr) !kg/s
kr=max(0.D0,alp_r) !s-1
where(kr>small.and.abs(kr-ks)>small) Cl=Wr + (Qrunf-Qr0)/kr*(1.D0-exp(-kr*dt)) + (Qrunf-Qs0)/(kr-ks)*(exp(-kr*dt)-exp(-ks*dt))
where(kr>small.and.abs(kr-ks)<=small) Cl=Wr + (Qrunf-Qr0)/kr*(1.D0-exp(-kr*dt)) - (Qrunf-Qs0)*dt*exp(-kr*dt)
where(kr<=small.and.ks>small) Cl=Wr + (Qrunf-Qr0)*dt - (Qrunf-Qs0)/ks*(1.D0-exp(-ks*dt))
where(kr<=small.and.ks<=small) Cl=Wr + (Qs0-Qr0)*dt
Al=Qs+Wr/dt-Cl/dt

!begin to river routing process
C1=0.D0
C2=0.D0
Qin=0.D0
Qout=0.D0
A1=0.D0
P=0.D0
B1=0.D0
B2=0.D0
!$omp end workshare
!$omp end parallel

do n=nmax,0,-1

  !$OMP PARALLEL default(shared) private(i,j,ui,co1,co2,co3) 
  !$OMP DO
  do i=1,nc
    if(nts(i)==n)then 

      if(nup(i)>=1)then
        do j=1,nup(i)
          ui=upID(j,i)
          if(ui==-1)exit
          if(kr(i)>small)then
            co1=max(0.D0,(1.D0-exp(-kr(i)*dt))/kr(i))
          else
            co1=dt
          endif
          C1(i)=C1(i)+co1*B1(ui)
          if(abs(kr(i)-kr(ui))>small)then
            co2=-(exp(-kr(i)*dt)-exp(-kr(ui)*dt))/(kr(i)-kr(ui))
          else
            co2=dt*exp(-kr(i)*dt)
          endif
          C2(i)=C2(i)+co2*B2(ui)
          Qin(i)=Qin(i)+Qout(ui)
        enddo
      endif
      Wr(i)=max(0.D0,Cl(i)+C1(i)+C2(i)) 
      A1(i)=Qin(i)-C1(i)/dt-C2(i)/dt 
      Qout(i)=max(0.D0,Al(i)+A1(i))

      if(kr(i)>small.and.Qin(i)+Qrunf(i)>small)then
        co3=max(0.D0,(1.D0-exp(-kr(i)*dt))/kr(i))
        P(i)=(dt*Qout(i)-co3*Qr0(i))/((Qin(i)+Qrunf(i))*(dt-co3))  
        if(P(i)>0.5D0.and.P(i)<1.5D0)then
          B1(i)=P(i)*(Qin(i)+Qrunf(i))
          B2(i)=-P(i)*(Qin(i)+Qrunf(i))+Qr0(i)
        else
          B1(i)=Qout(i)
          B2(i)=0.D0
        endif
      else
        B1(i)=Qout(i)
        B2(i)=0.D0
        P(i)=-9999.        
      endif

    endif
  enddo
  !$OMP END DO
  !$OMP END PARALLEL

enddo

open(88,file=trim(outputdir)//"/Pfaf_Qr_Sari_M"//trim(M_string)//"_mm"//trim(mm_string)//"_"//trim(yr_s)//trim(mon_s)//trim(day_s)//"_OL7000.txt")
do i=1,nc
  write(88,*)Qout(i)/1.D3 !m^3/s
enddo

open(88,file=trim(outputdir)//"/Pfaf_Ws_Sari_M"//trim(M_string)//"_mm"//trim(mm_string)//"_"//trim(yr_s)//trim(mon_s)//trim(day_s)//"_OL7000.txt")
do i=1,nc
  write(88,*)Ws(i)!/1.D9 !kg
enddo
open(88,file=trim(outputdir)//"/Pfaf_Wr_Sari_M"//trim(M_string)//"_mm"//trim(mm_string)//"_"//trim(yr_s)//trim(mon_s)//trim(day_s)//"_OL7000.txt")
do i=1,nc
  write(88,*)Wr(i)!/1.D9 !kg
enddo

print *,sum(Wr)/1.D12


ENDDO

end