#!/usr/bin/env python
# coding: utf-8

# In[5]:


###################################################################################
# set the parameters and press Ctrl-ENTER

spacecraft = 'A'
start = "2016-01-01" # included
end = "2016-02-01" # not included
#sampling = "PT60S" # set to None to use the original sampling
sampling = None
models = [
    "CHAOS-6-Core",
    "CHAOS-6-Static",
    "CHAOS-6-MMA-Primary",
    "CHAOS-6-MMA-Secondary"
]
variables=[
    "B_NEC", 
    "B_NEC_CHAOS-6-Core",
    "B_NEC_CHAOS-6-Static",
    "B_NEC_CHAOS-6-MMA-Primary",
    "B_NEC_CHAOS-6-MMA-Secondary",
    'Flags_B', 'Flags_q', 'Flags_Platform', 
    'Kp', 'Dst', 'F10_INDEX', 'OrbitNumber', 'QDLat', 'QDLon', 'QDBasis', 'MLT',
]

output_filename_template = "VIRES_MAG{spacecraft}_LR_1B_{start}_{end}.cdf"


# In[2]:


import sys
from os import remove, rename
from os.path import exists
from viresclient._wps.wps_vires import ViresWPS10Service
from viresclient._wps.environment import JINJA2_ENVIRONMENT
from viresclient._wps.time_util import parse_datetime, parse_duration, encode_duration, timedelta, datetime
from viresclient._client import ProgressBarProcessing, ProgressBarDownloading, ClientRequest

wps_proxy = ViresWPS10Service(
    url="https://staging.viresdisc.vires.services/openows"
)

request_template = JINJA2_ENVIRONMENT.get_template(
    "vires_fetch_filtered_data_async.xml"
)

spcectraft_to_collections_mapping = {
    "A": "SW_OPER_MAGA_LR_1B",
    "B": "SW_OPER_MAGB_LR_1B",
    "C": "SW_OPER_MAGC_LR_1B",
}

def chop_time_range_by_days(start_time, end_time):
    year, month, day = start_time.year, start_time.month, start_time.day

    chunk_offset= datetime(year, month, day)
    chunk_duration = timedelta(days=1, microseconds=-1)
    chunk_step = timedelta(days=1)

    while chunk_offset < end_time:
        yield (chunk_offset, chunk_offset + chunk_duration)
        chunk_offset += chunk_step


# In[6]:


start_time = parse_datetime(start)
end_time = parse_datetime(end)

# split time into daily chunks
time_windows = list(chop_time_range_by_days(start_time, end_time))
output_filenames = [
    output_filename_template.format(
        spacecraft=spacecraft,
        start=chunk_start_time.strftime("%Y%m%d-%H%M%S"),
        end=chunk_end_time.strftime("%Y%m%d-%H%M%S"),
    ) for chunk_start_time, chunk_end_time in time_windows
]

# validate sampling step
if sampling:
    parse_duration(sampling)
    
# prepare collection identifiers
collections = {
    spacecraft: [spcectraft_to_collections_mapping[spacecraft]]
}

# print request summary
print("request summary:")
print("\toutput_filenames:")
for filename in output_filenames:
    print("\t\t%s" % filename)
print("\tstart time: %sZ" % start_time.isoformat())
print("\tend time:   %sZ" % end_time.isoformat())
print("\tsampling:   %s" % (sampling or "original data sampling"))
print("\tcollections:")
for spacecraft, colls in collections.items():
    print("\t%s: %s" % (spacecraft, ", ".join(colls)))
print("\tmodels:")
for model in models:
    print("\t\t%s" % model)
print("\tvariables:")
for variable in variables:
    print("\t\t%s" % variable)
print()

# download data by chunks
for (chunk_start_time, chunk_end_time), output_filename in zip(time_windows, output_filenames):
    print("processing %s" % output_filename)
    sys.stdout.flush()

    # XML WPS request
    request = request_template.render(
        begin_time=chunk_start_time,
        end_time=chunk_end_time,
        sampling_step=sampling,
        variables=variables,
        collection_ids=collections,
        model_ids=models,
        response_type="application/x-cdf",
    )
    
    # download the data
    if not exists(output_filename):
        tmp_filename = output_filename + ".tmp"
        try:
            with open(tmp_filename, "wb") as fout:
                response_handler=ClientRequest._response_handler(fout)
                with ProgressBarProcessing() as progress_bar:
                    wps_proxy.retrieve_async(
                        request.encode('UTF-8'),
                        handler=response_handler,
                        status_handler=progress_bar.update
                    )
        except Exception as error:
            if exists(tmp_filename):
                remove(tmp_filename)
            print("ERROR: %s faled: %s" % (output_filename, error), file=sys.stderr)
        else:
            if exists(tmp_filename):
                rename(tmp_filename, output_filename)
    else:
        print("%s exists, skipped ... " % output_filename)