import pandas as pd
from pandas import Timestamp

from tadas.domain import features


def test_combine_data_sources():
    data_df = pd.DataFrame({
        'report_date': ['2025-09-20', '2025-09-20'],
        'isrc_cd': ['ZAB7A1600105', 'ZAB7A1800292'],
        'geo_country': ['US', 'US'],
        'pfn_geo': ['US_ZAB7A1600105', 'US_ZAB7A1800292'],
        'wow_amazon_search_streams': [1, 2],
        'wow_tt_creations': [3, 4],
    })
    days_trending_data = [
        {
            'report_date': '2025-09-20',
            'isrc_cd': 'ZAB7A1600105',
            'consecutive_trend': 10.0,
            'pfn_geo': 'US_ZAB7A1600105',
            'combined_forecast': 0.2,
            'vs_forecast_lift': 0.2},
        {
            'report_date': '2025-09-19',
            'isrc_cd': 'ZAB7A1600105',
            'consecutive_trend': 9.0,
            'pfn_geo': 'US_ZAB7A1600105',
            'combined_forecast': 0.2,
            'vs_forecast_lift': 0.2},
        {
            'report_date': '2025-09-18',
            'isrc_cd': 'ZAB7A1600105',
            'consecutive_trend': 8.0,
            'pfn_geo': 'US_ZAB7A1600105',
            'combined_forecast': 0.2,
            'vs_forecast_lift': 0.2},

        {
            # missing trending data for US_ZAB7A1800292 on '2025-09-20'
            'report_date': '2025-09-19',
            'isrc_cd': 'ZAB7A1800292',
            'consecutive_trend': 11.0,
            'pfn_geo': 'US_ZAB7A1800292',
            'combined_forecast': 0.2,
            'vs_forecast_lift': 0.2
        },
    ]
    days_trending_df = pd.DataFrame(days_trending_data)
    days_trending_df["report_date"] = pd.to_datetime(days_trending_df["report_date"], format="%Y-%m-%d")

    result = features.combine_data_sources(
        data_df=data_df,
        days_trending_df=days_trending_df,
        report_date='2025-09-20'
    )

    assert result.to_dict(orient='records') == [
        {
            'report_date': '2025-09-20',
            'isrc_cd': 'ZAB7A1600105',
            'consecutive_trend': 10.0,
            'combined_forecast': 0.2,
            'pfn_geo': 'US_ZAB7A1600105',
            'geo_country': 'US',
            'wow_amazon_search_streams': 1,
            'vs_forecast_lift': 0.2,
            'wow_tt_creations': 3},
        {
            'report_date': '2025-09-20',
            'isrc_cd': 'ZAB7A1800292',
            'consecutive_trend': 0,
            'combined_forecast': 0,
            'pfn_geo': 'US_ZAB7A1800292',
            'geo_country': 'US',
            'wow_amazon_search_streams': 2,
            'vs_forecast_lift': 0,
            'wow_tt_creations': 4},
    ]


def test_apply_adj_to_dod():
    data = [
        {'report_date': '2025-09-20',
         'isrc_cd': 'ZAB7A1600105',
         'dod_col_1': 10,
         'dod_col_2': 3},
        {'report_date': '2025-09-21',
         'isrc_cd': 'ZAB7A16001192',
         'dod_col_1': 11,
         'dod_col_2': 32},
    ]

    df = pd.DataFrame(data)
    df["report_date"] = pd.to_datetime(df["report_date"], format="%Y-%m-%d")

    result = features.apply_adj_to_dod(data_df=df)

    assert result.to_dict(orient='records') == [
        {'adj_dod_col_1': -1.551765520576064,
         'adj_dod_col_2': -0.4655296561728193,
         'dod_col_1': 10,
         'dod_col_2': 3,
         'isrc_cd': 'ZAB7A1600105',
         'pct_change_vs_prev': -6.444272583326691,
         'report_date': Timestamp('2025-09-20 00:00:00'),
         'weekday': 'Saturday',
         'weekday_num': 5},
        {'adj_dod_col_1': -0.9139593989016872,
         'adj_dod_col_2': -2.6587909786230903,
         'dod_col_1': 11,
         'dod_col_2': 32,
         'isrc_cd': 'ZAB7A16001192',
         'pct_change_vs_prev': -12.035545575896252,
         'report_date': Timestamp('2025-09-21 00:00:00'),
         'weekday': 'Sunday',
         'weekday_num': 6}]