"""Module for parsing EnergyPlus SQLite result files into DataFrames."""
import logging
from datetime import timedelta
from sqlite3 import connect
from typing import List, Optional, Sequence, Union

import numpy as np
import pandas as pd
from energy_pandas import EnergyDataFrame
from pandas import to_datetime
from path import Path
from typing_extensions import Literal

from archetypal.utils import log

_REPORTING_FREQUENCIES = Literal[
    "HVAC System Timestep",
    "Zone Timestep",
    "Hourly",
    "Daily",
    "Monthly",
    "Run Period",
]


class SqlOutput:
    """Represents a single output from the Sql file."""

    __slots__ = (
        "_file_path",
        "output_name",
        "reporting_frequency",
    )

    def __init__(self, file_path, output_name, reporting_frequency):
        self._file_path = file_path
        self.output_name = output_name
        self.reporting_frequency = reporting_frequency

    def values(self, environment_type: int = 3, units: str = None) -> EnergyDataFrame:
        """Get the time series values as an EnergyDataFrame.

        Args:
            environment_type (int): The environment type (1 = Design Day, 2 = Design
                Run Period, 3 = Weather Run Period) for the series.
            units (str): Convert original values to another unit. The original unit
                is detected automatically and a dimensionality check is performed.

        Returns:
            (EnergyDataFrame): The time series as an EnergyDataFrame.
        """
        cols = (
            "ReportDataDictionaryIndex, IndexGroup, KeyValue, Name, "
            "Units, ReportingFrequency"
        )
        query = f"""
            SELECT {cols} 
            FROM ReportDataDictionary 
            WHERE Name=@output_name 
            AND ReportingFrequency=@reporting_frequency;
        """
        with connect(self._file_path) as conn:
            header_rows = pd.read_sql(
                query,
                conn,
                params={
                    "output_name": self.output_name,
                    "reporting_frequency": self.reporting_frequency,
                },
            )
            header_rows.set_index("ReportDataDictionaryIndex", inplace=True)
            # extract all data of the relevant type from ReportData
            rel_indices = tuple(header_rows.index.to_list())
            data = _extract_timeseries(conn, environment_type, header_rows, rel_indices)

            if units is not None:
                data = data.to_units(units)

            return data


class _SqlOutputs:
    """Represents all the available outputs from the Sql file."""

    def __init__(self, file_path: str, available_outputs: List[tuple]):
        self._available_outputs = available_outputs
        self._properties = {}

        for output, reporting_frequency in self._available_outputs:
            name = (
                output.replace(":", "__").replace(" ", "_") + f"_{reporting_frequency}"
            )
            self._properties[name] = SqlOutput(file_path, output, reporting_frequency)
            setattr(self, name, self._properties[name])

    def __getitem__(self, meter_name):
        """Get item by key."""
        return self._properties[meter_name]


