# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# -*- coding: utf-8 -*-
import dataiku
import pandas as pd, numpy as np
from dataiku import pandasutils as pdu
import plotly.graph_objects as go
from dataiku import insights
import io
from datetime import datetime

# Read recipe inputs
waterfall = dataiku.Dataset("waterfall")
df = waterfall.get_dataframe()

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
now = datetime.now()
current_time = now.strftime("%Y-%m-%d_%H:%M:%S")

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
total_row = {'configs_stage': '5. closing_balance', 'ECL_change': df['ECL_change'].sum()}
df = df.append(total_row, ignore_index=True)

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
def format_currency(number):
    # Check if the number is negative
    is_negative = False
    if number < 0:
        is_negative = True
        number = abs(number)

    # Define the suffixes for each scale (thousand, million, billion, etc.)
    suffixes = ['', 'K', 'M', 'B', 'T']

    # Determine the appropriate scale to use
    scale_index = 0
    while number >= 1000 and scale_index < len(suffixes) - 1:
        number /= 1000
        scale_index += 1

    # Format the number with two decimal places
    formatted_number = "{:.2f}".format(number)

    # Remove trailing zeros and decimal point if the number is an integer
    formatted_number = formatted_number.rstrip('0').rstrip('.')

    # Add the appropriate suffix
    formatted_number += suffixes[scale_index]

    # Add the currency prefix
    formatted_number = f"USD {formatted_number}"

    # Add the negative sign if applicable
    if is_negative:
        formatted_number = "-" + formatted_number

    return formatted_number

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
def waterfall_plotly(categories, values):
    # Calculate cumulative values
    #cumsum = [sum(values[:i+1]) for i in range(len(values))]

    # Create the waterfall trace
    trace = go.Waterfall(
        x=categories,
        y=values,
        text=[format_currency(val) for val in values],
        measure=['relative']*(len(values)-1) + ['total'],  # Sets the bar type to relative
        increasing=dict(marker_color='#2ca02c'),  # Sets the color for positive changes to green
        decreasing=dict(marker_color='#d62728'),  # Sets the color for negative changes to red
        totals=dict(marker_color='#1f77b4', marker_line_color='#1f77b4'),  # Sets the color for totals to blue
    )

    # Create the layout
    layout = go.Layout(
        yaxis=dict(title='Change Amount'),
        showlegend=False,
        margin=dict(l=0, r=0, t=0, b=0),
    )

    # Create the figure and add the trace
    fig = go.Figure(data=trace, layout=layout)

    # Show the plot
    fig.show()
    return fig

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
waterfall = waterfall_plotly(df['configs_stage'].tolist(), df['ECL_change'].tolist())

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
img = waterfall.to_image('PNG')

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
stream = io.BytesIO(img)

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
folder_handle = dataiku.Folder("gLrPKkLQ")
folder_handle.upload_stream("/waterfall_" + current_time + ".png", stream.getvalue())

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
insights.save_plotly("waterfall", waterfall)