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

__all__ = ['AutoTS']

# Cell
import json
import os
import requests
import sys
import time
from requests.auth import HTTPBasicAuth
from typing import Dict, List, Optional, Union

import boto3
import botocore
import pandas as pd

# Cell
class AutoTS:
    """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) -> 'AutoTS':
        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 _parse_response(self, response: str) -> Dict:
        """Parses responde."""
        try:
            resp = json.loads(response)
        except Exception as e:
            raise Exception(response)

        return resp

    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, force_upload: bool = False) -> str:
        """Uploads file to s3.

        Parameters
        ----------
        file: str
            Local file.
        force_upload: bool
            If True forces uploading even if the file already exists.
        """
        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
        )
        # Check if exists
        if not force_upload:
            try:
                session = boto3.Session(
                    aws_access_key_id=self.aws_access_key_id,
                    aws_secret_access_key=self.aws_secret_access_key
                )

                session.resource('s3').Object(self.bucket_name, filename).load()
            except botocore.exceptions.ClientError as e:
                if e.response['Error']['Code'] != "404":
                    raise e
            else:
                # The object does exist
                return filename

        # Progress bar
        size = os.stat(file).st_size
        progress_bar = self._progress(size)

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

        return filename

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

        Parameters
        ----------
        filename: str
            File to download.
        filename_output: str
            Filename output.
        """
        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`'
            )

        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=self.bucket_name, Key=filename)
        size = int(meta_data.get('ContentLength', 0))
        progress_bar = self._progress(size)

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

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

        Parameters
        ----------
        filename: str
            Filename 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(
            s3_args=dict(
                s3_url=f's3://{self.bucket_name}',
            ),
            args = dict(
                filename=filename,
                kind=kind,
                freq=freq,
                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._parse_response(resp.text)

    def tsfeatures_temporal(self, filename: 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
        ----------
        filename: str
            Filename 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._tsfeatures(s3_uri=s3_uri,
                                kind='temporal',
                                freq=freq,
                                unique_id_column=unique_id_column,
                                ds_column=ds_column,
                                y_column=y_column)

    def tsfeatures_static(self, filename: str, freq: int,
                          unique_id_column: str = 'unique_id',
                          ds_column: str = 'ds',
                          y_column: str = 'y') -> Dict:
        """Calculates static features from S3 URL.

        Parameters
        ----------
        filename: str
            Filename 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._tsfeatures(s3_uri=s3_uri,
                                kind='static',
                                freq=freq,
                                unique_id_column=unique_id_column,
                                ds_column=ds_column,
                                y_column=y_column)

    def calendartsfeatures(self, filename: str, country: str,
                           events: Optional[Union[str, Dict[str, List[str]]]] = None,
                           unique_id_column: str = 'unique_id',
                           ds_column: str = 'ds',
                           y_column: str = 'y') -> Dict:
        """Calculates calendar features from S3 URL.

        Parameters
        ----------
        filename: str
            filename of the dataset.
            The dataset should contain at least three columns:
                - Time series identifier.
                - Time identifier.
                - Target identifier.
        country: str
            Country to calculate calendar features from.
        events: Union[str, Dict[str, List[str]]]
            - str file
                String file of events with the format
                event1=fecha1,fecha2/event2=fecha1,fecha2,fecha3
                allocated in the bucket.
            - Dictionary of events.
                The key defines the name of the event and
                the value is a list of the dates to consider.
                Ej. {'important_event_1': ['2010-02-05', '2011-02-07', '2012-02-08', '2013-02-09'],
                     'important_event_2': ['2010-03-14', '2011-03-12', '2012-03-16', '2013-03-18'],}
                The format of each date should be 'YYYY-MM-DD'.
        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(
            s3_args=dict(
                s3_url=f's3://{self.bucket_name}',
            ),
            args = dict(
                filename=filename,
                country=country,
                unique_id_column=unique_id_column,
                ds_column=ds_column,
                y_column=y_column,
            )
        )

        if events is not None:
            if isinstance(events, str):
                query['args']['events'] = events
            else:
                str_events = [f'{key}={",".join(value)}' for key, value in events.items()]
                str_events = '/'.join(str_events)
                query['args']['events'] = str_events

        resp = requests.post(f'{self.invoke_url}/calendartsfeatures',
                             headers={'x-api-key': self.api_key},
                             data=json.dumps(query))

        return self._parse_response(resp.text)

    def tsforecast(self,
                   filename_target: str,
                   freq: str,
                   horizon: int,
                   filename_static: str,
                   filename_temporal: str,
                   objective: Optional[str] = 'tweedie',
                   metric: Optional[str] = 'tweedie',
                   learning_rate: Optional[float] = 0.1,
                   n_estimators: Optional[int] = 100,
                   num_leaves: Optional[int] = 128,
                   min_data_in_leaf: Optional[int] = 20,
                   bagging_freq: Optional[int] = 0,
                   bagging_fraction: Optional[float] = 1.,
                   unique_id_column: str = 'unique_id',
                   ds_column: str = 'ds',
                   y_column: str = 'y') -> Dict:
        """Calculates forecast.

        Parameters
        ----------
        filename_target: str
            Filename of the target dataset.
            The dataset should contain at least three columns:
                - Time series identifier.
                - Time identifier.
                - Target identifier.
        freq: str
            Frequency of the data.
            Only 'D' supported.
        horizon: int
            Forcast horizon.
        filename_static: str
            Filename of the dataset with static features.
        filename_temporal: str
            Filename of the dataset with temporal features.
        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(
            s3_args=dict(
                s3_url=f's3://{self.bucket_name}',
            ),
            data_args = dict(
                filename=filename_target,
                freq=freq,
                horizon=horizon,
                filename_static=filename_static,
                filename_temporal=filename_temporal,
                unique_id_column=unique_id_column,
                ds_column=ds_column,
                y_column=y_column,
            ),
            model_args = dict(
                objective=objective,
                metric=metric,
                learning_rate=learning_rate,
                n_estimators=n_estimators,
                num_leaves=num_leaves,
                min_data_in_leaf=min_data_in_leaf,
                bagging_freq=bagging_freq,
                bagging_fraction=bagging_fraction,
            )
        )

        resp = requests.post(f'{self.invoke_url}/tsforecast',
                             headers={'x-api-key': self.api_key},
                             data=json.dumps(query))

        return self._parse_response(resp.text)

    def tsbenchmarks(self,
                     filename: str,
                     dataset: str,
                     unique_id_column: str = 'unique_id',
                     ds_column: str = 'ds',
                     y_column: str = 'y') -> Dict:
        """Calculates performance based on filename and dataset.

        Parameters
        ----------
        filename: str
            Filename of the dataset.
            The dataset should contain at least three columns:
                - Time series identifier.
                - Time identifier.
                - Target identifier.
        dataset: str
            Target dataset.
        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(
            s3_args=dict(
                s3_url=f's3://{self.bucket_name}',
            ),
            args = dict(
                filename=filename,
                dataset=dataset,
                unique_id_column=unique_id_column,
                ds_column=ds_column,
                y_column=y_column,
            )
        )

        resp = requests.post(f'{self.invoke_url}/tsbenchmarks',
                             headers={'x-api-key': self.api_key},
                             data=json.dumps(query))

        return self._parse_response(resp.text)


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

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

        output = self._parse_response(resp.text)

        return output

    def check_progress_logs(self, job_id: str, sleep_time: int = 45) -> Dict:
        """Waits untile `job_id` is done and prints logs."""
        sys.stdout.write(
            """
            @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
            @@@@.####.@@(.@@/((((.@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
            @@@@@.####.&(@.,((((.@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
            @@@@@@*####....((((.@@@@@@@@@@..*@@@@@@@.@@@@@@@@/.@@@@@@@@..@@@@@..@@@@@@..........@@@@@@..@@@@@@@@@@@@@@@@@.../@@@@@@@
            @@@@@@@(###.%@.(((.@@@@@@@@@@@.@*.(@@@@@.@@@@@@@@/.@@@@@@@@@@..@..@@@@@@@@@@@@..@@@@@@@@@@..@@@@@@@@@@@@@@@@..@@.,@@@@@@
            @@@@@@@@/#.....,(.@@@@@@@@@@@@.@@@,./@@@.@@@@@@@@/.@@@@@@@@@@@...@@@@@@@@@@@@@..@@@@@@@@@@..@@@@@@@@@@@@@@@.*@@@@..@@@@@
            @@@@@@@(###(.%*(((.@@@@@@@@@@@.@@@@@/.(@.@@@@@@@@/.@@@@@@@@@*.@@&.,@@@@@@@@@@@..@@@@@@@@@@..@@@@@@@@@@@@@@..........@@@@
            @@@@@@,####....((((.@@@@@@@@@@.@@@@@@@...@@@@@@@@/.@@@@@@@(./@@@@@..&@@@@@@@@@..@@@@@@@@@@......@@@@@@@@/.@@@@@@@@@,.@@@
            @@@@@.####.@,&.,((((.@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
            @@@@.####.@@@*@@/((((.@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

            Have some coffee while you wait...
               ( (
                ) )
              ........
              |      |]
              \      /
               `----'
            """
        )
        sys.stdout.write(f'\nSetting up Nixtla infrastructure (this will take up 5 mins)...\n')

        idx_logs = 0
        in_progress = True
        while in_progress:
            resp = self.get_status(job_id)
            status = resp['status']
            logs = json.loads(resp['logs'])

            if status != 'InProgress' and not logs:
                time.sleep(30)
                resp = self.get_status(job_id)
                status = resp['status']
                logs = json.loads(resp['logs'])

            if logs:
                #if logs != latest_logs:
                for log in logs[idx_logs:]:
                    sys.stdout.write(f'{log}\n')
                #latest_logs = logs
                idx_logs = len(logs)

            in_progress = status == 'InProgress'

            time.sleep(sleep_time)

        return status