# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_download.ipynb (unless otherwise specified).

__all__ = ['Ladsweb', 'run_parallel']

# Cell
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from pathlib import Path
import pandas as pd
from rasterio.coords import BoundingBox, disjoint_bounds
import numpy as np
import json
import requests
import warnings
import re
import os
from fastprogress.fastprogress import progress_bar
from nbdev.imports import test_eq
from functools import partial
from concurrent.futures import ThreadPoolExecutor
from time import sleep

from .external import geturl

# Cell
class Ladsweb():
    def __init__(self, product:str, collection:str, tstart:str, tend:str,
                 bbox:list, bands:list=None, coordsOrTiles:str="coords", daynight:str="DNB",
                 repName:str='GEO', repPixSize:float=0.01, repResample:str='bilinear',
                 doMosaic:str='False', **kwargs):
        self.product, self.collection = product, collection
        self.tstart, self.tend, self.bbox, self.bands = tstart, tend, bbox, bands
        self.coordsOrTiles, self.daynight, self.repName = coordsOrTiles, daynight, repName
        self.repPixSize, self.repResample, self.doMosaic = repPixSize, repResample, doMosaic
        self._maxOrderSize, self.max_parallel = 1800, 4

    def search_files(self):
        "Search for files for the product, region and time span given."
        url = (f"https://modwebsrv.modaps.eosdis.nasa.gov/axis2/services/MODAPSservices/" +
            f"searchForFiles?product={self.product}&collection={self.collection}&" +
            f"start={self.tstart}&stop={self.tend}&north={self.bbox[3]}&south={self.bbox[1]}" +
            f"&west={self.bbox[0]}&east={self.bbox[2]}&coordsOrTiles={self.coordsOrTiles}" +
            f"&dayNightBoth={self.daynight}")
        return re.findall('<return>(.*?)</return>', requests.get(url).text)

    def order_size(self):
        "Calculates the number of files in the order."
        if self.bands is None:
            raise Exception("`bands` list required to calculate order_size.")
        return len(self.search_files())*len(self.bands)

    def split_times(self):
        "Split a single order into multiple orders if the order size is too large."
        order_size = self.order_size()
        n_splits = order_size // self._maxOrderSize + 1
        times = pd.date_range(self.tstart, self.tend)
        splits = np.split(times, n_splits)
        times = [(str(t[0]), str(t[-1]+pd.Timedelta(days=1-1e-5).round('s'))) for t in splits]
        tstart, tend = zip(*times)
        kwargs = self.__dict__
        group = []
        for ti, tf in zip(tstart, tend):
            kwargs['tstart'], kwargs['tend'] = ti, tf
            group.append(Ladsweb(**kwargs))
        return group

    def send_order(self, ids, email=None):
        "Send order for a set of ids obtained with `search_files` method."
        if email is None: raise Exception("`email` is not defined")
        ids = ','.join(ids)
        bands = ','.join([self.product + f'___{b}' for b in self.bands])
        url = (f"http://modwebsrv.modaps.eosdis.nasa.gov/axis2/services/MODAPSservices/" +
            f"orderFiles?fileIds={ids}" +
            f"&subsetDataLayer={bands}" +
            f"&geoSubsetNorth={self.bbox[3]}" +
            f"&geoSubsetSouth={self.bbox[1]}" +
            f"&geoSubsetEast={self.bbox[2]}" +
            f"&geoSubsetWest={self.bbox[0]}" +
            f"&reprojectionName={self.repName}" +
            f"&reprojectionOutputPixelSize={self.repPixSize}" +
            f"&reprojectionResampleType={self.repResample}" +
            f"&doMosaic={self.doMosaic}" +
            f"&email={email}")
        return re.findall('<return>(.*?)</return>', requests.get(url).text)[0]

    def order_status(self, orderId):
        url = (f"http://modwebsrv.modaps.eosdis.nasa.gov/axis2/services/MODAPSservices/" +
                f"getOrderStatus?orderId={orderId}")
        return re.findall('<return>(.*?)</return>', requests.get(url).text)[0]

    def download_files(self, orderId, path_save, auth=None):
        "Download files if the order is Available."
        if auth is None: raise Exception("`auth` code is not defined")
        status = self.order_status(orderId)
        if status != 'Available':
            msg = f"Order is not Available, current status is {status}"
            warnings.warn(msg, UserWarning)
            return
        url = f'https://ladsweb.modaps.eosdis.nasa.gov/archive/orders/{orderId}'
        files = pd.DataFrame(json.loads(geturl(url + '.json', auth)))
        checksums = geturl(url + f'/checksums_{orderId}', auth)
        hdfs = re.findall('(.*?.hdf)', checksums)
        ch = [tuple([k for k in h.split(' ') if k != '']) for h in hdfs]
        check_df = pd.DataFrame(ch, columns=['checksum', 'size', 'name'])
        check_df = check_df.drop('size', axis=1)
        files = pd.merge(files, check_df, how='left', on='name')
        files['verified'] = False

        for i in progress_bar(range(len(files))):
            file, checksum = files.loc[i, ['name', 'checksum']]
            n_tries = 0
            while ~files.loc[i, 'verified'] and n_tries<5:
                #print(f'Downloading {file}')
                with open(path_save/f'{file}', mode='w+b') as fh:
                    geturl(f'{url}/{file}', auth, fh)
                csum = os.popen(f'cksum {str(path_save)}/{file}').read().split(' ')[0]
                if str(checksum) == 'nan': checksum = csum
                files.loc[i, 'verified'] = checksum == csum
                n_tries += 1
        log_file = f'download_log_{orderId}.csv'
        files.to_csv(path_save/log_file)
        not_verified = np.sum(~files.verified)
        if not_verified > 0:
            msg = f"Checksum failed for {not_verified} files. Check the {log_file}."
            warnings.warn(msg, UserWarning)
        return not_verified

    def release_order(self, orderId, email=None):
        "To release order after download the files."
        if email is None: raise Exception("`email` is not defined")
        url = (f"http://modwebsrv.modaps.eosdis.nasa.gov/axis2/services/MODAPSservices/" +
        f"releaseOrder?orderId={orderId}&email={email}")
        status = re.findall('<return>(.*?)</return>', requests.get(url).text)[0]
        return status == '1'

    def run(self, path_save, email, auth):
        "Run a complete request from search_files to release_order."
        ids = self.search_files()
        if len(ids) == 0:
            warnings.warn("No files found", UserWarning)
            return
        if email is None: raise Exception("The `email` is required to place an order")
        if self.bands is None: raise Exception("A list of `bands` is not defined")
        orderId = self.send_order(ids, email)
        try:
            status = None
            while status != 'Available':
                status = self.order_status(orderId)
                print(f"Current status: {status}")
                sleep(60)
        except KeyboardInterrupt:
            raise Exception("KeaybordInterrupt")
        if auth is None: raise Exception("The `auth` key is required to download the files")
        if path_save is None: raise Exception("The `path_save` is not defined")
        path_save = Path(path_save)
        path_save.mkdir(exist_ok=True, parents=True)
        print("Waiting 5 min to start download...")
        sleep(60*5)
        status = self.download_files(orderId, path_save, auth)
        if status==0: result = self.release_order(orderId, email)
        if result: print(f'Process complete, files saved at {path_save}.')

    def __repr__(self):
        s = ''
        for k in self.__dict__:
            s += f'{k}: {self.__dict__[k]}, '
        return s + '\n'

# Cell
def run_parallel(obj_list, path_save, email, auth, max_workers=8):
    "Run a list of requests in parallel."
    obj_list = [partial(obj.run, path_save=path_save, email=email, auth=auth)
                for obj in obj_list]
    with ThreadPoolExecutor(max_workers) as e:
        list(progress_bar(e.map(lambda x:x(), obj_list), total=len(obj_list)))