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

__all__ = ['ProgressPercentage', 'TSFeatures']

# Cell
import json
import os
import requests
import sys
from requests.auth import HTTPBasicAuth
from typing import Optional

import boto3
import pandas as pd

# Cell
class ProgressPercentage(object):
    def __init__(self, filename):
        self._filename = filename
        self._size = float(os.path.getsize(filename))
        self._seen_so_far = 0
        self._lock = threading.Lock()

    def __call__(self, bytes_amount):
        # To simplify we'll assume this is hooked up
        # to a single filename.
        with self._lock:
            self._seen_so_far += bytes_amount
            percentage = round((self._seen_so_far / self._size) * 100,2)
            LoggingFile('{} is the file name. {} out of {} done. The percentage completed is {} %'.format(str(self._filename), str(self._seen_so_far), str(self._size),str(percentage)))
            sys.stdout.flush()

# Cell
class TSFeatures:
    """Compute time series features at scale.
    Send an email to fede.garza.ramirez@gmail.com to request access.

    Parameters
    ----------
    bucket_name: str
        Bucket name.
    api_id: str
        API identifier.
    api_key: str
        API key.
    aws_access_key_id: str
        AWS ACCESS KEY ID.
    aws_secret_access_key: str
        AWS SECRET ACCESS KEY.
    """

    def __init__(self, bucket_name: str,
                 api_id: str, api_key: str,
                 aws_access_key_id: str,
                 aws_secret_access_key: str) -> 'TSFeatures':
        self.bucket_name = bucket_name
        self.invoke_url = f'https://{api_id}.execute-api.us-east-1.amazonaws.com/main'
        self.api_key = api_key
        self.aws_access_key_id = aws_access_key_id
        self.aws_secret_access_key = aws_secret_access_key

    def _from_api_to_df(self, response: str) -> pd.DataFrame:
        """Transforms api output to df."""
        resp = json.loads(response)

        logs = resp.pop('logs', None)
        if logs is not None:
            logs = json.loads(logs)

            for log in logs:
                sys.stdout.write(f'{log}\n')

        df = pd.DataFrame(resp, index=[0])

        return df

    def _progress(self, size: float):
        """Progress bar based on size."""
        downloaded = 0

        def progress(chunk):
            nonlocal downloaded
            downloaded += chunk
            done = int(50 * downloaded / size)
            sys.stdout.write(f'\r[{"=" * done}{" " * (50-done)}]')
            sys.stdout.flush()

        return progress

    def upload_to_s3(self, file: str) -> str:
        """Uploads file to s3.

        Parameters
        ----------
        file: str
            Local file.
        """
        if self.aws_access_key_id is None:
            raise Exception(
                'To use `upload_to_s3` you need to pass '
                '`aws_access_key_id` and '
                '`aws_secret_access_key`'
            )

        filename = file.split('/')[-1]

        s3 = boto3.client('s3',
                          aws_access_key_id=self.aws_access_key_id,
                          aws_secret_access_key=self.aws_secret_access_key)
        # Progress bar
        size = os.stat(filename).st_size
        progress_bar = self._progress(size)

        # Uploading file
        s3.upload_file(file, self.bucket_name, filename, Callback=progress_bar)

        s3_uri = f's3://{self.bucket_name}/{filename}'

        return s3_uri

    def download_from_s3(self, s3_url: str, file: Optional[str] = None) -> str:
        """Downloads file from s3.

        Parameters
        ----------
        s3_url: str
            S3 url.
        filename: str
            Destination file
        """
        if self.aws_access_key_id is None:
            raise Exception(
                'To use `download_from_s3` you need to pass '
                '`aws_access_key_id` and '
                '`aws_secret_access_key`'
            )

        bucket_name = s3_url.split('/')[2]
        s3_file = s3_url.split('/')[-1]

        s3 = boto3.client('s3',
                          aws_access_key_id=self.aws_access_key_id,
                          aws_secret_access_key=self.aws_secret_access_key)

        # Progress bar
        meta_data = s3.head_object(Bucket=bucket_name, Key=s3_file)
        size = int(meta_data.get('ContentLength', 0))
        progress_bar = self._progress(size)

        # Downloading file
        s3.download_file(bucket_name, s3_file, s3_file if file is None else file, Callback=progress_bar)

    def _calculate_features_from_s3_uri(self, s3_uri: str, freq: int,
                                       kind: str = 'static',
                                       unique_id_column: str = 'unique_id',
                                       ds_column: str = 'ds',
                                       y_column: str = 'y') -> pd.DataFrame:
        """Calculates features from S3 URL.

        Parameters
        ----------
        s3_uri: str
            S3 uri of the dataset.
            The dataset should contain at least three columns:
                - Time series identifier.
                - Time identifier.
                - Target identifier.
        freq: int
            Frequency of the time series.
        kind: str
            Kind of features. Static returns features for each time series.
            temporal returns for each ds and each time series.
        unique_id_column: str
            Column name identifying each time series.
        ds_column: str
            Column name identifying each time stamp.
        y_column: str
            Column name identifying the target variable.

        Notes
        -----
        [1] The dataset should contain time series of the same frequency.
        """
        query = dict(url=s3_uri, freq=freq, kind=kind,
                     unique_id_column=unique_id_column,
                     ds_column=ds_column,
                     y_column=y_column)
        resp = requests.post(f'{self.invoke_url}/tsfeatures',
                             headers={'x-api-key': self.api_key},
                             data=json.dumps(query))

        return self._from_api_to_df(resp.text)

    def calculate_temporal_features_from_s3_uri(self, s3_uri: str, freq: int,
                                                 unique_id_column: str = 'unique_id',
                                                 ds_column: str = 'ds',
                                                 y_column: str = 'y') -> pd.DataFrame:
        """Calculates temporal features from S3 URL.

        Parameters
        ----------
        s3_uri: str
            S3 uri of the dataset.
            The dataset should contain at least three columns:
                - Time series identifier.
                - Time identifier.
                - Target identifier.
        freq: int
            Frequency of the time series.
        unique_id_column: str
            Column name identifying each time series.
        ds_column: str
            Column name identifying each time stamp.
        y_column: str
            Column name identifying the target variable.

        Notes
        -----
        [1] The dataset should contain time series of the same frequency.
        """

        return self._calculate_features_from_s3_uri(s3_uri=s3_uri,
                                                    kind='temporal',
                                                    freq=freq,
                                                    unique_id_column=unique_id_column,
                                                    ds_column=ds_column,
                                                    y_column=y_column)

    def calculate_static_features_from_s3_uri(self, s3_uri: str, freq: int,
                                              unique_id_column: str = 'unique_id',
                                              ds_column: str = 'ds',
                                              y_column: str = 'y') -> pd.DataFrame:
        """Calculates static features from S3 URL.

        Parameters
        ----------
        s3_uri: str
            S3 uri of the dataset.
            The dataset should contain at least three columns:
                - Time series identifier.
                - Time identifier.
                - Target identifier.
        freq: int
            Frequency of the time series.
        unique_id_column: str
            Column name identifying each time series.
        ds_column: str
            Column name identifying each time stamp.
        y_column: str
            Column name identifying the target variable.

        Notes
        -----
        [1] The dataset should contain time series of the same frequency.
        """

        return self._calculate_features_from_s3_uri(s3_uri=s3_uri,
                                                    kind='static',
                                                    freq=freq,
                                                    unique_id_column=unique_id_column,
                                                    ds_column=ds_column,
                                                    y_column=y_column)

    def calculate_calendar_features_from_s3_uri(self, s3_uri: str, country: str,
                                                unique_id_column: str = 'unique_id',
                                                ds_column: str = 'ds',
                                                y_column: str = 'y') -> pd.DataFrame:
        """Calculates static features from S3 URL.

        Parameters
        ----------
        s3_uri: str
            S3 uri of the dataset.
            The dataset should contain at least three columns:
                - Time series identifier.
                - Time identifier.
                - Target identifier.
        country: int
            Country to calculate calendar features from.
        unique_id_column: str
            Column name identifying each time series.
        ds_column: str
            Column name identifying each time stamp.
        y_column: str
            Column name identifying the target variable.
        """
        query = dict(url=s3_uri, country=country,
                     unique_id_column=unique_id_column,
                     ds_column=ds_column,
                     y_column=y_column)
        resp = requests.post(f'{self.invoke_url}/calendartsfeatures',
                             headers={'x-api-key': self.api_key},
                             data=json.dumps(query))

        return self._from_api_to_df(resp.text)

    def get_status(self, job_id: str) -> pd.DataFrame:
        """Gets job status.

        Parameters
        ----------
        job_id: str
            ID job from `calculate_features_from_s3_uri`.
        """
        resp = requests.get(f'{self.invoke_url}/tsfeatures/jobs/',
                            params={'job_id': job_id},
                            headers={'x-api-key': self.api_key})

        return self._from_api_to_df(resp.text)