sunnyportal2file

#!/usr/bin/env python3

# Copyright (c) 2020 Joel Berglund <joebe975@protonmail.com>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
# USA

from datetime import date, datetime, timedelta
from getpass import getpass

import argparse
import configparser
import logging
import sunnyportal.client, sunnyportal.responses
import re
import os
import pandas as pd
import numpy as np


def valid_date_type(arg_date_str):
    try:
        d = datetime.strptime(arg_date_str, "%Y-%m-%d").date()
        if d >= date.today():
            # Sunny Portal will return values for the whole day, while setting
            # future times as blank (interpreted as 0)
            logging.debug(
                f"Cannot look into the future ({d}), changing to yesterday..."
            )
            return date.today() - timedelta(days=1)
        return d
    except ValueError:
        msg = f"Given Date {arg_date_str} not valid! Expected format, YYYY-MM-DD!"
        raise argparse.ArgumentTypeError(msg)


class DataFrameHandler:
    def __init__(self, plant, start_date, end_date, format):
        self._plant = plant
        self._plant_name = self._fix_plant_name(plant.name)
        self._file_extension = self._file_extension_dict[format]
        self._file_name = f"{self._plant_name}.{self._file_extension}"
        self._start_date = start_date
        self._end_date = end_date
        self._old_df = pd.DataFrame()
        self._new_df = pd.DataFrame()

    _file_extension_dict = {
        "json": "json",
        "pickle": "pkl",
        "csv": "csv",
        "feather": "feather",
        "parquet": "parquet",
        "excel": "xlsx",
        "sqlite": "db",
    }

    # Assuming the power in W is always an integer
    _dtype_dict = {
        "timestamp": np.datetime64,
        "mean_power": np.uint32,
        "min_power": np.uint32,
        "max_power": np.uint32,
    }

    def _get_power_measurements_for_day(self, date):
        day = self._plant.day_overview(date, include_all=True)
        return day.power_measurements

    def _set_dtype(self, df):
        dtypes = {k: self._dtype_dict[k] for k in df.columns}
        return df.astype(dtypes)

    def download_new_data(self):
        """Extracts data for the period between self._start_date and self._end_date"""

        start_date = self._start_date

        if start_date > self._end_date:
            logging.debug(
                f"start_date {start_date} cannot be larger than end_date {self._end_date}"
            )
            self._new_df = pd.DataFrame()
            return

        data = []

        while start_date <= self._end_date:
            res = self._get_power_measurements_for_day(start_date)
            if res:
                data.append(
                    pd.DataFrame(res)
                )  # The column names will be taken from the namedtuple fields

            start_date = start_date + timedelta(days=1)

        # Concatenate the dataframes, remove columns without values (eg. missing min, max),
        # replace NaN with 0 and rename the columns, reset the index and set the resulting
        # dtypes according to _dtype_dict.
        translation_dict = {
            "power": "mean_power",
            "min": "min_power",
            "max": "max_power",
        }

        if data:
            df = (
                pd.concat(data)
                .dropna(axis=1, how="all")
                .fillna(0)
                .rename(columns=translation_dict)
                .reset_index(drop=True)
            )

            self._new_df = df.astype({k: self._dtype_dict[k] for k in df.columns})

        else:
            logging.debug(
                f"No data found between {self._start_date} and {self._end_date}"
            )
            self._new_df = pd.DataFrame()

    def prepare(self):
        """Reads old data, if needed, and sets new start date accordingly.
        Some subclasses of DataFrameHandler might need to override this method."""

        if os.path.isfile(self._file_name):
            logging.debug(f"Reading old dataframe from file {self._file_name}")
            self._old_df = self._read_df()

            old_end_date = np.amax(self._old_df["timestamp"]).date()

            if old_end_date > self._start_date:
                self._start_date = old_end_date + timedelta(days=1)
                logging.debug(
                    f"New start_date {self._start_date} was extracted from file {self._file_name}"
                )

    def _fix_plant_name(self, old_plant_name):
        """Replaces potentially illegal characters in name with underscore (assumes worst case OS)"""
        replace_pattern = '[<>:"/|?*\\\\]'
        plant_name = re.sub(replace_pattern, "_", old_plant_name)

        if plant_name != old_plant_name:
            logging.info(
                f"Plant name '{old_plant_name}' was converted to '{plant_name}'"
            )
        return plant_name

    def _merge_dataframes(self):
        """Combines old and new data, if applicable"""

        if (not self._old_df.empty) & (not self._new_df.empty):
            logging.debug(f"Merging old and new data")
            self._df = pd.concat([self._old_df, self._new_df], axis=0).reset_index(
                drop=True
            )  # .sort_values('timestamp', ignore_index=True)
        elif not self._new_df.empty:
            logging.debug(f"Only new data, no need to merge")
            self._df = self._new_df
        else:
            self._df = pd.DataFrame()

    def _read_df(self):
        pass

    def _write_df(self):
        pass

    def save_data(self):
        """Saves data to file/database"""

        # Merge the old and new data (if applicable)
        self._merge_dataframes()

        if self._df.empty:
            logging.debug("Empty DataFrame. Nothing to save...")
            return

        logging.info(f"Saving to file {self._file_name}")
        self._write_df()


