#===================================================================
# Simple python script to download TIMPS data by the month
# Requires Python 3, os and Pandas (on python), and lftp (on system)
# James Russell 2021

#===================================================================
# Import libaries

import pandas as pd
import os

#===================================================================
# Namelist

getTIMPSfiles = True # Indicate whether to download TIMPS files
getTrackfiles = True # Indicate whether to download Track files

yearstart  = 2015 # First year for date range
yearend    = 2015 # Last year for date range
monthstart = 1    # First month
monthend   = 2    # Last month

url = "https://home.chpc.utah.edu/~u0816744/TIMPS/" # TIMPS URL

procs = 100 # parallel processes for lftp

#===================================================================
# Generate date list

months = pd.period_range(
start=f"{str(yearstart).zfill(4)}-{str(monthstart).zfill(2)}-01",
end=f"{str(yearend).zfill(4)}-{str(monthend).zfill(2)}-01",
freq="M").strftime('%Y%m').tolist()

#===================================================================
# Loop over months to download data using lftp

for m in months:

  # Get TIMPS files
  if getTIMPSfiles:

    # Download month directory
    currurl = f"{url}TIMPS_files/{m[:4]}/{m[4:]}/"
    print(f"Downloading files in {currurl}")
    os.system(f"lftp -c 'mirror --parallel={procs} {currurl} ;exit'")

    # Move data to new directory and organize
    newdir = f"TIMPS/{m[:4]}/{m[4:]}"
    print(f"Moving TIMPS files downloaded to {newdir}")
    try: os.mkdir("TIMPS")
    except: print("TIMPS directory already exists")
    try: os.mkdir(f"TIMPS/{m[:4]}/")
    except: print(f"TIMPS/{m[:4]}/ directory already exists")
    os.rename(f"{m[4:]}",f"{newdir}")

  # Get tracking files
  if getTrackfiles:

    # Download month directory
    currurl = f"{url}Track_files/{m[:4]}/{m[4:]}/"
    print(f"Downloading files in {currurl}")
    os.system(f"lftp -c 'mirror --parallel={procs} {currurl} ;exit'")

    # Move data to new directory and organize
    newdir = f"Track/{m[:4]}/{m[4:]}"
    print(f"Moving Track files downloaded to {newdir}")
    try: os.mkdir("Track")
    except: print("Track directory already exists")
    try: os.mkdir(f"Track/{m[:4]}/")
    except: print(f"Track/{m[:4]}/ directory already exists")
    os.rename(f"{m[4:]}",f"{newdir}")

  print("All files downloaded and organized")

#===================================================================
# End
#===================================================================