import pandas as pd
from .type_conversions import (convert_date_to_datetime,
                               from_pandas_timestamp_to_datetime_date,
                               from_datetime_date_to_pandas_timestamp,
                               from_datetime_to_dss_string_date
                               )
from .computations import compute_future_date, compute_antecedent_date


def generate_time_dimension_dataframe(start_date: str, end_date:str , time_granularity: str,
                                      output_format: str="pandas_timestamp"):
    """
    Generates a time dimension dataframe based on a time granularity.
    
    :param start_date: str: The start date of the time dimension.
        Accepted formats are:
            - yyyy-MM-dd.
            - Dataiku non-ambiguous dates format, also known as 'ISO-8601'.
    :param end_date: str: The end date of the time dimension.
    :param time_granularity: str: The granularity of time dimension.
        Accepted time granularities are ["days", "weeks", "months", "years"].
    :param output_format: str: The dates format expected in dimension. Available choices are:
        - pandas_timestamp [Default value]: a pandas timestamp
        - dss_string_date: Dataiku non-ambiguous dates format, also known as 'ISO-8601'

    :returns: time_dimension_dataframe: pandas.core.frame.DataFrame: The time dimension dataframe.
    """
    ALLOWED_TIME_GRANULARITIES = ["days", "weeks", "months", "years"]
    if time_granularity not in ALLOWED_TIME_GRANULARITIES:
        log_message = f"Time granularity '{time_granularity}' is not accepted in this function."\
        f"Please choose a time granularity that is in '{ALLOWED_TIME_GRANULARITIES}'."
        raise ValueError(log_message)
    
    start_date = convert_date_to_datetime(start_date)
    end_date = convert_date_to_datetime(end_date)

    if start_date > end_date:
        raise ValueError("You set wrong parameters: "\
                         f"start_date ({start_date}) > end_date ({end_date})!")

    if time_granularity in ["weeks", "months"]:
        start_date = compute_antecedent_date(start_date, time_granularity, 1)

    if time_granularity == "days":
        date_range = pd.date_range(start=start_date, end=end_date, freq="D")
        period_numbers = date_range.dayofyear
    
    elif time_granularity == "weeks":
        date_range = pd.date_range(start=start_date, end=end_date, freq="W-MON")
        period_numbers = date_range.weekofyear
    
    elif time_granularity == "months":
        date_range = pd.date_range(start=start_date, end=end_date, freq="MS")
        period_numbers = date_range.month
    
    elif time_granularity == "years":
        date_range = pd.date_range(start=start_date, end=end_date, freq="YS")
        period_numbers = date_range.year
    
    else:
        raise ValueError('Invalid time_granularity. Choose from "days", "weeks", "months", or "years".')
    
    period_years = date_range.year
    date_range_datetime = [
            from_pandas_timestamp_to_datetime_date(date) for date in date_range
        ]
    if time_granularity == "days":
        end_dates = date_range
    else:
        end_dates = []
        end_dates_datetime = []
        for start_date in date_range_datetime:
            next_period_start_date = compute_future_date(start_date, time_granularity, 1)
            end_date = compute_antecedent_date(next_period_start_date, "days", 1)
            end_dates.append(from_datetime_date_to_pandas_timestamp(end_date))
            end_dates_datetime.append(end_date)       

    time_dimension_dataframe = pd.DataFrame({
        "year": period_years,
        "period_number": period_numbers,
        "start_date": date_range,
        "end_date": end_dates
    })
    time_dimension_dataframe = time_dimension_dataframe[["year", "period_number", "start_date", "end_date"]]
    
    if time_granularity == "days":
        time_dimension_dataframe = time_dimension_dataframe.rename({"period_number": "day_of_year"}, axis=1)
    elif time_granularity == "weeks":
        time_dimension_dataframe = time_dimension_dataframe.rename({"period_number": "week_of_year"}, axis=1)
    elif time_granularity == "months":
        time_dimension_dataframe = time_dimension_dataframe.rename({"period_number": "month_of_year"}, axis=1)
    else:
        time_dimension_dataframe.drop("period_number", axis=1, inplace=True)
    
    if output_format == "dss_string_date":
        new_start_dates = []
        new_end_dates = []
        for start_date, end_date in zip(date_range_datetime, end_dates_datetime):
            new_start_dates.append(from_datetime_to_dss_string_date(start_date))
            new_end_dates.append(from_datetime_to_dss_string_date(end_date))
        time_dimension_dataframe["start_date"] = new_start_dates
        time_dimension_dataframe["end_date"] = new_end_dates
    return time_dimension_dataframe