class Sql:
    """Object for parsing EnergyPlus SQLite result files into DataFrames.

    Args:
        file_path: Full path to an SQLite file that was generated by EnergyPlus.

    Properties:
        * file_path
        * available_outputs
        * zone_info
        * environment_periods
    """

    _reporting_frequencies = (
        "HVAC System Timestep",
        "Zone Timestep",
        "Hourly",
        "Daily",
        "Monthly",
        "Run Period",
    )

    def __init__(self, file_path):
        """Initialize SQLiteResult"""
        assert Path(file_path).exists(), "No file was found at {}".format(file_path)
        self._file_path = file_path

        # values to be computed as soon as they are requested
        self._available_outputs = None
        self._zone_info = None
        self._environment_periods = None
        self._tabular_data_keys = None
        self._outputs = None
        self._surfaces_table = None
        self._constructions_table = None

    @property
    def file_path(self):
        """Get the path to the .sql file."""
        return self._file_path

    @property
    def tabular_data_keys(self):
        """Get tuples of (ReportName, TableName, ReportForString) from tabular data."""
        if self._tabular_data_keys is None:
            with connect(self.file_path) as conn:
                query = "SELECT DISTINCT ReportName, TableName, ReportForString FROM TabularDataWithStrings"
                data = conn.execute(query)
                self._tabular_data_keys = data.fetchall()
        return self._tabular_data_keys

    @property
    def available_outputs(self) -> List[tuple]:
        """Get tuples (OutputName, ReportingFrequency) that can be requested.

        Any of these outputs when input to data_collections_by_output_name will
        yield a result with data collections.
        """
        if self._available_outputs is None:
            self._available_outputs = self._extract_available_outputs()
        return self._available_outputs

    @property
    def outputs(self):
        if self._outputs is None:
            self._outputs = _SqlOutputs(self.file_path, self.available_outputs)
        return self._outputs

    @property
    def zone_info(self):
        """Get a list of strings for available timeseries outputs that can be requested.

        Any of these outputs when input to data_collections_by_output_name will
        yield a result with data collections.
        """
        if self._zone_info is None:
            self._zone_info = self._extract_zone_info()
        return self._zone_info

    @property
    def surfaces_table(self):
        if self._surfaces_table is None:
            self._surfaces_table = self._extract_surfaces_table()
        return self._surfaces_table

    @property
    def constructions_table(self):
        if self._constructions_table is None:
            self._constructions_table = self._extract_constructions_table()
        return self._constructions_table

    @property
    def environment_periods(self):
        """Get a list of environment periods for the simulation run periods.

        EnvironmentType: An enumeration of the environment type. (1 = Design Day,
            2 = Design Run Period, 3 = Weather Run Period).
        """
        if self._environment_periods is None:
            self._environment_periods = self._extract_environment_periods()
        return self._environment_periods

    def full_html_report(self):
        """Get the html report as a dictionary of DataFrames.

        The dictionary keys are tuples of ("ReportName", "TableName",
        "ReportForString").
        """
        with connect(self.file_path) as conn:
            cols = (
                "ReportName, TableName, ReportForString, ColumnName, RowName, "
                "Units, Value"
            )
            query = f"SELECT {cols} FROM TabularDataWithStrings"
            data = pd.read_sql(query, conn)

        data.RowName = data.RowName.replace({"": np.NaN, "-": np.NaN})
        data.dropna(subset=["RowName"], inplace=True)

        all_df = {}
        for name, df in data.groupby(["ReportName", "TableName", "ReportForString"]):
            try:
                pivoted = df.pivot(
                    columns=["ColumnName", "Units"], index="RowName", values="Value"
                )
            except ValueError:
                # Cannot pivot; return long form
                pivoted = df
            # apply to_numeric column-wise
            pivoted = pivoted.apply(pd.to_numeric, errors="ignore")

            # insert in final dict.
            all_df[name] = pivoted

        return all_df

    def timeseries_by_name(
        self,
        variable_or_meter: Union[str, Sequence],
        reporting_frequency: Union[_REPORTING_FREQUENCIES] = "Hourly",
        environment_type: Union[Literal[1, 2, 3]] = 3,
    ) -> EnergyDataFrame:
        """Get an EnergyDataFrame for specified meters and/or variables.

        The returned DataFrame has a column MultiIndex with levels ["IndexGroup",
        "KeyValue", "Name"]. KeyValue corresponds to the zone name for variables
        while KeyValue is None for meters.

        Note that if an output name is not not available for the reporting
        frequency or the environment type, the DataFrame can be empty.

        Args:
            variable_or_meter (str or list): The name of an EnergyPlus output meter or
                variable to be retrieved from the SQLite result file. This can also be an
                array of output names for which all data collections should be retrieved.
            reporting_frequency (str): The reporting interval. One of ("HVAC System
                Timestep", "Zone Timestep", "Hourly", "Daily", "Monthly" or "Run
                Period"
            environment_type (int): The environment type. (1 = Design Day, 2 = Design
                Run Period, 3 = Weather Run Period). Default = 3.

        Returns:
            EnergyDataFrame: An EnergyDataFrame with the variable_or_meter as columns.
        """
        reporting_frequency = reporting_frequency.title()
        assert (
            reporting_frequency in Sql._reporting_frequencies
        ), f"reporting_frequency is not one of {Sql._reporting_frequencies}"
        with connect(self.file_path) as conn:
            cols = "ReportDataDictionaryIndex, IndexGroup, KeyValue, Name, Units, ReportingFrequency"
            if isinstance(variable_or_meter, str):  # assume it's a single output
                if (
                    variable_or_meter,
                    reporting_frequency,
                ) not in self.available_outputs:
                    log(
                        f"{(variable_or_meter, reporting_frequency)} not "
                        f"an available output in the Sql file.",
                        level=logging.WARNING,
                    )
                query = f"""
                        SELECT {cols} 
                        FROM ReportDataDictionary 
                        WHERE Name=@output_name 
                        AND ReportingFrequency=@reporting_frequency;
                        """
                header_rows = pd.read_sql(
                    query,
                    conn,
                    params={
                        "output_name": variable_or_meter,
                        "reporting_frequency": reporting_frequency,
                    },
                )
            elif len(variable_or_meter) == 1:  # assume it's a list
                query = f"""
                        SELECT {cols} 
                        FROM ReportDataDictionary 
                        WHERE Name=@output_name 
                        AND ReportingFrequency=@reporting_frequency;
                        """
                header_rows = pd.read_sql(
                    query,
                    conn,
                    params={
                        "output_name": variable_or_meter[0],
                        "reporting_frequency": reporting_frequency,
                    },
                )
            else:  # assume it is a list of outputs
                query = f"""
                        SELECT {cols} 
                        FROM ReportDataDictionary 
                        WHERE Name IN {tuple(variable_or_meter)}
                        AND ReportingFrequency=@reporting_frequency;"""
                header_rows = pd.read_sql(
                    query,
                    conn,
                    params={
                        "reporting_frequency": reporting_frequency,
                    },
                )
            # if nothing was found, return an empty DataFrame
            if len(header_rows) == 0:
                return EnergyDataFrame([])
            else:
                header_rows.set_index("ReportDataDictionaryIndex", inplace=True)

            # extract all data of the relevant type from ReportData
            rel_indices = tuple(header_rows.index.to_list())
            data = _extract_timeseries(conn, environment_type, header_rows, rel_indices)
        log(f"collected data for {variable_or_meter}")
        return data

    def tabular_data_by_name(
        self, report_name: str, table_name: str, report_for_string: Optional[str] = None
    ) -> pd.DataFrame:
        """Get (ReportName, TableName) data as DataFrame.

        Args:
            report_name (str): The name of the report.
            table_name (str): The name of the table in the report.
            report_for_string (str): The “For” string.

        Returns:
            (pd.DataFrame): A DataFrame.
        """
        with connect(self.file_path) as conn:
            cols = "RowName, ColumnName, Value, Units"
            query = f"""
                SELECT {cols} FROM TabularDataWithStrings 
                WHERE 
                    (@report_name IS NULL OR ReportName=@report_name)
                AND 
                    (@table_name IS NULL OR TableName=@table_name)
                AND 
                    (@report_for_string IS NULL OR ReportForString=@report_for_string);
            """
            data = pd.read_sql(
                query,
                conn,
                params={
                    "report_name": report_name,
                    "table_name": table_name,
                    "report_for_string": report_for_string,
                },
            )
            try:
                pivoted = data.pivot(
                    index="RowName", columns=["ColumnName", "Units"], values="Value"
                )
            except ValueError:
                # Cannot pivot; return long-form DataFrame
                pivoted = data
                log(
                    f"{(report_name, table_name, report_name)} cannot be "
                    f"pivoted as RowName and ColumnName. The long-form "
                    f"DataFrame has been returned.",
                    level=logging.WARNING,
                )
            pivoted = pivoted.apply(pd.to_numeric, errors="ignore")
        return pivoted

    def _extract_available_outputs(self) -> List:
        """Extract the list of all available outputs from the SQLite file."""
        with connect(self.file_path) as conn:
            cols = "Name, ReportingFrequency"
            query = f"SELECT DISTINCT {cols} FROM ReportDataDictionary"
            data = conn.execute(query)
            return data.fetchall()

    def _extract_zone_info(self):
        """Extract the Zones table from the SQLite file."""
        with connect(self.file_path) as conn:
            query = "SELECT * from Zones"
            df = pd.read_sql(query, conn).set_index("ZoneIndex")
        return df

    def _extract_surfaces_table(self):
        """Extract the Zones table from the SQLite file."""
        with connect(self.file_path) as conn:
            query = "SELECT * from Surfaces"
            df = pd.read_sql(query, conn).set_index(["ZoneIndex", "SurfaceIndex"])
        return df

    def _extract_constructions_table(self):
        with connect(self.file_path) as conn:
            query = "SELECT * from Constructions"
            df = pd.read_sql(query, conn).set_index("ConstructionIndex")
        return df

    def _extract_environment_periods(self):
        """Extract the EnvironmentPeriods table from the SQLite file."""
        with connect(self.file_path) as conn:
            query = "SELECT * from EnvironmentPeriods"
            df = pd.read_sql(query, conn).set_index("EnvironmentPeriodIndex")
        return df


