from openpyxl.worksheet.worksheet import Worksheet
from pandas import DataFrame

# Define a function to write a pandas dataframe to excel

def write_to_excel(ws: Worksheet, start_row: int,
                   start_column: int, df: DataFrame,
                   index: bool = True, columns: bool = True) -> None:
    """
    Function which takes a worksheet and writes a table of data to it.

    Parameters
    ----------
    ws : Worksheet
        The worksheet to modify
    start_row : int
        The row to begin on
    start_column : int
        The column to start on
    df : DataFrame
        The data to write
    index : bool, optional
        Whether to write the index labels. The default is True.
    columns : bool, optional
        Whether to write the column labels. The default is True.

    """
    
    # If we want to include the index, reset the index on the dataframe first
    if index:
        write_data = df.reset_index().values.tolist()
    else:
        write_data = df.values.tolist()
    
    # If we want to include columns we need to add the columns as the first in the nested lists
    if columns:
        # If we already added an index, then we are going to want also include a blank space since
        # we don't want a label above the index
        if index:
            write_data = [[""]+list(df.columns)]+write_data
        else:
            write_data = [list(df.columns)]+write_data
    
    # For every value iterate through finding the location to place it in
    for row in range(len(write_data)):
        for column in range(len(write_data[0])):
            ws.cell(row=start_row+row,column=start_column+column).value = write_data[row][column]

wb = openpyxl.Workbook()
ws = wb.active

#Examples of writing to excel
write_to_excel(ws, 1, 1, yield_curve_table)
write_to_excel(ws, 16, 1, yield_curve_table,index=False)
write_to_excel(ws, 31, 1, yield_curve_table,columns=False)
write_to_excel(ws, 46, 1, yield_curve_table,index=False,columns=False)

wb.save("Workbooks/Example 2.xlsx")

from openpyxl.styles import PatternFill

# Create a workbook
wb = openpyxl.Workbook()
ws = wb.active

# Write the table in
write_to_excel(ws, 3, 1, yield_curve_table)

# PatternFill is a way to get colors, here we pass the start and end color as the same and use a fill type solid
# The color used here is red
my_fill_color = PatternFill(start_color='FF0000',
                   end_color='FF0000',
                   fill_type='solid')

# The fill attribute of a cell is the color to use
ws.cell(row=1,column=1).fill = my_fill_color

wb.save("Workbooks/Example 3.xlsx")

#Create a function to map red and green colors based on whether a field is positive or negative
def red_green_map(ws: Worksheet, df: DataFrame, start_row: int, start_column: int) -> None:
    """
    Function which colors cells green or red depending on if
    they are positive or negative

    Parameters
    ----------
    ws : Worksheet
        The worksheet to modify
    df : DataFrame
        The dataframe of data
    start_row : int
        The row to start coloring on
    start_column : int
        The color to start coloring on

    """
    
    # Get the values
    write_data = df.values.tolist()
    
    # Iterate over the rows and columns
    for row in range(len(write_data)):
        for column in range(len(write_data[0])):
            # Set the fill color to green for positive numbers
            if write_data[row][column] > 0:
                fill = fill = PatternFill(start_color='008000',
                   end_color='008000',
                   fill_type='solid')
            # Set the fill color to red for negative numbers
            elif write_data[row][column] < 0:
                fill = PatternFill(start_color='FF0000',
                   end_color='FF0000',
                   fill_type='solid')
                
            # Change the color for the cell
            ws.cell(row=start_row+row,column=start_column+column).fill = fill

# Create the excel workbook
wb = openpyxl.Workbook()
ws = wb.active

#Write the data to excel
write_to_excel(ws, 3, 1, yield_curve_table)

# We'll deal with only the two change columns
# So pass the data for only those and also pass their starting columns/rows
red_green_map(ws, yield_curve_table[["1 Year Change"]], 4, 4)
red_green_map(ws, yield_curve_table[["5 Year Change"]], 4, 6)

wb.save("Workbooks/Example 4.xlsx")

#Put it all together for a function that takes a worksheet and yield curve data then writes the sheet
def write_treasury_sheet(ws: Worksheet, yield_curve_table: DataFrame):
    """
    Function to fill out the treasury curve data

    Parameters
    ----------
    ws : Worksheet
        The worksheet to modify
    yield_curve_table : DataFrame
        The yield curve data


    """
    # Write the yield curve table
    write_to_excel(ws, 3, 1, yield_curve_table)
    
    # Color the yield curve changes red/green
    red_green_map(ws, yield_curve_table[["1 Year Change"]], 4, 4)
    red_green_map(ws, yield_curve_table[["5 Year Change"]], 4, 6)
    
wb = openpyxl.Workbook()
ws = wb.active
write_treasury_sheet(ws, yield_curve_table)
wb.save("Workbooks/Economic Report.xlsx")