# BSD 2-CLAUSE LICENSE

# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:

# Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# #ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# Original author: Reza Hosseini

"""The goal is to get "Expt" data which is experiment assignment data.
It also joins several expts if needed."""

import sys
from typing import List, Optional

from abvelocity.core.get_data.cursor import Cursor
from abvelocity.core.get_data.data_container import DataContainer
from abvelocity.core.get_data.get_expt_stats import get_expt_stats
from abvelocity.core.param.expt_info import ExptInfo, MultiExptInfo


def get_asgmnt_data(
    expt_info: ExptInfo,
    cursor: Optional[Cursor] = None,
    condition: Optional[str] = None,
    materialize_to_df: bool = True,
) -> DataContainer:
    """Returns experiment assignment data as a DataContainer.
    The query needs to return a schema with:

        - expt_unit_id
        - variant

    Optionally it might include

        - minimum trigger time and other aggregate info.

    It is important to make sure that the data is aggregated at the expt_unit_id and variant level

    Args:
        expt_info: Experiment info.
            See `~abvelocity.expt_info.ExptInfo` for more details.
        cursor: A Presto Cursor.
        condition: A SQL condition.
        materialize_to_df: A bool to determine if we want to materialize to pandas df.

    Returns:
        result: DataContainer with experiment assignment data.

    Raises:
        None.
    """
    query = expt_info.query
    if query is None:
        expt_info.gen_query(condition=condition)
        query = expt_info.query
        print("\n***: In `get_asgmnt_data` query was generated by calling `gen_query`," f"generated query for expt_name: {expt_info.name} :\n {query}")
    else:
        print(
            f"\n***: User has passed a direct query to get experiment data. "
            f"\n***: It is important to make sure that the data is aggregated at the expt_unit_id and variant level."
            f"\n***: query: {query}"
        )

    if not materialize_to_df:
        return DataContainer(query=query, is_pandas_df=False)

    result = cursor.get_df(query=query)

    return DataContainer(df=result.df, is_pandas_df=True)


def get_asgmnt_data_multi(
    multi_expt_info: MultiExptInfo,
    cursor: Optional[Cursor] = None,
    condition: Optional[str] = None,
    drop_unit_duplicates: bool = True,
    materialize_to_df: bool = True,
) -> List[DataContainer]:
    """This function queries expt data and returns a list of DataContainers.

    Args:
        multi_expt_info : `MultiExptInfo`
            A MultiExptInfo object with experiment information.
        cursor: A databse Cursor.
        condition: A SQL condition.
        drop_unit_duplicates : If True, this will only keep one row per unit (`expt_info.expt_unit_col`) for each experiment.
            It keeps the last row, but note that this is a rather arbitrary choice.
            Best practice is to pass data without duplicates in `expt_info.expt_unit_col` to this function.
            This argument is only supported when we do materialize to df.
        materialize_to_df: A bool to determine if we want to materialize to pandas df.

    Returns
        expt_dc_list : A list of DataContainers, one for each experiment in `multi_expt_info.expt_info_list`.

    Alters:
        expt_info_list : This function will amend each `expt_info` in the list with
        it's derived stats (`DerivedExptState`)
    """

    # Queries expt data for each expt in `multi_expt_info.expt_info_list`.
    expt_dc_list = []
    for expt_info in multi_expt_info.expt_info_list:
        dc = get_asgmnt_data(
            cursor=cursor,
            expt_info=expt_info,
            condition=condition,
            materialize_to_df=materialize_to_df,
        )
        print(f"***\n Retrieved expt assignment data for experiment dc: {dc}")
        if materialize_to_df:
            expt_df = dc.df

            if drop_unit_duplicates and expt_info.expt_unit_col is not None:
                expt_df.drop_duplicates(subset=[expt_info.expt_unit_col], keep="first", inplace=True)
            # Get statistics about one individual experiment
            expt_info.derived_stats = get_expt_stats(expt_df)

            expt_dc_list.append(DataContainer(pandas_df=expt_df, is_pandas_df=True))
        else:
            expt_dc_list.append(dc)

    size_in_megabytes = sys.getsizeof(expt_dc_list) / 10**6
    print(f"`expt_data` size in MB: {size_in_megabytes}")

    return expt_dc_list