def _extract_timeseries(
    conn, environment_type, header_rows, rel_indices
) -> EnergyDataFrame:
    """Extract time series given indices."""
    if len(rel_indices) == 1:
        data = pd.read_sql(
            """SELECT rd.Value,
                      rd.ReportDataDictionaryIndex, 
                      t.Month,
                      t.Day,
                      t.Hour,
                      t.Minute,
                      t.Interval
            FROM ReportData as rd 
                    LEFT JOIN Time As t ON rd.TimeIndex = t.TimeIndex
                    LEFT JOIN EnvironmentPeriods as p ON t.EnvironmentPeriodIndex = p.EnvironmentPeriodIndex
            WHERE ReportDataDictionaryIndex=@rel_indices
            AND (IFNULL(t.WarmupFlag, 0) = @warmup_flag)
            AND p.EnvironmentType = @environment_type
            ORDER BY t.TimeIndex;""",
            conn,
            params={
                "rel_indices": rel_indices[0],
                "warmup_flag": 0,
                "environment_type": environment_type,
            },
        )
    else:
        data = pd.read_sql(
            f"""SELECT rd.Value,
                          rd.ReportDataDictionaryIndex,
                          t.Month,
                          t.Day,
                          t.Hour,
                          t.Minute,
                          t.Interval
                FROM ReportData as rd
                        LEFT JOIN Time As t ON rd.TimeIndex = t.TimeIndex
                        LEFT JOIN EnvironmentPeriods as p ON t.EnvironmentPeriodIndex = p.EnvironmentPeriodIndex
                WHERE ReportDataDictionaryIndex IN {tuple(rel_indices)}
                AND (IFNULL(t.WarmupFlag, 0) = @warmup_flag)
                AND p.EnvironmentType = @environment_type
                ORDER BY rd.ReportDataDictionaryIndex, t.TimeIndex;""",
            conn,
            params={"warmup_flag": 0, "environment_type": environment_type},
        )
    # Join the header_rows on ReportDataDictionaryIndex
    data = data.join(
        header_rows[["IndexGroup", "KeyValue", "Name"]],
        on="ReportDataDictionaryIndex",
    )
    # Pivot the data so that ["Name", "KeyValue"] becomes the column MultiIndex.
    data = data.pivot(
        index=["Month", "Day", "Hour", "Minute", "Interval"],
        columns=["IndexGroup", "KeyValue", "Name"],
        values="Value",
    )
    # reset the index to prepare the DatetimeIndex
    date_time_names = data.index.names
    data.reset_index(inplace=True)
    index = to_datetime(
        {
            "year": 2018,
            "month": data.Month,
            "day": data.Day,
            "hour": data.Hour,
            "minute": data.Minute,
        }
    )
    # Adjust timeindex by timedelta
    index -= data["Interval"].apply(lambda x: timedelta(minutes=x))
    index = pd.DatetimeIndex(index, freq="infer")
    # get data
    data = data.drop(columns=date_time_names, level="IndexGroup")
    data.index = index
    # Create the EnergyDataFrame and set the units using dict
    data = EnergyDataFrame(data)
    data.units = header_rows.set_index(["IndexGroup", "KeyValue", "Name"])[
        "Units"
    ].to_dict()
    return data