class PickleHandler(DataFrameHandler):
    def __init__(self, plant, start_date, end_date):
        super().__init__(plant, start_date, end_date, format="pickle")
        self.kwargs = {}

    def _read_df(self):
        return pd.read_pickle(self._file_name, **self.kwargs)

    def _write_df(self):
        self._df.to_pickle(self._file_name, **self.kwargs)


class JSONHandler(DataFrameHandler):
    def __init__(self, plant, start_date, end_date):
        super().__init__(plant, start_date, end_date, format="json")
        # Change json format below
        self.kwargs = {"orient": "table"}

    def _read_df(self):
        return self._set_dtype(pd.read_json(self._file_name, **self.kwargs))

    def _write_df(self):
        self._df.to_json(self._file_name, index=False, **self.kwargs)


class CSVHandler(DataFrameHandler):
    def __init__(self, plant, start_date, end_date):
        super().__init__(plant, start_date, end_date, format="csv")
        self.kwargs = {}

    def _read_df(self):
        return self._set_dtype(
            pd.read_csv(self._file_name, parse_dates=["timestamp"], **self.kwargs)
        )

    def _write_df(self):
        self._df.to_csv(self._file_name, index=False, **self.kwargs)


class ParquetHandler(DataFrameHandler):
    def __init__(self, plant, start_date, end_date):
        super().__init__(plant, start_date, end_date, format="parquet")
        self.kwargs = {}

    def _read_df(self):
        return self._set_dtype(pd.read_parquet(self._file_name, **self.kwargs))

    def _write_df(self):
        self._df.to_parquet(self._file_name, **self.kwargs)


class SQLiteHandler(DataFrameHandler):
    def __init__(self, plant, start_date, end_date):
        import sqlite3

        super().__init__(plant, start_date, end_date, format="sqlite")
        self._file_name = "sunny_portal.db"
        self._conn = sqlite3.connect(self._file_name)
        self._c = self._conn.cursor()

    def __table_exists(self):
        return self._c.execute(
            f"""SELECT count(name) FROM sqlite_master WHERE type='table' AND name='{self._plant_name}' """
        ).fetchone()[0]

    def prepare(self):
        """No need to read the old data as dataframe in SQL. Just extract the latest date"""

        if os.path.isfile(self._file_name) & self.__table_exists():
            logging.debug(f"Extracting latest date from database {self._file_name}")
            old_end_date = datetime.strptime(
                self._conn.execute(
                    f'SELECT MAX(timestamp) FROM "{self._plant_name}"'
                ).fetchone()[0],
                "%Y-%m-%d %H:%M:%S",
            ).date()

            if old_end_date > self._start_date:
                self._start_date = old_end_date + timedelta(days=1)
                logging.debug(
                    f"New start_date {self._start_date} was extracted from file {self._file_name} and table {self._plant_name}"
                )

    def _read_df(self):
        return self._set_dtype(
            pd.read_sql(f"SELECT * FROM {self._plant_name}", con=self._conn)
        )

    def _write_df(self):
        self._df.to_sql(
            self._plant_name, con=self._conn, index=False, if_exists="append"
        )


