# -*- coding: utf-8 -*-
"""
seasonality-table: Seasonality Returns Table Generator
A library to generate monthly returns tables for financial analysis.
"""

import pandas as pd
import yfinance as yf
from datetime import datetime

__version__ = "0.1.1"
__author__ = "cagrigungor"
__email__ = "iletisim@cagrigungor.com"

__all__ = ["get_returns_table"]


def get_returns_table(
    ticker="BTC-USD",
    is_colored=True,
    start_date=None,
    end_date=None,
    column="Close",
    date="index",
    df=None
):
    """
    Generate a monthly returns table (seasonality table) for financial data.
    
    Parameters
    ----------
    ticker : str, optional
        The ticker symbol to fetch data for. Default is "BTC-USD".
    is_colored : bool, optional
        If True, returns a styled DataFrame with color coding. Default is True.
    start_date : str, optional
        Start date in "YYYY-MM-DD" format. Default is "2010-01-01".
    end_date : str, optional
        End date in "YYYY-MM-DD" format. Default is today.
    column : str, optional
        The column to use for calculations. Default is "Close".
    date : str, optional
        The date column name or "index" if date is in index. Default is "index".
    df : pd.DataFrame, optional
        Custom DataFrame to use instead of fetching from yfinance.
    
    Returns
    -------
    pd.DataFrame or pd.io.formats.style.Styler
        Monthly returns pivot table. If is_colored=True, returns a styled DataFrame.
    
    Examples
    --------
    >>> from sessionalty_table import get_returns_table
    >>> table = get_returns_table("BTC-USD", is_colored=False)
    >>> table = get_returns_table("AAPL", start_date="2020-01-01")
    """
    
    # 1. Determine data source
    if df is None:
        if start_date is None:
            start_date = "2010-01-01"
        if end_date is None:
            end_date = datetime.today().strftime('%Y-%m-%d')
        
        data = yf.download(ticker, start=start_date, end=end_date, progress=False)
        
        # MultiIndex fix
        if isinstance(data.columns, pd.MultiIndex):
            data.columns = data.columns.get_level_values(0)
    else:
        data = df.copy()

    if data.empty:
        return "No data found!"

    # 2. Set date column
    if date != "index":
        data = data.set_index(date)
    
    # Convert index to datetime (required for resample)
    data.index = pd.to_datetime(data.index)

    # 3. Calculate monthly returns
    monthly_data = data[column].resample('ME').last()
    monthly_returns = monthly_data.pct_change() * 100
    
    # 4. Convert to matrix format
    df_returns = monthly_returns.reset_index()
    df_returns.columns = ['Date', 'Returns']
    df_returns['Year'] = df_returns['Date'].dt.year
    df_returns['Month'] = df_returns['Date'].dt.month
    
    pivot_table = df_returns.pivot(index='Year', columns='Month', values='Returns')
    
    # Assign month names to columns
    month_names = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
    pivot_table.columns = [month_names[m-1] for m in pivot_table.columns]

    # 5. Display settings
    if is_colored:
        def color_logic(val):
            if pd.isna(val):
                return ''
            color = '#2ecc71' if val > 0 else '#e74c3c'
            return f'background-color: {color}; color: white; font-weight: bold; border: 1px solid white;'
        
        return pivot_table.style.map(color_logic).format("{:.2f}%", na_rep="-")
    else:
        return pivot_table.round(2).fillna("-")