class FeatherHandler(DataFrameHandler):
    def __init__(self, plant, start_date, end_date):
        super().__init__(plant, start_date, end_date, format="feather")
        self.kwargs = {}

    def _read_df(self):
        return pd.read_feather(self._file_name, **self.kwargs)

    def _write_df(self):
        self._df.to_feather(self._file_name, **self.kwargs)


class ExcelHandler(DataFrameHandler):
    def __init__(self, plant, start_date, end_date):
        super().__init__(plant, start_date, end_date, format="excel")
        self.kwargs = {}

    def _read_df(self):
        return pd.read_excel(self._file_name, dtype=self._dtype_dict, **self.kwargs)

    def _write_df(self):
        self._df.to_excel(self._file_name, index=False, **self.kwargs)


df_Handler_dict = {
    "json": JSONHandler,
    "pickle": PickleHandler,
    "csv": CSVHandler,
    "feather": FeatherHandler,
    "parquet": ParquetHandler,
    "excel": ExcelHandler,
    "sqlite": SQLiteHandler,
}


def get_df_handler(plant, format, start_date, end_date):
    df_handler = df_Handler_dict[format](plant, start_date, end_date)
    logging.debug(f"Initialized {df_handler.__class__.__name__} object")

    return df_handler


def main():
    logging.basicConfig(
        format="%(asctime)s %(levelname)s: %(message)s", level=logging.DEBUG
    )

    parser = argparse.ArgumentParser(
        description="Save information from Sunny Portal to file"
    )
    parser.add_argument("config", help="Configuration file to use")

    parser.add_argument(
        "-f",
        "--format",
        choices=["json", "csv", "pickle", "feather", "parquet", "excel", "sqlite"],
        required=True,
        help="Format for which the data is to be saved",
    )

    parser.add_argument(
        "-s",
        "--start-date",
        type=valid_date_type,
        default=date.today() - timedelta(days=1),
        required=False,
        help="The start date of data to be saved in the format YYYY-MM-DD (default yesterday)",
    )
    parser.add_argument(
        "-e",
        "--end-date",
        type=valid_date_type,
        default=date.today() - timedelta(days=1),
        required=False,
        help="The end date of data to be saved in the format YYYY-MM-DD (default yesterday)",
    )

    parser.add_argument(
        "-i",
        "--include-filter",
        type=str,
        default="",
        required=False,
        help="A string used to filter which plants to include (default includes all plants)",
    )

    parser.add_argument("-q", "--quiet", help="Silence output", action="store_true")

    args = parser.parse_args()

    if args.quiet:
        logging.disable(logging.DEBUG)

    config = configparser.ConfigParser()
    config["sunnyportal"] = {}
    config.read(args.config)

    modified = False
    if not config["sunnyportal"].get("email"):
        config["sunnyportal"]["email"] = input("Sunny Portal e-mail: ")
        modified = True
    if not config["sunnyportal"].get("password"):
        config["sunnyportal"]["password"] = getpass("Sunny Portal password: ")
        modified = True

    if modified:
        with open(args.config, "w") as configfile:
            config.write(configfile)

    client = sunnyportal.client.Client(
        config["sunnyportal"]["email"], config["sunnyportal"]["password"]
    )

    for plant in client.get_plants():
        if plant.name.startswith(args.include_filter):

            df_handler = get_df_handler(
                plant, args.format, args.start_date, args.end_date
            )

            # Prepare the df_handler where start date might be adjusted if
            # data already exists. For some formats, old data will be read.
            df_handler.prepare()

            # Downloads new data
            df_handler.download_new_data()

            # Save the data to file/databse
            df_handler.save_data()
        else:
            logging.debug(
                f'Plant {plant.name} does not match include filter "{args.include_filter}"'
            )
    client.logout()


if __name__ == "__main__":
    main()